Torchdata 0.5.0 is seemingly a lot slower and more memory usage than 0.4.1

jacobbieker: <p>We are heavily using torchdata and datapipes at Open Climate Fix in our <a href="https://github.com/openclimatefix/ocf_datapipes" class="inline-onebox" rel="noopener nofollow ugc">GitHub - openclimatefix/ocf_datapipes: OCF's DataPipe based dataloader for training and inference</a> repo. We have it working well with torchdata 0.4.1, but when upgrading to the newer PyTorch and torchdata 0.5.0, our CI tests timeout and fail because of memory issues. When running the same tests locally, with torchdata 0.4.1, 107 tests pass in 157 seconds, while with torchdata 0.5.0 its still running 10 minutes later with only 11 tests passed so far. On GitHub Actions, the tests timeout and run out of memory when updating to the newer version.</p>
<p>Looking through the release notes, there does not seem to be any real reason we can see why this is happening. We are not using any of the removed features, for example. Is there any changes between 0.4.1 and 0.5.0 not in the release notes that might have caused this?</p>

weiji14: <p>There are lots of differences in library versions in <a href="https://github.com/openclimatefix/ocf_datapipes/pull/86" class="inline-onebox" rel="noopener nofollow ugc">Update Requirements by jacobbieker · Pull Request #86 · openclimatefix/ocf_datapipes · GitHub</a>, and I see some libraries like <code>rioxarray</code>, <code>xarray</code> and <code>fsspec</code> has been downgraded between the main branch at <a href="https://github.com/openclimatefix/ocf_datapipes/actions/runs/3533593852/jobs/5929417140#step:6:83" class="inline-onebox" rel="noopener nofollow ugc">Merge branch 'issue/datamodule' · openclimatefix/ocf_datapipes@807ed6a · GitHub</a> and in the PR at <a href="https://github.com/openclimatefix/ocf_datapipes/actions/runs/3392566007/jobs/5638910278#step:6:73" class="inline-onebox" rel="noopener nofollow ugc">[pre-commit.ci] auto fixes from pre-commit.com hooks · openclimatefix/ocf_datapipes@f454c8b · GitHub</a>. Can you run a <code>pip list</code> for each environment (with torchdata 0.4.1 and 0.5.0) so that we know what the actual differences in library versions are and better isolate the cause of the timeouts?</p>
<p>Also, It seems like there are <code>rioxarray.open_rasterio</code> calls (e.g. at <a href="https://github.com/openclimatefix/ocf_datapipes/blob/eb08124fdfb4deee23438984562dd7d1dd605bd5/ocf_datapipes/load/topographic.py#L25-L27" class="inline-onebox" rel="noopener nofollow ugc">ocf_datapipes/topographic.py at eb08124fdfb4deee23438984562dd7d1dd605bd5 · openclimatefix/ocf_datapipes · GitHub</a>) that are not made in a context manager, which might make it prone to memory leaks. There was a recent fix at <a href="https://github.com/corteva/rioxarray/releases/tag/0.13.1" class="inline-onebox" rel="noopener nofollow ugc">Release Release 0.13.1 · corteva/rioxarray · GitHub</a>, but not sure if it’s related to your issue here. Oh, and in addition to that, you might want to consider using <a href="https://pytorch.org/data/0.5/generated/torchdata.datapipes.utils.StreamWrapper.html" class="inline-onebox" rel="noopener nofollow ugc">StreamWrapper — TorchData 0.5.0 (beta) documentation</a> for some of your datapipes to close the files properly. See e.g. <a href="https://github.com/weiji14/zen3geo/blob/a71639f9476dcf423dc806f4270f04938aa62739/zen3geo/datapipes/rioxarray.py#L72-L74" class="inline-onebox" rel="noopener nofollow ugc">zen3geo/rioxarray.py at a71639f9476dcf423dc806f4270f04938aa62739 · weiji14/zen3geo · GitHub</a> for an example.</p>

