datago 2025.12.2__tar.gz → 2026.1.2__tar.gz

{datago-2025.12.2 → datago-2026.1.2}/.github/workflows/ci-cd.yml

@@ -22,7 +22,7 @@ jobs:
         platform:
           - runner: ubuntu-latest
             target: x86_64
-        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
+        python-version: ['3.10', '3.11', '3.12', '3.13']
 
     environment:
       name: release
@@ -48,7 +48,7 @@ jobs:
 
       - name: Build the package
         run: |
-          maturin build -i python${{ matrix.python-version }} --release --out dist --target "x86_64-unknown-linux-gnu" --manylinux 2014 --zig
+          maturin build -i python${{ matrix.python-version }} --release --out dist --target "x86_64-unknown-linux-gnu" --manylinux 2_31 --zig
 
       - name: Test package
         env:

{datago-2025.12.2 → datago-2026.1.2}/.github/workflows/rust.yml

@@ -51,4 +51,4 @@ jobs:
           DATAROOM_TEST_SOURCE: ${{ secrets.DATAROOM_TEST_SOURCE }}
           DATAROOM_API_URL: ${{ secrets.DATAROOM_API_URL }}
 
-        run: cargo test --verbose
+        run: RUST_BACKTRACE=1 cargo test --verbose

{datago-2025.12.2 → datago-2026.1.2}/Cargo.lock

