chore(hf): migrate to new XetSession API from xet-core PR #747

Replace low-level xet-data/xet-client usage with the new high-level
session API from xet_pkg (hf-xet crate, lib name `xet`):

- Single hf-xet dependency replaces xet-client + xet-data
- ONE cached XetSession per HfCore (write token, serves both reads
  and writes); removes XetTokenRefresher and TranslatorConfig handling
- Writer: XetUploadCommit + XetStreamUpload replace FileUploadSession
  + SingleFileCleaner; no Mutex needed since oio::Write is &mut self
- Reader: XetDownloadStream replaces DownloadStream; range converted
  from BytesRange to Option<Range<u64>> for new download_stream() API

Author: kszucs

core/services/hf/Cargo.toml

@@ -45,8 +45,7 @@ reqwest = { version = "0.12", default-features = false, features = [
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 tokio = { workspace = true, features = ["sync"] }
-xet-client = { git = "https://github.com/huggingface/xet-core.git", rev = "101837f6" }
-xet-data = { git = "https://github.com/huggingface/xet-core.git", rev = "101837f6" }
+hf-xet = { git = "https://github.com/huggingface/xet-core.git", rev = "60a916f74df506a5d10c0656ee3e86a8411ab4f6" }
 
 [dev-dependencies]
 opendal-core = { path = "../../core", version = "0.55.0", features = [

core/services/hf/src/core.rs

@@ -27,11 +27,7 @@ use http::Response;
 use http::header;
 use serde::Deserialize;
 
-use xet_client::cas_client::auth::TokenRefresher;
-use xet_data::processing::FileDownloadSession;
-use xet_data::processing::FileUploadSession;
-use xet_data::processing::XetFileInfo;
-use xet_data::processing::configurations::TranslatorConfig;
+use xet::xet_session::{XetFileInfo, XetSession, XetSessionBuilder};
 
 use super::error::parse_error;
 use super::uri::HfRepo;
@@ -182,34 +178,6 @@ pub(super) struct XetToken {
     pub exp: u64,
 }
 
-pub(super) struct XetTokenRefresher {
-    core: HfCore,
-    token_type: &'static str,
-}
-
-impl XetTokenRefresher {
-    pub(super) fn new(core: &HfCore, token_type: &'static str) -> Self {
-        Self {
-            core: core.clone(),
-            token_type,
-        }
-    }
-}
-
-#[async_trait::async_trait]
-impl TokenRefresher for XetTokenRefresher {
-    async fn refresh(
-        &self,
-    ) -> std::result::Result<(String, u64), xet_client::cas_client::auth::AuthError> {
-        let token = self
-            .core
-            .xet_token(self.token_type)
-            .await
-            .map_err(xet_client::cas_client::auth::AuthError::token_refresh_failure)?;
-        Ok((token.access_token, token.exp))
-    }
-}
-
 // Core HuggingFace client that manages API interactions, authentication
 // and shared logic for reader/writer/lister.
 
@@ -226,8 +194,7 @@ pub struct HfCore {
     /// inspect headers on 302 responses.
     pub no_redirect_client: HttpClient,
 
-    upload_config: Arc<OnceCell<Arc<TranslatorConfig>>>,
-    download_config: Arc<OnceCell<Arc<TranslatorConfig>>>,
+    xet_session: Arc<OnceCell<XetSession>>,
 }
 
 impl Debug for HfCore {
@@ -256,8 +223,7 @@ impl HfCore {
             token,
             endpoint,
             no_redirect_client,
-            upload_config: Arc::new(OnceCell::new()),
-            download_config: Arc::new(OnceCell::new()),
+            xet_session: Arc::new(OnceCell::new()),
         }
     }
 
@@ -352,52 +318,22 @@ impl HfCore {
         Ok(token)
     }
 
-    async fn xet_config(&self, token_type: &'static str) -> Result<Arc<TranslatorConfig>> {
-        let cas_url = self.xet_token(token_type).await?.cas_url;
-        let refresher = Arc::new(XetTokenRefresher::new(self, token_type));
-        xet_data::processing::data_client::default_config(
-            cas_url,
-            None,
-            Some(refresher as Arc<dyn TokenRefresher>),
-            None,
-        )
-        .map(Arc::new)
-        .map_err(|err| {
-            Error::new(ErrorKind::Unexpected, "failed to create xet config").set_source(err)
-        })
-    }
-
-    async fn upload_config(&self) -> Result<Arc<TranslatorConfig>> {
-        let this = self.clone();
-        self.upload_config
-            .get_or_try_init(|| async move { this.xet_config("write").await })
-            .await
-            .map(Arc::clone)
-    }
-
-    async fn download_config(&self) -> Result<Arc<TranslatorConfig>> {
+    pub(super) async fn xet_session(&self) -> Result<XetSession> {
         let this = self.clone();
-        self.download_config
-            .get_or_try_init(|| async move { this.xet_config("read").await })
-            .await
-            .map(Arc::clone)
-    }
-
-    pub(super) async fn xet_upload_session(&self) -> Result<Arc<FileUploadSession>> {
-        FileUploadSession::new(self.upload_config().await?)
-            .await
-            .map_err(|err| {
-                Error::new(ErrorKind::Unexpected, "failed to create upload session").set_source(err)
+        self.xet_session
+            .get_or_try_init(|| async move {
+                let token = this.xet_token("write").await?;
+                XetSessionBuilder::new()
+                    .with_endpoint(token.cas_url)
+                    .with_token_info(token.access_token, token.exp)
+                    .build()
+                    .map_err(|err| {
+                        Error::new(ErrorKind::Unexpected, "failed to create xet session")
+                            .set_source(err)
+                    })
             })
-    }
-
-    pub(super) async fn xet_download_session(&self) -> Result<Arc<FileDownloadSession>> {
-        FileDownloadSession::new(self.download_config().await?)
             .await
-            .map_err(|err| {
-                Error::new(ErrorKind::Unexpected, "failed to create download session")
-                    .set_source(err)
-            })
+            .cloned()
     }
 
     /// Issue a HEAD request and extract XET file info (hash and size).

core/services/hf/src/reader.rs

@@ -19,10 +19,7 @@ use http::Response;
 use http::StatusCode;
 use http::header;
 
-use xet_client::ClientError;
-use xet_data::file_reconstruction::FileReconstructionError;
-use xet_data::processing::DownloadStream;
-use xet_data::processing::XetFileInfo;
+use xet::xet_session::{SessionError, XetDownloadStream, XetFileInfo};
 
 use super::core::HfCore;
 use super::uri::RepoType;
@@ -31,7 +28,7 @@ use opendal_core::*;
 
 pub enum HfReader {
     Http(HttpBody),
-    Xet(DownloadStream),
+    Xet(XetDownloadStream),
 }
 
 impl HfReader {
@@ -86,32 +83,28 @@ impl HfReader {
         file_info: &XetFileInfo,
         range: BytesRange,
     ) -> Result<Self> {
-        let session = core.xet_download_session().await?;
-
-        let (_id, stream) = session
-            .download_stream_range(file_info, range.to_range())
+        let session = core.xet_session().await?;
+        let xet_range = if range.is_full() {
+            None
+        } else {
+            let start = range.offset();
+            let end = range.size().map(|s| start + s).unwrap_or(u64::MAX);
+            Some(start..end)
+        };
+        let mut stream = session
+            .download_stream(file_info.clone(), xet_range)
             .await
             .map_err(|err| {
-                Error::new(
-                    ErrorKind::Unexpected,
-                    "failed to create xet download stream",
-                )
-                .set_source(err)
+                Error::new(ErrorKind::Unexpected, "failed to create xet download stream")
+                    .set_source(err)
             })?;
+        stream.start();
         Ok(Self::Xet(stream))
     }
 }
 
-fn map_reconstruction_error(e: FileReconstructionError) -> Error {
-    let kind = match &e {
-        FileReconstructionError::ClientError(arc) => match arc.as_ref() {
-            ClientError::FileNotFound(_) | ClientError::XORBNotFound(_) => ErrorKind::NotFound,
-            ClientError::InvalidRange => ErrorKind::RangeNotSatisfied,
-            _ => ErrorKind::Unexpected,
-        },
-        _ => ErrorKind::Unexpected,
-    };
-    Error::new(kind, "xet read error").set_source(e)
+fn map_session_error(e: SessionError) -> Error {
+    Error::new(ErrorKind::Unexpected, "xet read error").set_source(e)
 }
 
 impl oio::Read for HfReader {
@@ -121,7 +114,7 @@ impl oio::Read for HfReader {
             Self::Xet(stream) => match stream.next().await {
                 Ok(Some(bytes)) => Ok(Buffer::from(bytes)),
                 Ok(None) => Ok(Buffer::new()),
-                Err(e) => Err(map_reconstruction_error(e)),
+                Err(e) => Err(map_session_error(e)),
             },
         }
     }

core/services/hf/src/writer.rs

@@ -16,17 +16,14 @@
 // under the License.
 
 use std::sync::Arc;
-use std::sync::Mutex;
 
 use base64::Engine;
 
 use super::core::{BucketOperation, CommitFile, HfCore, LfsFile};
 use super::uri::RepoType;
 use opendal_core::raw::*;
 use opendal_core::*;
-use xet_data::processing::FileUploadSession;
-use xet_data::processing::Sha256Policy;
-use xet_data::processing::SingleFileCleaner;
+use xet::xet_session::{Sha256Policy, XetStreamUpload, XetUploadCommit};
 
 /// Writer that handles both regular (small) and XET (large) file uploads.
 pub enum HfWriter {
@@ -40,8 +37,8 @@ pub enum HfWriter {
     Xet {
         core: Arc<HfCore>,
         path: String,
-        session: Arc<FileUploadSession>,
-        cleaner: Mutex<Option<SingleFileCleaner>>,
+        commit: XetUploadCommit,
+        stream: Option<XetStreamUpload>,
     },
 }
 
@@ -54,17 +51,23 @@ impl HfWriter {
             || core.determine_upload_mode(&path).await? == "lfs";
 
         let writer = if use_xet {
-            let session = core.xet_upload_session().await?;
-            let (_id, cleaner) = session
-                .start_clean(None, None, Sha256Policy::Compute)
+            let session = core.xet_session().await?;
+            let commit = session.new_upload_commit().await.map_err(|err| {
+                Error::new(ErrorKind::Unexpected, "failed to create xet upload commit")
+                    .set_source(err)
+            })?;
+            let stream = commit
+                .upload_stream(None, Sha256Policy::Compute)
+                .await
                 .map_err(|err| {
-                    Error::new(ErrorKind::Unexpected, "failed to start xet upload").set_source(err)
+                    Error::new(ErrorKind::Unexpected, "failed to start xet upload stream")
+                        .set_source(err)
                 })?;
             HfWriter::Xet {
                 core,
                 path,
-                session,
-                cleaner: Mutex::new(Some(cleaner)),
+                commit,
+                stream: Some(stream),
             }
         } else {
             HfWriter::Regular {
@@ -93,14 +96,9 @@ impl oio::Write for HfWriter {
                 buf.push(bs);
                 Ok(())
             }
-            HfWriter::Xet { cleaner, .. } => {
-                // get_mut() required here: lock() would hold a MutexGuard across .await
-                let cleaner = cleaner
-                    .get_mut()
-                    .expect("mutex poisoned")
-                    .as_mut()
-                    .expect("cleaner missing");
-                cleaner.add_data(&bs.to_bytes()).await.map_err(|err| {
+            HfWriter::Xet { stream, .. } => {
+                let stream = stream.as_ref().expect("stream missing");
+                stream.write(bs.to_bytes()).await.map_err(|err| {
                     Error::new(ErrorKind::Unexpected, "failed to write xet chunk").set_source(err)
                 })
             }
@@ -130,65 +128,51 @@ impl oio::Write for HfWriter {
             HfWriter::Xet {
                 core,
                 path,
-                session,
-                cleaner,
+                commit,
+                stream,
             } => {
-                let cleaner = cleaner
-                    .get_mut()
-                    .expect("mutex poisoned")
+                let stream = stream
                     .take()
-                    .ok_or_else(|| {
-                        Error::new(ErrorKind::Unexpected, "xet writer already closed")
-                    })?;
+                    .ok_or_else(|| Error::new(ErrorKind::Unexpected, "xet writer already closed"))?;
 
-                let (file_info, _metrics) = cleaner.finish().await.map_err(|err| {
+                let file_meta = stream.finish().await.map_err(|err| {
                     Error::new(ErrorKind::Unexpected, "failed to finish xet upload").set_source(err)
                 })?;
-
+                let file_info = file_meta.xet_info;
                 let content_length = file_info
                     .file_size()
                     .expect("file_size must be set after finish()");
                 let meta = Metadata::default().with_content_length(content_length);
 
+                commit.commit().await.map_err(|err| {
+                    Error::new(ErrorKind::Unexpected, "failed to commit xet upload")
+                        .set_source(err)
+                })?;
+
                 if core.repo.repo_type == RepoType::Bucket {
                     let xet_hash = file_info.hash().to_string();
-                    session.clone().finalize().await.map_err(|err| {
-                        Error::new(
-                            ErrorKind::Unexpected,
-                            "failed to finalize xet upload session",
-                        )
-                        .set_source(err)
-                    })?;
-                    let operation = BucketOperation::AddFile {
+                    core.bucket_batch(vec![BucketOperation::AddFile {
                         path: path.clone(),
                         xet_hash,
-                    };
-                    core.bucket_batch(vec![operation]).await?;
+                    }])
+                    .await?;
                     Ok(meta)
                 } else {
                     let sha256 = file_info.sha256().ok_or_else(|| {
                         Error::new(ErrorKind::Unexpected, "xet upload returned no sha256")
                     })?;
-                    session.clone().finalize().await.map_err(|err| {
-                        Error::new(
-                            ErrorKind::Unexpected,
-                            "failed to finalize xet upload session",
-                        )
-                        .set_source(err)
-                    })?;
                     let lfs_file = LfsFile {
                         path: path.clone(),
                         oid: sha256.to_string(),
                         algo: "sha256".to_string(),
                         size: content_length,
                     };
                     let resp = core.commit_files(vec![], vec![lfs_file], vec![]).await?;
-
-                    if let Some(commit_oid) = resp.commit_oid {
-                        Ok(meta.with_version(commit_oid))
+                    Ok(if let Some(oid) = resp.commit_oid {
+                        meta.with_version(oid)
                     } else {
-                        Ok(meta)
-                    }
+                        meta
+                    })
                 }
             }
         }
@@ -199,8 +183,10 @@ impl oio::Write for HfWriter {
             HfWriter::Regular { buf, .. } => {
                 buf.clear();
             }
-            HfWriter::Xet { cleaner, .. } => {
-                cleaner.get_mut().expect("mutex poisoned").take();
+            HfWriter::Xet { stream, .. } => {
+                if let Some(s) = stream.take() {
+                    s.abort();
+                }
             }
         }
         Ok(())