jacobbieker: <p>Hello,</p>
<p>Thanks for all the detail! I’ll start making some of those changes, and I’ve updated that branch to match main now. At the moment, for torchdata 0.4.1 the <code>pip list</code> is</p>
<pre><code class="lang-auto">Package                 Version       Editable project location
----------------------- ------------- -------------------------------------
affine                  2.3.1
aiohttp                 3.8.3
aiosignal               1.3.1
alembic                 1.8.1
asciitree               0.3.3
async-timeout           4.0.2
attrs                   22.1.0
bokeh                   2.4.3
Bottleneck              1.3.5
branca                  0.6.0
brotlipy                0.7.0
cached-property         1.5.2
Cartopy                 0.21.0
certifi                 2022.9.24
cffi                    1.15.1
charset-normalizer      2.1.1
click                   8.1.3
click-plugins           1.1.1
cligj                   0.7.2
cloudpickle             2.2.0
configobj               5.0.6
contourpy               1.0.6
cramjam                 2.6.2
cryptography            38.0.3
cycler                  0.11.0
cytoolz                 0.12.0
dask                    2022.11.1
distributed             2022.11.1
einops                  0.6.0
entrypoints             0.4
exceptiongroup          1.0.4
fasteners               0.18
fastparquet             2022.11.0
Fiona                   1.8.22
fire                    0.4.0
folium                  0.13.0
fonttools               4.38.0
freezegun               1.2.2
frozenlist              1.3.3
fsspec                  2022.11.0
GDAL                    3.5.3
geopandas               0.12.1
gitdb                   4.0.10
GitPython               3.1.29
greenlet                2.0.1
h5netcdf                0.0.0
h5py                    3.7.0
HeapDict                1.0.1
idna                    3.4
imagecodecs             2022.9.26
iniconfig               1.1.1
Jinja2                  3.1.2
joblib                  1.2.0
jpeg-xl-float-with-nans 0.0.4
kiwisolver              1.4.4
lightning-utilities     0.3.0
locket                  1.0.0
lz4                     4.0.2
Mako                    1.2.4
mapclassify             2.4.3
MarkupSafe              2.1.1
matplotlib              3.6.2
msgpack                 1.0.4
multidict               6.0.2
munch                   2.5.0
munkres                 1.1.4
networkx                2.8.8
nowcasting-datamodel    1.1.54
numcodecs               0.10.2
numpy                   1.23.5
ocf-datapipes           0.5.30        /home/jacob/Development/ocf_datapipes
packaging               21.3
pandas                  1.5.2
partd                   1.3.0
pathy                   0.10.0
Pillow                  9.2.0
pip                     22.3.1
pluggy                  1.0.0
portalocker             2.6.0
protobuf                3.20.1
psutil                  5.9.4
psycopg2-binary         2.9.5
pvlib                   0.9.3
pvlive-api              0.11
pyaml-env               1.2.0
pycparser               2.21
pydantic                1.10.2
pykdtree                1.3.6
pyOpenSSL               22.1.0
pyparsing               3.0.9
pyproj                  3.4.0
pyresample              1.25.1
pyshp                   2.3.1
PySocks                 1.7.1
pytest                  7.2.0
python-dateutil         2.8.2
pytorch-lightning       1.8.3.post0
pytz                    2022.6
PyYAML                  5.4.1
rasterio                1.3.3
requests                2.28.1
rioxarray               0.13.1
Rtree                   1.0.1
scikit-learn            1.1.3
scipy                   1.9.3
setuptools              65.5.1
Shapely                 1.8.5.post1
six                     1.16.0
smart-open              5.2.1
smmap                   5.0.0
snuggs                  1.4.7
sortedcontainers        2.4.0
SQLAlchemy              1.4.44
tblib                   1.7.0
tensorboardX            2.5.1
termcolor               2.1.1
threadpoolctl           3.1.0
tomli                   2.0.1
toolz                   0.12.0
torch                   1.12.0
torchdata               0.4.1+f9ecd8b
torchmetrics            0.10.3
torchvision             0.13.0
tornado                 6.1
tqdm                    4.64.1
typer                   0.7.0
typing_extensions       4.4.0
unicodedata2            15.0.0
urllib3                 1.26.13
wheel                   0.38.4
xarray                  2022.11.0
xyzservices             2022.9.0
yarl                    1.8.1
zarr                    2.13.3
zict                    2.2.0
</code></pre>
<p>And for torchdata 0.5.0 is:</p>
<pre><code class="lang-auto">Package                 Version     Editable project location
----------------------- ----------- -------------------------------------
affine                  2.3.1
aiohttp                 3.8.3
aiosignal               1.3.1
alembic                 1.8.1
asciitree               0.3.3
async-timeout           4.0.2
attrs                   22.1.0
bokeh                   2.4.3
Bottleneck              1.3.5
branca                  0.6.0
brotlipy                0.7.0
cached-property         1.5.2
Cartopy                 0.21.0
certifi                 2022.9.24
cffi                    1.15.1
charset-normalizer      2.1.1
click                   8.1.3
click-plugins           1.1.1
cligj                   0.7.2
cloudpickle             2.2.0
configobj               5.0.6
contourpy               1.0.6
cramjam                 2.6.2
cryptography            38.0.3
cycler                  0.11.0
cytoolz                 0.12.0
dask                    2022.11.1
distributed             2022.11.1
einops                  0.6.0
entrypoints             0.4
exceptiongroup          1.0.4
execnet                 1.9.0
fasteners               0.18
fastparquet             2022.11.0
Fiona                   1.8.22
fire                    0.4.0
folium                  0.13.0
fonttools               4.38.0
freezegun               1.2.2
frozenlist              1.3.3
fsspec                  2022.11.0
GDAL                    3.5.3
geopandas               0.12.1
gitdb                   4.0.9
GitPython               3.1.29
greenlet                2.0.1
h5netcdf                0.0.0
h5py                    3.7.0
HeapDict                1.0.1
idna                    3.4
imagecodecs             2022.9.26
iniconfig               1.1.1
Jinja2                  3.1.2
joblib                  1.2.0
jpeg-xl-float-with-nans 0.0.4
kiwisolver              1.4.4
lightning-utilities     0.3.0
locket                  1.0.0
lz4                     4.0.2
Mako                    1.2.4
mapclassify             2.4.3
MarkupSafe              2.1.1
matplotlib              3.6.2
msgpack                 1.0.4
multidict               6.0.2
munch                   2.5.0
munkres                 1.1.4
networkx                2.8.8
nowcasting-datamodel    1.1.54
numcodecs               0.10.2
numpy                   1.23.5
ocf-datapipes           0.5.30      /home/jacob/Development/ocf_datapipes
packaging               21.3
pandas                  1.5.2
partd                   1.3.0
pathy                   0.9.0
Pillow                  9.2.0
pip                     22.3.1
pluggy                  1.0.0
portalocker             2.6.0
protobuf                3.20.1
psutil                  5.9.4
psycopg2-binary         2.9.5
pvlib                   0.9.3
pvlive-api              0.11
pyaml-env               1.2.0
pycparser               2.21
pydantic                1.10.2
pykdtree                1.3.6
pyOpenSSL               22.1.0
pyparsing               3.0.9
pyproj                  3.4.0
pyresample              1.25.1
pyshp                   2.3.1
PySocks                 1.7.1
pytest                  7.2.0
pytest-timeout          2.1.0
pytest-xdist            3.0.2
python-dateutil         2.8.2
pytorch-lightning       1.8.3.post0
pytz                    2022.6
PyYAML                  5.4.1
rasterio                1.3.3
requests                2.28.1
rioxarray               0.13.1
Rtree                   1.0.1
scikit-learn            1.1.3
scipy                   1.9.3
setuptools              65.5.1
Shapely                 1.8.5.post1
six                     1.16.0
smart-open              5.2.1
smmap                   5.0.0
snuggs                  1.4.7
sortedcontainers        2.4.0
SQLAlchemy              1.4.44
tblib                   1.7.0
tensorboardX            2.5.1
termcolor               2.1.1
threadpoolctl           3.1.0
tomli                   2.0.1
toolz                   0.12.0
torch                   1.13.0
torchdata               0.5.0
torchmetrics            0.10.3
torchvision             0.14.0
tornado                 6.1
tqdm                    4.64.1
typer                   0.7.0
typing_extensions       4.4.0
unicodedata2            15.0.0
urllib3                 1.26.12
wheel                   0.38.4
xarray                  2022.11.0
xyzservices             2022.9.0
yarl                    1.8.1
zarr                    2.13.3
zict                    2.2.0
</code></pre>