@@ -623,7 +623,7 @@ dependencies = [
 
 [[package]]
 name = "datago"
-version = "2025.12.2"
+version = "2026.1.2"
 dependencies = [
  "async-compression",
  "async-tar",

{datago-2025.12.2 → datago-2026.1.2}/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "datago"
 edition = "2021"
-version = "2025.12.2"
+version = "2026.1.2"
 readme = "README.md"
 
 [lib]

{datago-2025.12.2 → datago-2026.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datago
-Version: 2025.12.2
+Version: 2026.1.2
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
@@ -133,7 +133,7 @@ client_config = {
     "source_config": {
         "url": url,
         "random_sampling": False,
-        "max_concurrency": 8, # The number of TarballSamples which should be handled concurrently
+        "concurrent_downloads": 8, # The number of TarballSamples which should be handled concurrently
         "rank": 0,
         "world_size": 1,
     },
@@ -267,18 +267,46 @@ Create a new tag and a new release in this repo, a new package will be pushed au
 <details> <summary><strong>Benchmarks</strong></summary>
 As usual, benchmarks are a tricky game, and you shouldn't read too much into the following plots but do your own tests. Some python benchmark examples are provided in the [python](./python/) folder.
 
-In general, Datago will be impactful if you want to load a lot of images very fast, but if you consume them as you go at a more leisury pace then it's not really needed. The more CPU work there is with the images and the higher quality they are, the more Datago will shine. The following benchmarks are using ImageNet 1k, which is very low resolution and thus kind of a worst case scenario. Data is served from cache (i.e. the OS cache) and the images are not pre-processed. In this case the receiving python process is typically the bottleneck, and caps at around 3000 images per second.
+In general, Datago will be impactful if you want to load a lot of images very fast, but if you consume them as you go at a more leisury pace then it's not really needed. The more CPU work there is with the images and the higher quality they are, the more Datago will shine.
 
-### AMD Zen3 laptop - IN1k - disk
+## From disk: ImageNet
+
+The following benchmarks are using ImageNet 1k, which is very low resolution and thus kind of a worst case scenario. Data is served from cache (i.e. the OS cache) and the images are not pre-processed. In this case the receiving python process is typically the bottleneck, and caps at around 3000 images per second.
+
+### AMD Zen3 laptop - IN1k - disk - no processing
 ![AMD Zen3 laptop & M2 SSD](assets/zen3_ssd.png)
 
-### AMD EPYC 9454 - IN1k - disk
+### AMD EPYC 9454 - IN1k - disk - no processing
 ![AMD EPYC 9454](assets/epyc_vast.png)
 
-This benchmark is using the PD12M dataset, which hosts high resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), so given enough time the two results would look closer.
+## Webdataset: FakeIN
+
+This benchmark is using low resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), which emphasizes throughput differences depending on how long you test it for.
+
+Of note is also that this can be bottlenecked by your external bandwidth to the remote storage where WDS is hosted, in which case both solution would yield comparable numbers.
+
+### AMD Zen3 laptop - webdataset - no processing
+![AMD EPYC 9454](assets/zen3_wds_fakein.png)
+
+
+## Webdataset: PD12M
+
+This benchmark is using high resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), which emphasizes throughput differences depending on how long you test it for.
+
+Of note is also that this can be bottlenecked by your external bandwidth to the remote storage where WDS is hosted, in which case both solution would yield comparable numbers.
+
+### AMD Zen3 laptop - webdataset - no processing
+![AMD Zen3 laptop](assets/zen3_wds_pd12m.png)
+
+
+### AMD EPYC 9454 - pd12m - webdataset - no processing
+![AMD EPYC 9454](assets/epyc_wds_pd12m.png)
+
+
+### AMD Zen3 laptop - webdataset - processing
+Adding image processing (crop and resize to Transformer compatible size buckets) to the equation changes the picture, as the work spread becomes more important. If you're training a diffusion model or an image encoder from a diverse set of images, this is likely to be the most realistic micro-benchmark.
 
-### AMD EPYC 9454 - pd12m - webdataset
-![AMD EPYC 9454](assets/epyc_wds.png)
+![AMD Zen3 laptop](assets/zen3_wds_pd12m_processing.png)
 
 </details>
 

{datago-2025.12.2 → datago-2026.1.2}/README.md

@@ -116,7 +116,7 @@ client_config = {
     "source_config": {
         "url": url,
         "random_sampling": False,
-        "max_concurrency": 8, # The number of TarballSamples which should be handled concurrently
+        "concurrent_downloads": 8, # The number of TarballSamples which should be handled concurrently
         "rank": 0,
         "world_size": 1,
     },
@@ -250,18 +250,46 @@ Create a new tag and a new release in this repo, a new package will be pushed au
 <details> <summary><strong>Benchmarks</strong></summary>
 As usual, benchmarks are a tricky game, and you shouldn't read too much into the following plots but do your own tests. Some python benchmark examples are provided in the [python](./python/) folder.
 
-In general, Datago will be impactful if you want to load a lot of images very fast, but if you consume them as you go at a more leisury pace then it's not really needed. The more CPU work there is with the images and the higher quality they are, the more Datago will shine. The following benchmarks are using ImageNet 1k, which is very low resolution and thus kind of a worst case scenario. Data is served from cache (i.e. the OS cache) and the images are not pre-processed. In this case the receiving python process is typically the bottleneck, and caps at around 3000 images per second.
+In general, Datago will be impactful if you want to load a lot of images very fast, but if you consume them as you go at a more leisury pace then it's not really needed. The more CPU work there is with the images and the higher quality they are, the more Datago will shine.
 
-### AMD Zen3 laptop - IN1k - disk
+## From disk: ImageNet
+
+The following benchmarks are using ImageNet 1k, which is very low resolution and thus kind of a worst case scenario. Data is served from cache (i.e. the OS cache) and the images are not pre-processed. In this case the receiving python process is typically the bottleneck, and caps at around 3000 images per second.
+
+### AMD Zen3 laptop - IN1k - disk - no processing
 ![AMD Zen3 laptop & M2 SSD](assets/zen3_ssd.png)
 
-### AMD EPYC 9454 - IN1k - disk
+### AMD EPYC 9454 - IN1k - disk - no processing
 ![AMD EPYC 9454](assets/epyc_vast.png)
 
-This benchmark is using the PD12M dataset, which hosts high resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), so given enough time the two results would look closer.
+## Webdataset: FakeIN
+
+This benchmark is using low resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), which emphasizes throughput differences depending on how long you test it for.
+
+Of note is also that this can be bottlenecked by your external bandwidth to the remote storage where WDS is hosted, in which case both solution would yield comparable numbers.
+
+### AMD Zen3 laptop - webdataset - no processing
+![AMD EPYC 9454](assets/zen3_wds_fakein.png)
+
+
+## Webdataset: PD12M
+
+This benchmark is using high resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), which emphasizes throughput differences depending on how long you test it for.
+
+Of note is also that this can be bottlenecked by your external bandwidth to the remote storage where WDS is hosted, in which case both solution would yield comparable numbers.
+
+### AMD Zen3 laptop - webdataset - no processing
+![AMD Zen3 laptop](assets/zen3_wds_pd12m.png)
+
+
+### AMD EPYC 9454 - pd12m - webdataset - no processing
+![AMD EPYC 9454](assets/epyc_wds_pd12m.png)
+
+
+### AMD Zen3 laptop - webdataset - processing
+Adding image processing (crop and resize to Transformer compatible size buckets) to the equation changes the picture, as the work spread becomes more important. If you're training a diffusion model or an image encoder from a diverse set of images, this is likely to be the most realistic micro-benchmark.
 
-### AMD EPYC 9454 - pd12m - webdataset
-![AMD EPYC 9454](assets/epyc_wds.png)
+![AMD Zen3 laptop](assets/zen3_wds_pd12m_processing.png)
 
 </details>
 

