hf-xet 0.1.0__tar.gz

hf_xet-0.1.0/Cargo.toml

@@ -0,0 +1,31 @@
+[workspace]
+resolver = "2"
+
+members = [
+    "cas_client",
+    "cas_object",
+    "cas_types",
+    "data",
+    "error_printer",
+    "file_utils",
+    "merklehash",
+    "mdb_shard",
+    "parutils",
+    "progress_reporting",
+    "utils",
+    "cas_object",
+    "cas_types",
+    "chunk_cache",
+]
+
+exclude = ["hf_xet", "chunk_cache_bench"]
+
+[profile.release]
+opt-level = 3
+lto = true
+debug = 1
+
+[profile.opt-test]
+inherits = "dev"
+opt-level = 1
+debug = 1

hf_xet-0.1.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: hf-xet
+Version: 0.1.0
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Dist: pytest ; extra == 'tests'
+Provides-Extra: tests
+License: Apache-2.0
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+
+# Development Notes
+
+* `pip install maturin`
+* from this directory: `maturin develop`
+

hf_xet-0.1.0/cas_client/Cargo.toml

@@ -0,0 +1,41 @@
+[package]
+name = "cas_client"
+version = "0.14.5"
+edition = "2021"
+readme = "README.md"
+
+[features]
+strict = []
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+cas_types = { version = "0.1.0", path = "../cas_types" }
+cas_object = { path = "../cas_object" }
+chunk_cache = { path = "../chunk_cache" }
+error_printer = { path = "../error_printer" }
+file_utils = { path = "../file_utils" }
+utils = { path = "../utils" }
+merkledb = { path = "../merkledb" }
+mdb_shard = { path = "../mdb_shard" }
+merklehash = { path = "../merklehash" }
+thiserror = "2.0"
+tokio = { version = "1.41", features = ["full"] }
+async-trait = "0.1.9"
+anyhow = "1"
+http = "1.1.0"
+tempfile = "3.13.0"
+tracing = "0.1.31"
+bytes = "1"
+itertools = "0.10"
+reqwest = { version = "0.12.7", features = ["json", "stream"] }
+reqwest-middleware = "0.3.3"
+url = "2.5.2"
+reqwest-retry = "0.6.1"
+heed = "0.11"
+futures = "0.3.31"
+
+[dev-dependencies]
+rand = "0.8.5"
+httpmock = "0.7.0"
+tracing-test = { version = "0.2.5", features = ["no-env-filter"] }

hf_xet-0.1.0/cas_client/Dockerfile

