datago 2025.4.1__tar.gz → 2025.4.3__tar.gz

{datago-2025.4.1 → datago-2025.4.3}/Cargo.lock

@@ -247,7 +247,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02"
 dependencies = [
  "smallvec",
- "target-lexicon",
+ "target-lexicon 0.12.16",
 ]
 
 [[package]]
@@ -387,7 +387,7 @@ dependencies = [
 
 [[package]]
 name = "datago"
-version = "2025.4.1"
+version = "2025.4.3"
 dependencies = [
  "clap",
  "image",
@@ -1381,9 +1381,9 @@ checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e"
 
 [[package]]
 name = "openssl"
-version = "0.10.71"
+version = "0.10.72"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5e14130c6a98cd258fdcb0fb6d744152343ff729cbfcb28c656a9d12b999fbcd"
+checksum = "fedfea7d58a1f73118430a55da6a286e7b044961736ce96a16a17068ea25e5da"
 dependencies = [
  "bitflags 2.9.0",
  "cfg-if",
@@ -1422,9 +1422,9 @@ dependencies = [
 
 [[package]]
 name = "openssl-sys"
-version = "0.9.106"
+version = "0.9.107"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8bb61ea9811cc39e3c2069f40b8b8e2e70d8569b361f879786cc7ed48b777cdd"
+checksum = "8288979acd84749c744a9014b4382d42b8f7b2592847b5afb2ed29e5d16ede07"
 dependencies = [
  "cc",
  "libc",
@@ -1560,9 +1560,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3"
-version = "0.23.5"
+version = "0.24.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
+checksum = "17da310086b068fbdcefbba30aeb3721d5bb9af8db4987d6735b2183ca567229"
 dependencies = [
  "cfg-if",
  "indoc",
@@ -1578,19 +1578,19 @@ dependencies = [
 
 [[package]]
 name = "pyo3-build-config"
-version = "0.23.5"
+version = "0.24.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
+checksum = "e27165889bd793000a098bb966adc4300c312497ea25cf7a690a9f0ac5aa5fc1"
 dependencies = [
  "once_cell",
- "target-lexicon",
+ "target-lexicon 0.13.2",
 ]
 
 [[package]]
 name = "pyo3-ffi"
-version = "0.23.5"
+version = "0.24.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
+checksum = "05280526e1dbf6b420062f3ef228b78c0c54ba94e157f5cb724a609d0f2faabc"
 dependencies = [
  "libc",
  "pyo3-build-config",
@@ -1598,9 +1598,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros"
-version = "0.23.5"
+version = "0.24.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
+checksum = "5c3ce5686aa4d3f63359a5100c62a127c9f15e8398e5fdeb5deef1fed5cd5f44"
 dependencies = [
  "proc-macro2",
  "pyo3-macros-backend",
@@ -1610,9 +1610,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros-backend"
-version = "0.23.5"
+version = "0.24.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
+checksum = "f4cf6faa0cbfb0ed08e89beb8103ae9724eb4750e3a78084ba4017cbe94f3855"
 dependencies = [
  "heck",
  "proc-macro2",
@@ -2211,6 +2211,12 @@ version = "0.12.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
 
+[[package]]
+name = "target-lexicon"
+version = "0.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a"
+
 [[package]]
 name = "tempfile"
 version = "3.17.1"
@@ -2288,9 +2294,9 @@ dependencies = [
 
 [[package]]
 name = "tokio"
-version = "1.43.0"
+version = "1.43.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3d61fa4ffa3de412bfea335c6ecff681de2b609ba3c77ef3e00e521813a9ed9e"
+checksum = "492a604e2fd7f814268a378409e6c92b5525d747d10db9a229723f55a417958c"
 dependencies = [
  "backtrace",
  "bytes",

{datago-2025.4.1 → datago-2025.4.3}/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "datago"
 edition = "2021"
-version = "2025.4.1"
+version = "2025.4.3"
 
 [lib]
 # exposed by pyo3
@@ -21,9 +21,9 @@ serde_json = "1.0"
 url = "2.5.4"
 kanal = "0.1"
 clap = { version = "4.5.27", features = ["derive"] }
-tokio = { version = "1.43.0", features = ["rt-multi-thread", "macros"] }
+tokio = { version = "1.43.1", features = ["rt-multi-thread", "macros"] }
 prettytable-rs = "0.10.0"
-pyo3 = { version = "0.23.4", features = ["extension-module"] }
+pyo3 = { version = "0.24.1", features = ["extension-module"] }
 threadpool = "1.8.1"
 openssl = { version = "0.10", features = ["vendored"] }
 walkdir = "2.5.0"

{datago-2025.4.1 → datago-2025.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datago
-Version: 2025.4.1
+Version: 2025.4.3
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{datago-2025.4.1 → datago-2025.4.3}/src/client.rs

@@ -198,7 +198,6 @@ impl DatagoClient {
         if !self.is_started {
             self.start();
         }
-        const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
 
         // If no more samples and workers are closed, then wrap it up
         if self.samples_rx.is_closed() {
@@ -208,6 +207,10 @@ impl DatagoClient {
         }
 
         // Try to fetch a new sample from the queue
+        // The client will timeout if zero sample is received in 5 minutes
+        // At this point it will stop and wrap everything up
+        const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
+
         match self.samples_rx.recv_timeout(TIMEOUT) {
             Ok(sample) => match sample {
                 Some(sample) => Some(sample),

{datago-2025.4.1 → datago-2025.4.3}/src/main.rs

@@ -1,7 +1,7 @@
 use clap::{Arg, Command};
 use prettytable::{row, Table};
 use serde_json::json;
-
+use std::collections::HashMap;
 mod client;
 mod generator_files;
 mod generator_http;
@@ -118,7 +118,7 @@ fn main() {
     let mut client = client::DatagoClient::new(config.to_string());
 
     // -----------------------------------------------------------------
-    let mut size_buckets: std::collections::HashMap<String, i32> = std::collections::HashMap::new();
+    let mut size_buckets: HashMap<String, i32> = HashMap::new();
     let start_time = std::time::Instant::now();
     let mut rolling_time = std::time::Instant::now();
 

{datago-2025.4.1 → datago-2025.4.3}/src/structs.rs

@@ -1,6 +1,7 @@
 use crate::image_processing::ImageTransformConfig;
 use pyo3::prelude::*;
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 
 #[derive(Deserialize)]
 #[serde(rename_all = "lowercase")]
@@ -64,7 +65,7 @@ pub struct Sample {
     pub source: String,
 
     #[doc(hidden)]
-    pub attributes: std::collections::HashMap<String, serde_json::Value>,
+    pub attributes: HashMap<String, serde_json::Value>,
 
     #[pyo3(get, set)]
     pub duplicate_state: i32,
@@ -73,13 +74,13 @@ pub struct Sample {
     pub image: ImagePayload,
 
     #[pyo3(get, set)]
-    pub masks: std::collections::HashMap<String, ImagePayload>,
+    pub masks: HashMap<String, ImagePayload>,
 
     #[pyo3(get, set)]
-    pub additional_images: std::collections::HashMap<String, ImagePayload>,
+    pub additional_images: HashMap<String, ImagePayload>,
 
     #[pyo3(get, set)]
-    pub latents: std::collections::HashMap<String, LatentPayload>,
+    pub latents: HashMap<String, LatentPayload>,
 
     #[pyo3(get, set)]
     pub coca_embedding: Vec<f32>,

{datago-2025.4.1 → datago-2025.4.3}/src/worker_files.rs

@@ -2,6 +2,7 @@ use crate::image_processing;
 use crate::structs::{ImagePayload, Sample};
 use std::cmp::min;
 use std::collections::HashMap;
+use std::collections::VecDeque;
 use std::sync::Arc;
 
 async fn image_from_path(path: &str) -> Result<image::DynamicImage, image::ImageError> {
@@ -78,7 +79,7 @@ async fn pull_sample(
 }
 
 pub async fn consume_oldest_task(
-    tasks: &mut std::collections::VecDeque<tokio::task::JoinHandle<Result<(), ()>>>,
+    tasks: &mut VecDeque<tokio::task::JoinHandle<Result<(), ()>>>,
 ) -> Result<(), ()> {
     match tasks.pop_front().unwrap().await {
         Ok(_) => Ok(()),
@@ -100,7 +101,7 @@ async fn async_pull_samples(
     // We use async-await here, to better use IO stalls
     // We'll issue N async tasks in parallel, and wait for them to finish
     let max_tasks = min(num_cpus::get() * 2, limit);
-    let mut tasks = std::collections::VecDeque::new();
+    let mut tasks = VecDeque::new();
     let mut count = 0;
     let shareable_img_tfm = Arc::new(image_transform);
 

{datago-2025.4.1 → datago-2025.4.3}/src/worker_http.rs

@@ -5,6 +5,7 @@ use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
 use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
 use serde::{Deserialize, Serialize};
 use std::cmp::min;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 // We'll share a single connection pool across all worker threads
@@ -37,7 +38,7 @@ pub fn new_shared_client(max_connections: usize) -> SharedClient {
 struct SampleMetadata {
     id: String,
     source: String,
-    attributes: std::collections::HashMap<String, serde_json::Value>,
+    attributes: HashMap<String, serde_json::Value>,
     duplicate_state: Option<i32>,
     image_direct_url: Option<String>,
     latents: Option<Vec<UrlLatent>>,
@@ -76,6 +77,28 @@ async fn image_from_url(
     )))
 }
 
+async fn payload_from_url(
+    client: &SharedClient,
+    url: &str,
+    retries: i32,
+) -> Result<Vec<u8>, std::io::Error> {
+    // Retry on the fetch and decode a few times, could happen that we get a broken packet
+    for _ in 0..retries {
+        match bytes_from_url(client, url).await {
+            Some(bytes) => {
+                return Ok(bytes);
+            }
+            None => {
+                println!("Failed to get bytes from URL: {}. Retrying", url);
+            }
+        }
+    }
+    Err(std::io::Error::new(
+        std::io::ErrorKind::Other,
+        "Failed to fetch bytes buffer",
+    ))
+}
+
 async fn image_payload_from_url(
     client: &SharedClient,
     url: &str,
@@ -143,12 +166,9 @@ async fn pull_sample(
     }
 
     // Same for the latents, mask and masked images, if they exist
-    let mut masks: std::collections::HashMap<String, ImagePayload> =
-        std::collections::HashMap::new();
-    let mut additional_images: std::collections::HashMap<String, ImagePayload> =
-        std::collections::HashMap::new();
-    let mut latents: std::collections::HashMap<String, LatentPayload> =
-        std::collections::HashMap::new();
+    let mut masks: HashMap<String, ImagePayload> = HashMap::new();
+    let mut additional_images: HashMap<String, ImagePayload> = HashMap::new();
+    let mut latents: HashMap<String, LatentPayload> = HashMap::new();
 
     if let Some(exposed_latents) = &sample.latents {
         for latent in exposed_latents {
@@ -203,8 +223,8 @@ async fn pull_sample(
                 }
             } else {
                 // Vanilla latents, pure binary payloads
-                match bytes_from_url(&client, &latent.file_direct_url).await {
-                    Some(latent_payload) => {
+                match payload_from_url(&client, &latent.file_direct_url, 5).await {
+                    Ok(latent_payload) => {
                         latents.insert(
                             latent.latent_type.clone(),
                             LatentPayload {
@@ -213,8 +233,8 @@ async fn pull_sample(
                             },
                         );
                     }
-                    None => {
-                        println!("Error fetching latent: {}", latent.file_direct_url);
+                    Err(e) => {
+                        println!("Error fetching latent: {} {}", latent.file_direct_url, e);
                         return Err(());
                     }
                 }