torch-loader(example): use persistent workers to reduce test time (#694)

Author: skshetry

examples/get_started/torch-loader.py

@@ -19,7 +19,7 @@
 from datachain.torch import label_to_int
 
 STORAGE = "gs://datachain-demo/dogs-and-cats/"
-NUM_EPOCHS = os.getenv("NUM_EPOCHS", "3")
+NUM_EPOCHS = int(os.getenv("NUM_EPOCHS", "3"))
 
 # Define transformation for data preprocessing
 transform = v2.Compose(
@@ -68,7 +68,8 @@ def forward(self, x):
     train_loader = DataLoader(
         ds.to_pytorch(transform=transform),
         batch_size=25,
-        num_workers=4,
+        num_workers=max(4, os.cpu_count() or 2),
+        persistent_workers=True,
         multiprocessing_context=multiprocessing.get_context("spawn"),
     )
 
@@ -77,7 +78,7 @@ def forward(self, x):
     optimizer = optim.Adam(model.parameters(), lr=0.001)
 
     # Train the model
-    for epoch in range(int(NUM_EPOCHS)):
+    for epoch in range(NUM_EPOCHS):
         with tqdm(
             train_loader, desc=f"epoch {epoch + 1}/{NUM_EPOCHS}", unit="batch"
         ) as loader:

noxfile.py

@@ -81,6 +81,8 @@ def examples(session: nox.Session) -> None:
     session.install(".[examples]")
     session.run(
         "pytest",
+        "--durations=0",
+        "tests/examples",
         "-m",
         "examples",
         *session.posargs,

tests/examples/test_examples.py

@@ -52,7 +52,7 @@ def smoke_test(example: str, env: Optional[dict] = None):
 @pytest.mark.get_started
 @pytest.mark.parametrize("example", get_started_examples)
 def test_get_started_examples(example):
-    smoke_test(example, {"NUM_EPOCHS": "1"})
+    smoke_test(example)
 
 
 @pytest.mark.examples