xet_data/processing/migration_tool/
migrate.rs

1use std::sync::Arc;
2
3use http::header;
4use tracing::{Instrument, Span, info_span, instrument};
5use xet_client::cas_client::auth::TokenRefresher;
6use xet_client::hub_client::{BearerCredentialHelper, CredentialHelper, HubClient, Operation, RepoInfo};
7use xet_core_structures::metadata_shard::file_structs::MDBFileInfo;
8use xet_runtime::core::XetContext;
9use xet_runtime::core::par_utils::run_constrained;
10
11use super::super::data_client::{clean_file, default_config};
12use super::super::{FileUploadSession, Sha256Policy, XetFileInfo};
13use super::hub_client_token_refresher::HubClientTokenRefresher;
14use crate::error::{DataError, Result};
15
16const USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
17
18/// Migrate files to the Hub with external async runtime.
19/// How to use:
20/// ```no_run
21/// let file_paths = vec!["/path/to/file1".to_string(), "/path/to/file2".to_string()];
22/// let hub_endpoint = "https://huggingface.co";
23/// let hub_token = "your_token";
24/// let repo_type = "model";
25/// let repo_id = "your_repo_id";
26/// migrate_with_external_runtime(file_paths, hub_endpoint, hub_token, repo_type, repo_id).await?;
27/// ```
28pub async fn migrate_with_external_runtime(
29    file_paths: Vec<String>,
30    sha256s: Option<Vec<String>>,
31    hub_endpoint: &str,
32    cas_endpoint: Option<String>,
33    hub_token: &str,
34    repo_type: &str,
35    repo_id: &str,
36) -> Result<()> {
37    let cred_helper = BearerCredentialHelper::new(hub_token.to_owned(), "");
38    let mut headers = header::HeaderMap::new();
39    headers.insert(header::USER_AGENT, header::HeaderValue::from_static(USER_AGENT));
40    let ctx = XetContext::default()?;
41    let hub_client = HubClient::new(
42        ctx.clone(),
43        hub_endpoint,
44        RepoInfo::try_from(repo_type, repo_id)?,
45        Some("main".to_owned()),
46        "",
47        Some(cred_helper as Arc<dyn CredentialHelper>),
48        Some(headers),
49    )?;
50
51    migrate_files_impl(&ctx, file_paths, sha256s, false, hub_client, cas_endpoint, false).await?;
52
53    Ok(())
54}
55
56/// mdb file info (if dryrun), cleaned file info, total bytes uploaded
57pub type MigrationInfo = (Vec<MDBFileInfo>, Vec<(XetFileInfo, u64)>, u64);
58
59#[instrument(skip_all, name = "migrate_files", fields(session_id = tracing::field::Empty, num_files = file_paths.len()))]
60pub async fn migrate_files_impl(
61    ctx: &XetContext,
62    file_paths: Vec<String>,
63    sha256s: Option<Vec<String>>,
64    sequential: bool,
65    hub_client: HubClient,
66    cas_endpoint: Option<String>,
67    dry_run: bool,
68) -> Result<MigrationInfo> {
69    let operation = Operation::Upload;
70    let jwt_info = hub_client.get_cas_jwt(operation).await?;
71    let token_refresher = Arc::new(HubClientTokenRefresher {
72        operation,
73        client: Arc::new(hub_client),
74    }) as Arc<dyn TokenRefresher>;
75    let cas = cas_endpoint.unwrap_or(jwt_info.cas_url);
76
77    // Create headers with USER_AGENT
78    let mut headers = http::HeaderMap::new();
79    headers.insert(http::header::USER_AGENT, http::HeaderValue::from_static(USER_AGENT));
80
81    let config = default_config(
82        ctx,
83        cas,
84        Some((jwt_info.access_token, jwt_info.exp)),
85        Some(token_refresher),
86        Some(Arc::new(headers)),
87    )?;
88    Span::current().record("session_id", &config.session.session_id);
89
90    let num_workers = if sequential {
91        1
92    } else {
93        ctx.runtime.num_worker_threads()
94    };
95    let processor = if dry_run {
96        FileUploadSession::dry_run(config.into()).await?
97    } else {
98        FileUploadSession::new(config.into()).await?
99    };
100
101    let sha256_policies: Vec<Sha256Policy> = match sha256s {
102        Some(v) => {
103            if v.len() != file_paths.len() {
104                return Err(DataError::ParameterError(
105                    "mismatched length of the file list and the sha256 list".to_string(),
106                ));
107            }
108            v.iter().map(|s| Sha256Policy::from_hex(s)).collect()
109        },
110        None => vec![Sha256Policy::Compute; file_paths.len()],
111    };
112
113    let clean_futs = file_paths.into_iter().zip(sha256_policies).map(|(file_path, policy)| {
114        let proc = processor.clone();
115        async move {
116            let (pf, metrics) = clean_file(proc, file_path, policy).await?;
117            Ok::<(XetFileInfo, u64), DataError>((pf, metrics.new_bytes))
118        }
119        .instrument(info_span!("clean_file"))
120    });
121    let clean_ret = run_constrained(clean_futs, num_workers).await?;
122
123    if dry_run {
124        let (metrics, all_file_info) = processor.finalize_with_file_info().await?;
125        Ok((all_file_info, clean_ret, metrics.total_bytes_uploaded))
126    } else {
127        let metrics = processor.finalize().await?;
128        Ok((vec![], clean_ret, metrics.total_bytes_uploaded as u64))
129    }
130}
xet_data/processing/migration_tool/migrate.rs

xet_data/processing/migration_tool/
migrate.rs
