Plumb endpoint and access token (#19)

* Basic plumbing of the endpoint and access token through to calls to xetcas

Author: bpronan

cas_client/src/remote_client.rs

@@ -4,7 +4,7 @@ use anyhow::anyhow;
 use bytes::Buf;
 use cas::key::Key;
 use cas_types::{QueryChunkResponse, QueryReconstructionResponse, UploadXorbResponse};
-use reqwest::{StatusCode, Url};
+use reqwest::{header::{HeaderMap, HeaderValue}, StatusCode, Url};
 use serde::{de::DeserializeOwned, Serialize};
 
 use bytes::Bytes;
@@ -82,9 +82,9 @@ impl Client for RemoteClient {
 }
 
 impl RemoteClient {
-    pub async fn from_config(endpoint: String) -> Self {
-        Self {
-            client: CASAPIClient::new(&endpoint),
+    pub async fn from_config(endpoint: String, token: Option<String>) -> Self {
+        Self { 
+            client: CASAPIClient::new(&endpoint, token) 
         }
     }
 }
@@ -93,20 +93,22 @@ impl RemoteClient {
 pub struct CASAPIClient {
     client: reqwest::Client,
     endpoint: String,
+    token: Option<String>,
 }
 
 impl Default for CASAPIClient {
     fn default() -> Self {
-        Self::new(CAS_ENDPOINT)
+        Self::new(CAS_ENDPOINT, None)
     }
 }
 
 impl CASAPIClient {
-    pub fn new(endpoint: &str) -> Self {
+    pub fn new(endpoint: &str, token: Option<String>) -> Self {
         let client = reqwest::Client::builder().build().unwrap();
         Self {
             client,
             endpoint: endpoint.to_string(),
+            token
         }
     }
 
@@ -222,12 +224,17 @@ impl CASAPIClient {
     /// Reconstruct the file
     async fn reconstruct_file(&self, file_id: &MerkleHash) -> Result<QueryReconstructionResponse> {
         let url = Url::parse(&format!(
-            "{}/reconstruction/{}",
-            self.endpoint,
+            "{}/reconstruction/{}", 
+            self.endpoint, 
             file_id.hex()
         ))?;
 
-        let response = self.client.get(url).send().await?;
+        let mut headers = HeaderMap::new();
+        if let Some(tok) = &self.token {
+            headers.insert("Authorization", HeaderValue::from_str(&format!("Bearer {}", tok)).unwrap());
+        }
+
+        let response = self.client.get(url).headers(headers).send().await?;
         let response_body = response.bytes().await?;
         let response_parsed: QueryReconstructionResponse =
             serde_json::from_reader(response_body.reader())?;
@@ -290,7 +297,7 @@ mod tests {
     #[tokio::test]
     async fn test_basic_put() {
         // Arrange
-        let rc = RemoteClient::from_config(CAS_ENDPOINT.to_string()).await;
+        let rc = RemoteClient::from_config(CAS_ENDPOINT.to_string(), None).await;
         let prefix = PREFIX_DEFAULT;
         let (hash, data, chunk_boundaries) = gen_dummy_xorb(3, 10248, true);
 

cas_types/src/lib.rs

@@ -22,7 +22,9 @@ pub struct CASReconstructionTerm {
     pub hash: HexMerkleHash,
     pub unpacked_length: u32,
     pub range: Range,
-    pub range_start_offset: u32,
+    // TODO: disabling until https://github.com/huggingface-internal/xetcas/pull/31/files
+    // is merged.
+    // pub range_start_offset: u32,
     pub url: String,
 }
 

data/src/bin/example.rs

@@ -74,8 +74,7 @@ fn default_clean_config() -> Result<TranslatorConfig> {
         cas_storage_config: StorageConfig {
             endpoint: Endpoint::FileSystem(path.join("xorbs")),
             auth: Auth {
-                user_id: "".into(),
-                login_id: "".into(),
+                token: None,
             },
             prefix: "default".into(),
             cache_config: Some(CacheConfig {
@@ -88,8 +87,7 @@ fn default_clean_config() -> Result<TranslatorConfig> {
         shard_storage_config: StorageConfig {
             endpoint: Endpoint::FileSystem(path.join("xorbs")),
             auth: Auth {
-                user_id: "".into(),
-                login_id: "".into(),
+                token: None,
             },
             prefix: "default-merkledb".into(),
             cache_config: Some(CacheConfig {
@@ -123,8 +121,7 @@ fn default_smudge_config() -> Result<TranslatorConfig> {
         cas_storage_config: StorageConfig {
             endpoint: Endpoint::FileSystem(path.join("xorbs")),
             auth: Auth {
-                user_id: "".into(),
-                login_id: "".into(),
+                token: None,
             },
             prefix: "default".into(),
             cache_config: Some(CacheConfig {
@@ -137,8 +134,7 @@ fn default_smudge_config() -> Result<TranslatorConfig> {
         shard_storage_config: StorageConfig {
             endpoint: Endpoint::FileSystem(path.join("xorbs")),
             auth: Auth {
-                user_id: "".into(),
-                login_id: "".into(),
+                token: None,
             },
             prefix: "default-merkledb".into(),
             cache_config: Some(CacheConfig {
@@ -234,7 +230,7 @@ async fn smudge(mut reader: impl Read, mut writer: impl Write) -> Result<()> {
     let translator = PointerFileTranslator::new(default_smudge_config()?).await?;
 
     translator
-        .smudge_file_from_pointer(&pointer_file, &mut writer, None)
+        .smudge_file_from_pointer(&pointer_file, &mut writer, None, None, None)
         .await?;
 
     Ok(())

data/src/cas_interface.rs

@@ -36,9 +36,6 @@ pub(crate) async fn create_cas_client(
         unreachable!();
     };
 
-    // Auth info.
-    let _user_id = &cas_storage_config.auth.user_id;
-    let _auth = &cas_storage_config.auth.login_id;
 
     // Usage tracking.
     let _repo_paths = maybe_repo_info
@@ -49,7 +46,7 @@ pub(crate) async fn create_cas_client(
 
     // Raw remote client.
     let remote_client = Arc::new(
-        RemoteClient::from_config(endpoint.to_string()).await,
+        RemoteClient::from_config(endpoint.to_string(), cas_storage_config.auth.token.clone()).await,
     );
 
     // Try add in caching capability.

data/src/configurations.rs

@@ -11,8 +11,7 @@ pub enum Endpoint {
 
 #[derive(Debug)]
 pub struct Auth {
-    pub user_id: String,
-    pub login_id: String,
+    pub token: Option<String>,
 }
 
 #[derive(Debug)]

data/src/data_processing.rs

@@ -313,8 +313,10 @@ impl PointerFileTranslator {
         pointer: &PointerFile,
         writer: &mut impl std::io::Write,
         range: Option<(usize, usize)>,
+        endpoint: Option<String>,
+        token: Option<String>,
     ) -> Result<()> {
-        self.smudge_file_from_hash(&pointer.hash()?, writer, range)
+        self.smudge_file_from_hash(&pointer.hash()?, writer, range, endpoint, token)
             .await
     }
 
@@ -323,13 +325,21 @@ impl PointerFileTranslator {
         file_id: &MerkleHash,
         writer: &mut impl std::io::Write,
         _range: Option<(usize, usize)>,
+        endpoint: Option<String>,
+        token: Option<String>,
     ) -> Result<()> {
         let endpoint = match &self.config.cas_storage_config.endpoint {
-            Endpoint::Server(endpoint) => endpoint.clone(),
+            Endpoint::Server(config_endpoint) => {
+                if let Some(endpoint) = endpoint {
+                    endpoint
+                } else {
+                config_endpoint.clone()
+                }
+            },
             Endpoint::FileSystem(_) => panic!("aaaaaaaa no server"),
         };
 
-        let rc = CASAPIClient::new(&endpoint);
+        let rc = CASAPIClient::new(&endpoint, token);
 
         rc.write_file(file_id, writer).await?;
 

data/src/shard_interface.rs

@@ -39,7 +39,7 @@ pub async fn create_shard_client(
 ) -> Result<Arc<dyn ShardClientInterface>> {
     info!("Shard endpoint = {:?}", shard_storage_config.endpoint);
     let client: Arc<dyn ShardClientInterface> = match &shard_storage_config.endpoint {
-        Server(endpoint) => Arc::new(HttpShardClient::new(endpoint)),
+        Server(endpoint) => Arc::new(HttpShardClient::new(endpoint, shard_storage_config.auth.token.clone())),
         FileSystem(path) => Arc::new(LocalShardClient::new(path).await?),
     };
 

hf_xet/Cargo.lock

@@ -5,7 +5,7 @@ use data::{DEFAULT_BLOCK_SIZE, errors};
 
 pub const SMALL_FILE_THRESHOLD: usize = 1;
 
-pub fn default_config(endpoint: String) -> errors::Result<TranslatorConfig> {
+pub fn default_config(endpoint: String, token: Option<String>) -> errors::Result<TranslatorConfig> {
     let path = current_dir()?.join(".xet");
     fs::create_dir_all(&path)?;
 
@@ -14,8 +14,7 @@ pub fn default_config(endpoint: String) -> errors::Result<TranslatorConfig> {
         cas_storage_config: StorageConfig {
             endpoint: Endpoint::Server(endpoint.clone()),
             auth: Auth {
-                user_id: "".into(),
-                login_id: "".into(),
+                token: token.clone(),
             },
             prefix: "default".into(),
             cache_config: Some(CacheConfig {
@@ -28,8 +27,7 @@ pub fn default_config(endpoint: String) -> errors::Result<TranslatorConfig> {
         shard_storage_config: StorageConfig {
             endpoint: Endpoint::Server(endpoint),
             auth: Auth {
-                user_id: "".into(),
-                login_id: "".into(),
+                token: token,
             },
             prefix: "default-merkledb".into(),
             cache_config: Some(CacheConfig {

hf_xet/src/config.rs

@@ -12,16 +12,17 @@ use crate::config::default_config;
 pub const MAX_CONCURRENT_UPLOADS: usize = 8; // TODO
 pub const MAX_CONCURRENT_DOWNLOADS: usize = 8; // TODO
 
-const DEFAULT_CAS_ENDPOINT: &str = "https://cas-server.us.dev.moon.huggingface.tech"; 
+const DEFAULT_CAS_ENDPOINT: &str = "http://localhost:8080";
 const READ_BLOCK_SIZE: usize = 1024 * 1024;
 
-pub async fn upload_async(file_paths: Vec<String>) -> errors::Result<Vec<PointerFile>> {
+pub async fn upload_async(file_paths: Vec<String>, endpoint: Option<String>, token: Option<String>) -> errors::Result<Vec<PointerFile>> {
     // chunk files
     // produce Xorbs + Shards
     // upload shards and xorbs
     // for each file, return the filehash
+    let endpoint = endpoint.unwrap_or(DEFAULT_CAS_ENDPOINT.to_string());
 
-    let config = default_config(DEFAULT_CAS_ENDPOINT.to_string())?;
+    let config = default_config(endpoint, token)?;
     let processor = Arc::new(PointerFileTranslator::new(config).await?);
     let processor = &processor;
     // for all files, clean them, producing pointer files.
@@ -45,16 +46,20 @@ pub async fn upload_async(file_paths: Vec<String>) -> errors::Result<Vec<Pointer
     Ok(pointers)
 }
 
-pub async fn download_async(pointer_files: Vec<PointerFile>) -> errors::Result<Vec<String>> {
-    let config = default_config(DEFAULT_CAS_ENDPOINT.to_string())?;
+pub async fn download_async(pointer_files: Vec<PointerFile>, endpoint: Option<String>, token: Option<String>) -> errors::Result<Vec<String>> {
+    let config = default_config(endpoint.clone().unwrap_or(DEFAULT_CAS_ENDPOINT.to_string()), token.clone())?;
     let processor = Arc::new(PointerFileTranslator::new(config).await?);
     let processor = &processor;
     let paths = tokio_par_for_each(
         pointer_files,
         MAX_CONCURRENT_DOWNLOADS,
-        |pointer_file, _| async move {
-            let proc = processor.clone();
-            smudge_file(&proc, &pointer_file).await
+        |pointer_file, _| {
+            let tok = token.clone();
+            let end = endpoint.clone();
+            async move {
+                let proc = processor.clone();
+                smudge_file(&proc, &pointer_file, end.clone(), tok.clone()).await
+            }
         },
     ).await.map_err(|e| match e {
         ParallelError::JoinError => {
@@ -87,13 +92,13 @@ async fn clean_file(processor: &PointerFileTranslator, f: String) -> errors::Res
     Ok(pf)
 }
 
-async fn smudge_file(proc: &PointerFileTranslator, pointer_file: &PointerFile) -> errors::Result<String> {
+async fn smudge_file(proc: &PointerFileTranslator, pointer_file: &PointerFile, endpoint: Option<String>, token: Option<String>) -> errors::Result<String> {
     let path = PathBuf::from(pointer_file.path());
     if let Some(parent_dir) = path.parent() {
         fs::create_dir_all(parent_dir)?;
     }
     let mut f = File::create(&path)?;
-    proc.smudge_file_from_pointer(&pointer_file, &mut f, None).await?;
+    proc.smudge_file_from_pointer(&pointer_file, &mut f, None, endpoint, token).await?;
     Ok(pointer_file.path().to_string())
 }
 
@@ -122,7 +127,7 @@ mod tests {
         let pointers = vec![
             PointerFile::init_from_info("/tmp/foo.rs", "6999733a46030e67f6f020651c91442ace735572458573df599106e54646867c", 4203),
         ];
-        let paths = download_async(pointers).await.unwrap();
+        let paths = download_async(pointers, "http://localhost:8080", "12345").await.unwrap();
         println!("paths: {paths:?}");
     }
 }