@@ -0,0 +1,37 @@
+FROM rust:1.58 as builder
+
+RUN USER=root rustup component add rustfmt
+RUN USER=root cargo new --bin cas_client
+
+WORKDIR ./cas_client
+ADD ./utils ../utils
+COPY ./cas_client/Cargo.toml ./Cargo.toml
+RUN cargo build --release
+RUN rm src/*.rs
+
+COPY ./cas_client .
+RUN rm ./target/release/deps/cas_client*
+
+RUN cargo build --release
+
+FROM debian:buster-slim
+ARG APP=/usr/src/app
+
+RUN apt-get update \
+    && apt-get install -y ca-certificates tzdata \
+    && rm -rf /var/lib/apt/lists/*
+
+ENV TZ=Etc/UTC \
+    APP_USER=appuser
+
+RUN groupadd $APP_USER \
+    && useradd -g $APP_USER $APP_USER \
+    && mkdir -p ${APP}
+
+COPY --from=builder /cas_client/target/release/cas_client ${APP}/cas_client
+RUN mkdir ${APP}/config
+
+RUN chown -R $APP_USER:$APP_USER ${APP}
+
+USER $APP_USER
+WORKDIR ${APP}

hf_xet-0.1.0/cas_client/README.md

@@ -0,0 +1,35 @@
+# CAS client
+
+This package is responsible for handling all communication with the CAS services.
+
+## Layout
+
+Check out the traits published by this crate to understand how it is intended to be used. These are stored in [src/interface.rs].
+
+### Main impl of Client trait
+
+- [src/remote_client.rs]: This is the main impl of Client - and is responsible for communicating with a remote CAS.
+- [src/local_client.rs]: This is an impl of Client for local filesystem usage. It is currently only used for testing.
+
+### Caching
+
+Caching happens locally using the [chunk_cache](../chunk_cache) crate, specifically using the ChunkCache. When RemoteClient is provided a ChunkCache then it will use this on download calls (the ReconstructionClient trait `get_file` and `get_file_byte_range`).
+
+## Overall CAS Communication Design
+
+### Authentication
+
+Authentication is done using AuthMiddleware, which sets an Authorization Header and refreshes it periodically with CAS. See [src/http_client.rs].
+
+### Retry
+
+HTTP operations are retried using a RetryPolicy defined in [src/http_client.rs]. This is implemented as Middleware for the reqwest HTTP clients.
+
+### Operations
+
+CAS offers a set of services used by the client to upload and download user files. These files are stored using two different
+types of storage objects, Xorbs and Shards. Xorbs contain chunks and Shards contain mappings of files to Xorb chunks.
+
+### Logging / Tracing
+
+Logging & Tracing is done through the tracing crate, with info!, warn!, and debug! macros widely used in the code.

hf_xet-0.1.0/cas_client/src/error.rs

@@ -0,0 +1,75 @@
+use merklehash::MerkleHash;
+use thiserror::Error;
+
+#[non_exhaustive]
+#[derive(Error, Debug)]
+pub enum CasClientError {
+    #[error("ChunkCache Error: {0}")]
+    ChunkCache(#[from] chunk_cache::error::ChunkCacheError),
+
+    #[error("Cas Object Error: {0}")]
+    CasObjectError(#[from] cas_object::error::CasObjectError),
+
+    #[error("Configuration Error: {0} ")]
+    ConfigurationError(String),
+
+    #[error("Invalid Range")]
+    InvalidRange,
+
+    #[error("Invalid Arguments")]
+    InvalidArguments,
+
+    #[error("File not found for hash: {0}")]
+    FileNotFound(MerkleHash),
+
+    #[error("IO Error: {0}")]
+    IOError(#[from] std::io::Error),
+
+    #[error("Invalid Shard Key: {0}")]
+    InvalidShardKey(String),
+
+    #[error("Other Internal Error: {0}")]
+    InternalError(#[from] anyhow::Error),
+
+    #[error("MerkleDB Shard Error : {0}")]
+    MDBShardError(#[from] mdb_shard::error::MDBShardError),
+
+    #[error("Error : {0}")]
+    Other(String),
+
+    #[error("Parse Error: {0}")]
+    ParseError(#[from] url::ParseError),
+
+    #[error("ReqwestMiddleware Error: {0}")]
+    ReqwestMiddlewareError(#[from] reqwest_middleware::Error),
+
+    #[error("Reqwest Error: {0}")]
+    ReqwestError(#[from] reqwest::Error),
+
+    #[error("LMDB Error: {0}")]
+    ShardDedupDBError(String),
+
+    #[error("CAS object not found for hash: {0}")]
+    XORBNotFound(MerkleHash),
+}
+
+// Define our own result type here (this seems to be the standard).
+pub type Result<T> = std::result::Result<T, CasClientError>;
+
+impl PartialEq for CasClientError {
+    fn eq(&self, other: &CasClientError) -> bool {
+        match (self, other) {
+            (CasClientError::XORBNotFound(a), CasClientError::XORBNotFound(b)) => a == b,
+            (e1, e2) => std::mem::discriminant(e1) == std::mem::discriminant(e2),
+        }
+    }
+}
+
+impl From<utils::errors::SingleflightError<CasClientError>> for CasClientError {
+    fn from(value: utils::singleflight::SingleflightError<CasClientError>) -> Self {
+        match value {
+            utils::singleflight::SingleflightError::InternalError(e) => e,
+            e => CasClientError::Other(format!("single flight error: {e}")),
+        }
+    }
+}

hf_xet-0.1.0/cas_client/src/global_dedup_table.rs

@@ -0,0 +1,230 @@
+use std::collections::HashMap;
+use std::path::Path;
+use std::sync::Arc;
+
+use heed::types::*;
+use heed::EnvOpenOptions;
+use itertools::Itertools;
+use merkledb::aggregate_hashes::with_salt;
+use merklehash::MerkleHash;
+use tokio::sync::RwLock;
+use tracing::{info, warn};
+
+use crate::error::{CasClientError, Result};
+
+type DB = heed::Database<OwnedType<MerkleHash>, OwnedType<MerkleHash>>;
+
+pub struct DiskBasedGlobalDedupTable {
+    env: heed::Env,
+    table: RwLock<HashMap<String, Arc<DB>>>, // map of chunk_hash -> shard_hash
+}
+
+// Annoyingly, heed::Error is not Send/Sync, so convert to string.
+fn map_db_error(e: heed::Error) -> CasClientError {
+    let msg = format!("Global shard dedup database error: {e:?}");
+    warn!("{msg}");
+    CasClientError::ShardDedupDBError(msg)
+}
+
+impl DiskBasedGlobalDedupTable {
+    pub fn open_or_create(path: impl AsRef<Path>) -> Result<Self> {
+        let db_path = path.as_ref().join("global_shard_dedup.db");
+        info!("Using {db_path:?} as path to global shard dedup database.");
+
+        std::fs::create_dir_all(&db_path)?;
+        let env = EnvOpenOptions::new()
+            .max_dbs(32)
+            .max_readers(32)
+            .open(&db_path)
+            .map_err(map_db_error)?;
+
+        Ok(Self {
+            env,
+            table: RwLock::new(HashMap::new()),
+        })
+    }
+
+    async fn get_db(&self, prefix: &str) -> Result<Arc<DB>> {
+        if let Some(db) = self.table.read().await.get(prefix).cloned() {
+            return Ok(db);
+        }
+
+        let mut write_lock = self.table.write().await;
+
+        match write_lock.entry(prefix.to_owned()) {
+            std::collections::hash_map::Entry::Occupied(db) => Ok(db.get().clone()),
+            std::collections::hash_map::Entry::Vacant(entry_ref) => {
+                let db = Arc::new(self.env.create_database(Some(prefix)).map_err(map_db_error)?);
+                entry_ref.insert(db.clone());
+                Ok(db)
+            },
+        }
+    }
+
+    pub async fn batch_add(
+        &self,
+        chunk_hashes: &[MerkleHash],
+        shard_hash: &MerkleHash,
+        prefix: &str,
+        salt: &[u8; 32],
+    ) -> Result<()> {
+        let db = self.get_db(prefix).await?;
+
+        let mut write_txn = self.env.write_txn().map_err(map_db_error)?;
+
+        chunk_hashes.iter().for_each(|chunk| {
+            let maybe_salted_chunk_hash = with_salt(chunk, salt).ok();
+            if let Some(salted_chunk_hash) = maybe_salted_chunk_hash {
+                let _ = db.put(&mut write_txn, &salted_chunk_hash, shard_hash).map_err(map_db_error);
+                // Prints warning for error, otherwise ignores.
+            }
+        });
+        write_txn.commit().map_err(map_db_error)?;
+
+        Ok(())
+    }
+
+    pub async fn query(&self, salted_chunk_hash: &[MerkleHash], prefix: &str) -> Vec<MerkleHash> {
+        let Ok(db) = self.get_db(prefix).await else {
+            return vec![];
+        };
+
+        let Ok(read_txn) = self.env.read_txn().map_err(|e| {
+            warn!("Error starting read transaction for prefix {prefix}: {e:?}");
+            e
+        }) else {
+            return vec![];
+        };
+
+        salted_chunk_hash
+            .iter()
+            .filter_map(|chunk| db.get(&read_txn, chunk).unwrap_or(None))
+            .collect_vec()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use itertools::Itertools;
+    use mdb_shard::shard_format::test_routines::rng_hash;
+    use merkledb::aggregate_hashes::with_salt;
+    use rand::{thread_rng, Rng};
+    use tempfile::TempDir;
+
+    use super::DiskBasedGlobalDedupTable;
+
+    #[tokio::test]
+    async fn test_basic_insert_retrieval() -> anyhow::Result<()> {
+        let tempdir = TempDir::new()?;
+
+        let db_file = tempdir.path().join("db");
+
+        let db = DiskBasedGlobalDedupTable::open_or_create(&db_file)?;
+
+        let mut rng = thread_rng();
+
+        let prefix = "default";
+        let chunk_hash = rng_hash(rng.gen());
+        let shard_hash = rng_hash(rng.gen());
+        let salt: [u8; 32] = rng.gen();
+
+        db.batch_add(&[chunk_hash], &shard_hash, "default", &salt).await?;
+
+        let query_shard = db.query(&[with_salt(&chunk_hash, &salt)?], prefix).await;
+
+        assert_eq!(query_shard.len(), 1);
+        assert_eq!(query_shard.first(), Some(&shard_hash));
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_multithread_insert_retrieval() -> anyhow::Result<()> {
+        let tempdir = TempDir::new()?;
+
+        let db_file = tempdir.path().join("db");
+
+        let db = Arc::new(DiskBasedGlobalDedupTable::open_or_create(&db_file)?);
+
+        let mut rng = thread_rng();
+        let prefix = "default";
+        let chunk_hashes = (0..10).map(|_| rng_hash(rng.gen())).collect_vec();
+        let shard_hashes = (0..10).map(|_| rng_hash(rng.gen())).collect_vec();
+        let salt: [u8; 32] = rng.gen();
+
+        // insert to the db concurrently
+        let handles = (0..10)
+            .map(|i| {
+                let chunk_hash = chunk_hashes[i];
+                let shard_hash = shard_hashes[i];
+                let db = db.clone();
+
+                tokio::spawn(async move { db.batch_add(&[chunk_hash], &shard_hash, prefix, &salt).await })
+            })
+            .collect_vec();
+
+        for h in handles {
+            let _ = h.await?;
+        }
+
+        // now examine that inserts succeeded
+        for i in 0..10 {
+            let chunk_hash = chunk_hashes[i];
+            let shard_hash = shard_hashes[i];
+            let query_shard = db.query(&[with_salt(&chunk_hash, &salt)?], prefix).await;
+
+            assert_eq!(query_shard.len(), 1);
+            assert_eq!(query_shard.first(), Some(&shard_hash));
+        }
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_multi_db_instance_insert_retrieval() -> anyhow::Result<()> {
+        let tempdir = TempDir::new()?;
+
+        let db_file = tempdir.path().join("db");
+
+        let mut rng = thread_rng();
+        let prefix = "default";
+        let chunk_hashes = (0..1000).map(|_| rng_hash(rng.gen())).collect_vec();
+        let shard_hashes = (0..10).map(|_| rng_hash(rng.gen())).collect_vec();
+        let salt: [u8; 32] = rng.gen();
+
+        // insert to the db concurrently
+        let handles = (0..10)
+            .map(|i| {
+                let chunk_hashes = chunk_hashes[i * 100..(i + 1) * 100].to_vec();
+                let shard_hash = shard_hashes[i];
+                let db_file = db_file.clone();
+
+                tokio::spawn(async move {
+                    let db = DiskBasedGlobalDedupTable::open_or_create(&db_file).unwrap();
+                    db.batch_add(&chunk_hashes, &shard_hash, prefix, &salt).await
+                })
+            })
+            .collect_vec();
+
+        for h in handles {
+            let _ = h.await?;
+        }
+
+        // now examine that inserts succeeded
+        let db = DiskBasedGlobalDedupTable::open_or_create(&db_file)?;
+        for i in 0..10 {
+            let shard_hash = shard_hashes[i];
+
+            for chunk_hash in &chunk_hashes[i * 100..(i + 1) * 100] {
+                let query_shard = db.query(&[with_salt(chunk_hash, &salt)?], prefix).await;
+
+                assert_eq!(query_shard.len(), 1);
+                assert_eq!(query_shard.first(), Some(&shard_hash));
+            }
+        }
+
+        Ok(())
+    }
+}