{datago-2025.12.2 → datago-2026.1.2}/python/benchmark_webdataset.py

@@ -1,6 +1,7 @@
 import json
 import os
 import time
+from typing import Any
 
 import typer
 from benchmark_defaults import IMAGE_CONFIG
@@ -9,45 +10,123 @@ from tqdm import tqdm
 
 
 def benchmark(
-    limit: int = typer.Option(10, help="The number of samples to test on"),
+    limit: int = typer.Option(1000, help="The number of samples to test on"),
     crop_and_resize: bool = typer.Option(
-        True, help="Crop and resize the images on the fly"
+        False, help="Crop and resize the images on the fly"
     ),
     compare_wds: bool = typer.Option(True, help="Compare against torch dataloader"),
+    num_downloads: int = typer.Option(
+        32,
+        help="Number of concurrent downloads",
+    ),
     num_workers: int = typer.Option(
-        16,
-        help="Number of processes to use",
+        8,
+        help="Number of CPU workers",
+    ),
+    sweep: bool = typer.Option(False, help="Sweep over the number of workers"),
+    plot: bool = typer.Option(
+        False, help="Whether to save a plot at the end of the run"
     ),
-    sweep: bool = typer.Option(False, help="Sweep over the number of processes"),
 ):
-    if sweep:
-        results = {}
-        for num_workers in range(2, max(64, (os.cpu_count() or 1)), 8):
-            results[num_workers] = benchmark(limit, crop_and_resize, compare_wds, num_workers, False)
-
-        # Save results to a json file
-        with open("benchmark_results_wds.json", "w") as f:
-            json.dump(results, f, indent=2)
-
-        return results
+    results: dict[Any, Any] = {}
+    if plot and not sweep:
+        print("Plot option only makes sense if we sweeped results, will not be used since sweep is False")
+        plot = False
 
     # URL of the test bucket
     # bucket = "https://storage.googleapis.com/webdataset/fake-imagenet"
     # dataset = "/imagenet-train-{000000..001281}.tar"
+    # source = "FakeIN"
 
     bucket = "https://huggingface.co/datasets/sayakpaul/pd12m-full/resolve/"
     dataset = "main/{00155..02480}.tar"
+    source = "PD12M"
+
     url = bucket + dataset
 
     print(
-        f"Benchmarking Datago WDS path on {url}.\nRunning benchmark for {limit} samples"
+        f"Benchmarking Datago WDS path on {url}.\nRunning benchmark for {limit} samples. Source {source}"
     )
+
+    if sweep:
+        max_cpus = os.cpu_count() or 16
+
+        num_workers = 1
+        while num_workers < max_cpus:
+            results[num_workers] = benchmark(
+                limit,
+                crop_and_resize,
+                compare_wds,
+                num_downloads,
+                num_workers,
+                False,
+                False,
+            )
+            num_workers *= 2
+
+        # Save results to a json file
+        with open("benchmark_results_wds.json", "w") as f:
+            json.dump(results, f, indent=2)
+
+        if plot:
+            import matplotlib.pyplot as plt
+            import pandas as pd
+
+            # Convert to a DataFrame for plotting
+            df = pd.DataFrame(
+                {
+                    "Thread Count": [int(k) for k in results.keys()],
+                    "Datago FPS": [results[k]["datago"]["fps"] for k in results.keys()],
+                    "Webdataset FPS": [
+                        results[k]["webdataset"]["fps"] for k in results.keys()
+                    ],
+                }
+            )
+
+            # Plotting with vertical axis starting at 0
+            plt.figure(figsize=(10, 6))
+            plt.plot(
+                df["Thread Count"],
+                df["Datago FPS"],
+                marker="o",
+                label="Datago",
+            )
+            plt.plot(
+                df["Thread Count"],
+                df["Webdataset FPS"],
+                marker="o",
+                label="Webdataset",
+            )
+            plt.xlabel("Thread Count")
+            plt.ylabel("Frames Per Second (FPS)")
+            plt.title(f"Throughput: Datago vs Webdataset. Source: {source}")
+            plt.ylim(
+                0,
+                max(df["Datago FPS"].max(), df["Webdataset FPS"].max()) + 20,
+            )
+            plt.legend()
+            plt.grid(True)
+            plt.xticks(df["Thread Count"])
+            plt.tight_layout()
+            plt.savefig(
+                "bench_datago_webdataset.png",
+                format="PNG",
+                dpi=200,
+                bbox_inches="tight",
+            )
+            plt.close()
+
+        return results
+
+    # This setting is not exposed in the config, but an env variable can be used instead
+    os.environ["DATAGO_MAX_TASKS"] = str(num_workers)
+
     client_config = {
         "source_type": "webdataset",
         "source_config": {
             "url": url,
             "shuffle": True,
-            "max_concurrency": num_workers,  # Number of concurrent TarballSample downloads and dispatch
+            "concurrent_downloads": num_downloads,  # Number of concurrent TarballSample downloads and dispatch
             "auth_token": os.environ.get("HF_TOKEN", default=""),
         },
         "prefetch_buffer_size": 256,
@@ -98,7 +177,7 @@ def benchmark(
                 ]
             )
             if crop_and_resize