VitalyFedyunin: <p>CC <a class="mention" href="/u/nivek">@nivek</a> <a class="mention" href="/u/ejguan">@ejguan</a> PTAL</p>

weiji14: <p>Ok, I see you’ve made some force pushes… The main diffs seem to be on <code>torch</code>, <code>torchdata</code> and <code>torchvision</code> now:</p>
<pre><code class="lang-diff">diff --git a/requirements.txt b/requirements.txt
index 1269b7be..09bebd84 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -31,6 +31,7 @@ distributed             2022.11.1
 einops                  0.6.0
 entrypoints             0.4
 exceptiongroup          1.0.4
+execnet                 1.9.0
 fasteners               0.18
 fastparquet             2022.11.0
 Fiona                   1.8.22
@@ -42,7 +43,7 @@ frozenlist              1.3.3
 fsspec                  2022.11.0
 GDAL                    3.5.3
 geopandas               0.12.1
-gitdb                   4.0.10
+gitdb                   4.0.9
 GitPython               3.1.29
 greenlet                2.0.1
 h5netcdf                0.0.0
@@ -74,7 +75,7 @@ ocf-datapipes           0.5.30        /home/jacob/Development/ocf_datapipes
:...skipping...
diff --git a/requirements.txt b/requirements.txt
index 1269b7be..09bebd84 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -31,6 +31,7 @@ distributed             2022.11.1
 einops                  0.6.0
 entrypoints             0.4
 exceptiongroup          1.0.4
