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

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

* fix: Rust version for Neuron

* fix: PR comments, use rust-toolchain.toml
This commit is contained in:
Hugo Larcher 2025-03-04 16:43:50 +01:00 committed by GitHub
parent e88f6f6ee9
commit d8ff7f2623
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 48 additions and 37 deletions

42
Cargo.lock generated
View File

@ -1007,9 +1007,9 @@ dependencies = [
[[package]] [[package]]
name = "cxx" name = "cxx"
version = "1.0.140" version = "1.0.141"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc49567e08c72902f4cbc7242ee8d874ec9cbe97fbabf77b4e0e1f447513e13a" checksum = "8bc580dceb395cae0efdde0a88f034cfd8a276897e40c693a7b87bed17971d33"
dependencies = [ dependencies = [
"cc", "cc",
"cxxbridge-cmd", "cxxbridge-cmd",
@ -1021,9 +1021,9 @@ dependencies = [
[[package]] [[package]]
name = "cxx-build" name = "cxx-build"
version = "1.0.140" version = "1.0.141"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fe46b5309c99e9775e7a338c98e4097455f52db5b684fd793ca22848fde6e371" checksum = "49d8c1baedad72a7efda12ad8d7ad687b3e7221dfb304a12443fd69e9de8bb30"
dependencies = [ dependencies = [
"cc", "cc",
"codespan-reporting", "codespan-reporting",
@ -1035,9 +1035,9 @@ dependencies = [
[[package]] [[package]]
name = "cxxbridge-cmd" name = "cxxbridge-cmd"
version = "1.0.140" version = "1.0.141"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4315c4ce8d23c26d87f2f83698725fd5718d8e6ace4a9093da2664d23294d372" checksum = "e43afb0e3b2ef293492a31ecd796af902112460d53e5f923f7804f348a769f9c"
dependencies = [ dependencies = [
"clap 4.5.30", "clap 4.5.30",
"codespan-reporting", "codespan-reporting",
@ -1048,15 +1048,15 @@ dependencies = [
[[package]] [[package]]
name = "cxxbridge-flags" name = "cxxbridge-flags"
version = "1.0.140" version = "1.0.141"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f55d69deb3a92f610a60ecc524a72c7374b6dc822f8fb7bb4e5d9473f10530c4" checksum = "0257ad2096a2474fe877e9e055ab69603851c3d6b394efcc7e0443899c2492ce"
[[package]] [[package]]
name = "cxxbridge-macro" name = "cxxbridge-macro"
version = "1.0.140" version = "1.0.141"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5bee7a1d9b5091462002c2b8de2a4ed0f0fde011d503cc272633f66075bd5141" checksum = "b46cbd7358a46b760609f1cb5093683328e58ca50e594a308716f5403fdc03e5"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
@ -1660,30 +1660,27 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2b780635574b3d92f036890d8373433d6f9fc7abb320ee42a5c25897fc8ed732" checksum = "2b780635574b3d92f036890d8373433d6f9fc7abb320ee42a5c25897fc8ed732"
dependencies = [ dependencies = [
"dirs", "dirs",
"futures",
"indicatif", "indicatif",
"log", "log",
"native-tls", "native-tls",
"num_cpus",
"rand 0.8.5", "rand 0.8.5",
"reqwest 0.11.27",
"serde", "serde",
"serde_json", "serde_json",
"thiserror 1.0.69", "thiserror 1.0.69",
"tokio",
"ureq", "ureq",
] ]
[[package]] [[package]]
name = "hf-hub" name = "hf-hub"
version = "0.4.1" version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "112fa2f6ad4ab815b9e1b938b4b1e437032d055e2f92ed10fd6ab2e62d02c6b6" checksum = "cc03dcb0b0a83ae3f3363ec811014ae669f083e4e499c66602f447c4828737a1"
dependencies = [ dependencies = [
"dirs", "dirs",
"futures", "futures",
"http 1.2.0", "http 1.2.0",
"indicatif", "indicatif",
"libc",
"log", "log",
"native-tls", "native-tls",
"num_cpus", "num_cpus",
@ -1694,6 +1691,7 @@ dependencies = [
"thiserror 2.0.11", "thiserror 2.0.11",
"tokio", "tokio",
"ureq", "ureq",
"windows-sys 0.59.0",
] ]
[[package]] [[package]]
@ -4627,7 +4625,7 @@ dependencies = [
"cxx", "cxx",
"cxx-build", "cxx-build",
"hashbrown 0.15.2", "hashbrown 0.15.2",
"hf-hub 0.3.2", "hf-hub 0.4.2",
"pkg-config", "pkg-config",
"pyo3", "pyo3",
"text-generation-router", "text-generation-router",
@ -4645,7 +4643,7 @@ dependencies = [
"average", "average",
"clap 4.5.30", "clap 4.5.30",
"float-ord", "float-ord",
"hf-hub 0.3.2", "hf-hub 0.4.2",
"ratatui", "ratatui",
"serde", "serde",
"serde_json", "serde_json",
@ -4683,7 +4681,7 @@ dependencies = [
"clap 4.5.30", "clap 4.5.30",
"ctrlc", "ctrlc",
"float_eq", "float_eq",
"hf-hub 0.4.1", "hf-hub 0.4.2",
"nix 0.28.0", "nix 0.28.0",
"once_cell", "once_cell",
"pyo3", "pyo3",
@ -4712,7 +4710,7 @@ dependencies = [
"csv", "csv",
"futures", "futures",
"futures-util", "futures-util",
"hf-hub 0.3.2", "hf-hub 0.4.2",
"image", "image",
"init-tracing-opentelemetry", "init-tracing-opentelemetry",
"itertools 0.10.5", "itertools 0.10.5",
@ -4779,7 +4777,7 @@ dependencies = [
"futures", "futures",
"futures-util", "futures-util",
"grpc-metadata", "grpc-metadata",
"hf-hub 0.3.2", "hf-hub 0.4.2",
"image", "image",
"init-tracing-opentelemetry", "init-tracing-opentelemetry",
"jsonschema", "jsonschema",
@ -4829,7 +4827,7 @@ dependencies = [
"futures", "futures",
"futures-util", "futures-util",
"grpc-metadata", "grpc-metadata",
"hf-hub 0.3.2", "hf-hub 0.4.2",
"image", "image",
"init-tracing-opentelemetry", "init-tracing-opentelemetry",
"itertools 0.13.0", "itertools 0.13.0",

View File

@ -29,7 +29,7 @@ homepage = "https://github.com/huggingface/text-generation-inference"
[workspace.dependencies] [workspace.dependencies]
base64 = "0.22.0" base64 = "0.22.0"
tokenizers = { version = "0.20.0", features = ["http"] } 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 = { version = "0.23.0" }
metrics-exporter-prometheus = { version = "0.15.1", features = [] } metrics-exporter-prometheus = { version = "0.15.1", features = [] }
minijinja = { version = "2.2.0", features = ["json"] } minijinja = { version = "2.2.0", features = ["json"] }

View File

@ -18,7 +18,8 @@ RUN apt-get update -y \
&& rm -rf /var/lib/apt/lists/* \ && rm -rf /var/lib/apt/lists/* \
&& apt-get clean && 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}" ENV PATH="/root/.cargo/bin:${PATH}"
RUN cargo install cargo-chef --locked RUN cargo install cargo-chef --locked

View File

@ -22,7 +22,7 @@ homepage = "https://github.com/huggingface/text-generation-inference"
[workspace.dependencies] [workspace.dependencies]
base64 = "0.22.0" base64 = "0.22.0"
tokenizers = { version = "0.20.0", features = ["http"] } 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 = { version = "0.23.0" }
metrics-exporter-prometheus = { version = "0.15.1", features = [] } metrics-exporter-prometheus = { version = "0.15.1", features = [] }
minijinja = { version = "2.2.0", features = ["json"] } minijinja = { version = "2.2.0", features = ["json"] }

View File

@ -86,6 +86,10 @@ async fn get_tokenizer(tokenizer_name: &str, revision: Option<&str>) -> Option<T
builder = builder.with_cache_dir(cache_dir.into()); 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 builder
}; };

View File

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

View File

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

View File

@ -1719,6 +1719,10 @@ pub async fn run(
builder = builder.with_cache_dir(cache_dir.into()); 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 builder
}; };

View File

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

View File

@ -720,7 +720,7 @@ wheels = [
[[package]] [[package]]
name = "huggingface-hub" name = "huggingface-hub"
version = "0.28.1" version = "0.29.1"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
dependencies = [ dependencies = [
{ name = "filelock" }, { name = "filelock" },
@ -731,9 +731,9 @@ dependencies = [
{ name = "tqdm" }, { name = "tqdm" },
{ name = "typing-extensions" }, { name = "typing-extensions" },
] ]
sdist = { url = "https://files.pythonhosted.org/packages/e7/ce/a734204aaae6c35a22f9956ebcd8d8708ae5b842e15d6f42bd6f49e634a4/huggingface_hub-0.28.1.tar.gz", hash = "sha256:893471090c98e3b6efbdfdacafe4052b20b84d59866fb6f54c33d9af18c303ae", size = 387074 } sdist = { url = "https://files.pythonhosted.org/packages/22/37/797d6476f13e5ef6af5fc48a5d641d32b39c37e166ccf40c3714c5854a85/huggingface_hub-0.29.1.tar.gz", hash = "sha256:9524eae42077b8ff4fc459ceb7a514eca1c1232b775276b009709fe2a084f250", size = 389776 }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/ea/da/6c2bea5327b640920267d3bf2c9fc114cfbd0a5de234d81cda80cc9e33c8/huggingface_hub-0.28.1-py3-none-any.whl", hash = "sha256:aa6b9a3ffdae939b72c464dbb0d7f99f56e649b55c3d52406f49e0a5a620c0a7", size = 464068 }, { url = "https://files.pythonhosted.org/packages/ae/05/75b90de9093de0aadafc868bb2fa7c57651fd8f45384adf39bd77f63980d/huggingface_hub-0.29.1-py3-none-any.whl", hash = "sha256:352f69caf16566c7b6de84b54a822f6238e17ddd8ae3da4f8f2272aea5b198d5", size = 468049 },
] ]
[[package]] [[package]]
@ -2563,6 +2563,7 @@ dependencies = [
{ name = "grpcio-status" }, { name = "grpcio-status" },
{ name = "hf-kernels" }, { name = "hf-kernels" },
{ name = "hf-transfer" }, { name = "hf-transfer" },
{ name = "huggingface-hub" },
{ name = "loguru" }, { name = "loguru" },
{ name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
{ name = "numpy", version = "2.2.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, { name = "numpy", version = "2.2.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
@ -2627,6 +2628,7 @@ requires-dist = [
{ name = "grpcio-tools", marker = "extra == 'gen'", specifier = ">=1.69.0" }, { name = "grpcio-tools", marker = "extra == 'gen'", specifier = ">=1.69.0" },
{ name = "hf-kernels", specifier = ">=0.1.5" }, { name = "hf-kernels", specifier = ">=0.1.5" },
{ name = "hf-transfer", specifier = ">=0.1.8" }, { name = "hf-transfer", specifier = ">=0.1.8" },
{ name = "huggingface-hub", specifier = ">=0.29.0" },
{ name = "loguru", specifier = ">=0.7.3" }, { name = "loguru", specifier = ">=0.7.3" },
{ name = "mypy-protobuf", marker = "extra == 'gen'", specifier = ">=3.6.0" }, { name = "mypy-protobuf", marker = "extra == 'gen'", specifier = ">=3.6.0" },
{ name = "numpy", specifier = ">=1.26,<3" }, { name = "numpy", specifier = ">=1.26,<3" },