feat: add support for HF_HUB_USER_AGENT_ORIGIN to add user-agent Origin field in Hub requests.

Cargo.toml

@@ -29,7 +29,7 @@ homepage = "https://github.com/huggingface/text-generation-inference"
 [workspace.dependencies]
 base64 = "0.22.0"
 tokenizers = { version = "0.20.0", features = ["http"] }
-hf-hub = { version = "0.3.1", features = ["tokio"] }
+hf-hub = { version = "0.4.1", features = ["tokio"] }
 metrics = { version = "0.23.0" }
 metrics-exporter-prometheus = { version = "0.15.1", features = [] }
 minijinja = { version = "2.2.0", features = ["json"] }

Dockerfile.neuron

@@ -18,7 +18,8 @@ RUN apt-get update -y \
     && rm -rf /var/lib/apt/lists/* \
     && apt-get clean
 
-RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --default-toolchain 1.80.1 --profile minimal -y
+COPY rust-toolchain.toml rust-toolchain.toml
+RUN curl -sSf https://sh.rustup.rs | sh -s -- -y --no-modify-path --default-toolchain none
 ENV PATH="/root/.cargo/bin:${PATH}"
 RUN cargo install cargo-chef --locked
 

backends/neuron/Cargo.toml

@@ -22,7 +22,7 @@ homepage = "https://github.com/huggingface/text-generation-inference"
 [workspace.dependencies]
 base64 = "0.22.0"
 tokenizers = { version = "0.20.0", features = ["http"] }
-hf-hub = { version = "0.3.1", features = ["tokio"] }
+hf-hub = { version = "0.4.2", features = ["tokio"] }
 metrics = { version = "0.23.0" }
 metrics-exporter-prometheus = { version = "0.15.1", features = [] }
 minijinja = { version = "2.2.0", features = ["json"] }

backends/trtllm/src/main.rs

@@ -86,6 +86,10 @@ async fn get_tokenizer(tokenizer_name: &str, revision: Option<&str>) -> Option<T
             builder = builder.with_cache_dir(cache_dir.into());
         }
 
+        if let Ok(origin) = std::env::var("HF_HUB_USER_AGENT_ORIGIN") {
+            builder = builder.with_user_agent("origin", origin.as_str());
+        }
+
         builder
     };
 

launcher/Cargo.toml

@@ -9,7 +9,7 @@ homepage.workspace = true
 [dependencies]
 clap = { version = "4.4.5", features = ["derive", "env"] }
 ctrlc = { version = "3.4.1", features = ["termination"] }
-hf-hub = "0.4.1"
+hf-hub = "0.4.2"
 nix = { version = "0.28.0", features = ["signal"] }
 once_cell = "1.19.0"
 pyo3 = { workspace = true }

launcher/src/main.rs

@@ -1,8 +1,5 @@
 use clap::{Parser, ValueEnum};
-use hf_hub::{
-    api::sync::{Api, ApiBuilder},
-    Repo, RepoType,
-};
+use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
 use nix::sys::signal::{self, Signal};
 use nix::unistd::Pid;
 use serde::Deserialize;
@@ -100,12 +97,16 @@ fn get_config(
     let filename = if !path.exists() {
         // Assume it's a hub id
 
-        let api = if let Ok(token) = std::env::var("HF_TOKEN") {
+        let mut builder = if let Ok(token) = std::env::var("HF_TOKEN") {
             // env variable has precedence over on file token.
-            ApiBuilder::new().with_token(Some(token)).build()?
+            ApiBuilder::new().with_token(Some(token))
         } else {
-            Api::new()?
+            ApiBuilder::new()
         };
+        if let Ok(origin) = env::var("HF_HUB_USER_AGENT_ORIGIN") {
+            builder = builder.with_user_agent("origin", origin.as_str());
+        }
+        let api = builder.build()?;
         let repo = if let Some(ref revision) = revision {
             api.repo(Repo::with_revision(
                 model_id,

router/src/server.rs

@@ -1719,6 +1719,10 @@ pub async fn run(
             builder = builder.with_cache_dir(cache_dir.into());
         }
 
+        if let Ok(origin) = std::env::var("HF_HUB_USER_AGENT_ORIGIN") {
+            builder = builder.with_user_agent("origin", origin.as_str());
+        }
+
         builder
     };
 

server/pyproject.toml

@@ -31,7 +31,8 @@ dependencies = [
     "sentencepiece>=0.2.0",
     "tokenizers>=0.20.3",
     "typer>=0.15.1",
-    "transformers>=4.48.0"
+    "transformers>=4.48.0",
+    "huggingface-hub>=0.29.0",
 ]
 
 [build-system]