+execnet                 1.9.0
 fasteners               0.18
 fastparquet             2022.11.0
 Fiona                   1.8.22
@@ -42,7 +43,7 @@ frozenlist              1.3.3
 fsspec                  2022.11.0
 GDAL                    3.5.3
 geopandas               0.12.1
-gitdb                   4.0.10
+gitdb                   4.0.9
 GitPython               3.1.29
 greenlet                2.0.1
 h5netcdf                0.0.0
@@ -74,7 +75,7 @@ ocf-datapipes           0.5.30        /home/jacob/Development/ocf_datapipes
 packaging               21.3
 pandas                  1.5.2
 partd                   1.3.0
-pathy                   0.10.0
+pathy                   0.9.0
 Pillow                  9.2.0
 pip                     22.3.1
 pluggy                  1.0.0
@@ -95,6 +96,8 @@ pyresample              1.25.1
 pyshp                   2.3.1
 PySocks                 1.7.1
 pytest                  7.2.0
+pytest-timeout          2.1.0
+pytest-xdist            3.0.2
 python-dateutil         2.8.2
 pytorch-lightning       1.8.3.post0
 pytz                    2022.6
@@ -119,16 +122,16 @@ termcolor               2.1.1
 threadpoolctl           3.1.0
 tomli                   2.0.1
 toolz                   0.12.0
-torch                   1.12.0
-torchdata               0.4.1+f9ecd8b
+torch                   1.13.0
+torchdata               0.5.0
 torchmetrics            0.10.3
-torchvision             0.13.0
+torchvision             0.14.0
 tornado                 6.1
 tqdm                    4.64.1
 typer                   0.7.0
 typing_extensions       4.4.0
 unicodedata2            15.0.0
-urllib3                 1.26.13
+urllib3                 1.26.12
 wheel                   0.38.4
 xarray                  2022.11.0
 xyzservices             2022.9.0
</code></pre>
<p>Can you disable <code>pytest-xdist</code> temporarily at <a href="https://github.com/openclimatefix/ocf_datapipes/blob/86f2a1397a6696c536b692d64cff16dc25e8afe9/.github/workflows/workflows.yaml#L12" class="inline-onebox" rel="noopener nofollow ugc">ocf_datapipes/workflows.yaml at 86f2a1397a6696c536b692d64cff16dc25e8afe9 · openclimatefix/ocf_datapipes · GitHub</a> (i.e. run in serial) to see if it’s possible to isolate what test it’s crashing on? Also, I think your <code>yield</code> statement in the <code>__iter__</code> function is still outside of the <code>with rioxarray.open_rasterio(...)</code> block, though I’m not 100% sure if that’s the main issue anymore. It might actually be something in <code>pytorch</code>/<code>torchdata</code> itself.</p>
<p>As an aside, have you considered using something like <a href="https://github.com/xarray-contrib/xbatcher" rel="noopener nofollow ugc"><code>xbatcher</code></a> to read smaller slices of data at a time? Depends on what your data looks like of course <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=12" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:" loading="lazy" width="20" height="20"></p>

