torch-loader(example): use prefetch and try to run example in linux

iterative · Dec 11, 2024 · 37403b9 · 37403b9
1 parent 6ca3c98
commit 37403b9
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 10 deletions.
diff --git a/examples/get_started/torch-loader.py b/examples/get_started/torch-loader.py
@@ -54,6 +54,7 @@ def forward(self, x):
 if __name__ == "__main__":
     ds = (
         DataChain.from_storage(STORAGE, type="image")
+        .settings(prefetch=25, cache=True)
         .filter(C("file.path").glob("*.jpg"))
         .map(
             label=lambda path: label_to_int(basename(path)[:3], CLASSES),
@@ -64,7 +65,7 @@ def forward(self, x):
 
     train_loader = DataLoader(
         ds.to_pytorch(transform=transform),
-        batch_size=16,
+        batch_size=25,
         num_workers=2,
     )
 

diff --git a/src/datachain/asyn.py b/src/datachain/asyn.py
@@ -165,7 +165,7 @@ async def _break_iteration(self) -> None:
 
     def iterate(self, timeout=None) -> Generator[ResultT, None, None]:
         init = asyncio.run_coroutine_threadsafe(self.init(), self.loop)
-        init.result(timeout=1)
+        init.result()
         async_run = asyncio.run_coroutine_threadsafe(self.run(), self.loop)
         try:
             while True:

diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py
@@ -6,14 +6,7 @@
 
 import pytest
 
-get_started_examples = sorted(
-    [
-        filename
-        for filename in glob.glob("examples/get_started/**/*.py", recursive=True)
-        # torch-loader will not finish within an hour on Linux runner
-        if "torch" not in filename or os.environ.get("RUNNER_OS") != "Linux"
-    ]
-)
+get_started_examples = sorted(glob.glob("examples/get_started/**/*.py", recursive=True))
 
 llm_and_nlp_examples = sorted(glob.glob("examples/llm_and_nlp/**/*.py", recursive=True))