-            else None
+            else lambda x: x
         )
 
         def custom_transform(sample):
@@ -117,16 +196,17 @@ def benchmark(
             # .to_tuple("png", "cls")  # Map keys to output tuple
         )
 
-        dataloader = DataLoader(
+        dataloader = DataLoader(  #  type:ignore
             dataset,
             batch_size=1,
             num_workers=num_workers,
-            prefetch_factor=2,
+            prefetch_factor=8,  # Didn't sweep on that, but probably not super impactful
             collate_fn=lambda x: x,
         )
 
         # Iterate over the DataLoader
         start = time.time()
+        n_images = 0
         for n_images, _ in enumerate(tqdm(dataloader, desc="WDS", dynamic_ncols=True)):
             if n_images > limit:
                 break
@@ -136,5 +216,6 @@ def benchmark(
         results["webdataset"] = {"fps": fps, "count": n_images}
         return results
 
+
 if __name__ == "__main__":
     typer.run(benchmark)

datago-2026.1.2/python/test_datago_wds.py

@@ -0,0 +1,250 @@
+"""
+Test suite for WebDataset (WDS) functionality in Datago.
+
+This module tests that Datago correctly serves images and attributes from WebDataset sources.
+"""
+
+import os
+
+from dataset import DatagoIterDataset
+from PIL import Image
+
+# Test buckets - using the same ones as benchmark_webdataset.py
+TEST_BUCKETS = {
+    "pd12m": {
+        "url": "https://huggingface.co/datasets/sayakpaul/pd12m-full/resolve/main/{00155..02480}.tar",
+        "source": "PD12M",
+    },
+    "fakein": {
+        "url": "https://storage.googleapis.com/webdataset/fake-imagenet/imagenet-train-{000000..001281}.tar",
+        "source": "FakeIN",
+    },
+}
+
+
+def test_wds_basic_functionality():
+    """Test basic WDS functionality - that we can get samples with proper structure."""
+    limit = 5  # Small limit for quick testing
+
+    # Use the PD12M bucket for testing
+    bucket_config = TEST_BUCKETS["pd12m"]
+
+    client_config = {
+        "source_type": "webdataset",
+        "source_config": {
+            "url": bucket_config["url"],
+            "shuffle": True,
+            "concurrent_downloads": 4,  # Reduced for testing
+            "auth_token": os.environ.get("HF_TOKEN", default=""),
+        },
+        "prefetch_buffer_size": 32,
+        "samples_buffer_size": 32,
+        "limit": limit,
+    }
+
+    # Test with return_python_types=True to get proper Python objects
+    dataset = DatagoIterDataset(client_config, return_python_types=True)
+
+    count = 0
+    for sample in dataset:
+        count += 1
+
+        # Basic structure checks
+        assert "id" in sample, "Sample should contain 'id' field"
+        assert sample["id"] != "", "Sample ID should not be empty"
+
+        # Check that we have an image
+        assert "image" in sample, "Sample should contain 'image' field"
+        assert sample["image"] is not None, "Image should not be None"
+
+        # If it's a PIL Image, check its properties
+        if isinstance(sample["image"], Image.Image):
+            assert sample["image"].width > 0, "Image should have positive width"
+            assert sample["image"].height > 0, "Image should have positive height"
+            assert sample["image"].mode in ["RGB", "RGBA", "L"], (
+                f"Image should have valid mode, got {sample['image'].mode}"
+            )
+
+        # Check for attributes if present
+        if "attributes" in sample:
+            assert isinstance(sample["attributes"], dict), "Attributes should be a dictionary"
+            # Attributes should be non-empty if present
+            if sample["attributes"]:
+                assert len(sample["attributes"]) > 0, "Attributes dictionary should not be empty"
+
+        # We should get at least the basic fields
+        assert len(sample) >= 2, "Sample should contain at least id and image"
+
+        if count >= limit:
+            break
+
+    assert count == limit, f"Expected {limit} samples, got {count}"
+
+
+def test_wds_image_properties():
+    """Test that images from WDS have proper properties and can be processed."""
+    limit = 3
+
+    bucket_config = TEST_BUCKETS["pd12m"]
+
+    client_config = {
+        "source_type": "webdataset",
+        "source_config": {
+            "url": bucket_config["url"],
+            "shuffle": True,
+            "concurrent_downloads": 4,
+            "auth_token": os.environ.get("HF_TOKEN", default=""),
+        },
+        "prefetch_buffer_size": 32,
+        "samples_buffer_size": 32,
+        "limit": limit,
+    }
+
+    dataset = DatagoIterDataset(client_config, return_python_types=True)
+
+    for sample in dataset:
+        if "image" in sample and sample["image"] is not None:
+            image = sample["image"]
+
+            # Test that we can get image properties
+            if isinstance(image, Image.Image):
+                width, height = image.size
+                assert width > 0 and height > 0, "Image should have valid dimensions"
+
+                # Test that we can convert to different modes
+                rgb_image = image.convert("RGB")
+                assert rgb_image.mode == "RGB", "Image should convert to RGB mode"
+
+                # Test that we can get thumbnail
+                thumbnail = image.copy()
+                thumbnail.thumbnail((100, 100))
+                assert thumbnail.size[0] <= 100 and thumbnail.size[1] <= 100, "Thumbnail should be resized"
+
+                # Test that image data is valid by trying to get pixel data
+                pixels = image.get_flattened_data()
+                assert len(pixels) > 0, "Image should have pixel data"
+
+                break  # Just test one image
+
+
+def test_wds_with_image_processing():
+    """Test WDS with image processing configuration (crop and resize)."""
+    limit = 3
+
+    bucket_config = TEST_BUCKETS["pd12m"]
+
+    client_config = {
+        "source_type": "webdataset",
+        "source_config": {
+            "url": bucket_config["url"],
+            "shuffle": True,
+            "concurrent_downloads": 4,
+            "auth_token": os.environ.get("HF_TOKEN", default=""),
+        },
+        "image_config": {
+            "crop_and_resize": True,
+            "default_image_size": 256,
+            "downsampling_ratio": 16,
+            "min_aspect_ratio": 0.5,
+            "max_aspect_ratio": 2.0,
+        },
+        "prefetch_buffer_size": 32,
+        "samples_buffer_size": 32,
+        "limit": limit,
+    }
+
+    dataset = DatagoIterDataset(client_config, return_python_types=True)
+
+    for sample in dataset:
+        if "image" in sample and sample["image"] is not None:
+            image = sample["image"]
+
+            if isinstance(image, Image.Image):
+                # With crop_and_resize=True, images should be processed
+                width, height = image.size
+                assert width > 0 and height > 0, "Processed image should have valid dimensions"
+
+                # The processed image should be in RGB mode
+                assert image.mode == "RGB", f"Processed image should be RGB, got {image.mode}"
+
+                break  # Just test one image
+
+
+def test_wds_attributes_structure():
+    """Test that WDS attributes are properly structured when present."""
+    limit = 5
+
+    bucket_config = TEST_BUCKETS["pd12m"]
+
+    client_config = {
+        "source_type": "webdataset",
+        "source_config": {
+            "url": bucket_config["url"],
+            "shuffle": True,
+            "concurrent_downloads": 4,
+            "auth_token": os.environ.get("HF_TOKEN", default=""),
+        },
+        "prefetch_buffer_size": 32,
+        "samples_buffer_size": 32,
+        "limit": limit,
+    }
+
+    dataset = DatagoIterDataset(client_config, return_python_types=True)
+
+    for sample in dataset:
+        if "attributes" in sample and sample["attributes"]:
+            attributes = sample["attributes"]
+
+            # Attributes should be a dictionary
+            assert isinstance(attributes, dict), "Attributes should be a dictionary"
+
+            # Check that we can access attribute values
+            for key, value in attributes.items():
+                # Values should be JSON-serializable types
+                assert isinstance(key, str), "Attribute keys should be strings"
+                assert isinstance(value, (str, int, float, bool, list, dict)), (
+                    f"Attribute values should be JSON-serializable, got {type(value)}"
+                )
+
+            break  # Just test one sample with attributes
+
+    # Note: Not all samples may have attributes
+
+
+def test_wds_sample_consistency():
+    """Test that WDS samples have consistent structure across multiple samples."""
+    limit = 10
+
+    bucket_config = TEST_BUCKETS["pd12m"]
+
+    client_config = {
+        "source_type": "webdataset",
+        "source_config": {
+            "url": bucket_config["url"],
+            "shuffle": True,
+            "concurrent_downloads": 4,
+            "auth_token": os.environ.get("HF_TOKEN", default=""),
+        },
+        "prefetch_buffer_size": 32,
+        "samples_buffer_size": 32,
+        "limit": limit,
+    }
+
+    dataset = DatagoIterDataset(client_config, return_python_types=True)
+
+    first_sample = True
+
+    for sample in dataset:
+        current_keys = set(sample.keys())
+
+        if first_sample:
+            first_sample = False
+        else:
+            # All samples should have at least the core fields (id, image)
+            required_keys = {"id", "image"}
+            assert required_keys.issubset(current_keys), \
+                f"Sample missing required keys. Expected at least {required_keys}, got {current_keys}"
+
+        # Check that we don't have any unexpected None values for core fields
+        assert sample.get("id") != "", "Sample ID should not be empty"
+        assert sample.get("image") is not None, "Sample image should not be None"

datago-2026.1.2/requirements-dev.txt

@@ -0,0 +1,3 @@
+-r requirements.txt
+matplotlib
+pandas

{datago-2025.12.2 → datago-2026.1.2}/src/client.rs

@@ -217,14 +217,16 @@ impl DatagoClient {
             debug!("Sample pipe closed...");
 
             if let Some(feeder) = engine.feeder.take() {
-                if feeder.join().is_err() {
-                    error!("Failed to join feeder thread");
+                match feeder.join() {
+                    Ok(_) => debug!("Feeder thread joined successfully"),
+                    Err(e) => error!("Failed to join feeder thread: {:?}", e),
                 }
             }
 
             if let Some(worker) = engine.worker.take() {
-                if worker.join().is_err() {
-                    error!("Failed to join worker thread");
+                match worker.join() {
+                    Ok(_) => debug!("Worker thread joined successfully"),
+                    Err(e) => error!("Failed to join worker thread: {:?}", e),
                 }
             }
             self.is_started = false;

{datago-2025.12.2 → datago-2026.1.2}/src/generator_wds.rs

@@ -35,7 +35,7 @@ pub struct SourceWebDatasetConfig {
     pub random_sampling: bool,
 
     #[serde(default)]
-    pub max_concurrency: usize,
+    pub concurrent_downloads: usize,
 
     #[serde(default)]
     pub auth_token: String,
@@ -81,13 +81,23 @@ async fn pull_tarballs(
 
     // Grab an async byte stream from the request, we'll try to untar the results on the fly
     let response = request_builder.send().await;
-    if response.is_err() {
-        return Err("Failed to send request".into());
-    }
-    let response = response.unwrap();
+    let response = match response {
+        Ok(resp) => resp,
+        Err(e) => {
+            error!("Failed to send request for {}: {}", url, e);
+            return Err(format!("Failed to send request: {}", e));
+        }
+    };
+
     if !response.status().is_success() {
+        error!(
+            "Failed to download TarballSample {}: HTTP {}",
+            url,
+            response.status()
+        );
         return Err(format!(
-            "Failed to download TarballSample: {}",
+            "Failed to download TarballSample {}: HTTP {}",
+            url,
             response.status()
         ));
     }
@@ -158,7 +168,7 @@ async fn pull_tarballs(
             if samples_metadata_tx.send(current_files_for_sample).is_err() {
                 debug!("dispatch_shards (streaming): samples_metadata_tx channel closed.");
                 let _ = samples_metadata_tx.close(); // Make sure that we close on both ends
-                return Err("Channel closed".into());
+                return Ok(());
             }
 
             // Start a new sample
@@ -184,9 +194,9 @@ async fn pull_tarballs(
     // Send the last collected sample if any
     if !current_files_for_sample.content.is_empty()
         && samples_metadata_tx.send(current_files_for_sample).is_err()
+        && !samples_metadata_tx.is_closed()
     {
-        debug!("dispatch_shards (streaming): samples_metadata_tx channel closed for last sample.");
-        return Err("Channel closed".into());
+        return Err("Failed to send last sample".into());
     }
 
     debug!("dispatch_shards (streaming): finished processing TarballSample {url}");
@@ -308,7 +318,18 @@ async fn tasks_from_shards(
             let mut count = 0;
             let mut join_error: Option<String> = None;
 
+            info!("WDS: Using {} download tasks", config.concurrent_downloads);
+
             for url in task_list {
+                // Escape out if the channel is closed
+                if samples_metadata_tx.is_closed() {
+                    debug!(
+                        "dispatch_shards: channel is closed, enough samples probably. Bailing out"
+                    );
+                    break;
+                }
+
+                // All good, submit a new async task
                 tasks.spawn(pull_tarballs_task(
                     shared_client.clone(),
                     url,
@@ -319,7 +340,7 @@ async fn tasks_from_shards(
 
                 // Some bookkeeping, to limit the number of tasks in flight
                 // we'll wait for the first one to finish before adding a new one
-                if tasks.len() >= config.max_concurrency {
+                if tasks.len() >= config.concurrent_downloads {
                     match tasks.join_next().await {
                         Some(res) => {
                             match res.unwrap() {
@@ -334,7 +355,6 @@ async fn tasks_from_shards(
                                     break;
                                 }
                             }
-                            debug!("dispatch_shards: task completed successfully");
                         }
                         None => {
                             // Task was cancelled or panicked
@@ -407,8 +427,10 @@ fn query_shards_and_dispatch(
         .and_then(|v| v.parse::<u8>().ok())
         .unwrap_or(3);
 
+    // Use more threads for the download runtime to handle increased concurrency
+    let download_threads = std::cmp::max(4, source_config.concurrent_downloads);
     tokio::runtime::Builder::new_multi_thread()
-        .worker_threads(source_config.max_concurrency)
+        .worker_threads(download_threads)
         .enable_all()
         .build()
         .unwrap()
@@ -434,7 +456,8 @@ fn query_shards_and_dispatch(
 // ---- Global orchestration ---------
 pub fn orchestrate(client: &DatagoClient) -> DatagoEngine {
     // Allocate all the message passing pipes
-    let (samples_metadata_tx, samples_metadata_rx) = bounded::<TarballSample>(32);
+    let metadata_buffer_size = std::cmp::max(128, client.samples_buffer * 2);
+    let (samples_metadata_tx, samples_metadata_rx) = bounded::<TarballSample>(metadata_buffer_size);
     let (samples_tx, samples_rx) = bounded(client.samples_buffer);
 
     info!("Using webdataset as source");
@@ -444,9 +467,9 @@ pub fn orchestrate(client: &DatagoClient) -> DatagoEngine {
         serde_json::from_value(client.source_config.clone()).unwrap();
     let extension_reference_image_type: String = source_config.reference_image_type.clone();
 
-    if source_config.max_concurrency == 0 {
-        info!("WDS: Defaulting to 8 max_concurrency");
-        source_config.max_concurrency = 8;
+    if source_config.concurrent_downloads == 0 {
+        info!("WDS: Defaulting to 8 concurrent_downloads");
+        source_config.concurrent_downloads = 8;
     }
 
     // List the contents of the bucket and feed the workers
@@ -518,7 +541,7 @@ mod tests {
                 auth_token: "".into(),
                 reference_image_type: "jpg".into(),
                 random_sampling: s,
-                max_concurrency: 2,
+                concurrent_downloads: 2,
                 rank: 0,
                 world_size: 1,
             };
@@ -569,7 +592,7 @@ mod tests {
                 "source_config": {
                     "url": "https://storage.googleapis.com/storage/v1/b/webdataset/o?prefix=fake-imagenet/",
                     "random_sampling": do_random_sampling,
-                    "max_concurrency": 2
+                    "concurrent_downloads": 2
                 },
                 "limit": n_samples,
                 "num_threads": 1,
@@ -632,7 +655,7 @@ mod tests {
                 "source_config": {
                     "url": "https://storage.googleapis.com/storage/v1/b/webdataset/o?prefix=fake-imagenet/",
                     "random_sampling": false,
-                    "max_concurrency": 2,
+                    "concurrent_downloads": 2,
                     "rank": rank,
                     "world_size": world_size,
                 },

{datago-2025.12.2 → datago-2026.1.2}/src/worker_wds.rs

@@ -1,6 +1,6 @@
 use crate::image_processing;
 use crate::structs::{to_python_image_payload, ImagePayload, Sample, TarballSample};
-use log::{debug, error, info, warn};
+use log::{debug, error, info};
 use std::cmp::min;
 use std::collections::HashMap;
 use std::path::Path;
@@ -77,30 +77,56 @@ async fn process_sample(
 
                                 if ext == extension_reference_image {
                                     // If this is the reference image, we store it in the main image field
-                                    final_sample = Some(Sample {
-                                        id: String::from(sample_id.to_str().unwrap_or("unknown")),
-                                        source: sample.name.clone(),
-                                        image,
-                                        attributes: attributes.clone(),
-                                        coca_embedding: vec![],
-                                        tags: vec![],
-                                        masks: HashMap::new(),
-                                        latents: HashMap::new(),
-                                        additional_images: HashMap::new(),
-                                        duplicate_state: 0,
-                                    });
+                                    if let Some(mut_final_sample) = &mut final_sample {
+                                        mut_final_sample.image = image;
+                                    } else {
+                                        // Init the sample
+                                        final_sample = Some(Sample {
+                                            id: String::from(
+                                                sample_id.to_str().unwrap_or("unknown"),
+                                            ),
+                                            source: sample.name.clone(),
+                                            image,
+                                            attributes: HashMap::new(),
+                                            coca_embedding: vec![],
+                                            tags: vec![],
+                                            masks: HashMap::new(),
+                                            latents: HashMap::new(),
+                                            additional_images: HashMap::new(),
+                                            duplicate_state: 0,
+                                        });
+                                    }
                                 } else {
                                     // Otherwise, we store it in the additional images
-                                    match final_sample {
-                                        Some(ref mut final_sample_ref) => {
-                                            final_sample_ref
-                                                .additional_images
-                                                .insert(item.filename.clone(), image.clone());
-                                        }
-                                        None => {
-                                            // If final_sample is not initialized, we create it
-                                            panic!( "Final sample should be initialized before adding additional images");
-                                        }
+                                    if let Some(mut_final_sample) = &mut final_sample {
+                                        mut_final_sample
+                                            .additional_images
+                                            .insert(item.filename.clone(), image);
+                                    } else {
+                                        // Init the sample
+                                        final_sample = Some(Sample {
+                                            id: String::from(
+                                                sample_id.to_str().unwrap_or("unknown"),
+                                            ),
+                                            source: sample.name.clone(),
+                                            image: to_python_image_payload(ImagePayload {
+                                                data: vec![],
+                                                width: 0,
+                                                height: 0,
+                                                original_height: 0,
+                                                original_width: 0,
+                                                bit_depth: 0,
+                                                channels: 0,
+                                                is_encoded: false,
+                                            }),
+                                            attributes: HashMap::new(),
+                                            coca_embedding: vec![],
+                                            tags: vec![],
+                                            masks: HashMap::new(),
+                                            latents: HashMap::new(),
+                                            additional_images: HashMap::new(),
+                                            duplicate_state: 0,
+                                        });
                                     }
                                 }
                                 debug!("wds_worker: unpacked {}", item.filename);
@@ -114,15 +140,25 @@ async fn process_sample(
                         // Load the file in to a string
                         let class_file = String::from_utf8_lossy(&item.buffer).to_string();
                         attributes.insert(ext.to_string(), serde_json::json!(class_file));
-                        debug!("wds_worker: unpacked {}", item.filename);
+                        debug!("wds_worker: unpacked {} {}", item.filename, class_file);
                     }
                 }
 
-                if samples_tx.send(final_sample).is_err() {
-                    debug!("wds_worker: stream already closed, wrapping up");
-                    return Err(());
+                // Make sure that the sample has the attributes we decoded
+                if let Some(ref mut final_sample_ref) = final_sample {
+                    final_sample_ref.attributes = attributes;
+                    match samples_tx.send(final_sample) {
+                        Ok(_) => (),
+                        Err(e) => {
+                            if !samples_tx.is_closed() {
+                                debug!("wds_worker: error dispatching sample: {e}");
+                                return Err(());
+                            }
+                        }
+                    }
+                    return Ok(());
                 }
-                return Ok(());
+                return Err(());
             }
             None => {
                 debug!("wds_worker: unpacking sample with no ID");
@@ -147,10 +183,10 @@ async fn async_deserialize_samples(
     let default_max_tasks = std::env::var("DATAGO_MAX_TASKS")
         .unwrap_or_else(|_| "0".to_string())
         .parse::<usize>()
-        .unwrap_or(num_cpus::get() * 4);
-    let max_tasks = min(default_max_tasks, limit);
+        .unwrap_or(num_cpus::get());
+    let max_tasks = min(num_cpus::get() * 4, default_max_tasks); // Ensure minimum of 8 processing tasks
 
-    info!("Using {max_tasks} tasks in the async threadpool");
+    info!("WDS: Using {max_tasks} processing tasks in worker threadpool");
     let mut tasks = tokio::task::JoinSet::new();
     let mut count = 0;
     let shareable_channel_tx: Arc<kanal::Sender<Option<Sample>>> = Arc::new(samples_tx);
@@ -159,7 +195,7 @@ async fn async_deserialize_samples(
 
     while let Ok(sample) = samples_metadata_rx.recv() {
         if sample.is_empty() {
-            warn!("wds_worker: end of stream received, stopping there");
+            info!("wds_worker: end of stream received, stopping there");
             let _ = samples_metadata_rx.close();
             break;
         }
@@ -228,7 +264,7 @@ pub fn deserialize_samples(
     extension_reference_image: String,
 ) {
     tokio::runtime::Builder::new_multi_thread()
-        .worker_threads(num_cpus::get())
+        .worker_threads(num_cpus::get()) // Tasks in flight are limited by DATAGO_MAX_TASKS env
         .enable_all()
         .build()
         .unwrap()