jacobbieker: <p>Yeah, I force pushed to rebase the branch on the newest version so there is less changes between it and the main branch. And yeah, the yield still is, I can move it inside though. I have heard of xbatcher, and will look into it, our data is a mix of different xarray datasets with different spatial coordinates systems loaded all at once, and we have to compute the temporal and spatial overlaps and then select batches of data based off those intersections, which seems like it might be a bit more complicated than what xbatcher is more designed for.</p>

Peter_Dudfield: <p>I think a found the issue in <a href="https://github.com/pytorch/data/issues/865" class="inline-onebox" rel="noopener nofollow ugc">Dataloader does not stop if one of the zipped DataPipes has a perpetual cycle · Issue #865 · pytorch/data · GitHub</a><br>
Essentially the Zipper datapipe has changed behaviour a bit</p>

weiji14: <p>Right, so unclosed file streams are messing up at the <code>ZipperIterDataPipe</code> in <code>torchdata=0.5.0</code> according to <a href="https://github.com/pytorch/data/issues/865#issuecomment-1294013285" class="inline-onebox" rel="noopener nofollow ugc">Dataloader does not stop if one of the zipped DataPipes has a perpetual cycle · Issue #865 · pytorch/data · GitHub</a>. So either wait for a bugfix, pin to <code>torchdata=0.4.1</code>, or handle yield statements properly <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=12" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="jacobbieker" data-post="6" data-topic="166743">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jacobbieker/48/55111_2.png" class="avatar"> jacobbieker:</div>
<blockquote>
<p>I have heard of xbatcher, and will look into it, our data is a mix of different xarray datasets with different spatial coordinates systems loaded all at once, and we have to compute the temporal and spatial overlaps and then select batches of data based off those intersections, which seems like it might be a bit more complicated than what xbatcher is more designed for.</p>
</blockquote>
</aside>
<p>Hmm, <a href="https://xbatcher.readthedocs.io/en/latest/api.html#xbatcher.BatchGenerator" rel="noopener nofollow ugc"><code>xbatcher.BatchGenerator</code></a> does have an <code>input_overlap</code> parameter, but I’m guessing you’re after something a bit more advanced. There <em>is</em> a possibly related issue open at <a href="https://github.com/xarray-contrib/xbatcher/issues/30" class="inline-onebox" rel="noopener nofollow ugc">Handling of overlapping samples and shuffling · Issue #30 · xarray-contrib/xbatcher · GitHub</a>, and you’re welcome to articulate your use-case there or open another issue. Just want to save you from having to reinvent the wheel.</p>

ejguan: <p>Thank you guys for keeping us posted about the bug you found. I will send a patch ASAP and let you know when the new nightly release becomes available for you to test.</p>

ejguan: <p>After <a href="https://github.com/pytorch/pytorch/pull/89974" class="inline-onebox" rel="noopener nofollow ugc">[4/4][DataPipe] Remove iterator depletion in Zipper by ejguan · Pull Request #89974 · pytorch/pytorch · GitHub</a> is landed, this issue should have been fixed. I have tested the <a href="https://github.com/pytorch/data/pull/912" rel="noopener nofollow ugc">newly created tests</a> locally.</p>
<p>You should be able to find the nightly releases for both <code>torch</code> and <code>torchdata</code> tomorrow. In terms of how to install nightly releases, pls check the <a href="https://pytorch.org/get-started/locally/" rel="noopener nofollow ugc">official guidance</a></p>
<pre><code class="lang-auto">pip3 install --pre torch torchdata  --extra-index-url https://download.pytorch.org/whl/nightly/cpu
</code></pre>
<p>You can replace <code>cpu</code> by <code>cu116</code>, <code>cu117</code> or <code>rocm5.2</code> to install <code>torch</code> based on your requirement.</p>