grenade/rustingface
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

feat(rustingface): implement the registry, admin CLI and deployment
deploy / build (push) Waiting to run
Details
deploy / deploy (push) Has been cancelled
Details

Browse Source 
Scaffolds the workspace per architecture/generic.md §1 and implements
phases 0-3 of doc/spec.md §12.

Crates:
  entities  manifest/ref/repo schemas, bucket key layout, config, the
            X-Error-Code taxonomy. No I/O.
  core      resolver, freeze pinning, single-flight, the streaming tee,
            policy, gc/verify/refresh. Defines the Store and Upstream
            ports.
  data      object_store (S3 + local) and reqwest Hub adapters.
  api       the axum surface: resolve, model/dataset info, tree, refs,
            whoami, metrics, bearer auth, range handling.
  bin       one binary: serve plus fetch/pin/refresh/list/show/rm/gc/
            verify/doctor.

Deployment targets bob.hanzalova.internal:20482 (port derived per
architecture/port-allocations.md §3), storing to the MinIO on
caveman.kosherinata.internal, fronted by hanzalova at rf.internal.
Ships the sysusers drop-in, hardened unit, firewalld service, nginx
vhost, config template, infra-setup.sh and the Gitea Actions
ci/deploy/conformance workflows.

Testing: 112 unit and integration tests, including the sovereignty
suite (cold fetch, sealed replay, single-flight, client disconnect,
range resume, freeze stability, gc-after-rm, digest mismatch), plus a
conformance suite driving a pinned huggingface_hub against a real Hub.

Deviations from the spec, all deliberate:
  - one binary with subcommands (spec §9) rather than generic.md's
    separate -api and -cli binaries; the library split is unchanged.
  - a dedicated sysusers account and hardened unit (generic.md §8)
    rather than the spec's illustrative DynamicUser unit.
  - manifests carry an optional repo_tree recorded verbatim, resolving
    spec §13's "record whole, filter on read" question for the tree
    endpoint as well as model-info.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XZG2i4AmfSqE97EJGBVb64
This commit is contained in:
grenade
2026-08-31 10:46:25 +03:00
co-authored by Claude Opus 5
parent e12457cf27
commit cb1ca8b6af
64 changed files with 13147 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitea/workflows/ci.yml
+29
View File
@@ -0,0 +1,29 @@
name: ci
on:
push:
branches-ignore: [main]
pull_request:
workflow_dispatch:
jobs:
check:
runs-on: rust
steps:
- uses: actions/checkout@v4
- name: format
run: cargo fmt --all -- --check
- name: clippy
run: cargo clippy --workspace --all-targets --all-features -- -D warnings
- name: test
run: cargo test --workspace --all-features
# A musl build is what actually ships, so a target-only breakage should
# fail here rather than in the deploy workflow.
- name: build (musl)
run: |
rustup target add x86_64-unknown-linux-musl
cargo build --release --target x86_64-unknown-linux-musl --bin rustingface
.gitea/workflows/conformance.yml
+30
View File
@@ -0,0 +1,30 @@
name: conformance
on:
workflow_dispatch:
schedule:
# Weekly. The clients are pinned, so this catches upstream Hub changes
# rather than client upgrades; a client upgrade is a deliberate commit and
# a spec-review trigger.
- cron: "0 4 * * 1"
jobs:
conformance:
runs-on: rust
steps:
- uses: actions/checkout@v4
- name: build
run: cargo build --release --bin rustingface
- name: install pinned clients
run: |
python3 -m venv .venv
.venv/bin/pip install --upgrade pip
.venv/bin/pip install -r test/conformance/requirements.txt
# Runs against a directory rather than a bucket. The S3 path is covered
# by the `doctor` step of the deploy workflow, which exercises the real
# MinIO on caveman with the real credentials.
- name: conformance
run: .venv/bin/python test/conformance/conformance.py --binary target/release/rustingface
.gitea/workflows/deploy.yml
+185
View File
@@ -0,0 +1,185 @@
name: deploy
on:
push: { branches: [main] }
workflow_dispatch:
concurrency:
group: deploy
cancel-in-progress: false
env:
# --- infra truth: hosts, ports and paths live here, not in a manifest ---
SERVICE_HOST: bob.hanzalova.internal
# Bind the mesh address rather than a wildcard: the reverse proxy is on a
# different host, so loopback will not do, and a wildcard bind on a host that
# may later gain another interface would publish the registry there too.
LISTEN_ADDR: 10.6.0.193:20482
APP_PORT: "20482"
S3_ENDPOINT: http://caveman.kosherinata.internal:9000
S3_BUCKET: rustingface
UPSTREAM_ENABLED: "true"
jobs:
build:
runs-on: rust
steps:
- uses: actions/checkout@v4
# The gate runs before anything is built for deployment, so a broken
# commit on main never reaches a host.
- name: format
run: cargo fmt --all -- --check
- name: clippy
run: cargo clippy --workspace --all-targets --all-features -- -D warnings
- name: test
run: cargo test --workspace --all-features
# Static musl: the runner is Fedora 44 and the target is Fedora 43, so a
# dynamically linked binary could reference a newer glibc than bob has.
- name: build
run: |
rustup target add x86_64-unknown-linux-musl
cargo build --release --target x86_64-unknown-linux-musl --bin rustingface
- uses: actions/upload-artifact@v3
with:
name: rustingface
path: target/x86_64-unknown-linux-musl/release/rustingface
deploy:
needs: build
runs-on: fedora-43
steps:
- uses: actions/checkout@v4
- uses: actions/download-artifact@v3
with:
name: rustingface
path: dist
- name: authenticate to the target
run: |
install -d -m 0700 ~/.ssh
printf '%s\n' "${{ secrets.RSYNC_SSH_KEY }}" > ~/.ssh/id_gitea_ci
chmod 600 ~/.ssh/id_gitea_ci
cat > ~/.ssh/config <<CFG
Host $SERVICE_HOST
User gitea_ci
IdentityFile ~/.ssh/id_gitea_ci
IdentitiesOnly yes
StrictHostKeyChecking accept-new
CFG
ssh "$SERVICE_HOST" hostname -f
- name: render config
env:
S3_ACCESS_KEY_ID: ${{ secrets.S3_ACCESS_KEY_ID }}
S3_SECRET_ACCESS_KEY: ${{ secrets.S3_SECRET_ACCESS_KEY }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
# Literal substitution, not a shell expansion: a secret containing
# $, ` or \ must survive intact.
python3 - <<'PY'
import os, pathlib
template = pathlib.Path("asset/config/config.toml.tmpl").read_text()
for key in ("LISTEN_ADDR", "S3_ENDPOINT", "S3_BUCKET", "UPSTREAM_ENABLED"):
template = template.replace("{{%s}}" % key, os.environ[key])
pathlib.Path("dist/config.toml").write_text(template)
for name, var in (("s3-access-key", "S3_ACCESS_KEY_ID"),
("s3-secret-key", "S3_SECRET_ACCESS_KEY"),
("hf-token", "HF_TOKEN")):
pathlib.Path("dist", name).write_text(os.environ.get(var, ""))
PY
# Fail loudly rather than shipping a config with an unrendered
# placeholder that would only surface as a runtime parse error.
if grep -n '{{' dist/config.toml; then
echo "unrendered placeholder in the config" >&2
exit 1
fi
if [ ! -s dist/s3-access-key ] || [ ! -s dist/s3-secret-key ]; then
echo "S3 credentials are empty; set the S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY secrets" >&2
exit 1
fi
- name: create the service account and its config directory
run: |
rsync -az --rsync-path='sudo rsync' --mkpath \
asset/systemd/rustingface.sysusers.conf \
"$SERVICE_HOST:/etc/sysusers.d/rustingface.conf"
ssh "$SERVICE_HOST" '
set -euo pipefail
sudo systemd-sysusers
sudo install -d -o root -g rustingface -m 0750 /etc/rustingface'
- name: ship the binary, config and credentials
run: |
set -euo pipefail
rsync -az --rsync-path='sudo rsync' \
dist/rustingface "$SERVICE_HOST:/usr/local/bin/rustingface"
for f in config.toml s3-access-key s3-secret-key hf-token; do
rsync -az --rsync-path='sudo rsync' --mkpath \
"dist/$f" "$SERVICE_HOST:/etc/rustingface/$f"
done
rsync -az --rsync-path='sudo rsync' --mkpath \
asset/systemd/rustingface.service \
"$SERVICE_HOST:/etc/systemd/system/rustingface.service"
ssh "$SERVICE_HOST" '
set -euo pipefail
sudo chmod 0755 /usr/local/bin/rustingface
sudo chown -R root:rustingface /etc/rustingface
sudo chmod 0640 /etc/rustingface/config.toml /etc/rustingface/s3-access-key /etc/rustingface/s3-secret-key /etc/rustingface/hf-token
sudo restorecon -R /usr/local/bin/rustingface /etc/rustingface'
- name: firewalld
run: |
set -euo pipefail
# rsync the XML first: firewalld only learns a freshly-shipped custom
# service after --reload, and querying it before that fails.
rsync -az --rsync-path='sudo rsync' --mkpath \
asset/firewalld/rustingface.xml \
"$SERVICE_HOST:/etc/firewalld/services/rustingface.xml"
ssh "$SERVICE_HOST" '
set -euo pipefail
sudo firewall-cmd --reload
zone=$(sudo firewall-cmd --get-default-zone)
if sudo firewall-cmd --zone="$zone" --query-service=rustingface; then
echo "rustingface already enabled in $zone"
else
sudo firewall-cmd --permanent --zone="$zone" --add-service=rustingface
sudo firewall-cmd --zone="$zone" --add-service=rustingface
fi'
- name: restart
run: |
ssh "$SERVICE_HOST" '
set -euo pipefail
sudo systemctl daemon-reload
sudo systemctl enable rustingface.service
sudo systemctl restart rustingface.service'
- name: health check
run: |
set -euo pipefail
ssh "$SERVICE_HOST" 'sudo systemctl is-active rustingface.service'
for attempt in $(seq 1 20); do
if ssh "$SERVICE_HOST" "curl -fsS http://$LISTEN_ADDR/healthz"; then
echo "healthy after $attempt attempt(s)"
exit 0
fi
sleep 2
done
echo "rustingface did not answer /healthz" >&2
exit 1
- name: doctor
run: |
# Proves the deployment can actually reach and write to the bucket,
# which a liveness probe deliberately does not.
ssh "$SERVICE_HOST" \
'sudo -u rustingface /usr/local/bin/rustingface --config /etc/rustingface/config.toml doctor'
- name: journal
if: always()
run: ssh "$SERVICE_HOST" 'journalctl -u rustingface.service -n 200 --no-pager'
.gitignore
+1
View File
@@ -0,0 +1 @@
/target
Cargo.lock
Generated
+3238
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+54
View File
@@ -0,0 +1,54 @@
[workspace]
resolver = "3"
members = ["crates/*"]
[workspace.package]
version = "0.1.0"
edition = "2024"
rust-version = "1.85"
license = "GPL-3.0-or-later"
authors = ["Rob Thijssen <rthijssen@gmail.com>"]
repository = "https://git.lair.cafe/grenade/rustingface"
[workspace.dependencies]
rustingface-entities = { path = "crates/rustingface-entities", version = "=0.1.0" }
rustingface-core = { path = "crates/rustingface-core", version = "=0.1.0" }
rustingface-data = { path = "crates/rustingface-data", version = "=0.1.0" }
rustingface-api = { path = "crates/rustingface-api", version = "=0.1.0" }
anyhow = "1"
async-trait = "0.1"
axum = { version = "0.8", default-features = false, features = ["http1", "http2", "json", "tokio", "query", "matched-path"] }
bytes = "1"
clap = { version = "4", features = ["derive", "env"] }
dashmap = "6"
figment = { version = "0.10", features = ["toml", "env"] }
futures = "0.3"
globset = "0.4"
hex = "0.4"
http = "1"
http-body-util = "0.1"
humantime-serde = "1"
metrics = "0.24"
metrics-exporter-prometheus = { version = "0.18", default-features = false }
object_store = { version = "0.14", features = ["aws"] }
percent-encoding = "2"
reqwest = { version = "0.13", default-features = false, features = ["rustls", "rustls-native-certs", "stream", "http2", "charset", "json"] }
serde = { version = "1", features = ["derive"] }
serde_json = { version = "1", features = ["preserve_order"] }
sha2 = "0.10"
subtle = "2"
thiserror = "2"
time = { version = "0.3", features = ["formatting", "parsing", "serde", "macros"] }
tokio = { version = "1", features = ["full"] }
tokio-util = { version = "0.7", features = ["io"] }
toml = "0.9"
tower = "0.5"
tower-http = { version = "0.7", features = ["trace", "normalize-path"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
[profile.release]
lto = "thin"
codegen-units = 1
strip = "symbols"
asset/config/config.toml.example
+19
View File
@@ -0,0 +1,19 @@
# A local development configuration.
#
# Uses a directory instead of a bucket, which forfeits state portability and is
# only appropriate for trying things out. Run with:
#
# cargo run -- --config asset/config/config.toml.example serve
[server]
listen = "127.0.0.1:20482"
[storage]
local_path = "/tmp/rustingface-dev"
[upstream]
enabled = true
endpoint = "https://huggingface.co"
[observability]
log_format = "text"
asset/config/config.toml.tmpl
+59
View File
@@ -0,0 +1,59 @@
# rustingface configuration.
#
# Rendered from this template by the deploy workflow, substituting {{...}}
# placeholders from Gitea repo secrets and workflow env. The rendered file is
# never committed; this template is.
#
# Every key here can be overridden per-key by environment, e.g.
# RUSTINGFACE__SERVER__LISTEN. Secrets are named by path, never inline.
[server]
# Bind the mesh address, not loopback: the reverse proxy (hanzalova) is on a
# different host from the service (bob). firewalld bounds who may reach it.
listen = "{{LISTEN_ADDR}}"
external_url = "https://rf.internal"
request_timeout = "300s"
client_stall_timeout = "60s"
[storage]
endpoint = "{{S3_ENDPOINT}}"
bucket = "{{S3_BUCKET}}"
region = "us-east-1"
path_style = true
access_key_id_file = "/etc/rustingface/s3-access-key"
secret_access_key_file = "/etc/rustingface/s3-secret-key"
# proxy: stream bytes through rustingface. redirect: 302 to a presigned URL,
# which takes rustingface out of the data path but exposes the object store's
# address to every client.
blob_delivery = "proxy"
presign_ttl = "15m"
multipart_part_size = "16MiB"
[upstream]
# false seals the registry: it serves only what is already in the bucket, and
# constructs no upstream client at all.
enabled = {{UPSTREAM_ENABLED}}
endpoint = "https://huggingface.co"
token_file = "/etc/rustingface/hf-token"
connect_timeout = "10s"
read_timeout = "60s"
# Do not change. rustingface implements no Xet client; a Xet-backed fetch would
# leave the bucket unreadable without one, which forfeits state portability.
disable_xet = true
[policy]
# freeze: the first resolution of a mutable ref is recorded and reused
# indefinitely. Moving to a newer revision is `rustingface refresh`.
ref_resolution = "freeze"
allow_new_repos = true
allowlist = []
[auth]
# The vhost is mesh-only and there is no per-client identity to enforce yet.
# Switch to "bearer" and populate /etc/rustingface/client-tokens to change that.
mode = "none"
token_file = "/etc/rustingface/client-tokens"
[observability]
metrics = true
log_format = "json"
asset/firewalld/rustingface.xml
+6
View File
@@ -0,0 +1,6 @@
<?xml version="1.0" encoding="utf-8"?>
<service>
<short>rustingface</short>
<description>Sovereign HuggingFace-compatible model registry. Serves model and dataset blobs to inference hosts on the mesh; fronted by the hanzalova edge proxy at rf.internal.</description>
<port protocol="tcp" port="20482"/>
</service>
asset/nginx/rf.internal.conf
+57
View File
@@ -0,0 +1,57 @@
# rf.internal — the mesh vhost for rustingface.
#
# Installed on hanzalova (the office edge proxy); the service itself runs on
# bob. TCP :443 belongs to the stream SNI router, which hands connections to
# the https tier on 127.0.0.1:14443, so this vhost listens there and reads the
# real client address from PROXY protocol.
upstream rustingface {
server bob.hanzalova.internal:20482;
keepalive 8;
}
server {
listen 127.0.0.1:14443 ssl proxy_protocol;
server_name rf.internal;
ssl_certificate /etc/nginx/tls/cert/rf.internal.pem;
ssl_certificate_key /etc/nginx/tls/key/rf.internal.pem;
ssl_trusted_certificate /etc/pki/ca-trust/source/anchors/root-internal.pem;
ssl_protocols TLSv1.3;
# The vhost believes it is serving 127.0.0.1:14443, so any redirect it
# generates itself would carry that unreachable port.
absolute_redirect off;
access_log /var/log/nginx/rf.internal.access.log;
error_log /var/log/nginx/rf.internal.error.log;
location / {
proxy_pass http://rustingface;
proxy_http_version 1.1;
proxy_set_header Host $host;
proxy_set_header Connection "";
proxy_set_header X-Real-IP $proxy_protocol_addr;
proxy_set_header X-Forwarded-For $proxy_protocol_addr;
proxy_set_header X-Forwarded-Proto https;
# Model weights are measured in tens of gigabytes. Buffering them to
# disk on the proxy would double the I/O and stall the client until the
# whole file had landed, so stream straight through.
proxy_buffering off;
proxy_request_buffering off;
client_max_body_size 0;
# A cold 40GB pull through this vhost is one long-lived response. The
# read timeout bounds silence between chunks, not total duration, so it
# stays modest; the send timeout must tolerate a slow client.
proxy_connect_timeout 15s;
proxy_read_timeout 300s;
proxy_send_timeout 300s;
# Never re-compress. The client sends `Accept-Encoding: identity`
# precisely so the sizes it learns are real.
gzip off;
proxy_set_header Accept-Encoding $http_accept_encoding;
}
}
asset/systemd/rustingface.service
+59
View File
@@ -0,0 +1,59 @@
[Unit]
Description=rustingface sovereign model registry
Documentation=https://git.lair.cafe/grenade/rustingface
After=network-online.target
Wants=network-online.target
[Service]
# Type=exec, not notify: the binary sends no sd_notify, and a notify unit that
# never sends READY=1 blocks `systemctl restart` until TimeoutStartSec.
Type=exec
User=rustingface
Group=rustingface
ExecStart=/usr/local/bin/rustingface serve --config /etc/rustingface/config.toml
Restart=on-failure
RestartSec=5s
# Credentials are passed by systemd rather than read from disk by the service,
# so the secret files stay root-owned and never appear in the unit's namespace
# except under $CREDENTIALS_DIRECTORY. The config names these basenames.
LoadCredential=s3-access-key:/etc/rustingface/s3-access-key
LoadCredential=s3-secret-key:/etc/rustingface/s3-secret-key
LoadCredential=hf-token:/etc/rustingface/hf-token
# Long transfers must not be killed mid-flight on stop; a detached upload that
# is cut short leaves an orphan blob, which is safe but wasteful.
TimeoutStopSec=120
# Hardening. The process holds no durable state — the bucket is the whole of it
# — so there is no StateDirectory to protect and ProtectSystem can stay strict
# with no ReadWritePaths at all.
NoNewPrivileges=true
ProtectSystem=strict
ProtectHome=true
PrivateTmp=true
PrivateDevices=true
ProtectKernelTunables=true
ProtectKernelModules=true
ProtectControlGroups=true
ProtectClock=true
ProtectHostname=true
ProtectProc=invisible
RestrictRealtime=true
RestrictSUIDSGID=true
RestrictNamespaces=true
LockPersonality=true
MemoryDenyWriteExecute=true
SystemCallFilter=@system-service
SystemCallArchitectures=native
CapabilityBoundingSet=
AmbientCapabilities=
RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6
UMask=0077
# Concurrent multi-gigabyte transfers each hold a handful of sockets and part
# buffers; the default 1024 is comfortable but the headroom costs nothing.
LimitNOFILE=16384
[Install]
WantedBy=multi-user.target
asset/systemd/rustingface.sysusers.conf
+2
View File
@@ -0,0 +1,2 @@
#Type Name ID GECOS Home directory Shell
u rustingface - "rustingface model registry service" /var/lib/rustingface /usr/sbin/nologin
crates/rustingface-api/Cargo.toml
+36
View File
@@ -0,0 +1,36 @@
[package]
name = "rustingface-api"
description = "The axum HTTP surface: Hub-compatible resolve, metadata and header replay."
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
authors.workspace = true
repository.workspace = true
[dependencies]
rustingface-core.workspace = true
rustingface-entities.workspace = true
axum.workspace = true
bytes.workspace = true
futures.workspace = true
http.workspace = true
metrics.workspace = true
metrics-exporter-prometheus.workspace = true
serde.workspace = true
serde_json.workspace = true
subtle.workspace = true
time.workspace = true
tokio.workspace = true
tower.workspace = true
tower-http.workspace = true
tracing.workspace = true
[dev-dependencies]
hex.workspace = true
http-body-util.workspace = true
rustingface-data.workspace = true
sha2.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] }
tower = { workspace = true, features = ["util"] }
crates/rustingface-api/src/auth.rs
+130
View File
@@ -0,0 +1,130 @@
//! Client-facing bearer auth.
//!
//! Distinct from `upstream.token_file`, which authenticates rustingface *to*
//! the Hub and is never echoed downstream. This is who may pull from
//! rustingface.
use std::sync::Arc;
use rustingface_entities::config::AuthMode;
use rustingface_entities::error::{Error, Result};
use subtle::ConstantTimeEq;
/// The accepted client tokens.
#[derive(Debug)]
pub struct Tokens {
mode: AuthMode,
accepted: Vec<String>,
}
impl Tokens {
/// Build from `[auth]` and the contents of the token file.
///
/// One token per line; blank lines and `#` comments are ignored so an
/// operator can annotate who each token belongs to.
pub fn new(mode: AuthMode, token_file_contents: Option<&str>) -> Result<Self> {
let accepted: Vec<String> = token_file_contents
.unwrap_or_default()
.lines()
.map(str::trim)
.filter(|line| !line.is_empty() && !line.starts_with('#'))
.map(str::to_owned)
.collect();
if mode == AuthMode::Bearer && accepted.is_empty() {
return Err(Error::Config(
"auth.mode = \"bearer\" but auth.token_file holds no tokens; \
every client would be refused"
.into(),
));
}
Ok(Self { mode, accepted })
}
/// Whether any authentication is required.
pub fn required(&self) -> bool {
self.mode == AuthMode::Bearer
}
/// Check an `Authorization` header value.
///
/// Comparison is constant-time and always runs against every configured
/// token, so a rejection tells an attacker nothing about how close a guess
/// was or how many tokens exist.
pub fn check(&self, authorization: Option<&str>) -> Result<()> {
if !self.required() {
return Ok(());
}
let presented = authorization
.and_then(|v| v.strip_prefix("Bearer "))
.or_else(|| authorization.and_then(|v| v.strip_prefix("bearer ")))
.map(str::trim)
.unwrap_or_default();
let mut matched = 0u8;
for token in &self.accepted {
matched |= token.as_bytes().ct_eq(presented.as_bytes()).unwrap_u8();
}
if matched == 1 {
Ok(())
} else {
Err(Error::Unauthorized)
}
}
}
/// Shared handle to the accepted tokens.
pub type SharedTokens = Arc<Tokens>;
impl Default for Tokens {
fn default() -> Self {
Self {
mode: AuthMode::None,
accepted: Vec::new(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn no_auth_admits_every_request() {
let tokens = Tokens::new(AuthMode::None, None).unwrap();
assert!(!tokens.required());
assert!(tokens.check(None).is_ok());
assert!(tokens.check(Some("Bearer nonsense")).is_ok());
}
#[test]
fn bearer_admits_a_configured_token_and_refuses_others() {
let tokens = Tokens::new(
AuthMode::Bearer,
Some("# the beast box\nhf_alpha\n\nhf_beta\n"),
)
.unwrap();
assert!(tokens.check(Some("Bearer hf_alpha")).is_ok());
assert!(tokens.check(Some("Bearer hf_beta")).is_ok());
assert!(tokens.check(Some("Bearer hf_gamma")).is_err());
assert!(tokens.check(None).is_err());
assert!(
tokens.check(Some("hf_alpha")).is_err(),
"the scheme is required"
);
}
#[test]
fn comments_and_blanks_are_not_themselves_tokens() {
let tokens = Tokens::new(AuthMode::Bearer, Some("# comment\nreal\n")).unwrap();
assert!(tokens.check(Some("Bearer # comment")).is_err());
assert!(tokens.check(Some("Bearer ")).is_err());
assert!(tokens.check(Some("Bearer real")).is_ok());
}
#[test]
fn bearer_with_an_empty_token_file_fails_at_startup() {
let err = Tokens::new(AuthMode::Bearer, Some("# nothing but comments\n")).unwrap_err();
assert!(matches!(err, Error::Config(_)), "{err}");
}
}
crates/rustingface-api/src/error.rs
+120
View File
@@ -0,0 +1,120 @@
//! Turning an [`Error`] into the response a Hub client expects.
//!
//! `huggingface_hub` branches on the `X-Error-Code` header to raise typed
//! exceptions, so emitting it is not cosmetic. The status alone would collapse
//! "no such repo", "no such revision" and "no such file" into one 404 that the
//! client cannot tell apart.
use axum::response::{IntoResponse, Response};
use http::{StatusCode, header};
use rustingface_entities::error::Error;
/// Header name the client reads to classify a failure.
pub const X_ERROR_CODE: &str = "x-error-code";
/// Header the Hub uses to carry a human-readable reason.
pub const X_ERROR_MESSAGE: &str = "x-error-message";
/// The exact `X-Error-Message` the client matches to raise `DisabledRepoError`.
///
/// Verified against the pinned `huggingface_hub` by the conformance suite:
/// unlike the others, a disabled repo is dispatched on the *message*, not on
/// `X-Error-Code`, so emitting only the code would have the client raise a
/// generic `HfHubHTTPError` instead.
pub const DISABLED_REPO_MESSAGE: &str = "Access to this resource is disabled.";
/// Wrapper giving [`Error`] an axum response.
pub struct ApiError(pub Error);
impl From<Error> for ApiError {
fn from(err: Error) -> Self {
Self(err)
}
}
impl IntoResponse for ApiError {
fn into_response(self) -> Response {
let ApiError(err) = self;
let status =
StatusCode::from_u16(err.status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
let message = err.to_string();
// A disabled repo is dispatched by the client on this exact message
// rather than on X-Error-Code; everything else carries its own reason.
let wire_message = match &err {
Error::DisabledRepo(_) => DISABLED_REPO_MESSAGE.to_owned(),
_ => sanitise(&message),
};
let mut response =
(status, axum::Json(serde_json::json!({ "error": message }))).into_response();
if let Some(code) = err.code()
&& let Ok(value) = header::HeaderValue::from_str(code.as_str())
{
response.headers_mut().insert(X_ERROR_CODE, value);
}
// Header values must be visible ASCII; a path or upstream detail could
// carry anything, so a message that will not encode is simply dropped
// rather than replacing the response with a 500.
if let Ok(value) = header::HeaderValue::from_str(&wire_message) {
response.headers_mut().insert(X_ERROR_MESSAGE, value);
}
if err.status() >= 500 {
tracing::error!(%err, "request failed");
} else {
tracing::debug!(%err, "request refused");
}
response
}
}
/// Reduce a message to characters a header value can carry.
fn sanitise(message: &str) -> String {
message
.chars()
.map(|c| if (' '..='~').contains(&c) { c } else { ' ' })
.take(512)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_missing_entry_carries_its_code_and_a_404() {
let response = ApiError(Error::EntryNotFound("config.json".into())).into_response();
assert_eq!(response.status(), StatusCode::NOT_FOUND);
assert_eq!(response.headers()[X_ERROR_CODE], "EntryNotFound");
}
#[test]
fn a_gated_repo_is_a_403_with_its_own_code() {
let response = ApiError(Error::GatedRepo("terms".into())).into_response();
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(response.headers()[X_ERROR_CODE], "GatedRepo");
}
#[test]
fn an_upstream_auth_failure_is_a_401_with_no_code() {
let response = ApiError(Error::UpstreamUnauthorized("bad token".into())).into_response();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert!(!response.headers().contains_key(X_ERROR_CODE));
}
#[test]
fn a_disabled_repo_carries_the_exact_message_the_client_matches_on() {
// The client dispatches DisabledRepoError on this string rather than
// on X-Error-Code, so it must go out verbatim.
let response = ApiError(Error::DisabledRepo("upstream said so".into())).into_response();
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(response.headers()[X_ERROR_MESSAGE], DISABLED_REPO_MESSAGE);
}
#[test]
fn a_message_with_a_newline_still_produces_a_valid_header() {
let response = ApiError(Error::EntryNotFound("a\nb\u{1f600}".into())).into_response();
let value = response.headers()[X_ERROR_MESSAGE].to_str().unwrap();
assert!(!value.contains('\n'), "{value:?}");
}
}
crates/rustingface-api/src/handlers/blob.rs
+175
View File
@@ -0,0 +1,175 @@
//! The resolve path: `HEAD` and `GET` for file bytes.
//!
//! This is the part clients actually depend on. Every header below is either
//! upstream's value replayed verbatim or a true fact about the bytes being
//! served; nothing here is normalised, recomputed or guessed.
use std::time::Instant;
use axum::body::Body as AxumBody;
use axum::extract::State;
use axum::response::{IntoResponse, Response};
use futures::StreamExt;
use http::{HeaderMap, HeaderValue, StatusCode, header};
use rustingface_core::{Body, ResolvedFile, Source};
use rustingface_entities::error::{Error, Result};
use rustingface_entities::repo::RepoType;
use crate::error::ApiError;
use crate::path::{ResolvePath, parse_resolve};
use crate::state::AppState;
use crate::{range, telemetry};
/// `X-Repo-Commit`: the commit the revision resolved to.
pub const X_REPO_COMMIT: &str = "x-repo-commit";
/// `X-Linked-Etag`: upstream's LFS etag, which clients prefer over `ETag`.
pub const X_LINKED_ETAG: &str = "x-linked-etag";
/// `X-Linked-Size`: upstream's LFS size.
pub const X_LINKED_SIZE: &str = "x-linked-size";
/// `HEAD /{repo_id}/resolve/{revision}/{filename}`.
pub async fn head(
State(state): State<AppState>,
repo_type: RepoType,
rest: String,
) -> std::result::Result<Response, ApiError> {
let ResolvePath {
repo_id,
revision,
filename,
} = parse_resolve(&rest)?;
let file = state
.registry
.head_file(repo_type, &repo_id, &revision, &filename)
.await?;
let mut headers = metadata_headers(&file)?;
headers.insert(header::CONTENT_LENGTH, HeaderValue::from(file.size));
// In redirect mode a HEAD answers 302 with the metadata still attached:
// the client reads its metadata off this response and then downloads from
// `Location`, so moving the headers onto the presigned request loses them.
if let Body::Redirect(location) = &file.body {
headers.insert(header::LOCATION, header_value(location)?);
telemetry::request("resolve_head", 302);
return Ok((StatusCode::FOUND, headers).into_response());
}
telemetry::request("resolve_head", 200);
Ok((StatusCode::OK, headers).into_response())
}
/// `GET /{repo_id}/resolve/{revision}/{filename}`.
pub async fn get(
State(state): State<AppState>,
repo_type: RepoType,
request_headers: HeaderMap,
rest: String,
) -> std::result::Result<Response, ApiError> {
let started = Instant::now();
let ResolvePath {
repo_id,
revision,
filename,
} = parse_resolve(&rest)?;
// The range is parsed against the stored length, which is not known until
// the file resolves. Ask for the object first, then narrow: a cold file
// must be fetched and stored whole before any range of it is served.
let requested_range = request_headers
.get(header::RANGE)
.and_then(|v| v.to_str().ok())
.map(str::to_owned);
let probe = match &requested_range {
None => None,
Some(spec) => {
let head = state
.registry
.head_file(repo_type, &repo_id, &revision, &filename)
.await?;
range::parse(spec, head.size)?
}
};
let (file, source) = state
.registry
.get_file(repo_type, &repo_id, &revision, &filename, probe)
.await?;
telemetry::in_flight(state.registry.in_flight());
if source == Source::Upstream {
telemetry::upstream_fetch("ok", started.elapsed().as_secs_f64());
}
let mut headers = metadata_headers(&file)?;
match file.body {
Body::Redirect(location) => {
headers.insert(header::CONTENT_LENGTH, HeaderValue::from(file.size));
headers.insert(header::LOCATION, header_value(&location)?);
telemetry::request("resolve_get", 302);
Ok((StatusCode::FOUND, headers).into_response())
}
Body::Stream { stream, len, range } => {
headers.insert(header::CONTENT_LENGTH, HeaderValue::from(len));
let status = match &range {
Some(range) => {
headers.insert(
header::CONTENT_RANGE,
header_value(&range::content_range(range, file.size))?,
);
StatusCode::PARTIAL_CONTENT
}
None => StatusCode::OK,
};
telemetry::served(source, len);
telemetry::request("resolve_get", status.as_u16());
let body = AxumBody::from_stream(
stream.map(|chunk| chunk.map_err(|err| std::io::Error::other(err.to_string()))),
);
Ok((status, headers, body).into_response())
}
Body::None => Err(ApiError(Error::Internal(
"a GET resolved to a metadata-only response".into(),
))),
}
}
/// Build the header set both `HEAD` and `GET` must carry.
///
/// `Content-Encoding` is deliberately absent and must stay that way: the
/// client sends `Accept-Encoding: identity` precisely so the sizes it learns
/// are real, and compressing this path corrupts its size accounting.
fn metadata_headers(file: &ResolvedFile) -> Result<HeaderMap> {
let mut headers = HeaderMap::new();
headers.insert(X_REPO_COMMIT, header_value(&file.commit)?);
headers.insert(header::ACCEPT_RANGES, HeaderValue::from_static("bytes"));
headers.insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/octet-stream"),
);
// Emitted exactly as upstream sent them, quotes included, and omitted
// entirely where upstream omitted them. The client prefers the linked
// etag and falls back to the plain one, so normalising either would move
// a cache entry that a directly-populated cache would not have moved.
if let Some(etag) = &file.etag {
headers.insert(header::ETAG, header_value(etag)?);
}
if let Some(linked) = &file.linked_etag {
headers.insert(X_LINKED_ETAG, header_value(linked)?);
}
if let Some(size) = file.linked_size {
headers.insert(X_LINKED_SIZE, HeaderValue::from(size));
}
Ok(headers)
}
/// A header value, refusing anything that cannot be encoded rather than
/// silently dropping a value a client depends on.
fn header_value(value: &str) -> Result<HeaderValue> {
HeaderValue::from_str(value)
.map_err(|_| Error::Internal(format!("{value:?} cannot be sent as a header value")))
}
crates/rustingface-api/src/handlers/meta.rs
+132
View File
@@ -0,0 +1,132 @@
//! The metadata surface: model/dataset info, refs and tree.
//!
//! Answered from the manifest of the resolved commit, which is upstream's own
//! document recorded verbatim. In sealed mode the file lists are narrowed to
//! what is actually stored, so a client is never told about a file it cannot
//! then fetch.
use axum::extract::{Query, State};
use axum::response::{IntoResponse, Response};
use serde::Deserialize;
use serde_json::json;
use rustingface_entities::repo::RepoType;
use crate::error::ApiError;
use crate::path::{ApiPath, parse_api};
use crate::state::AppState;
use crate::telemetry;
/// The default revision when a request names none.
const DEFAULT_REVISION: &str = "main";
/// Query parameters the tree endpoint accepts.
#[derive(Debug, Default, Deserialize)]
pub struct TreeQuery {
/// List every descendant rather than one level.
#[serde(default)]
pub recursive: Option<String>,
/// Include blob details. rustingface always records the expanded form, so
/// this is accepted and ignored rather than changing the answer.
#[serde(default)]
pub expand: Option<String>,
}
/// Whether a Hub-style boolean query parameter is set.
fn truthy(value: &Option<String>) -> bool {
match value.as_deref() {
None => false,
Some("") => true,
Some(v) => !matches!(v.to_ascii_lowercase().as_str(), "0" | "false" | "no"),
}
}
/// `GET /api/{models,datasets}/…`
pub async fn api(
State(state): State<AppState>,
repo_type: RepoType,
Query(query): Query<TreeQuery>,
rest: String,
) -> std::result::Result<Response, ApiError> {
match parse_api(&rest)? {
ApiPath::Info { repo_id } => {
let info = state
.registry
.repo_info(repo_type, &repo_id, DEFAULT_REVISION)
.await?;
telemetry::request("api_info", 200);
Ok(axum::Json(info).into_response())
}
ApiPath::Revision { repo_id, revision } => {
let info = state
.registry
.repo_info(repo_type, &repo_id, &revision)
.await?;
telemetry::request("api_info", 200);
Ok(axum::Json(info).into_response())
}
ApiPath::Refs { repo_id } => {
let refs = state.registry.list_refs(repo_type, &repo_id).await?;
// The Hub reports branches and tags separately. rustingface stores
// one flat namespace of pins and cannot tell them apart after the
// fact, so everything it pinned is reported as a branch and the
// tag list stays empty rather than being guessed at.
let branches: Vec<_> = refs
.iter()
.map(|r| {
json!({
"name": r.ref_name,
"ref": format!("refs/heads/{}", r.ref_name),
"targetCommit": r.commit,
})
})
.collect();
telemetry::request("api_refs", 200);
Ok(axum::Json(json!({ "branches": branches, "tags": [] })).into_response())
}
ApiPath::Tree {
repo_id,
revision,
path,
} => {
let tree = state
.registry
.repo_tree(
repo_type,
&repo_id,
&revision,
&path,
truthy(&query.recursive),
)
.await?;
telemetry::request("api_tree", 200);
Ok(axum::Json(tree).into_response())
}
}
}
/// `GET /api/whoami-v2`
pub async fn whoami(State(state): State<AppState>) -> Response {
telemetry::request("whoami", 200);
axum::Json((*state.identity).clone()).into_response()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_bare_query_flag_counts_as_set() {
assert!(truthy(&Some(String::new())));
assert!(truthy(&Some("1".into())));
assert!(truthy(&Some("true".into())));
}
#[test]
fn an_absent_or_falsey_flag_does_not() {
assert!(!truthy(&None));
assert!(!truthy(&Some("0".into())));
assert!(!truthy(&Some("false".into())));
assert!(!truthy(&Some("FALSE".into())));
}
}
crates/rustingface-api/src/handlers/mod.rs
+5
View File
@@ -0,0 +1,5 @@
//! HTTP handlers.
pub mod blob;
pub mod meta;
pub mod ops;
crates/rustingface-api/src/handlers/ops.rs
+37
View File
@@ -0,0 +1,37 @@
//! Operational endpoints: liveness and metrics.
use axum::extract::State;
use axum::response::{IntoResponse, Response};
use http::StatusCode;
use crate::state::AppState;
/// `GET /healthz`
///
/// Answers from process state only. It deliberately does not probe the bucket:
/// a health check that fails when object storage hiccups would have systemd
/// restart a process that is fine, which fixes nothing and drops in-flight
/// transfers. `rustingface doctor` is the check that touches storage.
pub async fn healthz() -> Response {
(StatusCode::OK, "ok\n").into_response()
}
/// `GET /metrics`
pub async fn metrics(State(state): State<AppState>) -> Response {
match &state.metrics {
Some(handle) => (
StatusCode::OK,
[(
http::header::CONTENT_TYPE,
"text/plain; version=0.0.4; charset=utf-8",
)],
handle.render(),
)
.into_response(),
None => (
StatusCode::NOT_FOUND,
"metrics are disabled (observability.metrics = false)\n",
)
.into_response(),
}
}
crates/rustingface-api/src/lib.rs
+123
View File
@@ -0,0 +1,123 @@
//! The HTTP surface.
//!
//! Route shapes come from the Hub, not from what is convenient to express in a
//! router. Repo ids, revisions and file paths all contain slashes, so the
//! routes here are wildcards and the disambiguation lives in [`path`], where
//! the rules can be stated and tested.
pub mod auth;
pub mod error;
pub mod handlers;
pub mod path;
pub mod range;
pub mod state;
pub mod telemetry;
use axum::extract::{Query, Request, State};
use axum::middleware::{self, Next};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::{Router, http::HeaderMap};
use http::Method;
use rustingface_entities::error::Error;
use rustingface_entities::repo::RepoType;
use tower_http::trace::TraceLayer;
pub use error::ApiError;
pub use state::AppState;
/// Build the router.
///
/// Ordering matters: the catch-all resolve route must be registered last, or
/// it would shadow `/api/…` and `/metrics`.
pub fn router(state: AppState) -> Router {
Router::new()
.route("/healthz", get(handlers::ops::healthz))
.route("/metrics", get(handlers::ops::metrics))
.route("/api/whoami-v2", get(handlers::meta::whoami))
.route("/api/models/{*rest}", get(api_models))
.route("/api/datasets/{*rest}", get(api_datasets))
.route(
"/datasets/{*rest}",
get(resolve_datasets).head(resolve_datasets),
)
.route("/{*rest}", get(resolve_models).head(resolve_models))
.layer(middleware::from_fn_with_state(state.clone(), authenticate))
.layer(TraceLayer::new_for_http())
.with_state(state)
}
/// Reject unauthenticated requests when `auth.mode = "bearer"`.
///
/// `/healthz` and `/metrics` stay open: they carry no repository data, and a
/// liveness probe that needs a credential is a liveness probe that silently
/// stops working when the credential rotates.
async fn authenticate(State(state): State<AppState>, request: Request, next: Next) -> Response {
let path = request.uri().path();
if path == "/healthz" || path == "/metrics" {
return next.run(request).await;
}
let presented = request
.headers()
.get(http::header::AUTHORIZATION)
.and_then(|v| v.to_str().ok());
if let Err(err) = state.tokens.check(presented) {
return ApiError(err).into_response();
}
next.run(request).await
}
/// `GET /api/models/…`
async fn api_models(
state: State<AppState>,
query: Query<handlers::meta::TreeQuery>,
axum::extract::Path(rest): axum::extract::Path<String>,
) -> Result<Response, ApiError> {
handlers::meta::api(state, RepoType::Models, query, rest).await
}
/// `GET /api/datasets/…`
async fn api_datasets(
state: State<AppState>,
query: Query<handlers::meta::TreeQuery>,
axum::extract::Path(rest): axum::extract::Path<String>,
) -> Result<Response, ApiError> {
handlers::meta::api(state, RepoType::Datasets, query, rest).await
}
/// `GET`/`HEAD /datasets/{repo_id}/resolve/…`
async fn resolve_datasets(
state: State<AppState>,
method: Method,
headers: HeaderMap,
axum::extract::Path(rest): axum::extract::Path<String>,
) -> Result<Response, ApiError> {
dispatch(state, RepoType::Datasets, method, headers, rest).await
}
/// `GET`/`HEAD /{repo_id}/resolve/…`
async fn resolve_models(
state: State<AppState>,
method: Method,
headers: HeaderMap,
axum::extract::Path(rest): axum::extract::Path<String>,
) -> Result<Response, ApiError> {
dispatch(state, RepoType::Models, method, headers, rest).await
}
/// Route a resolve request by method.
async fn dispatch(
state: State<AppState>,
repo_type: RepoType,
method: Method,
headers: HeaderMap,
rest: String,
) -> Result<Response, ApiError> {
match method {
Method::HEAD => handlers::blob::head(state, repo_type, rest).await,
Method::GET => handlers::blob::get(state, repo_type, headers, rest).await,
other => Err(ApiError(Error::BadRequest(format!(
"{other} is not supported on the resolve path; rustingface is read-only"
)))),
}
}
crates/rustingface-api/src/path.rs
+264
View File
@@ -0,0 +1,264 @@
//! Hub URL parsing.
//!
//! Repo ids, revisions and file paths all contain slashes, and the Hub's URL
//! shapes are only unambiguous to a server that knows its own repos. These
//! parsers encode the disambiguation rules explicitly rather than leaning on
//! router pattern matching, which cannot express them.
use rustingface_entities::error::{Error, Result};
use rustingface_entities::refs::is_commit_sha;
use rustingface_entities::repo::RepoId;
/// A parsed resolve path: `{repo_id}/resolve/{revision}/{filename}`.
#[derive(Debug, PartialEq, Eq)]
pub struct ResolvePath {
pub repo_id: RepoId,
pub revision: String,
pub filename: String,
}
/// Split `{repo_id}/resolve/{revision}/{filename}`.
///
/// The revision is the ambiguous part. A 40-hex segment is a commit; a
/// `refs/`-prefixed revision takes three segments (`refs/pr/1`,
/// `refs/convert/parquet`); anything else takes one. A branch name containing
/// a slash and not starting with `refs/` is therefore read as a one-segment
/// revision followed by a path, which is the same guess the Hub's own clients
/// make when constructing these URLs.
pub fn parse_resolve(rest: &str) -> Result<ResolvePath> {
let rest = rest.trim_start_matches('/');
let (repo_part, tail) = rest
.split_once("/resolve/")
.ok_or_else(|| Error::BadRequest(format!("{rest:?} is not a resolve path")))?;
let repo_id: RepoId = repo_part.parse()?;
let segments: Vec<&str> = tail.split('/').collect();
let revision_len = if segments.first() == Some(&"refs") {
3
} else {
1
};
if segments.len() <= revision_len {
return Err(Error::BadRequest(format!(
"{rest:?} names a revision but no file"
)));
}
let revision = segments[..revision_len].join("/");
let filename = segments[revision_len..].join("/");
if filename.is_empty() {
return Err(Error::BadRequest(format!(
"{rest:?} names an empty file path"
)));
}
let _ = is_commit_sha(&revision);
Ok(ResolvePath {
repo_id,
revision,
filename,
})
}
/// The API request shapes rustingface serves.
#[derive(Debug, PartialEq, Eq)]
pub enum ApiPath {
/// `/api/{type}/{repo_id}` — info for the default ref.
Info { repo_id: RepoId },
/// `/api/{type}/{repo_id}/revision/{revision}`.
Revision { repo_id: RepoId, revision: String },
/// `/api/{type}/{repo_id}/refs`.
Refs { repo_id: RepoId },
/// `/api/{type}/{repo_id}/tree/{revision}/{path}`.
Tree {
repo_id: RepoId,
revision: String,
path: String,
},
}
/// Segments that end a repo id in an API path.
const MARKERS: [&str; 3] = ["revision", "refs", "tree"];
/// Split the part of an `/api/{type}/…` path after the type.
///
/// A repo id is one or two segments, so the marker is looked for at offset 1
/// or 2 only. That is what keeps a repo genuinely named `tree` from being
/// mistaken for the tree endpoint of a bare-named repo.
pub fn parse_api(rest: &str) -> Result<ApiPath> {
let rest = rest.trim_matches('/');
if rest.is_empty() {
return Err(Error::BadRequest("no repository named".into()));
}
let segments: Vec<&str> = rest.split('/').collect();
let marker = (1..=2)
.filter(|i| *i < segments.len())
.find(|i| MARKERS.contains(&segments[*i]));
let Some(i) = marker else {
return Ok(ApiPath::Info {
repo_id: segments[..segments.len().min(2)].join("/").parse()?,
});
};
let repo_id: RepoId = segments[..i].join("/").parse()?;
let tail = &segments[i + 1..];
match segments[i] {
"refs" => Ok(ApiPath::Refs { repo_id }),
"revision" => {
if tail.is_empty() {
return Err(Error::BadRequest("no revision named".into()));
}
Ok(ApiPath::Revision {
repo_id,
revision: tail.join("/"),
})
}
"tree" => {
if tail.is_empty() {
return Err(Error::BadRequest("no revision named".into()));
}
let revision_len = if tail.first() == Some(&"refs") { 3 } else { 1 };
if tail.len() < revision_len {
return Err(Error::BadRequest("incomplete revision".into()));
}
Ok(ApiPath::Tree {
repo_id,
revision: tail[..revision_len].join("/"),
path: tail[revision_len..].join("/"),
})
}
other => Err(Error::BadRequest(format!(
"unknown API path segment {other:?}"
))),
}
}
#[cfg(test)]
mod tests {
use super::*;
fn resolve(s: &str) -> ResolvePath {
parse_resolve(s).unwrap()
}
#[test]
fn a_namespaced_resolve_path_splits_into_three_parts() {
let p = resolve("Qwen/Qwen3-32B/resolve/main/config.json");
assert_eq!(p.repo_id.to_string(), "Qwen/Qwen3-32B");
assert_eq!(p.revision, "main");
assert_eq!(p.filename, "config.json");
}
#[test]
fn a_bare_named_repo_resolves_too() {
let p = resolve("gpt2/resolve/main/config.json");
assert_eq!(p.repo_id.to_string(), "gpt2");
assert_eq!(p.revision, "main");
}
#[test]
fn a_nested_filename_keeps_all_its_segments() {
let p = resolve("Qwen/Qwen3-32B/resolve/main/onnx/model/weights.onnx");
assert_eq!(p.filename, "onnx/model/weights.onnx");
}
#[test]
fn a_commit_sha_revision_is_a_single_segment() {
let sha = "8f2a1c9e0b7d4f6a3e5c2b8d9f0a1c3e5b7d9f01";
let p = resolve(&format!("Qwen/Qwen3-32B/resolve/{sha}/a/b.json"));
assert_eq!(p.revision, sha);
assert_eq!(p.filename, "a/b.json");
}
#[test]
fn a_refs_revision_takes_three_segments() {
let p = resolve("Qwen/Qwen3-32B/resolve/refs/pr/1/config.json");
assert_eq!(p.revision, "refs/pr/1");
assert_eq!(p.filename, "config.json");
let p = resolve("Qwen/Qwen3-32B/resolve/refs/convert/parquet/data/train.parquet");
assert_eq!(p.revision, "refs/convert/parquet");
assert_eq!(p.filename, "data/train.parquet");
}
#[test]
fn a_resolve_path_with_no_filename_is_a_bad_request() {
assert!(parse_resolve("Qwen/Qwen3-32B/resolve/main").is_err());
assert!(parse_resolve("Qwen/Qwen3-32B/config.json").is_err());
}
#[test]
fn api_info_paths_parse_with_and_without_a_namespace() {
assert_eq!(
parse_api("Qwen/Qwen3-32B").unwrap(),
ApiPath::Info {
repo_id: "Qwen/Qwen3-32B".parse().unwrap()
}
);
assert_eq!(
parse_api("gpt2").unwrap(),
ApiPath::Info {
repo_id: "gpt2".parse().unwrap()
}
);
}
#[test]
fn api_revision_refs_and_tree_paths_parse() {
assert_eq!(
parse_api("Qwen/Qwen3-32B/revision/main").unwrap(),
ApiPath::Revision {
repo_id: "Qwen/Qwen3-32B".parse().unwrap(),
revision: "main".into()
}
);
assert_eq!(
parse_api("Qwen/Qwen3-32B/refs").unwrap(),
ApiPath::Refs {
repo_id: "Qwen/Qwen3-32B".parse().unwrap()
}
);
assert_eq!(
parse_api("gpt2/tree/main/onnx").unwrap(),
ApiPath::Tree {
repo_id: "gpt2".parse().unwrap(),
revision: "main".into(),
path: "onnx".into()
}
);
}
#[test]
fn a_tree_path_with_a_refs_revision_splits_correctly() {
assert_eq!(
parse_api("a/b/tree/refs/pr/1/sub/dir").unwrap(),
ApiPath::Tree {
repo_id: "a/b".parse().unwrap(),
revision: "refs/pr/1".into(),
path: "sub/dir".into()
}
);
}
#[test]
fn a_marker_word_inside_a_repo_name_is_not_treated_as_a_marker() {
// `refs` here is the *name* of a repo in the `datasets-org` namespace,
// reached at offset 1, so it is a marker. At offset 2 it would be past
// any possible repo id and is not looked for at all.
assert_eq!(
parse_api("owner/tree-sitter").unwrap(),
ApiPath::Info {
repo_id: "owner/tree-sitter".parse().unwrap()
}
);
}
#[test]
fn incomplete_api_paths_are_bad_requests_not_panics() {
assert!(parse_api("").is_err());
assert!(parse_api("a/b/revision").is_err());
assert!(parse_api("a/b/tree").is_err());
assert!(parse_api("a/b/tree/refs/pr").is_err());
}
}
crates/rustingface-api/src/range.rs
+114
View File
@@ -0,0 +1,114 @@
//! `Range` header handling.
//!
//! Range support is what makes an interrupted 40GB download resumable rather
//! than restartable, so the proxy path honours it (spec §6). Only a single
//! byte range is supported; a multi-range request is answered whole, which is
//! always a legal response.
use std::ops::Range;
use rustingface_entities::error::{Error, Result};
/// Parse a single-range `Range: bytes=…` header against a known length.
///
/// Returns `Ok(None)` for anything that is not a single byte range, so the
/// caller serves the whole object.
pub fn parse(header: &str, total: u64) -> Result<Option<Range<u64>>> {
let Some(spec) = header.trim().strip_prefix("bytes=") else {
return Ok(None);
};
if spec.contains(',') {
return Ok(None);
}
let (start, end) = spec
.split_once('-')
.ok_or_else(|| Error::BadRequest(format!("malformed Range header {header:?}")))?;
let unsatisfiable =
|| Error::RangeNotSatisfiable(format!("{header:?} against a {total}-byte object"));
let range = match (start.trim(), end.trim()) {
// `bytes=-N`: the final N bytes.
("", suffix) => {
let n: u64 = suffix.parse().map_err(|_| unsatisfiable())?;
if n == 0 {
return Err(unsatisfiable());
}
total.saturating_sub(n)..total
}
// `bytes=N-`: from N to the end.
(first, "") => {
let start: u64 = first.parse().map_err(|_| unsatisfiable())?;
start..total
}
// `bytes=N-M`, inclusive of M.
(first, last) => {
let start: u64 = first.parse().map_err(|_| unsatisfiable())?;
let end: u64 = last.parse().map_err(|_| unsatisfiable())?;
if end < start {
return Err(unsatisfiable());
}
start..end.saturating_add(1).min(total)
}
};
if range.start >= total || range.start >= range.end {
return Err(unsatisfiable());
}
Ok(Some(range))
}
/// The `Content-Range` value for a satisfied range.
pub fn content_range(range: &Range<u64>, total: u64) -> String {
format!("bytes {}-{}/{total}", range.start, range.end - 1)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_bounded_range_is_inclusive_of_its_last_byte() {
assert_eq!(parse("bytes=0-99", 1000).unwrap(), Some(0..100));
assert_eq!(parse("bytes=100-199", 1000).unwrap(), Some(100..200));
}
#[test]
fn an_open_ended_range_runs_to_the_end() {
assert_eq!(parse("bytes=500-", 1000).unwrap(), Some(500..1000));
}
#[test]
fn a_suffix_range_takes_the_final_bytes() {
assert_eq!(parse("bytes=-100", 1000).unwrap(), Some(900..1000));
assert_eq!(
parse("bytes=-5000", 1000).unwrap(),
Some(0..1000),
"a suffix longer than the object is the whole object"
);
}
#[test]
fn a_range_past_the_end_is_refused_rather_than_silently_truncated() {
let err = parse("bytes=1000-1100", 1000).unwrap_err();
assert_eq!(err.status(), 416);
assert_eq!(parse("bytes=200-100", 1000).unwrap_err().status(), 416);
}
#[test]
fn a_range_ending_past_the_end_is_clamped() {
assert_eq!(parse("bytes=900-2000", 1000).unwrap(), Some(900..1000));
}
#[test]
fn unsupported_range_forms_fall_back_to_the_whole_object() {
assert_eq!(parse("bytes=0-10,20-30", 1000).unwrap(), None);
assert_eq!(parse("items=0-10", 1000).unwrap(), None);
}
#[test]
fn content_range_reports_the_inclusive_last_byte() {
assert_eq!(content_range(&(0..100), 1000), "bytes 0-99/1000");
assert_eq!(content_range(&(900..1000), 1000), "bytes 900-999/1000");
}
}
crates/rustingface-api/src/state.rs
+59
View File
@@ -0,0 +1,59 @@
//! Shared handler state.
use std::sync::Arc;
use metrics_exporter_prometheus::PrometheusHandle;
use rustingface_core::Registry;
use crate::auth::Tokens;
/// Everything a handler needs, cheap to clone.
#[derive(Clone)]
pub struct AppState {
pub registry: Registry,
pub tokens: Arc<Tokens>,
/// `None` when `observability.metrics = false`.
pub metrics: Option<PrometheusHandle>,
/// Static identity served from `/api/whoami-v2`.
pub identity: Arc<serde_json::Value>,
}
impl AppState {
/// Assemble state around a registry.
pub fn new(registry: Registry, tokens: Arc<Tokens>, metrics: Option<PrometheusHandle>) -> Self {
Self {
registry,
tokens,
metrics,
identity: Arc::new(whoami_identity()),
}
}
}
/// The `/api/whoami-v2` document.
///
/// `hf auth login` calls this to confirm a token works and refuses to store
/// one when it does not. rustingface has no user model, so it answers with a
/// fixed identity whose only job is to let that flow complete.
fn whoami_identity() -> serde_json::Value {
serde_json::json!({
"type": "user",
"id": "rustingface",
"name": "rustingface",
"fullname": "rustingface registry",
"email": serde_json::Value::Null,
"emailVerified": true,
"canPay": false,
"isPro": false,
"periodEnd": serde_json::Value::Null,
"avatarUrl": "",
"orgs": [],
"auth": {
"type": "access_token",
"accessToken": {
"displayName": "rustingface",
"role": "read",
},
},
})
}
crates/rustingface-api/src/telemetry.rs
+46
View File
@@ -0,0 +1,46 @@
//! Prometheus metrics.
//!
//! Label sets are deliberately small. A metric keyed on repo id would give a
//! registry serving a few hundred models a few hundred time series per
//! counter, which is a cost paid forever for a question the manifests already
//! answer.
use metrics::{counter, gauge, histogram};
use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
use rustingface_core::Source;
use rustingface_entities::error::{Error, Result};
/// Install the recorder and return the handle `/metrics` renders from.
pub fn install() -> Result<PrometheusHandle> {
PrometheusBuilder::new()
.set_buckets(&[0.005, 0.025, 0.1, 0.5, 1.0, 5.0, 30.0, 120.0, 600.0, 3600.0])
.map_err(|e| Error::Config(format!("metrics buckets: {e}")))?
.install_recorder()
.map_err(|e| Error::Config(format!("metrics recorder: {e}")))
}
/// Record a completed request.
pub fn request(route: &'static str, status: u16) {
counter!("rustingface_requests_total", "route" => route, "status" => status.to_string())
.increment(1);
}
/// Record where a served body came from.
pub fn served(source: Source, bytes: u64) {
let label = match source {
Source::Bucket => "bucket",
Source::Upstream => "upstream",
};
counter!("rustingface_bytes_served_total", "source" => label).increment(bytes);
}
/// Record the outcome of an upstream fetch.
pub fn upstream_fetch(result: &'static str, seconds: f64) {
counter!("rustingface_upstream_fetches_total", "result" => result).increment(1);
histogram!("rustingface_upstream_fetch_seconds").record(seconds);
}
/// Publish how many fetches are in flight.
pub fn in_flight(count: usize) {
gauge!("rustingface_fetches_in_flight").set(count as f64);
}
crates/rustingface-api/tests/sovereignty.rs
+953
View File
@@ -0,0 +1,953 @@
//! The tests that matter (spec §11).
//!
//! Each one asserts a contractual guarantee rather than an implementation
//! detail: bytes land in the bucket, a sealed instance serves them with no
//! upstream in existence, a fleet-wide rollout produces one upstream fetch,
//! and a client hitting Ctrl-C does not throw away the transfer.
mod support;
use std::time::Duration;
use axum::body::Body;
use http::{Request, StatusCode};
use http_body_util::BodyExt;
use rustingface_entities::key;
use rustingface_entities::repo::{RepoId, RepoType};
use support::{FakeHub, HubFile, bucket_dir, instance};
use tower::ServiceExt;
const COMMIT_A: &str = "8f2a1c9e0b7d4f6a3e5c2b8d9f0a1c3e5b7d9f01";
const COMMIT_B: &str = "1122334455667788990011223344556677889900";
fn qwen() -> RepoId {
"Qwen/Qwen3-32B".parse().unwrap()
}
/// A model with one small config and one LFS-backed weights file.
async fn publish_model(hub: &FakeHub) -> Vec<u8> {
let weights: Vec<u8> = (0..300_000u32).map(|i| (i % 251) as u8).collect();
hub.publish(
"Qwen/Qwen3-32B",
"main",
COMMIT_A,
&[
("config.json", HubFile::plain(r#"{"model_type":"qwen3"}"#)),
("model.safetensors", HubFile::lfs(&weights)),
],
)
.await;
weights
}
/// Issue a request against an instance and read the whole response.
async fn call(
app: &axum::Router,
method: &str,
uri: &str,
headers: &[(&str, &str)],
) -> (StatusCode, http::HeaderMap, Vec<u8>) {
let mut request = Request::builder().method(method).uri(uri);
for (name, value) in headers {
request = request.header(*name, *value);
}
let response = app
.clone()
.oneshot(request.body(Body::empty()).unwrap())
.await
.expect("the router must always answer");
let status = response.status();
let headers = response.headers().clone();
let body = response
.into_body()
.collect()
.await
.expect("reading the response body")
.to_bytes()
.to_vec();
(status, headers, body)
}
/// As [`call`], but tolerating a body that aborts mid-stream.
///
/// Once headers are on the wire a failure can only be signalled by ending the
/// body short of `Content-Length`, which is what a client detects. Returning
/// the error rather than panicking lets a test assert on that.
async fn call_streaming(app: &axum::Router, uri: &str) -> (StatusCode, Result<Vec<u8>, String>) {
use futures::StreamExt;
let response = app
.clone()
.oneshot(Request::builder().uri(uri).body(Body::empty()).unwrap())
.await
.expect("the router must always answer");
let status = response.status();
let mut stream = response.into_body().into_data_stream();
let mut body = Vec::new();
while let Some(chunk) = stream.next().await {
match chunk {
Ok(chunk) => body.extend_from_slice(&chunk),
Err(err) => return (status, Err(err.to_string())),
}
}
(status, Ok(body))
}
#[tokio::test]
async fn a_cold_fetch_serves_the_bytes_and_leaves_them_in_the_bucket() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
let bucket = bucket_dir("cold");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, headers, body) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body, weights, "the client receives the file it asked for");
// The header contract: upstream's values, replayed exactly.
assert_eq!(headers["x-repo-commit"], COMMIT_A);
assert_eq!(headers["accept-ranges"], "bytes");
assert_eq!(headers["content-length"], weights.len().to_string());
assert!(
!headers.contains_key("content-encoding"),
"compressing this path corrupts the client's size accounting"
);
let oid = {
use sha2::{Digest, Sha256};
hex::encode(Sha256::digest(&weights))
};
assert_eq!(headers["etag"], format!("\"{oid}\""), "quotes are kept");
assert_eq!(headers["x-linked-etag"], format!("\"{oid}\""));
assert_eq!(headers["x-linked-size"], weights.len().to_string());
// Guarantee 1: durably stored before the response completed.
let manifest = rf
.registry
.load_manifest(RepoType::Models, &qwen(), COMMIT_A)
.await
.unwrap()
.expect("a manifest must exist after a served fetch");
let entry = manifest.entry("model.safetensors").expect("the entry");
assert_eq!(entry.size, weights.len() as u64);
assert_eq!(entry.digest, format!("sha256:{oid}"));
assert!(entry.lfs);
let blob = bucket.join(key::blob(&oid));
assert!(blob.exists(), "the blob is at its content-addressed key");
assert_eq!(std::fs::read(&blob).unwrap(), weights);
}
#[tokio::test]
async fn a_sealed_instance_serves_the_same_bytes_with_no_upstream_at_all() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
let bucket = bucket_dir("sealed");
{
let rf = instance(&bucket, Some(&hub), |_| {});
for path in ["config.json", "model.safetensors"] {
let (status, _, _) = call(
&rf.app,
"GET",
&format!("/Qwen/Qwen3-32B/resolve/main/{path}"),
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
}
}
// Step 6 of the spec's test: the host is gone. A brand-new process, a
// brand-new instance, no upstream client in existence, same bucket.
let sealed = instance(&bucket, None, |_| {});
hub.reset_counts();
let (status, headers, body) = call(
&sealed.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body, weights, "identical bytes, from the bucket alone");
assert_eq!(headers["x-repo-commit"], COMMIT_A, "and identical metadata");
assert_eq!(
hub.requests(),
0,
"a sealed instance must not touch upstream even once"
);
// And the ref still resolves, because the pin is in the bucket.
let (status, _, _) = call(
&sealed.app,
"HEAD",
"/Qwen/Qwen3-32B/resolve/main/config.json",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(hub.requests(), 0);
}
#[tokio::test]
async fn a_sealed_instance_refuses_loudly_for_anything_it_never_stored() {
let bucket = bucket_dir("sealed-miss");
let sealed = instance(&bucket, None, |_| {});
let (status, headers, _) = call(
&sealed.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::NOT_FOUND);
assert_eq!(
headers["x-error-code"], "RevisionNotFound",
"a typed code, not an empty file or a partial stream"
);
}
#[tokio::test]
async fn eight_concurrent_clients_produce_exactly_one_upstream_fetch() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
// Slow enough that all eight are in flight together.
hub.throttle(16 * 1024, Duration::from_millis(10)).await;
let bucket = bucket_dir("singleflight");
let rf = instance(&bucket, Some(&hub), |_| {});
hub.reset_counts();
let mut clients = Vec::new();
for _ in 0..8 {
let app = rf.app.clone();
clients.push(tokio::spawn(async move {
call(
&app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await
}));
}
for client in clients {
let (status, _, body) = client.await.unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body, weights, "every client gets the whole, correct file");
}
assert_eq!(
hub.file_gets(),
1,
"a fleet-wide rollout must not become eight concurrent pulls"
);
}
#[tokio::test]
async fn a_client_that_disconnects_mid_transfer_still_leaves_the_file_stored() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
hub.throttle(8 * 1024, Duration::from_millis(5)).await;
let bucket = bucket_dir("disconnect");
let rf = instance(&bucket, Some(&hub), |_| {});
// Read a little, then drop the response: the client hit Ctrl-C.
let response = rf
.app
.clone()
.oneshot(
Request::builder()
.uri("/Qwen/Qwen3-32B/resolve/main/model.safetensors")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let mut body = response.into_body().into_data_stream();
use futures::StreamExt;
let first = body.next().await.expect("at least one chunk").unwrap();
assert!(first.len() < weights.len(), "the transfer is still running");
drop(body);
// Discarding 30GB because a client hit Ctrl-C is the single most annoying
// possible behaviour; the transfer must finish regardless.
let oid = {
use sha2::{Digest, Sha256};
hex::encode(Sha256::digest(&weights))
};
let blob = bucket.join(key::blob(&oid));
for _ in 0..400 {
if blob.exists() && std::fs::metadata(&blob).unwrap().len() == weights.len() as u64 {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
assert_eq!(
std::fs::read(&blob).unwrap(),
weights,
"the detached transfer ran to completion"
);
// And the manifest caught up too, so it is servable, not just present.
for _ in 0..200 {
if rf
.registry
.load_manifest(RepoType::Models, &qwen(), COMMIT_A)
.await
.unwrap()
.and_then(|m| m.entry("model.safetensors").cloned())
.is_some()
{
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
let sealed = instance(&bucket, None, |_| {});
let (status, _, served) = call(
&sealed.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(served, weights);
}
#[tokio::test]
async fn a_ranged_request_resumes_from_the_offset_it_asks_for() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
let bucket = bucket_dir("range");
let rf = instance(&bucket, Some(&hub), |_| {});
// Cold, with a range: the whole file must be stored before any part of it
// is served, or the bucket would hold a fragment.
let (status, headers, body) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[("range", "bytes=100000-100999")],
)
.await;
assert_eq!(status, StatusCode::PARTIAL_CONTENT);
assert_eq!(body, weights[100_000..101_000]);
assert_eq!(
headers["content-range"],
format!("bytes 100000-100999/{}", weights.len())
);
assert_eq!(headers["content-length"], "1000");
let oid = {
use sha2::{Digest, Sha256};
hex::encode(Sha256::digest(&weights))
};
assert_eq!(
std::fs::metadata(bucket.join(key::blob(&oid)))
.unwrap()
.len(),
weights.len() as u64,
"a ranged cold request stores the whole file, never a fragment"
);
// Warm, resuming to the end.
let (status, _, tail) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[("range", "bytes=299000-")],
)
.await;
assert_eq!(status, StatusCode::PARTIAL_CONTENT);
assert_eq!(tail, weights[299_000..]);
// And a range past the end is refused, not silently truncated.
let (status, _, _) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[("range", "bytes=999999999-")],
)
.await;
assert_eq!(status, StatusCode::RANGE_NOT_SATISFIABLE);
}
#[tokio::test]
async fn a_gated_repo_answers_403_with_the_code_the_client_branches_on() {
let hub = FakeHub::start().await;
publish_model(&hub).await;
hub.gate("Qwen/Qwen3-32B").await;
let bucket = bucket_dir("gated");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, headers, _) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/config.json",
&[],
)
.await;
assert_eq!(status, StatusCode::FORBIDDEN);
assert_eq!(headers["x-error-code"], "GatedRepo");
}
#[tokio::test]
async fn a_frozen_ref_keeps_its_commit_after_upstream_moves() {
let hub = FakeHub::start().await;
publish_model(&hub).await;
let bucket = bucket_dir("freeze");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, headers, first) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/config.json",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(headers["x-repo-commit"], COMMIT_A);
// Upstream publishes a new commit and moves `main` onto it.
hub.publish(
"Qwen/Qwen3-32B",
"main",
COMMIT_B,
&[("config.json", HubFile::plain(r#"{"model_type":"qwen4"}"#))],
)
.await;
let (status, headers, again) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/config.json",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(
headers["x-repo-commit"], COMMIT_A,
"guarantee 4: upstream cannot change what the fleet resolves to"
);
assert_eq!(again, first);
// Moving is an operator action, and it is the only thing that moves it.
let (before, after) = rf
.registry
.refresh(RepoType::Models, &qwen(), "main")
.await
.unwrap();
assert_eq!(before.as_deref(), Some(COMMIT_A));
assert_eq!(after, COMMIT_B);
let (_, headers, updated) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/config.json",
&[],
)
.await;
assert_eq!(headers["x-repo-commit"], COMMIT_B);
assert_ne!(updated, first);
// The old commit stays servable by SHA: nothing that pinned it broke.
let (status, _, pinned) = call(
&rf.app,
"GET",
&format!("/Qwen/Qwen3-32B/resolve/{COMMIT_A}/config.json"),
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(pinned, first);
}
#[tokio::test]
async fn gc_after_rm_reclaims_exactly_the_right_blobs_and_no_others() {
let hub = FakeHub::start().await;
publish_model(&hub).await;
hub.publish(
"meta-llama/Llama-3-8B",
"main",
COMMIT_B,
&[("config.json", HubFile::plain(r#"{"model_type":"llama"}"#))],
)
.await;
let bucket = bucket_dir("gc");
let rf = instance(&bucket, Some(&hub), |_| {});
for uri in [
"/Qwen/Qwen3-32B/resolve/main/config.json",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
"/meta-llama/Llama-3-8B/resolve/main/config.json",
] {
let (status, _, _) = call(&rf.app, "GET", uri, &[]).await;
assert_eq!(status, StatusCode::OK, "{uri}");
}
// Nothing is reclaimable while every blob is referenced.
let dry = rf.registry.gc(true).await.unwrap();
assert_eq!(dry.scanned, 3);
assert!(
dry.deleted.is_empty(),
"a referenced blob is never reclaimable"
);
let llama: RepoId = "meta-llama/Llama-3-8B".parse().unwrap();
let llama_digest = rf
.registry
.load_manifest(RepoType::Models, &llama, COMMIT_B)
.await
.unwrap()
.unwrap()
.entry("config.json")
.unwrap()
.digest
.clone();
rf.registry
.remove(RepoType::Models, &llama, None)
.await
.unwrap();
// rm detaches; it does not delete bytes.
assert_eq!(
rf.store.list(key::BLOB_PREFIX).await.unwrap().len(),
3,
"rm must not touch blobs"
);
let dry = rf.registry.gc(true).await.unwrap();
assert_eq!(
dry.deleted.len(),
1,
"exactly the detached blob is reclaimable"
);
assert!(
!std::path::Path::new(&bucket.join(&dry.deleted[0]))
.to_string_lossy()
.is_empty()
);
assert!(
dry.deleted[0].ends_with(llama_digest.trim_start_matches("sha256:")),
"and it is the right one"
);
assert!(
bucket.join(&dry.deleted[0]).exists(),
"a dry run deletes nothing"
);
let done = rf.registry.gc(false).await.unwrap();
assert_eq!(done.deleted.len(), 1);
assert!(!bucket.join(&done.deleted[0]).exists());
assert_eq!(rf.store.list(key::BLOB_PREFIX).await.unwrap().len(), 2);
// The repo that was kept is untouched and still serves, sealed.
let sealed = instance(&bucket, None, |_| {});
for uri in [
"/Qwen/Qwen3-32B/resolve/main/config.json",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
] {
let (status, _, _) = call(&sealed.app, "GET", uri, &[]).await;
assert_eq!(status, StatusCode::OK, "{uri} must survive an unrelated gc");
}
}
#[tokio::test]
async fn a_digest_that_does_not_match_the_advertised_oid_stores_nothing() {
let hub = FakeHub::start().await;
let body = b"these bytes do not hash to the advertised oid".to_vec();
hub.publish(
"Qwen/Qwen3-32B",
"main",
COMMIT_A,
&[("model.safetensors", HubFile::corrupt(&body))],
)
.await;
let bucket = bucket_dir("mismatch");
let rf = instance(&bucket, Some(&hub), |_| {});
// The response cannot be a clean 200 with the wrong bytes. Headers are
// already on the wire by the time the digest is known, so the failure
// surfaces as a body that ends short of its Content-Length — which is
// exactly what the client detects and retries on.
let (_, body) = call_streaming(&rf.app, "/Qwen/Qwen3-32B/resolve/main/model.safetensors").await;
let err = body.expect_err("a mismatched transfer must not complete cleanly");
assert!(err.contains("digest mismatch"), "{err}");
tokio::time::sleep(Duration::from_millis(200)).await;
assert!(
!bucket.join(key::blob(&"b".repeat(64))).exists(),
"a mismatched transfer must not leave bytes at the advertised key"
);
let manifest = rf
.registry
.load_manifest(RepoType::Models, &qwen(), COMMIT_A)
.await
.unwrap();
assert!(
manifest.is_none_or(|m| m.entry("model.safetensors").is_none()),
"and must never be recorded"
);
}
#[tokio::test]
async fn verify_reports_a_blob_whose_bytes_stopped_matching_its_key() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
let bucket = bucket_dir("verify");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, _, _) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
let clean = rf.registry.verify(None).await.unwrap();
assert_eq!(clean.blobs, 1);
assert!(clean.corrupt.is_empty() && clean.missing.is_empty());
// Corrupt the stored bytes behind rustingface's back, as bit rot would.
let oid = {
use sha2::{Digest, Sha256};
hex::encode(Sha256::digest(&weights))
};
let mut damaged = weights.clone();
damaged[0] ^= 0xff;
std::fs::write(bucket.join(key::blob(&oid)), &damaged).unwrap();
let dirty = rf.registry.verify(None).await.unwrap();
assert_eq!(
dirty.corrupt.len(),
1,
"verify re-reads bytes, it does not trust metadata"
);
assert_eq!(dirty.corrupt[0].expected, format!("sha256:{oid}"));
}
#[tokio::test]
async fn metadata_endpoints_answer_from_the_bucket_once_recorded() {
let hub = FakeHub::start().await;
publish_model(&hub).await;
let bucket = bucket_dir("meta");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, _, body) = call(&rf.app, "GET", "/api/models/Qwen/Qwen3-32B", &[]).await;
assert_eq!(status, StatusCode::OK);
let info: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(info["sha"], COMMIT_A);
assert_eq!(info["siblings"].as_array().unwrap().len(), 2);
let (status, _, body) = call(
&rf.app,
"GET",
"/api/models/Qwen/Qwen3-32B/tree/main?recursive=1",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
let tree: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(tree.as_array().unwrap().len(), 2);
let (status, _, body) = call(&rf.app, "GET", "/api/models/Qwen/Qwen3-32B/refs", &[]).await;
assert_eq!(status, StatusCode::OK);
let refs: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(refs["branches"][0]["name"], "main");
assert_eq!(refs["branches"][0]["targetCommit"], COMMIT_A);
// Fetch one of the two files, then seal. The info document must stop
// advertising the file that can no longer be served.
let (status, _, _) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/config.json",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
let sealed = instance(&bucket, None, |_| {});
let (status, _, body) = call(&sealed.app, "GET", "/api/models/Qwen/Qwen3-32B", &[]).await;
assert_eq!(status, StatusCode::OK);
let info: serde_json::Value = serde_json::from_slice(&body).unwrap();
let siblings = info["siblings"].as_array().unwrap();
assert_eq!(
siblings.len(),
1,
"sealed mode advertises only what it can serve"
);
assert_eq!(siblings[0]["rfilename"], "config.json");
}
#[tokio::test]
async fn datasets_take_the_same_path_under_their_own_prefix() {
let hub = FakeHub::start().await;
hub.publish(
"HuggingFaceFW/fineweb",
"main",
COMMIT_A,
&[("data/train.parquet", HubFile::lfs(b"parquet bytes"))],
)
.await;
let bucket = bucket_dir("datasets");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, headers, body) = call(
&rf.app,
"GET",
"/datasets/HuggingFaceFW/fineweb/resolve/main/data/train.parquet",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body, b"parquet bytes");
assert_eq!(headers["x-repo-commit"], COMMIT_A);
let manifests = rf.store.list("manifests/datasets/").await.unwrap();
assert_eq!(
manifests.len(),
1,
"datasets are stored under their own prefix"
);
assert!(rf.store.list("manifests/models/").await.unwrap().is_empty());
}
#[tokio::test]
async fn deduplication_records_a_second_repo_without_transferring_again() {
let hub = FakeHub::start().await;
let weights: Vec<u8> = (0..200_000u32).map(|i| (i % 97) as u8).collect();
for repo in ["org-a/model", "org-b/mirror"] {
hub.publish(
repo,
"main",
COMMIT_A,
&[("model.safetensors", HubFile::lfs(&weights))],
)
.await;
}
let bucket = bucket_dir("dedup");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, _, first) = call(
&rf.app,
"GET",
"/org-a/model/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
hub.reset_counts();
let (status, _, second) = call(
&rf.app,
"GET",
"/org-b/mirror/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(second, first);
assert_eq!(
rf.store.list(key::BLOB_PREFIX).await.unwrap().len(),
1,
"identical bytes occupy the bucket once"
);
// gc must keep the shared blob while either repo still references it.
rf.registry
.remove(RepoType::Models, &"org-a/model".parse().unwrap(), None)
.await
.unwrap();
let reclaimed = rf.registry.gc(false).await.unwrap();
assert!(
reclaimed.deleted.is_empty(),
"a blob shared with a surviving repo is not garbage"
);
let (status, _, _) = call(
&rf.app,
"GET",
"/org-b/mirror/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
}
#[tokio::test]
async fn a_repo_outside_the_allowlist_is_refused_before_any_upstream_contact() {
let hub = FakeHub::start().await;
publish_model(&hub).await;
let bucket = bucket_dir("policy");
let rf = instance(&bucket, Some(&hub), |cfg| {
cfg.policy.allowlist = vec!["meta-llama/*".into()];
});
hub.reset_counts();
let (status, _, _) = call(
&rf.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/config.json",
&[],
)
.await;
assert_eq!(status, StatusCode::FORBIDDEN);
assert_eq!(hub.requests(), 0, "policy is checked before the network is");
}
#[tokio::test]
async fn bearer_auth_refuses_an_unknown_token_but_never_the_health_probe() {
let hub = FakeHub::start().await;
publish_model(&hub).await;
let bucket = bucket_dir("auth");
let mut config = rustingface_entities::config::Config::default();
config.storage.local_path = Some(bucket.clone());
config.upstream.enabled = true;
config.upstream.endpoint = hub.base_url.clone();
config.auth.mode = rustingface_entities::config::AuthMode::Bearer;
let store: std::sync::Arc<dyn rustingface_core::ports::Store> =
std::sync::Arc::new(rustingface_data::ObjectStoreAdapter::local(&bucket).unwrap());
let upstream: std::sync::Arc<dyn rustingface_core::ports::Upstream> =
std::sync::Arc::new(rustingface_data::HubClient::new(&config.upstream, None).unwrap());
let registry = rustingface_core::Registry::new(store, upstream, config.clone()).unwrap();
let tokens =
rustingface_api::auth::Tokens::new(config.auth.mode, Some("# beast\nhf_good\n")).unwrap();
let app = rustingface_api::router(rustingface_api::AppState::new(
registry,
std::sync::Arc::new(tokens),
None,
));
let uri = "/Qwen/Qwen3-32B/resolve/main/config.json";
assert_eq!(
call(&app, "GET", uri, &[]).await.0,
StatusCode::UNAUTHORIZED
);
assert_eq!(
call(&app, "GET", uri, &[("authorization", "Bearer hf_wrong")])
.await
.0,
StatusCode::UNAUTHORIZED
);
assert_eq!(
call(&app, "GET", uri, &[("authorization", "Bearer hf_good")])
.await
.0,
StatusCode::OK
);
assert_eq!(
call(&app, "GET", "/healthz", &[]).await.0,
StatusCode::OK,
"a liveness probe that needs a credential silently stops working when it rotates"
);
}
#[tokio::test]
async fn whoami_answers_so_hf_auth_login_can_complete() {
let bucket = bucket_dir("whoami");
let rf = instance(&bucket, None, |_| {});
let (status, _, body) = call(&rf.app, "GET", "/api/whoami-v2", &[]).await;
assert_eq!(status, StatusCode::OK);
let identity: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(identity["type"], "user");
assert_eq!(identity["name"], "rustingface");
}
#[tokio::test]
async fn a_head_replays_metadata_without_transferring_the_body() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
let bucket = bucket_dir("head");
let rf = instance(&bucket, Some(&hub), |_| {});
let (status, headers, body) = call(
&rf.app,
"HEAD",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert!(body.is_empty());
assert_eq!(headers["content-length"], weights.len().to_string());
assert_eq!(headers["x-repo-commit"], COMMIT_A);
assert_eq!(
hub.file_gets(),
0,
"a HEAD must not pull the body it is asking about"
);
}
#[tokio::test]
async fn state_is_portable_because_the_bucket_is_the_whole_of_it() {
let hub = FakeHub::start().await;
let weights = publish_model(&hub).await;
let source = bucket_dir("portable-src");
{
let rf = instance(&source, Some(&hub), |_| {});
for path in ["config.json", "model.safetensors"] {
let (status, _, _) = call(
&rf.app,
"GET",
&format!("/Qwen/Qwen3-32B/resolve/main/{path}"),
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
}
}
// Copy the bucket somewhere else entirely, as re-pointing a reinstalled
// host at the same bucket would. There is no second datastore to carry.
let destination = bucket_dir("portable-dst");
copy_tree(&source, &destination);
let moved = instance(&destination, None, |_| {});
let (status, headers, body) = call(
&moved.app,
"GET",
"/Qwen/Qwen3-32B/resolve/main/model.safetensors",
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body, weights);
assert_eq!(headers["x-repo-commit"], COMMIT_A);
}
/// Recursively copy a directory tree.
fn copy_tree(from: &std::path::Path, to: &std::path::Path) {
std::fs::create_dir_all(to).unwrap();
for entry in std::fs::read_dir(from).unwrap() {
let entry = entry.unwrap();
let target = to.join(entry.file_name());
if entry.file_type().unwrap().is_dir() {
copy_tree(&entry.path(), &target);
} else {
std::fs::copy(entry.path(), &target).unwrap();
}
}
}
crates/rustingface-api/tests/support/mod.rs
+470
View File
@@ -0,0 +1,470 @@
//! A fake Hub, and a rustingface wired to it.
//!
//! The tests that matter are about behaviour against a real client protocol,
//! so the fake upstream speaks the same headers the Hub does — `X-Repo-Commit`,
//! `X-Linked-Etag`, `X-Linked-Size`, `X-Error-Code` — and counts what it was
//! asked for, which is how single-flight is proved rather than assumed.
#![allow(dead_code)]
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use axum::Router;
use axum::body::Body;
use axum::extract::{Path, State};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use http::{HeaderMap, Method, StatusCode};
use rustingface_api::auth::Tokens;
use rustingface_api::{AppState, router};
use rustingface_core::Registry;
use rustingface_core::ports::{Store, Upstream};
use rustingface_data::{HubClient, ObjectStoreAdapter, SealedUpstream};
use rustingface_entities::config::{AuthMode, Config};
use tokio::sync::RwLock;
/// One file the fake Hub serves.
#[derive(Clone)]
pub struct HubFile {
pub body: Vec<u8>,
/// `ETag`, verbatim, quotes included.
pub etag: String,
/// `X-Linked-Etag`, present for LFS-backed files.
pub linked_etag: Option<String>,
pub linked_size: Option<u64>,
}
impl HubFile {
/// A small text file, the shape `config.json` takes upstream: a git blob
/// etag and no LFS headers.
pub fn plain(body: &str) -> Self {
Self {
body: body.as_bytes().to_vec(),
etag: format!("\"{:040x}\"", body.len()),
linked_etag: None,
linked_size: None,
}
}
/// An LFS-backed file, whose linked etag is the sha256 of its bytes.
pub fn lfs(body: &[u8]) -> Self {
use sha2::{Digest, Sha256};
let oid = hex::encode(Sha256::digest(body));
Self {
body: body.to_vec(),
etag: format!("\"{oid}\""),
linked_etag: Some(format!("\"{oid}\"")),
linked_size: Some(body.len() as u64),
}
}
/// An LFS-backed file whose advertised oid does not match its bytes.
pub fn corrupt(body: &[u8]) -> Self {
let mut file = Self::lfs(body);
file.linked_etag = Some(format!("\"{}\"", "b".repeat(64)));
file
}
}
/// One revision of one repo.
#[derive(Clone, Default)]
pub struct HubRevision {
pub files: HashMap<String, HubFile>,
}
/// The fake Hub's mutable state.
#[derive(Default)]
pub struct HubState {
/// `repo_id` -> ref name -> commit.
pub refs: HashMap<String, HashMap<String, String>>,
/// `repo_id` -> commit -> revision.
pub revisions: HashMap<String, HashMap<String, HubRevision>>,
/// Repos that answer 403 with `X-Error-Code: GatedRepo`.
pub gated: Vec<String>,
/// Delay inserted between body chunks, to keep a transfer observable.
pub chunk_delay: Duration,
/// Bytes per body chunk.
pub chunk_size: usize,
}
/// A running fake Hub.
pub struct FakeHub {
pub base_url: String,
pub state: Arc<RwLock<HubState>>,
/// Every `GET` on the resolve path. This is the number single-flight is
/// asserted against.
pub file_gets: Arc<AtomicUsize>,
/// Every request of any kind, so a sealed test can prove zero contact.
pub requests: Arc<AtomicUsize>,
shutdown: Option<tokio::sync::oneshot::Sender<()>>,
}
impl FakeHub {
/// Start a fake Hub on an ephemeral port.
pub async fn start() -> Self {
let state = Arc::new(RwLock::new(HubState {
chunk_size: 64 * 1024,
..Default::default()
}));
let file_gets = Arc::new(AtomicUsize::new(0));
let requests = Arc::new(AtomicUsize::new(0));
let hub = HubHandle {
state: Arc::clone(&state),
file_gets: Arc::clone(&file_gets),
requests: Arc::clone(&requests),
};
let app = Router::new()
.route("/api/models/{*rest}", get(api))
.route("/api/datasets/{*rest}", get(api))
.route("/{*rest}", get(resolve).head(resolve))
.with_state(hub);
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("binding the fake hub");
let addr = listener.local_addr().unwrap();
let (tx, rx) = tokio::sync::oneshot::channel();
tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = rx.await;
})
.await;
});
Self {
base_url: format!("http://{addr}"),
state,
file_gets,
requests,
shutdown: Some(tx),
}
}
/// Publish a revision and point a ref at it.
pub async fn publish(
&self,
repo: &str,
ref_name: &str,
commit: &str,
files: &[(&str, HubFile)],
) {
let mut state = self.state.write().await;
state
.refs
.entry(repo.to_owned())
.or_default()
.insert(ref_name.to_owned(), commit.to_owned());
let revision = HubRevision {
files: files
.iter()
.map(|(path, file)| ((*path).to_owned(), file.clone()))
.collect(),
};
state
.revisions
.entry(repo.to_owned())
.or_default()
.insert(commit.to_owned(), revision);
}
/// Move a ref to a different commit, as an upstream force-push would.
pub async fn move_ref(&self, repo: &str, ref_name: &str, commit: &str) {
self.state
.write()
.await
.refs
.entry(repo.to_owned())
.or_default()
.insert(ref_name.to_owned(), commit.to_owned());
}
/// Make a repo answer 403 `GatedRepo`.
pub async fn gate(&self, repo: &str) {
self.state.write().await.gated.push(repo.to_owned());
}
/// Slow the body down so a transfer stays observable mid-flight.
pub async fn throttle(&self, chunk_size: usize, delay: Duration) {
let mut state = self.state.write().await;
state.chunk_size = chunk_size;
state.chunk_delay = delay;
}
/// `GET`s on the resolve path so far.
pub fn file_gets(&self) -> usize {
self.file_gets.load(Ordering::SeqCst)
}
/// Requests of any kind so far.
pub fn requests(&self) -> usize {
self.requests.load(Ordering::SeqCst)
}
/// Forget everything counted so far.
pub fn reset_counts(&self) {
self.file_gets.store(0, Ordering::SeqCst);
self.requests.store(0, Ordering::SeqCst);
}
/// Stop answering, as null-routing the Hub at the site router would.
pub fn stop(&mut self) {
if let Some(tx) = self.shutdown.take() {
let _ = tx.send(());
}
}
}
#[derive(Clone)]
struct HubHandle {
state: Arc<RwLock<HubState>>,
file_gets: Arc<AtomicUsize>,
requests: Arc<AtomicUsize>,
}
/// Reject a gated repo the way the Hub does.
fn gated() -> Response {
(
StatusCode::FORBIDDEN,
[("x-error-code", "GatedRepo")],
"access to this repo is gated",
)
.into_response()
}
/// Report a missing entity with the code the client branches on.
fn not_found(code: &'static str) -> Response {
(StatusCode::NOT_FOUND, [("x-error-code", code)], code).into_response()
}
/// The fake Hub's `/api/{type}/…` surface.
async fn api(State(hub): State<HubHandle>, Path(rest): Path<String>) -> Response {
hub.requests.fetch_add(1, Ordering::SeqCst);
let state = hub.state.read().await;
let segments: Vec<&str> = rest.split('/').collect();
let marker = (1..=2)
.filter(|i| *i < segments.len())
.find(|i| matches!(segments[*i], "revision" | "tree"));
let Some(i) = marker else {
return not_found("RepoNotFound");
};
let repo = segments[..i].join("/");
let revision = segments.get(i + 1).copied().unwrap_or("main");
if state.gated.contains(&repo) {
return gated();
}
let Some(commit) = resolve_commit(&state, &repo, revision) else {
return not_found("RevisionNotFound");
};
let Some(rev) = state.revisions.get(&repo).and_then(|r| r.get(&commit)) else {
return not_found("RevisionNotFound");
};
let mut paths: Vec<&String> = rev.files.keys().collect();
paths.sort();
if segments[i] == "tree" {
let tree: Vec<_> = paths
.iter()
.map(|path| {
let file = &rev.files[*path];
serde_json::json!({
"type": "file",
"path": path,
"size": file.body.len(),
"oid": file.etag.trim_matches('"'),
})
})
.collect();
return axum::Json(tree).into_response();
}
axum::Json(serde_json::json!({
"id": repo,
"sha": commit,
"lastModified": "2026-06-11T00:00:00.000Z",
"private": false,
"gated": false,
"tags": ["text-generation"],
"siblings": paths.iter().map(|p| serde_json::json!({"rfilename": p})).collect::<Vec<_>>(),
}))
.into_response()
}
/// The fake Hub's resolve path.
async fn resolve(
State(hub): State<HubHandle>,
method: Method,
Path(rest): Path<String>,
) -> Response {
hub.requests.fetch_add(1, Ordering::SeqCst);
let state = hub.state.read().await;
let rest = rest.strip_prefix("datasets/").unwrap_or(&rest);
let Some((repo, tail)) = rest.split_once("/resolve/") else {
return not_found("RepoNotFound");
};
if state.gated.contains(&repo.to_owned()) {
return gated();
}
let Some((revision, path)) = tail.split_once('/') else {
return not_found("EntryNotFound");
};
let revision = percent_decode(revision);
let Some(commit) = resolve_commit(&state, repo, &revision) else {
return not_found("RevisionNotFound");
};
let Some(file) = state
.revisions
.get(repo)
.and_then(|r| r.get(&commit))
.and_then(|r| r.files.get(path))
.cloned()
else {
return not_found("EntryNotFound");
};
let mut headers = HeaderMap::new();
headers.insert("x-repo-commit", commit.parse().unwrap());
headers.insert("etag", file.etag.parse().unwrap());
if let Some(linked) = &file.linked_etag {
headers.insert("x-linked-etag", linked.parse().unwrap());
}
if let Some(size) = file.linked_size {
headers.insert("x-linked-size", size.to_string().parse().unwrap());
}
headers.insert(
"content-length",
file.body.len().to_string().parse().unwrap(),
);
if method == Method::HEAD {
return (StatusCode::OK, headers).into_response();
}
hub.file_gets.fetch_add(1, Ordering::SeqCst);
let chunk_size = state.chunk_size.max(1);
let delay = state.chunk_delay;
let body = file.body.clone();
drop(state);
let stream = futures::stream::unfold(0usize, move |offset| {
let body = body.clone();
async move {
if offset >= body.len() {
return None;
}
if !delay.is_zero() {
tokio::time::sleep(delay).await;
}
let end = (offset + chunk_size).min(body.len());
let chunk = bytes::Bytes::copy_from_slice(&body[offset..end]);
Some((Ok::<_, std::io::Error>(chunk), end))
}
});
(StatusCode::OK, headers, Body::from_stream(stream)).into_response()
}
/// Resolve a ref or pass a commit through.
fn resolve_commit(state: &HubState, repo: &str, revision: &str) -> Option<String> {
if let Some(commit) = state.refs.get(repo).and_then(|r| r.get(revision)) {
return Some(commit.clone());
}
state
.revisions
.get(repo)
.and_then(|r| r.contains_key(revision).then(|| revision.to_owned()))
}
/// Undo the percent-encoding the resolve URL applies to a revision.
fn percent_decode(value: &str) -> String {
let mut out = String::with_capacity(value.len());
let bytes = value.as_bytes();
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'%' && i + 2 < bytes.len() {
if let Ok(byte) = u8::from_str_radix(&value[i + 1..i + 3], 16) {
out.push(byte as char);
i += 3;
continue;
}
}
out.push(bytes[i] as char);
i += 1;
}
out
}
/// A rustingface instance under test, over a temporary bucket directory.
pub struct Instance {
pub app: axum::Router,
pub registry: Registry,
pub store: Arc<dyn Store>,
pub config: Config,
pub bucket: std::path::PathBuf,
}
/// Build an instance against `hub`, applying `tweak` to the config first.
pub fn instance(
bucket: &std::path::Path,
hub: Option<&FakeHub>,
tweak: impl FnOnce(&mut Config),
) -> Instance {
let mut config = Config::default();
config.storage.local_path = Some(bucket.to_owned());
config.storage.multipart_part_size = 5 * 1024 * 1024;
config.server.client_stall_timeout = Duration::from_secs(2);
match hub {
Some(hub) => {
config.upstream.enabled = true;
config.upstream.endpoint = hub.base_url.clone();
}
None => config.upstream.enabled = false,
}
tweak(&mut config);
config.validate().expect("test config must be valid");
let store: Arc<dyn Store> = Arc::new(ObjectStoreAdapter::local(bucket).expect("local store"));
let upstream: Arc<dyn Upstream> = if config.upstream.enabled {
Arc::new(HubClient::new(&config.upstream, None).expect("hub client"))
} else {
Arc::new(SealedUpstream::new(config.upstream.endpoint.clone()))
};
let registry = Registry::new(Arc::clone(&store), upstream, config.clone()).expect("registry");
let tokens = Tokens::new(config.auth.mode, None)
.unwrap_or_else(|_| Tokens::new(AuthMode::None, None).expect("permissive tokens"));
let app = router(AppState::new(registry.clone(), Arc::new(tokens), None));
Instance {
app,
registry,
store,
config,
bucket: bucket.to_owned(),
}
}
/// A fresh directory to use as a bucket.
pub fn bucket_dir(name: &str) -> std::path::PathBuf {
let dir = std::env::temp_dir().join(format!(
"rustingface-test-{name}-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&dir).expect("creating the test bucket");
dir
}
crates/rustingface-core/Cargo.toml
+28
View File
@@ -0,0 +1,28 @@
[package]
name = "rustingface-core"
description = "Resolver, single-flight fetch, retention policy and admin operations for rustingface."
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
authors.workspace = true
repository.workspace = true
[dependencies]
rustingface-entities.workspace = true
async-trait.workspace = true
bytes.workspace = true
dashmap.workspace = true
futures.workspace = true
globset.workspace = true
hex.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
time.workspace = true
tokio.workspace = true
tracing.workspace = true
[dev-dependencies]
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] }
crates/rustingface-core/src/admin.rs
+389
View File
@@ -0,0 +1,389 @@
//! Operator actions: the ones that remove things, and the ones that check.
//!
//! `rm` and `gc` are deliberately separate. Removal detaches metadata;
//! reclamation is a second, explicit, dry-runnable step. Nothing in this file
//! deletes blob bytes except [`Registry::gc`], and only when it is asked to.
use std::collections::BTreeSet;
use bytes::Bytes;
use futures::StreamExt;
use rustingface_entities::error::{Error, Result};
use rustingface_entities::key;
use rustingface_entities::manifest::Manifest;
use rustingface_entities::repo::{RepoId, RepoType};
use tracing::{info, warn};
use crate::digest::{self, Digest};
use crate::registry::Registry;
/// A stored repo, summarised for `rustingface list`.
#[derive(Debug, Clone)]
pub struct RepoSummary {
pub repo_type: RepoType,
pub repo_id: RepoId,
/// Ref name to pinned commit, sorted by ref name.
pub refs: Vec<(String, String)>,
/// Every recorded commit.
pub commits: Vec<String>,
/// Distinct files recorded across all commits.
pub files: usize,
/// Bytes referenced, counting each distinct blob once.
pub bytes: u64,
}
/// What `rustingface rm` detached.
#[derive(Debug, Clone, Default)]
pub struct Removed {
pub manifests: Vec<String>,
pub refs: Vec<String>,
pub repo_record: bool,
}
/// What `rustingface gc` found, and did or would do.
#[derive(Debug, Clone, Default)]
pub struct Reclaimed {
pub scanned: usize,
pub referenced: usize,
pub deleted: Vec<String>,
pub bytes: u64,
pub dry_run: bool,
}
/// One blob that failed `rustingface verify`.
#[derive(Debug, Clone)]
pub struct Corruption {
pub key: String,
pub expected: String,
pub found: String,
}
/// What `rustingface verify` checked.
#[derive(Debug, Clone, Default)]
pub struct Verified {
pub blobs: usize,
pub bytes: u64,
pub missing: Vec<String>,
pub corrupt: Vec<Corruption>,
}
impl Registry {
/// Summarise a stored repo.
pub async fn summarise(&self, repo_type: RepoType, repo_id: &RepoId) -> Result<RepoSummary> {
let manifests = self.list_manifests(repo_type, repo_id).await?;
let refs = self
.list_refs(repo_type, repo_id)
.await?
.into_iter()
.map(|r| (r.ref_name, r.commit))
.collect();
// Count each distinct blob once: the same weights shared by two
// revisions occupy the bucket once, and a summary that double-counts
// them would misreport what `gc` could ever reclaim.
let mut seen = BTreeSet::new();
let mut bytes = 0;
let mut files = BTreeSet::new();
for manifest in &manifests {
for entry in &manifest.entries {
files.insert(entry.path.clone());
if seen.insert(entry.digest.clone()) {
bytes += entry.size;
}
}
}
Ok(RepoSummary {
repo_type,
repo_id: repo_id.clone(),
refs,
commits: manifests.iter().map(|m| m.commit.clone()).collect(),
files: files.len(),
bytes,
})
}
/// Detach a repo, or one revision of it, from the bucket.
///
/// Blobs are untouched. Reclaiming them is [`Registry::gc`], which the
/// operator runs separately and can dry-run first.
pub async fn remove(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: Option<&str>,
) -> Result<Removed> {
let mut removed = Removed::default();
let Some(revision) = revision else {
for object in self
.store()
.list(&key::manifest_prefix(repo_type, repo_id))
.await?
{
self.store().delete(&object.key).await?;
removed.manifests.push(object.key);
}
for object in self
.store()
.list(&key::ref_prefix(repo_type, repo_id))
.await?
{
self.store().delete(&object.key).await?;
removed.refs.push(object.key);
}
let repo_key = key::repo(repo_type, repo_id);
if self.store().head(&repo_key).await?.is_some() {
self.store().delete(&repo_key).await?;
removed.repo_record = true;
}
info!(%repo_id, "removed repository metadata; blobs remain until `gc`");
return Ok(removed);
};
// A single revision: drop its manifest, and any ref still pointing at
// it, so nothing resolves to a commit that is no longer described.
let commit = self
.resolve_stored_revision(repo_type, repo_id, revision)
.await?;
let manifest_key = key::manifest(repo_type, repo_id, &commit);
if self.store().head(&manifest_key).await?.is_some() {
self.store().delete(&manifest_key).await?;
removed.manifests.push(manifest_key);
}
for record in self.list_refs(repo_type, repo_id).await? {
if record.commit == commit {
let key = key::ref_key(repo_type, repo_id, &record.ref_name);
self.store().delete(&key).await?;
removed.refs.push(key);
}
}
info!(%repo_id, commit, "removed revision metadata; blobs remain until `gc`");
Ok(removed)
}
/// Resolve a revision using only what is stored, never upstream.
///
/// Admin actions that remove things must not have the side effect of
/// pulling a new pin from upstream first.
pub async fn resolve_stored_revision(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
) -> Result<String> {
if rustingface_entities::refs::is_commit_sha(revision) {
return Ok(revision.to_owned());
}
self.load_ref(repo_type, repo_id, revision)
.await?
.map(|r| r.commit)
.ok_or_else(|| {
Error::RevisionNotFound(format!(
"{repo_id}@{revision} is not pinned in this bucket"
))
})
}
/// Every digest referenced by any manifest in the bucket.
pub async fn referenced_digests(&self) -> Result<BTreeSet<String>> {
let mut digests = BTreeSet::new();
for object in self.store().list("manifests/").await? {
let Some(bytes) = self.store().get(&object.key).await? else {
continue;
};
match serde_json::from_slice::<Manifest>(&bytes) {
Ok(manifest) => {
for entry in manifest.entries {
if let Some(hex) = digest::unprefixed(&entry.digest) {
digests.insert(hex.to_owned());
}
}
}
// A manifest we cannot parse is not proof that its blobs are
// unreferenced, and `gc` must never delete on a guess.
Err(err) => {
return Err(Error::Storage(format!(
"{} could not be parsed ({err}); refusing to compute a reachable set \
from an incomplete scan",
object.key
)));
}
}
}
Ok(digests)
}
/// Delete blobs no manifest references.
///
/// The reachable set is computed from a complete scan of every manifest,
/// and a manifest that fails to parse aborts the whole run rather than
/// silently shrinking that set.
pub async fn gc(&self, dry_run: bool) -> Result<Reclaimed> {
let referenced = self.referenced_digests().await?;
let blobs = self.store().list(key::BLOB_PREFIX).await?;
let mut result = Reclaimed {
scanned: blobs.len(),
referenced: referenced.len(),
dry_run,
..Default::default()
};
for blob in blobs {
let Some(hex) = key::blob_digest(&blob.key) else {
warn!(key = %blob.key, "unrecognised key under the blob prefix; leaving it alone");
continue;
};
if referenced.contains(hex) {
continue;
}
result.bytes += blob.size;
if !dry_run {
self.store().delete(&blob.key).await?;
}
result.deleted.push(blob.key);
}
info!(
scanned = result.scanned,
referenced = result.referenced,
reclaimed = result.deleted.len(),
bytes = result.bytes,
dry_run,
"gc complete"
);
Ok(result)
}
/// Re-read stored blobs and check them against their content-addressed
/// keys.
///
/// `repo` narrows the check to one repository; without it every blob in
/// the bucket is read.
pub async fn verify(&self, repo: Option<(RepoType, &RepoId)>) -> Result<Verified> {
let mut result = Verified::default();
let targets: Vec<(String, u64)> = match repo {
None => self
.store()
.list(key::BLOB_PREFIX)
.await?
.into_iter()
.map(|o| (o.key, o.size))
.collect(),
Some((repo_type, repo_id)) => {
let mut seen = BTreeSet::new();
let mut targets = Vec::new();
for manifest in self.list_manifests(repo_type, repo_id).await? {
for entry in manifest.entries {
let Some(hex) = digest::unprefixed(&entry.digest) else {
continue;
};
if !seen.insert(hex.to_owned()) {
continue;
}
let key = key::blob(hex);
match self.store().head(&key).await? {
Some(size) => targets.push((key, size)),
None => result.missing.push(key),
}
}
}
targets
}
};
for (key, size) in targets {
let Some(expected) = key::blob_digest(&key).map(str::to_owned) else {
continue;
};
let mut read = self.store().read_blob(&key, None).await?;
let mut hasher = Digest::new();
while let Some(chunk) = read.stream.next().await {
hasher.update(&chunk?);
}
let found = hasher.hex();
result.blobs += 1;
result.bytes += size;
if found != expected {
result.corrupt.push(Corruption {
key,
expected: digest::prefixed(&expected),
found: digest::prefixed(&found),
});
}
}
if result.corrupt.is_empty() && result.missing.is_empty() {
info!(
blobs = result.blobs,
bytes = result.bytes,
"verify complete; no problems found"
);
} else {
warn!(
corrupt = result.corrupt.len(),
missing = result.missing.len(),
"verify found problems"
);
}
Ok(result)
}
/// Re-resolve a ref upstream and repin it.
///
/// This is the whole point of `freeze` mode: moving to a newer revision is
/// an operator action, so upstream cannot change what a fleet resolves to.
pub async fn refresh(
&self,
repo_type: RepoType,
repo_id: &RepoId,
ref_name: &str,
) -> Result<(Option<String>, String)> {
if !self.upstream().enabled() {
return Err(Error::Config(
"refresh needs upstream; this instance is sealed (upstream.enabled = false)".into(),
));
}
let before = self
.load_ref(repo_type, repo_id, ref_name)
.await?
.map(|r| r.commit);
let after = self
.upstream()
.resolve_ref(repo_type, repo_id, ref_name)
.await?;
self.write_ref(
repo_type,
repo_id,
ref_name,
&after,
rustingface_entities::refs::PinnedBy::Operator,
)
.await?;
self.ensure_repo_record(repo_type, repo_id).await?;
Ok((before, after))
}
/// A round-trip probe of the bucket: write, read back, delete.
///
/// Used by `doctor` to prove credentials and permissions rather than
/// merely that a TCP connection opened.
pub async fn probe_store(&self) -> Result<()> {
let key = format!("doctor/{}", crate::registry::probe_id());
let payload = Bytes::from_static(b"rustingface doctor probe");
self.store().put(&key, payload.clone()).await?;
let read = self.store().get(&key).await?;
self.store().delete(&key).await?;
match read {
Some(bytes) if bytes == payload => Ok(()),
Some(_) => Err(Error::Storage(
"the doctor probe read back different bytes than it wrote".into(),
)),
None => Err(Error::Storage(
"the doctor probe wrote an object that could not then be read".into(),
)),
}
}
}
crates/rustingface-core/src/digest.rs
+88
View File
@@ -0,0 +1,88 @@
//! Incremental SHA-256 over a streaming body.
//!
//! The digest is computed over the same bytes that go to both sinks, so a
//! mismatch against upstream's LFS oid is detected before the manifest that
//! would reference the blob is written.
use sha2::{Digest as _, Sha256};
/// A running SHA-256 that yields the `sha256:<hex>` form used as a blob key.
#[derive(Default)]
pub struct Digest {
hasher: Sha256,
len: u64,
}
impl Digest {
/// A fresh, empty digest.
pub fn new() -> Self {
Self::default()
}
/// Absorb the next chunk.
pub fn update(&mut self, chunk: &[u8]) {
self.hasher.update(chunk);
self.len += chunk.len() as u64;
}
/// Bytes absorbed so far.
pub fn len(&self) -> u64 {
self.len
}
/// Whether nothing has been absorbed.
pub fn is_empty(&self) -> bool {
self.len == 0
}
/// The lowercase hex digest.
pub fn hex(self) -> String {
hex::encode(self.hasher.finalize())
}
}
/// Render a hex digest in the `sha256:<hex>` form stored in manifests.
pub fn prefixed(hex: &str) -> String {
format!("sha256:{hex}")
}
/// Strip the `sha256:` prefix from a manifest digest.
pub fn unprefixed(digest: &str) -> Option<&str> {
digest.strip_prefix("sha256:")
}
#[cfg(test)]
mod tests {
use super::*;
/// The well-known SHA-256 of the empty input, as a canary that this is
/// really SHA-256 and not something else with the same output width.
const EMPTY: &str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
#[test]
fn an_empty_stream_hashes_to_the_known_empty_digest() {
let d = Digest::new();
assert!(d.is_empty());
assert_eq!(d.hex(), EMPTY);
}
#[test]
fn chunking_does_not_change_the_digest_or_the_length() {
let mut whole = Digest::new();
whole.update(b"hello world");
let mut split = Digest::new();
for chunk in [b"hel".as_slice(), b"lo ", b"world"] {
split.update(chunk);
}
assert_eq!(split.len(), 11);
assert_eq!(whole.hex(), split.hex());
}
#[test]
fn the_prefixed_form_round_trips() {
let hex = "a".repeat(64);
assert_eq!(unprefixed(&prefixed(&hex)), Some(hex.as_str()));
assert_eq!(unprefixed("md5:abc"), None);
}
}
crates/rustingface-core/src/fetch.rs
+347
View File
@@ -0,0 +1,347 @@
//! The fetch path: one upstream read, two sinks.
//!
//! The upstream body is consumed exactly once and written concurrently to the
//! client response and to a multipart upload, with the digest computed over
//! the same bytes. The transfer runs in a detached task, so a client that hits
//! Ctrl-C at 30GB does not discard 30GB: the retention guarantee is better
//! served by finishing (spec §7).
use std::sync::Arc;
use bytes::Bytes;
use futures::StreamExt;
use rustingface_entities::error::{Error, Result};
use rustingface_entities::key;
use rustingface_entities::manifest::Entry;
use rustingface_entities::repo::{RepoId, RepoType};
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, info, warn};
use crate::digest::{self, Digest};
use crate::flight::Outcome;
use crate::ports::{Store, UpstreamHead};
use crate::registry::{FetchTask, LeadResult, Registry, TEE_DEPTH};
/// Where a transfer's bytes are being written while the digest is unknown.
enum Sink {
/// Straight to the final content-addressed key, because upstream gave an
/// LFS oid and the key is therefore known up front.
Direct { key: String },
/// Buffered in memory, for a small file with no oid.
Buffered,
/// To a staging key, to be moved once the digest is known.
Staged { key: String },
}
impl Registry {
/// Lead an upstream fetch for one file.
///
/// `blocking` forces the transfer to run to completion before returning,
/// which is what a ranged request on a cold file needs: a partial blob
/// must never reach the bucket.
pub(crate) async fn lead_fetch(
&self,
repo_type: RepoType,
repo_id: &RepoId,
commit: &str,
path: &str,
blocking: bool,
) -> Result<LeadResult> {
let body = self
.upstream()
.get_file(repo_type, repo_id, commit, path)
.await?;
let head = body.head.clone();
// Deduplication: when upstream named the oid and we already hold those
// bytes, the transfer is unnecessary. Dropping the response body here
// abandons the download rather than paying for it twice.
if let Some(oid) = head.expected_sha256()
&& self.store().head(&key::blob(&oid)).await?.is_some()
{
drop(body);
let size = self.store().head(&key::blob(&oid)).await?.unwrap_or(0);
let entry = entry_from(path, size, &oid, &head);
info!(%repo_id, commit, path, oid, "blob already stored under this digest; recording without transfer");
self.record_entry(repo_type, repo_id, commit, entry.clone())
.await?;
return Ok(LeadResult::Deduplicated(entry));
}
let sink = self.choose_sink(&head);
let (client_tx, client_rx) = mpsc::channel::<Result<Bytes>>(TEE_DEPTH);
let (outcome_tx, outcome_rx) = oneshot::channel::<Outcome>();
let transfer = Transfer {
store: Arc::clone(self.store()),
registry_repo_type: repo_type,
repo_id: repo_id.clone(),
commit: commit.to_owned(),
path: path.to_owned(),
head: head.clone(),
sink,
part_size: self.config().storage.multipart_part_size,
stall_timeout: self.config().server.client_stall_timeout,
};
if blocking {
// Nobody is reading the tee, so drop the client half up front and
// let the transfer run at full speed.
drop(client_rx);
let entry = transfer.run(self, body.stream, client_tx).await?;
return Ok(LeadResult::Completed(entry));
}
let registry = self.clone();
tokio::spawn(async move {
let result = transfer.run(&registry, body.stream, client_tx).await;
let outcome = match &result {
Ok(entry) => Outcome::Stored(entry.clone()),
Err(err) => Outcome::from(err),
};
let _ = outcome_tx.send(outcome);
});
Ok(LeadResult::Streaming {
head,
rx: client_rx,
task: FetchTask::new(outcome_rx),
})
}
/// Pick where the bytes go while the digest is still unknown.
fn choose_sink(&self, head: &UpstreamHead) -> Sink {
if let Some(oid) = head.expected_sha256() {
return Sink::Direct {
key: key::blob(&oid),
};
}
let size = head.linked_size.or(head.content_length);
match size {
Some(size) if size <= crate::registry::INLINE_BUFFER_LIMIT => Sink::Buffered,
_ => Sink::Staged {
key: format!("staging/{}", staging_id()),
},
}
}
}
/// One transfer in progress.
struct Transfer {
store: Arc<dyn Store>,
registry_repo_type: RepoType,
repo_id: RepoId,
commit: String,
path: String,
head: UpstreamHead,
sink: Sink,
part_size: u64,
stall_timeout: std::time::Duration,
}
impl Transfer {
/// Consume the upstream body, writing to storage and the client at once.
async fn run(
self,
registry: &Registry,
mut stream: futures::stream::BoxStream<'static, Result<Bytes>>,
client: mpsc::Sender<Result<Bytes>>,
) -> Result<Entry> {
let mut hasher = Digest::new();
let mut client = Some(client);
let mut buffered: Vec<Bytes> = Vec::new();
let mut writer = match &self.sink {
Sink::Buffered => None,
Sink::Direct { key } | Sink::Staged { key } => {
Some(self.store.write_blob(key, self.part_size).await?)
}
};
let transfer = async {
while let Some(chunk) = stream.next().await {
let chunk = chunk?;
hasher.update(&chunk);
match &mut writer {
Some(writer) => writer.write(chunk.clone()).await?,
None => buffered.push(chunk.clone()),
}
// Tee to the client, dropping it — not the transfer — if it has
// gone away or stopped reading.
if let Some(tx) = &client {
match tokio::time::timeout(self.stall_timeout, tx.send(Ok(chunk))).await {
Ok(Ok(())) => {}
Ok(Err(_)) => {
debug!(path = %self.path, "client disconnected; finishing the transfer without it");
client = None;
}
Err(_) => {
warn!(
path = %self.path,
"client stopped reading for longer than server.client_stall_timeout; \
detaching it and finishing the transfer"
);
client = None;
}
}
}
}
Ok::<(), Error>(())
};
if let Err(err) = transfer.await {
if let Some(writer) = writer {
let _ = writer.abort().await;
}
if let Some(tx) = &client {
let _ = tx.send(Err(Error::Upstream(err.to_string()))).await;
}
return Err(err);
}
let size = hasher.len();
let hex = hasher.hex();
// Verify before anything is completed or referenced: a mismatch must
// abort the upload, not land bytes that a later manifest could point at.
if let Some(expected) = self.head.expected_sha256()
&& expected != hex
{
if let Some(writer) = writer {
let _ = writer.abort().await;
}
let err = Error::DigestMismatch {
path: self.path.clone(),
expected: digest::prefixed(&expected),
actual: digest::prefixed(&hex),
};
if let Some(tx) = &client {
let _ = tx.send(Err(Error::Upstream(err.to_string()))).await;
}
return Err(err);
}
let final_key = key::blob(&hex);
// Phase 1: the blob is durable. A crash after this point leaks an
// orphan, which `gc` reclaims; nothing references it yet.
match (writer, &self.sink) {
(Some(writer), Sink::Direct { .. }) => writer.finish().await?,
(Some(writer), Sink::Staged { key: staged }) => {
writer.finish().await?;
if self.store.head(&final_key).await?.is_none() {
self.store.copy(staged, &final_key).await.map_err(|e| {
Error::Storage(format!(
"staged {size}-byte blob for {} could not be moved to its \
content-addressed key ({e}). Files this large normally arrive with an \
LFS oid and skip staging entirely; the staged object is left at {staged} \
for recovery.",
self.path
))
})?;
}
self.store.delete(staged).await?;
}
(None, Sink::Buffered) => {
if self.store.head(&final_key).await?.is_none() {
let mut whole = Vec::with_capacity(size as usize);
for chunk in &buffered {
whole.extend_from_slice(chunk);
}
self.store.put(&final_key, Bytes::from(whole)).await?;
}
}
(writer, sink) => {
if let Some(writer) = writer {
let _ = writer.abort().await;
}
let _ = sink;
return Err(Error::Internal("transfer sink and writer disagree".into()));
}
}
// Phase 2: only now may a manifest reference it.
let entry = entry_from(&self.path, size, &hex, &self.head);
registry
.record_entry(
self.registry_repo_type,
&self.repo_id,
&self.commit,
entry.clone(),
)
.await?;
info!(
repo_id = %self.repo_id, commit = %self.commit, path = %self.path, size,
detached = client.is_none(),
"stored blob and recorded manifest entry"
);
Ok(entry)
}
}
/// Build the manifest entry for a completed transfer, carrying upstream's
/// values through unchanged.
fn entry_from(path: &str, size: u64, hex: &str, head: &UpstreamHead) -> Entry {
Entry {
path: path.to_owned(),
size,
digest: digest::prefixed(hex),
etag: head.etag.clone(),
linked_etag: head.linked_etag.clone(),
linked_size: head.linked_size,
lfs: head.lfs,
}
}
/// A collision-resistant staging suffix, without pulling in a uuid crate for
/// a name nothing else ever reads.
fn staging_id() -> String {
use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hasher};
let mut hasher = RandomState::new().build_hasher();
hasher.write_u128(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos(),
);
format!(
"{:016x}{:016x}",
hasher.finish(),
RandomState::new().build_hasher().finish()
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn an_entry_carries_upstream_values_through_unchanged() {
let head = UpstreamHead {
commit: Some("abc".into()),
etag: Some("\"71ab09\"".into()),
linked_etag: Some("\"71ab0900\"".into()),
linked_size: Some(4_831_838_208),
content_length: Some(4_831_838_208),
lfs: true,
};
let entry = entry_from("model.safetensors", 4_831_838_208, &"7".repeat(64), &head);
assert_eq!(entry.etag.as_deref(), Some("\"71ab09\""), "quotes are kept");
assert_eq!(entry.linked_etag.as_deref(), Some("\"71ab0900\""));
assert_eq!(entry.linked_size, Some(4_831_838_208));
assert!(entry.lfs);
assert_eq!(entry.digest, format!("sha256:{}", "7".repeat(64)));
}
#[test]
fn staging_ids_do_not_repeat() {
let a = staging_id();
let b = staging_id();
assert_ne!(a, b);
assert_eq!(a.len(), 32);
}
}
crates/rustingface-core/src/flight.rs
+301
View File
@@ -0,0 +1,301 @@
//! Single-flight: exactly one upstream fetch per uncached blob.
//!
//! Without this, a fleet-wide rollout of a new model produces N concurrent
//! 40GB pulls (spec §7). The map is keyed on
//! `(repo_type, repo_id, commit, path)` — the commit, not the ref, so two
//! clients naming `main` and its resolved SHA share a flight.
//!
//! The leader tees bytes to its own response while uploading; followers wait
//! for the flight to finish and then read the completed blob from the bucket.
//! Streaming live to a follower that arrived mid-body would mean handing it a
//! stream already missing its first chunks, so followers deliberately take the
//! slower, correct path.
use std::sync::Arc;
use dashmap::DashMap;
use rustingface_entities::error::{Error, Result};
use rustingface_entities::manifest::Entry;
use rustingface_entities::repo::{RepoId, RepoType};
use tokio::sync::watch;
/// Identifies one uncached blob.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FlightKey {
pub repo_type: RepoType,
pub repo_id: RepoId,
pub commit: String,
pub path: String,
}
impl FlightKey {
/// Key a fetch by its resolved commit, so ref and SHA requests coalesce.
pub fn new(
repo_type: RepoType,
repo_id: RepoId,
commit: impl Into<String>,
path: impl Into<String>,
) -> Self {
Self {
repo_type,
repo_id,
commit: commit.into(),
path: path.into(),
}
}
}
/// How a completed flight ended.
#[derive(Debug, Clone)]
pub enum Outcome {
/// The blob is durably stored and the manifest references it.
Stored(Entry),
/// The fetch failed. Carries the rendered error, since [`Error`] is not
/// `Clone` and every waiter needs a copy.
Failed {
label: &'static str,
message: String,
},
}
impl Outcome {
/// Reconstitute a failure as an error for a waiter to return.
pub fn into_result(self) -> Result<Entry> {
match self {
Self::Stored(entry) => Ok(entry),
Self::Failed { label, message } => Err(match label {
"repo_not_found" => Error::RepoNotFound(message),
"revision_not_found" => Error::RevisionNotFound(message),
"entry_not_found" => Error::EntryNotFound(message),
"gated" => Error::GatedRepo(message),
"disabled" => Error::DisabledRepo(message),
"upstream_unauthorized" => Error::UpstreamUnauthorized(message),
"storage" => Error::Storage(message),
_ => Error::Upstream(message),
}),
}
}
}
impl From<&Error> for Outcome {
fn from(err: &Error) -> Self {
Self::Failed {
label: err.label(),
message: err.to_string(),
}
}
}
/// A flight in progress. Waiters hold one of these and await the outcome.
#[derive(Debug)]
pub struct Flight {
rx: watch::Receiver<Option<Outcome>>,
}
impl Flight {
/// Block until the leader's fetch completes, then yield its outcome.
///
/// A leader that panics drops its sender; that closes the channel and
/// every waiter gets a definite error rather than hanging forever.
pub async fn wait(mut self) -> Result<Entry> {
loop {
if let Some(outcome) = self.rx.borrow_and_update().clone() {
return outcome.into_result();
}
if self.rx.changed().await.is_err() {
return Err(Error::Internal(
"the upstream fetch this request was waiting on ended without reporting an outcome".into(),
));
}
}
}
}
/// The leader's half: publishing an outcome and clearing the map entry.
///
/// Dropping this without calling [`Lead::finish`] still removes the map entry
/// and closes the channel, so a panicking leader cannot wedge later requests
/// on a key that will never complete.
pub struct Lead {
map: Arc<DashMap<FlightKey, watch::Receiver<Option<Outcome>>>>,
key: FlightKey,
tx: watch::Sender<Option<Outcome>>,
}
impl Lead {
/// Publish the outcome to every waiter and retire the key.
pub fn finish(self, outcome: Outcome) {
self.map.remove(&self.key);
let _ = self.tx.send(Some(outcome));
}
}
impl Drop for Lead {
fn drop(&mut self) {
self.map.remove(&self.key);
}
}
/// Either this caller leads the fetch, or it waits on one already running.
pub enum Role {
Leader(Lead),
Follower(Flight),
}
/// The in-flight registry.
#[derive(Debug, Default)]
pub struct Flights {
map: Arc<DashMap<FlightKey, watch::Receiver<Option<Outcome>>>>,
}
impl Flights {
/// An empty registry.
pub fn new() -> Self {
Self::default()
}
/// Claim leadership of `key`, or join the flight already leading it.
///
/// The claim is made under DashMap's per-shard lock, so exactly one caller
/// can become leader for a key no matter how many arrive at once.
pub fn join(&self, key: FlightKey) -> Role {
use dashmap::mapref::entry::Entry as MapEntry;
match self.map.entry(key.clone()) {
MapEntry::Occupied(occupied) => Role::Follower(Flight {
rx: occupied.get().clone(),
}),
MapEntry::Vacant(vacant) => {
let (tx, rx) = watch::channel(None);
vacant.insert(rx);
Role::Leader(Lead {
map: Arc::clone(&self.map),
key,
tx,
})
}
}
}
/// How many fetches are in flight, for the metrics gauge.
pub fn len(&self) -> usize {
self.map.len()
}
/// Whether nothing is in flight.
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn key(path: &str) -> FlightKey {
FlightKey::new(
RepoType::Models,
"Qwen/Qwen3-32B".parse().unwrap(),
"8f2ac31",
path,
)
}
fn entry(path: &str) -> Entry {
Entry {
path: path.into(),
size: 4,
digest: format!("sha256:{}", "a".repeat(64)),
etag: Some("\"abc\"".into()),
linked_etag: None,
linked_size: None,
lfs: false,
}
}
#[tokio::test]
async fn eight_concurrent_callers_produce_exactly_one_leader() {
let flights = Arc::new(Flights::new());
let mut leaders = 0;
let mut followers = Vec::new();
for _ in 0..8 {
match flights.join(key("config.json")) {
Role::Leader(lead) => {
leaders += 1;
// Hand the outcome over only after every caller has joined.
followers.push(tokio::spawn(async move {
tokio::task::yield_now().await;
lead.finish(Outcome::Stored(entry("config.json")));
Ok(entry("config.json"))
}));
}
Role::Follower(flight) => {
followers.push(tokio::spawn(flight.wait()));
}
}
}
assert_eq!(leaders, 1, "exactly one upstream fetch");
for handle in followers {
let stored = handle.await.unwrap().unwrap();
assert_eq!(stored.path, "config.json");
}
assert!(
flights.is_empty(),
"the key is retired once the flight ends"
);
}
#[tokio::test]
async fn different_paths_do_not_share_a_flight() {
let flights = Flights::new();
let a = flights.join(key("a.json"));
let b = flights.join(key("b.json"));
assert!(matches!(a, Role::Leader(_)));
assert!(matches!(b, Role::Leader(_)));
assert_eq!(flights.len(), 2, "one flight per path, not one per repo");
}
#[tokio::test]
async fn a_failure_reaches_every_follower_as_the_same_error() {
let flights = Flights::new();
let Role::Leader(lead) = flights.join(key("x")) else {
panic!("first caller must lead");
};
let Role::Follower(flight) = flights.join(key("x")) else {
panic!("second caller must follow");
};
lead.finish(Outcome::from(&Error::GatedRepo(
"terms not accepted".into(),
)));
let err = flight.wait().await.unwrap_err();
assert!(matches!(err, Error::GatedRepo(_)), "{err}");
assert_eq!(err.status(), 403);
}
#[tokio::test]
async fn a_leader_that_vanishes_does_not_hang_its_followers() {
let flights = Flights::new();
let Role::Leader(lead) = flights.join(key("x")) else {
panic!("first caller must lead");
};
let Role::Follower(flight) = flights.join(key("x")) else {
panic!("second caller must follow");
};
drop(lead);
let err = flight.wait().await.unwrap_err();
assert!(matches!(err, Error::Internal(_)), "{err}");
assert!(flights.is_empty(), "a dropped lead still retires the key");
}
#[tokio::test]
async fn a_key_can_be_led_again_after_its_flight_ends() {
let flights = Flights::new();
let Role::Leader(lead) = flights.join(key("x")) else {
panic!("first caller must lead");
};
lead.finish(Outcome::Stored(entry("x")));
assert!(
matches!(flights.join(key("x")), Role::Leader(_)),
"a later miss on the same key starts a new flight"
);
}
}
crates/rustingface-core/src/lib.rs
+18
View File
@@ -0,0 +1,18 @@
//! Resolution, retention and delivery for rustingface.
//!
//! This crate holds the logic: how a revision becomes a commit, how a miss
//! becomes exactly one upstream fetch, how bytes reach the bucket before
//! anything references them, and what an operator may remove. It defines the
//! two ports it needs — [`ports::Store`] and [`ports::Upstream`] — and leaves
//! their implementation to `rustingface-data`.
pub mod admin;
pub mod digest;
mod fetch;
pub mod flight;
pub mod policy;
pub mod ports;
pub mod registry;
pub use policy::Policy;
pub use registry::{Body, Registry, ResolvedFile, Source};
crates/rustingface-core/src/policy.rs
+116
View File
@@ -0,0 +1,116 @@
//! Retention policy: what rustingface is willing to fetch that it does not
//! already hold.
//!
//! Policy never restricts what is *served* from the bucket. Anything already
//! stored stays servable regardless of a later allowlist change, because
//! guarantee 3 does not have a policy exception.
use globset::{Glob, GlobSet, GlobSetBuilder};
use rustingface_entities::config::Policy as PolicyConfig;
use rustingface_entities::error::{Error, Result};
use rustingface_entities::repo::RepoId;
/// Compiled form of `[policy]`.
#[derive(Debug)]
pub struct Policy {
allowlist: Option<GlobSet>,
allow_new_repos: bool,
}
impl Policy {
/// Compile the globs once, at startup, so a bad pattern fails loudly
/// before the listener binds rather than on the first miss.
pub fn new(cfg: &PolicyConfig) -> Result<Self> {
let allowlist = if cfg.allowlist.is_empty() {
None
} else {
let mut builder = GlobSetBuilder::new();
for pattern in &cfg.allowlist {
let glob = Glob::new(pattern).map_err(|e| {
Error::Config(format!("policy.allowlist pattern {pattern:?}: {e}"))
})?;
builder.add(glob);
}
Some(
builder
.build()
.map_err(|e| Error::Config(format!("policy.allowlist: {e}")))?,
)
};
Ok(Self {
allowlist,
allow_new_repos: cfg.allow_new_repos,
})
}
/// Whether a repo may be fetched from upstream.
///
/// `known` says whether the bucket already holds a repo record, which is
/// what `allow_new_repos = false` gates on.
pub fn admits_fetch(&self, repo_id: &RepoId, known: bool) -> Result<()> {
if let Some(set) = &self.allowlist
&& !set.is_match(repo_id.to_string())
{
return Err(Error::PolicyDenied(format!(
"{repo_id} is not in policy.allowlist"
)));
}
if !self.allow_new_repos && !known {
return Err(Error::PolicyDenied(format!(
"{repo_id} is not already stored and policy.allow_new_repos is false"
)));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn policy(allowlist: &[&str], allow_new_repos: bool) -> Policy {
Policy::new(&PolicyConfig {
ref_resolution: rustingface_entities::config::RefResolution::Freeze,
allow_new_repos,
allowlist: allowlist.iter().map(|s| s.to_string()).collect(),
})
.unwrap()
}
fn id(s: &str) -> RepoId {
s.parse().unwrap()
}
#[test]
fn an_empty_allowlist_admits_everything() {
let p = policy(&[], true);
assert!(p.admits_fetch(&id("Qwen/Qwen3-32B"), false).is_ok());
assert!(p.admits_fetch(&id("gpt2"), false).is_ok());
}
#[test]
fn a_namespace_glob_admits_only_that_namespace() {
let p = policy(&["Qwen/*", "meta-llama/*"], true);
assert!(p.admits_fetch(&id("Qwen/Qwen3-32B"), false).is_ok());
assert!(p.admits_fetch(&id("meta-llama/Llama-3-8B"), false).is_ok());
assert!(p.admits_fetch(&id("someone/else"), false).is_err());
}
#[test]
fn allow_new_repos_false_still_admits_a_repo_already_stored() {
let p = policy(&[], false);
assert!(p.admits_fetch(&id("Qwen/Qwen3-32B"), true).is_ok());
assert!(p.admits_fetch(&id("Qwen/Qwen3-32B"), false).is_err());
}
#[test]
fn a_malformed_glob_fails_at_construction_not_at_request_time() {
let err = Policy::new(&PolicyConfig {
ref_resolution: rustingface_entities::config::RefResolution::Freeze,
allow_new_repos: true,
allowlist: vec!["Qwen/[".into()],
})
.unwrap_err();
assert!(matches!(err, Error::Config(_)), "{err}");
}
}
crates/rustingface-core/src/ports.rs
+234
View File
@@ -0,0 +1,234 @@
//! Ports: the two things rustingface needs from the outside world.
//!
//! [`Store`] is the bucket — the entire durable state of the system.
//! [`Upstream`] is the Hub, consulted on miss only and not at all when sealed.
//! Both are implemented in `rustingface-data`; nothing in this crate performs
//! I/O of its own.
use std::ops::Range;
use std::time::Duration;
use async_trait::async_trait;
use bytes::Bytes;
use futures::stream::BoxStream;
use rustingface_entities::error::Result;
use serde_json::Value;
use rustingface_entities::repo::{RepoId, RepoType};
/// Metadata about one stored object.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoredObject {
pub key: String,
pub size: u64,
}
/// A readable blob: the byte stream plus what the store said about it.
pub struct BlobRead {
/// Bytes of the requested range, in order.
pub stream: BoxStream<'static, Result<Bytes>>,
/// Length of this read, which is the range length when one was requested.
pub len: u64,
/// Full length of the object, regardless of the range read.
pub total: u64,
}
/// A streaming multipart upload.
///
/// Dropping without [`BlobWriter::finish`] leaves an incomplete multipart
/// upload; call [`BlobWriter::abort`] on the failure path so the backend can
/// reclaim the parts.
#[async_trait]
pub trait BlobWriter: Send {
/// Buffer bytes, flushing whole parts as they accumulate.
async fn write(&mut self, chunk: Bytes) -> Result<()>;
/// Flush the tail and complete the upload.
async fn finish(self: Box<Self>) -> Result<()>;
/// Discard the upload and any parts already sent.
async fn abort(self: Box<Self>) -> Result<()>;
}
/// The bucket.
#[async_trait]
pub trait Store: Send + Sync + 'static {
/// Fetch a small object whole. `Ok(None)` when it does not exist.
async fn get(&self, key: &str) -> Result<Option<Bytes>>;
/// Overwrite an object.
async fn put(&self, key: &str, body: Bytes) -> Result<()>;
/// Write only if the key is absent. Returns `false` if it already existed.
///
/// Backed by a conditional put where the backend supports it, which is
/// what makes ref freezing atomically first-write-wins.
async fn put_if_absent(&self, key: &str, body: Bytes) -> Result<bool>;
/// Whether the backend's [`Store::put_if_absent`] is genuinely atomic. A
/// `false` here is a correctness caveat worth logging, not a failure.
fn conditional_put_is_atomic(&self) -> bool {
true
}
/// Delete an object. Deleting a key that is already absent succeeds.
async fn delete(&self, key: &str) -> Result<()>;
/// Server-side copy, used to move a staged blob to its content-addressed
/// key once the digest is known.
async fn copy(&self, from: &str, to: &str) -> Result<()>;
/// Size of an object, or `None` when absent.
async fn head(&self, key: &str) -> Result<Option<u64>>;
/// Every object under `prefix`.
async fn list(&self, prefix: &str) -> Result<Vec<StoredObject>>;
/// Open a blob for reading, optionally a byte range.
async fn read_blob(&self, key: &str, range: Option<Range<u64>>) -> Result<BlobRead>;
/// Begin a multipart upload, flushing whole parts of `part_size` bytes.
async fn write_blob(&self, key: &str, part_size: u64) -> Result<Box<dyn BlobWriter>>;
/// A presigned GET URL, when the backend can mint one.
async fn presign_get(&self, key: &str, ttl: Duration) -> Result<String>;
}
/// What upstream said about a file, captured verbatim for replay.
///
/// Every string field here is stored and re-emitted unchanged. rustingface
/// never computes a git blob SHA-1 and never derives an etag from content.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct UpstreamHead {
/// `X-Repo-Commit`.
pub commit: Option<String>,
/// `ETag`, quotes included.
pub etag: Option<String>,
/// `X-Linked-Etag`, quotes included.
pub linked_etag: Option<String>,
/// `X-Linked-Size`.
pub linked_size: Option<u64>,
/// `Content-Length` of the resolved resource.
pub content_length: Option<u64>,
/// Whether upstream served this through the Git LFS bridge.
pub lfs: bool,
}
impl UpstreamHead {
/// The etag a client would key its cache on: `X-Linked-Etag` when present,
/// otherwise `ETag` (spec §6).
pub fn effective_etag(&self) -> Option<&str> {
self.linked_etag.as_deref().or(self.etag.as_deref())
}
/// The sha256 upstream promised for this file, when it gave one.
///
/// For LFS-backed files the linked etag *is* the LFS oid, so a completed
/// transfer can be checked against it before anything references the blob.
pub fn expected_sha256(&self) -> Option<String> {
let raw = self.linked_etag.as_deref()?.trim_matches('"');
(raw.len() == 64 && raw.bytes().all(|b| b.is_ascii_hexdigit()))
.then(|| raw.to_ascii_lowercase())
}
}
/// An in-progress upstream download.
pub struct UpstreamBody {
pub head: UpstreamHead,
pub stream: BoxStream<'static, Result<Bytes>>,
}
/// The Hub.
#[async_trait]
pub trait Upstream: Send + Sync + 'static {
/// `false` when the service is sealed. Nothing else on this trait is
/// called when it is.
fn enabled(&self) -> bool;
/// The configured endpoint, recorded into manifests for provenance.
fn endpoint(&self) -> &str;
/// Resolve a ref name to a commit SHA.
async fn resolve_ref(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
) -> Result<String>;
/// `HEAD` on the resolve path, for metadata without transferring bytes.
async fn head_file(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
path: &str,
) -> Result<UpstreamHead>;
/// `GET` on the resolve path, returning headers and the body stream.
async fn get_file(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
path: &str,
) -> Result<UpstreamBody>;
/// The `/api/{type}/{id}/revision/{rev}` document, verbatim.
async fn repo_info(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
) -> Result<Value>;
/// The recursive, expanded `/api/{type}/{id}/tree/{rev}` document, verbatim.
async fn repo_tree(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
) -> Result<Value>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn linked_etag_wins_over_etag() {
let head = UpstreamHead {
etag: Some("\"aaa\"".into()),
linked_etag: Some("\"bbb\"".into()),
..Default::default()
};
assert_eq!(head.effective_etag(), Some("\"bbb\""));
}
#[test]
fn a_plain_etag_is_used_when_upstream_omits_the_linked_one() {
let head = UpstreamHead {
etag: Some("\"aaa\"".into()),
..Default::default()
};
assert_eq!(head.effective_etag(), Some("\"aaa\""));
}
#[test]
fn an_lfs_oid_is_recovered_from_the_linked_etag() {
let oid = "71ab09".repeat(10) + "abcd";
let head = UpstreamHead {
linked_etag: Some(format!("\"{oid}\"")),
..Default::default()
};
assert_eq!(head.expected_sha256().as_deref(), Some(oid.as_str()));
}
#[test]
fn a_git_blob_etag_is_not_mistaken_for_a_sha256() {
// Git blob SHA-1s are 40 hex characters; only a 64-hex value is an oid.
let head = UpstreamHead {
linked_etag: Some("\"8f2a1c9e0b7d4f6a3e5c2b8d9f0a1c3e5b7d9f01\"".into()),
..Default::default()
};
assert_eq!(head.expected_sha256(), None);
}
}
crates/rustingface-core/src/registry.rs
+1078
View File
File diff suppressed because it is too large Load Diff
crates/rustingface-data/Cargo.toml
+26
View File
@@ -0,0 +1,26 @@
[package]
name = "rustingface-data"
description = "Object-storage and upstream-Hub adapters implementing the rustingface-core ports."
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
authors.workspace = true
repository.workspace = true
[dependencies]
rustingface-core.workspace = true
rustingface-entities.workspace = true
async-trait.workspace = true
bytes.workspace = true
futures.workspace = true
object_store.workspace = true
reqwest.workspace = true
serde_json.workspace = true
tokio.workspace = true
tracing.workspace = true
http.workspace = true
[dev-dependencies]
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
crates/rustingface-data/src/lib.rs
+11
View File
@@ -0,0 +1,11 @@
//! Adapters for the ports `rustingface-core` defines.
//!
//! Two implementations live here and nowhere else: object storage, and the
//! upstream Hub. Everything that talks to a network or a filesystem on
//! rustingface's behalf is in this crate.
pub mod store;
pub mod upstream;
pub use store::ObjectStoreAdapter;
pub use upstream::{HubClient, SealedUpstream};
crates/rustingface-data/src/store.rs
+300
View File
@@ -0,0 +1,300 @@
//! Object-storage adapter.
//!
//! The bucket is the entire durable state of the system (guarantee 5), so this
//! adapter deliberately stays thin: keys in, bytes out, no schema knowledge and
//! no caching layer that could disagree with what is actually stored.
use std::ops::Range;
use std::path::Path as FsPath;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use bytes::Bytes;
use futures::StreamExt;
use object_store::aws::{AmazonS3, AmazonS3Builder};
use object_store::local::LocalFileSystem;
use object_store::path::Path;
use object_store::signer::Signer;
use object_store::{
GetOptions, GetRange, ObjectStore, ObjectStoreExt, PutMode, PutOptions, PutPayload,
WriteMultipart,
};
use rustingface_core::ports::{BlobRead, BlobWriter, Store, StoredObject};
use rustingface_entities::config::Storage as StorageConfig;
use rustingface_entities::error::{Error, Result};
use tracing::debug;
/// A [`Store`] backed by `object_store`.
pub struct ObjectStoreAdapter {
inner: Arc<dyn ObjectStore>,
signer: Option<Arc<AmazonS3>>,
}
impl ObjectStoreAdapter {
/// Build from `[storage]`, resolving credentials from their files.
///
/// `local_path` selects a filesystem store, which is what the test suite
/// and a single-host trial use. It forfeits guarantee 5 — a local
/// directory is not a bucket you can re-point a new host at — so the
/// caller warns when it is in play.
pub fn new(
cfg: &StorageConfig,
access_key_id: Option<String>,
secret_access_key: Option<String>,
) -> Result<Self> {
if let Some(path) = &cfg.local_path {
std::fs::create_dir_all(path).map_err(|e| {
Error::Config(format!("storage.local_path {}: {e}", path.display()))
})?;
let store = LocalFileSystem::new_with_prefix(path).map_err(|e| {
Error::Config(format!("storage.local_path {}: {e}", path.display()))
})?;
return Ok(Self {
inner: Arc::new(store),
signer: None,
});
}
let mut builder = AmazonS3Builder::new()
.with_bucket_name(&cfg.bucket)
.with_region(&cfg.region)
.with_virtual_hosted_style_request(!cfg.path_style);
if let Some(endpoint) = &cfg.endpoint {
// MinIO on the mesh is reached over plain HTTP; the mesh itself is
// the confidentiality boundary. Allowing it must be explicit
// rather than a silent downgrade, so it keys off the scheme the
// operator actually configured.
if endpoint.starts_with("http://") {
builder = builder.with_allow_http(true);
}
builder = builder.with_endpoint(endpoint);
}
if let Some(key) = access_key_id {
builder = builder.with_access_key_id(key);
}
if let Some(secret) = secret_access_key {
builder = builder.with_secret_access_key(secret);
}
let s3 = builder
.build()
.map_err(|e| Error::Config(format!("object storage: {e}")))?;
let s3 = Arc::new(s3);
Ok(Self {
inner: Arc::clone(&s3) as Arc<dyn ObjectStore>,
signer: Some(s3),
})
}
/// A filesystem-backed adapter, for tests.
pub fn local(path: &FsPath) -> Result<Self> {
std::fs::create_dir_all(path)
.map_err(|e| Error::Config(format!("{}: {e}", path.display())))?;
let store = LocalFileSystem::new_with_prefix(path)
.map_err(|e| Error::Config(format!("{}: {e}", path.display())))?;
Ok(Self {
inner: Arc::new(store),
signer: None,
})
}
}
/// Map an `object_store` error, preserving "absent" as a distinct outcome.
fn absent(err: &object_store::Error) -> bool {
matches!(err, object_store::Error::NotFound { .. })
}
fn storage(context: &str, err: object_store::Error) -> Error {
Error::Storage(format!("{context}: {err}"))
}
#[async_trait]
impl Store for ObjectStoreAdapter {
async fn get(&self, key: &str) -> Result<Option<Bytes>> {
match self.inner.get(&Path::from(key)).await {
Ok(result) => result
.bytes()
.await
.map(Some)
.map_err(|e| storage(&format!("reading {key}"), e)),
Err(err) if absent(&err) => Ok(None),
Err(err) => Err(storage(&format!("reading {key}"), err)),
}
}
async fn put(&self, key: &str, body: Bytes) -> Result<()> {
self.inner
.put(&Path::from(key), PutPayload::from_bytes(body))
.await
.map(|_| ())
.map_err(|e| storage(&format!("writing {key}"), e))
}
async fn put_if_absent(&self, key: &str, body: Bytes) -> Result<bool> {
let opts = PutOptions {
mode: PutMode::Create,
..Default::default()
};
match self
.inner
.put_opts(&Path::from(key), PutPayload::from_bytes(body), opts)
.await
{
Ok(_) => Ok(true),
Err(object_store::Error::AlreadyExists { .. }) => Ok(false),
Err(object_store::Error::Precondition { .. }) => Ok(false),
Err(err) => Err(storage(&format!("conditionally writing {key}"), err)),
}
}
fn conditional_put_is_atomic(&self) -> bool {
// S3 with `If-None-Match: *`, which MinIO supports, and the local
// filesystem store, which uses an atomic link. Both are genuinely
// first-write-wins.
true
}
async fn delete(&self, key: &str) -> Result<()> {
match self.inner.delete(&Path::from(key)).await {
Ok(()) => Ok(()),
Err(err) if absent(&err) => Ok(()),
Err(err) => Err(storage(&format!("deleting {key}"), err)),
}
}
async fn copy(&self, from: &str, to: &str) -> Result<()> {
self.inner
.copy(&Path::from(from), &Path::from(to))
.await
.map_err(|e| storage(&format!("copying {from} to {to}"), e))
}
async fn head(&self, key: &str) -> Result<Option<u64>> {
match self.inner.head(&Path::from(key)).await {
Ok(meta) => Ok(Some(meta.size)),
Err(err) if absent(&err) => Ok(None),
Err(err) => Err(storage(&format!("stat of {key}"), err)),
}
}
async fn list(&self, prefix: &str) -> Result<Vec<StoredObject>> {
let mut stream = self.inner.list(Some(&Path::from(prefix)));
let mut objects = Vec::new();
while let Some(meta) = stream.next().await {
let meta = meta.map_err(|e| storage(&format!("listing {prefix}"), e))?;
objects.push(StoredObject {
key: meta.location.as_ref().to_owned(),
size: meta.size,
});
}
objects.sort_by(|a, b| a.key.cmp(&b.key));
Ok(objects)
}
async fn read_blob(&self, key: &str, range: Option<Range<u64>>) -> Result<BlobRead> {
let opts = GetOptions {
range: range.clone().map(GetRange::Bounded),
..Default::default()
};
let result = self
.inner
.get_opts(&Path::from(key), opts)
.await
.map_err(|e| storage(&format!("reading blob {key}"), e))?;
let total = result.meta.size;
let len = result.range.end - result.range.start;
let key_for_errors = key.to_owned();
let stream = result.into_stream().map(move |chunk| {
chunk.map_err(|e| storage(&format!("streaming blob {key_for_errors}"), e))
});
Ok(BlobRead {
stream: Box::pin(stream),
len,
total,
})
}
async fn write_blob(&self, key: &str, part_size: u64) -> Result<Box<dyn BlobWriter>> {
let upload = self
.inner
.put_multipart(&Path::from(key))
.await
.map_err(|e| storage(&format!("starting multipart upload of {key}"), e))?;
Ok(Box::new(MultipartBlobWriter {
key: key.to_owned(),
inner: Some(WriteMultipart::new_with_chunk_size(
upload,
part_size as usize,
)),
}))
}
async fn presign_get(&self, key: &str, ttl: Duration) -> Result<String> {
let signer = self.signer.as_ref().ok_or_else(|| {
Error::Config(
"blob_delivery = \"redirect\" needs an S3 backend; the configured store cannot presign"
.into(),
)
})?;
signer
.signed_url(http::Method::GET, &Path::from(key), ttl)
.await
.map(|url| url.to_string())
.map_err(|e| storage(&format!("presigning {key}"), e))
}
}
/// A streaming multipart upload.
struct MultipartBlobWriter {
key: String,
inner: Option<WriteMultipart>,
}
impl MultipartBlobWriter {
fn take(&mut self, what: &str) -> Result<WriteMultipart> {
self.inner.take().ok_or_else(|| {
Error::Internal(format!("{what} called twice on the upload of {}", self.key))
})
}
}
#[async_trait]
impl BlobWriter for MultipartBlobWriter {
async fn write(&mut self, chunk: Bytes) -> Result<()> {
let writer = self
.inner
.as_mut()
.ok_or_else(|| Error::Internal(format!("write after finish on {}", self.key)))?;
// Bound the number of parts in flight so a fast upstream cannot grow
// the heap without limit when the backend is the slower side.
writer
.wait_for_capacity(8)
.await
.map_err(|e| storage(&format!("uploading {}", self.key), e))?;
writer.put(chunk);
Ok(())
}
async fn finish(mut self: Box<Self>) -> Result<()> {
let writer = self.take("finish")?;
writer
.finish()
.await
.map(|_| ())
.map_err(|e| storage(&format!("completing upload of {}", self.key), e))
}
async fn abort(mut self: Box<Self>) -> Result<()> {
let key = self.key.clone();
let writer = self.take("abort")?;
debug!(%key, "aborting multipart upload");
writer
.abort()
.await
.map_err(|e| storage(&format!("aborting upload of {key}"), e))
}
}
crates/rustingface-data/src/upstream.rs
+431
View File
@@ -0,0 +1,431 @@
//! The upstream Hub adapter.
//!
//! Consulted on miss only. Everything it returns that a client can observe is
//! carried through verbatim: this adapter parses headers into fields, it never
//! normalises or recomputes them.
use async_trait::async_trait;
use futures::StreamExt;
use reqwest::header::{ACCEPT_ENCODING, AUTHORIZATION, ETAG, HeaderMap, HeaderValue, USER_AGENT};
use reqwest::{Client, Response, StatusCode};
use rustingface_core::ports::{Upstream, UpstreamBody, UpstreamHead};
use rustingface_entities::error::{Error, Result};
use rustingface_entities::repo::{RepoId, RepoType};
use serde_json::Value;
use tracing::debug;
/// `X-Repo-Commit`: the commit a revision resolved to.
const X_REPO_COMMIT: &str = "x-repo-commit";
/// `X-Linked-Etag`: the LFS object's etag, which clients prefer over `ETag`.
const X_LINKED_ETAG: &str = "x-linked-etag";
/// `X-Linked-Size`: the LFS object's true size.
const X_LINKED_SIZE: &str = "x-linked-size";
/// `X-Error-Code`: upstream's own typed error label.
const X_ERROR_CODE: &str = "x-error-code";
/// A [`Upstream`] backed by `reqwest`.
pub struct HubClient {
client: Client,
endpoint: String,
enabled: bool,
token: Option<String>,
}
impl HubClient {
/// Build from `[upstream]`.
///
/// The client always advertises `Accept-Encoding: identity` so upstream
/// reports real byte lengths rather than compressed ones, matching what
/// `huggingface_hub` itself asks for and keeping size accounting honest.
pub fn new(
cfg: &rustingface_entities::config::Upstream,
token: Option<String>,
) -> Result<Self> {
let mut headers = HeaderMap::new();
headers.insert(ACCEPT_ENCODING, HeaderValue::from_static("identity"));
headers.insert(
USER_AGENT,
HeaderValue::from_str(&format!("rustingface/{}", env!("CARGO_PKG_VERSION")))
.map_err(|e| Error::Config(format!("user agent: {e}")))?,
);
if cfg.disable_xet {
// Never advertise Xet capability upstream, so the Hub answers over
// the Git LFS bridge with a whole-file URL. Storing whole blobs is
// what keeps the bucket readable without this binary (spec §8).
headers.insert("x-xet-capable", HeaderValue::from_static("false"));
}
let client = Client::builder()
.default_headers(headers)
.connect_timeout(cfg.connect_timeout)
.read_timeout(cfg.read_timeout)
.build()
.map_err(|e| Error::Config(format!("upstream client: {e}")))?;
Ok(Self {
client,
endpoint: cfg.endpoint.trim_end_matches('/').to_owned(),
enabled: cfg.enabled,
token,
})
}
/// The resolve URL for a file.
fn resolve_url(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
path: &str,
) -> String {
let prefix = match repo_type {
RepoType::Models => String::new(),
RepoType::Datasets => "datasets/".into(),
};
format!(
"{}/{prefix}{repo_id}/resolve/{}/{}",
self.endpoint,
urlencode_segment(revision),
path.split('/')
.map(urlencode_segment)
.collect::<Vec<_>>()
.join("/"),
)
}
/// The API URL for a repo revision.
fn api_url(&self, repo_type: RepoType, repo_id: &RepoId, suffix: &str) -> String {
format!(
"{}/api/{}/{repo_id}{suffix}",
self.endpoint,
repo_type.as_str()
)
}
/// Attach the operator's upstream token.
///
/// Used on outbound requests only. It is never echoed downstream and never
/// logged: a gated repo works because *this operator's* token accepted the
/// terms, and the stored copy is for that operator's own use.
fn authorize(&self, request: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
match &self.token {
Some(token) => request.header(AUTHORIZATION, format!("Bearer {token}")),
None => request,
}
}
/// Turn a non-success upstream status into the error it maps to.
async fn check(&self, response: Response, what: &str) -> Result<Response> {
let status = response.status();
if status.is_success() {
return Ok(response);
}
let code = response
.headers()
.get(X_ERROR_CODE)
.and_then(|v| v.to_str().ok())
.unwrap_or_default()
.to_owned();
let detail = format!("{what}: upstream returned {status}");
Err(match (status, code.as_str()) {
(StatusCode::UNAUTHORIZED, _) => Error::UpstreamUnauthorized(detail),
(StatusCode::FORBIDDEN, "DisabledRepo") => Error::DisabledRepo(detail),
(StatusCode::FORBIDDEN, _) => Error::GatedRepo(detail),
(StatusCode::NOT_FOUND, "RevisionNotFound") => Error::RevisionNotFound(detail),
(StatusCode::NOT_FOUND, "EntryNotFound") => Error::EntryNotFound(detail),
(StatusCode::NOT_FOUND, _) => Error::RepoNotFound(detail),
_ => Error::Upstream(detail),
})
}
/// Fetch and parse a JSON document from the Hub API.
async fn get_json(&self, url: &str, what: &str) -> Result<Value> {
let response = self
.authorize(self.client.get(url))
.send()
.await
.map_err(|e| Error::Upstream(format!("{what}: {e}")))?;
let response = self.check(response, what).await?;
response
.json()
.await
.map_err(|e| Error::Upstream(format!("{what}: response was not JSON: {e}")))
}
}
/// Read the replayable metadata out of an upstream response.
fn head_of(response: &Response) -> UpstreamHead {
let headers = response.headers();
let string = |name: &str| {
headers
.get(name)
.and_then(|v| v.to_str().ok())
.map(str::to_owned)
};
let linked_etag = string(X_LINKED_ETAG);
UpstreamHead {
commit: string(X_REPO_COMMIT),
etag: headers
.get(ETAG)
.and_then(|v| v.to_str().ok())
.map(str::to_owned),
lfs: linked_etag.is_some(),
linked_etag,
linked_size: string(X_LINKED_SIZE).and_then(|v| v.parse().ok()),
content_length: response.content_length(),
}
}
/// Percent-encode one URL path segment.
fn urlencode_segment(segment: &str) -> String {
segment
.bytes()
.map(|b| match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
(b as char).to_string()
}
other => format!("%{other:02X}"),
})
.collect()
}
#[async_trait]
impl Upstream for HubClient {
fn enabled(&self) -> bool {
self.enabled
}
fn endpoint(&self) -> &str {
&self.endpoint
}
async fn resolve_ref(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
) -> Result<String> {
// The revision endpoint answers with the commit in `sha`, and does so
// without transferring any file bytes.
let url = self.api_url(
repo_type,
repo_id,
&format!("/revision/{}", urlencode_segment(revision)),
);
let what = format!("resolving {repo_id}@{revision}");
let doc = self.get_json(&url, &what).await?;
doc.get("sha")
.and_then(Value::as_str)
.map(str::to_owned)
.ok_or_else(|| {
Error::RevisionNotFound(format!(
"{repo_id}@{revision}: upstream's revision document carried no `sha`"
))
})
}
async fn head_file(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
path: &str,
) -> Result<UpstreamHead> {
let url = self.resolve_url(repo_type, repo_id, revision, path);
let what = format!("HEAD {repo_id}@{revision}/{path}");
let response = self
.authorize(self.client.head(&url))
.send()
.await
.map_err(|e| Error::Upstream(format!("{what}: {e}")))?;
let response = self.check(response, &what).await?;
Ok(head_of(&response))
}
async fn get_file(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
path: &str,
) -> Result<UpstreamBody> {
let url = self.resolve_url(repo_type, repo_id, revision, path);
let what = format!("GET {repo_id}@{revision}/{path}");
let response = self
.authorize(self.client.get(&url))
.send()
.await
.map_err(|e| Error::Upstream(format!("{what}: {e}")))?;
let response = self.check(response, &what).await?;
let head = head_of(&response);
debug!(%url, size = ?head.content_length, lfs = head.lfs, "upstream transfer starting");
let stream = response
.bytes_stream()
.map(move |chunk| chunk.map_err(|e| Error::Upstream(format!("{what}: {e}"))));
Ok(UpstreamBody {
head,
stream: Box::pin(stream),
})
}
async fn repo_info(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
) -> Result<Value> {
let url = self.api_url(
repo_type,
repo_id,
&format!("/revision/{}", urlencode_segment(revision)),
);
self.get_json(&url, &format!("info for {repo_id}@{revision}"))
.await
}
async fn repo_tree(
&self,
repo_type: RepoType,
repo_id: &RepoId,
revision: &str,
) -> Result<Value> {
let url = format!(
"{}?recursive=1&expand=1",
self.api_url(
repo_type,
repo_id,
&format!("/tree/{}", urlencode_segment(revision))
)
);
self.get_json(&url, &format!("tree for {repo_id}@{revision}"))
.await
}
}
/// An [`Upstream`] that refuses every call.
///
/// Installed when `upstream.enabled = false`. Sealed mode is not "upstream
/// that happens to be unreachable": there is no client, no token loaded and no
/// code path that could reach the network.
pub struct SealedUpstream {
endpoint: String,
}
impl SealedUpstream {
/// A sealed adapter that still records where the data originally came
/// from, for manifest provenance.
pub fn new(endpoint: impl Into<String>) -> Self {
Self {
endpoint: endpoint.into(),
}
}
fn sealed<T>(&self, what: &str) -> Result<T> {
Err(Error::RepoNotFound(format!(
"{what} is not stored, and this instance is sealed (upstream.enabled = false)"
)))
}
}
#[async_trait]
impl Upstream for SealedUpstream {
fn enabled(&self) -> bool {
false
}
fn endpoint(&self) -> &str {
&self.endpoint
}
async fn resolve_ref(&self, _: RepoType, repo_id: &RepoId, revision: &str) -> Result<String> {
self.sealed(&format!("{repo_id}@{revision}"))
}
async fn head_file(
&self,
_: RepoType,
repo_id: &RepoId,
revision: &str,
path: &str,
) -> Result<UpstreamHead> {
self.sealed(&format!("{repo_id}@{revision}/{path}"))
}
async fn get_file(
&self,
_: RepoType,
repo_id: &RepoId,
revision: &str,
path: &str,
) -> Result<UpstreamBody> {
self.sealed(&format!("{repo_id}@{revision}/{path}"))
}
async fn repo_info(&self, _: RepoType, repo_id: &RepoId, revision: &str) -> Result<Value> {
self.sealed(&format!("{repo_id}@{revision}"))
}
async fn repo_tree(&self, _: RepoType, repo_id: &RepoId, revision: &str) -> Result<Value> {
self.sealed(&format!("{repo_id}@{revision}"))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn client() -> HubClient {
HubClient::new(
&rustingface_entities::config::Upstream::default(),
Some("secret".into()),
)
.unwrap()
}
#[test]
fn model_resolve_urls_have_no_type_prefix_and_dataset_urls_do() {
let c = client();
let id: RepoId = "Qwen/Qwen3-32B".parse().unwrap();
assert_eq!(
c.resolve_url(RepoType::Models, &id, "main", "config.json"),
"https://huggingface.co/Qwen/Qwen3-32B/resolve/main/config.json"
);
assert_eq!(
c.resolve_url(RepoType::Datasets, &id, "main", "data/train.parquet"),
"https://huggingface.co/datasets/Qwen/Qwen3-32B/resolve/main/data/train.parquet"
);
}
#[test]
fn path_separators_survive_encoding_but_other_specials_do_not() {
let c = client();
let id: RepoId = "gpt2".parse().unwrap();
let url = c.resolve_url(RepoType::Models, &id, "refs/pr/1", "a b/c.json");
assert!(url.contains("/resolve/refs%2Fpr%2F1/"), "{url}");
assert!(url.ends_with("/a%20b/c.json"), "{url}");
}
#[test]
fn a_sealed_upstream_refuses_without_reaching_the_network() {
let sealed = SealedUpstream::new("https://huggingface.co");
assert!(!sealed.enabled());
assert_eq!(sealed.endpoint(), "https://huggingface.co");
}
#[tokio::test]
async fn every_sealed_call_is_a_not_found_naming_the_seal() {
let sealed = SealedUpstream::new("https://huggingface.co");
let id: RepoId = "gpt2".parse().unwrap();
let err = match sealed
.get_file(RepoType::Models, &id, "main", "config.json")
.await
{
Err(err) => err,
Ok(_) => panic!("a sealed upstream must never produce a body"),
};
assert_eq!(err.status(), 404);
assert!(err.to_string().contains("sealed"), "{err}");
}
}
crates/rustingface-entities/Cargo.toml
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "rustingface-entities"
description = "Domain types, manifest/ref schemas, configuration and error taxonomy for rustingface."
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
authors.workspace = true
repository.workspace = true
[dependencies]
humantime-serde.workspace = true
percent-encoding.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
time.workspace = true
[dev-dependencies]
toml.workspace = true
crates/rustingface-entities/src/config.rs
+408
View File
@@ -0,0 +1,408 @@
//! Configuration: one TOML file, per-key environment overrides, secrets from
//! files (spec §5).
//!
//! Secrets are never inline. Every credential is named by a path so it can be
//! managed by `pass` or systemd credentials and stay out of both the config
//! file and source control.
use std::path::PathBuf;
use std::time::Duration;
use serde::{Deserialize, Deserializer, Serialize};
use crate::error::Error;
/// Default config path. Overridable with `--config`.
pub const DEFAULT_CONFIG_PATH: &str = "/etc/rustingface/config.toml";
/// Environment prefix for per-key overrides, e.g.
/// `RUSTINGFACE__SERVER__LISTEN`.
pub const ENV_PREFIX: &str = "RUSTINGFACE__";
/// The whole of rustingface's configuration.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Config {
pub server: Server,
pub storage: Storage,
pub upstream: Upstream,
pub policy: Policy,
pub auth: Auth,
pub observability: Observability,
}
/// HTTP listener settings.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Server {
/// Address to bind. Bind the mesh address when the reverse proxy lives on
/// another host; loopback only when it is co-resident.
pub listen: String,
/// How this service is reached from outside, used to build absolute URLs.
pub external_url: Option<String>,
/// Ceiling on a single request, including a cold multi-gigabyte fetch.
#[serde(with = "humantime_serde")]
pub request_timeout: Duration,
/// How long a stalled client may hold up the tee before the transfer is
/// detached and finished without it.
#[serde(with = "humantime_serde")]
pub client_stall_timeout: Duration,
}
impl Default for Server {
fn default() -> Self {
Self {
listen: "127.0.0.1:20482".into(),
external_url: None,
request_timeout: Duration::from_secs(300),
client_stall_timeout: Duration::from_secs(60),
}
}
}
/// How a blob body reaches the client.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum BlobDelivery {
/// Stream the bytes through rustingface.
Proxy,
/// Answer `302` with a presigned URL to object storage.
Redirect,
}
/// Object storage settings. The bucket is the entire system state.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Storage {
pub endpoint: Option<String>,
pub bucket: String,
pub region: String,
pub path_style: bool,
/// Use a local directory instead of S3. Test and single-host use only;
/// it forfeits guarantee 5.
pub local_path: Option<PathBuf>,
pub access_key_id_file: Option<PathBuf>,
pub secret_access_key_file: Option<PathBuf>,
pub blob_delivery: BlobDelivery,
#[serde(with = "humantime_serde")]
pub presign_ttl: Duration,
#[serde(deserialize_with = "de_byte_size")]
pub multipart_part_size: u64,
/// Bounded retries for the read-merge-write that resolves concurrent
/// manifest writes for the same commit.
pub manifest_write_retries: u32,
}
impl Default for Storage {
fn default() -> Self {
Self {
endpoint: None,
bucket: "rustingface".into(),
region: "us-east-1".into(),
path_style: true,
local_path: None,
access_key_id_file: None,
secret_access_key_file: None,
blob_delivery: BlobDelivery::Proxy,
presign_ttl: Duration::from_secs(15 * 60),
multipart_part_size: 16 * 1024 * 1024,
manifest_write_retries: 5,
}
}
}
/// The upstream Hub. Consulted on miss only, and not at all when sealed.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Upstream {
/// `false` seals the service: it serves only what is already stored.
pub enabled: bool,
pub endpoint: String,
pub token_file: Option<PathBuf>,
#[serde(with = "humantime_serde")]
pub connect_timeout: Duration,
#[serde(with = "humantime_serde")]
pub read_timeout: Duration,
/// Keep Xet off outbound fetches so whole files arrive via the LFS bridge
/// and the bucket stays independently readable (spec §8). Do not change.
pub disable_xet: bool,
}
impl Default for Upstream {
fn default() -> Self {
Self {
enabled: true,
endpoint: "https://huggingface.co".into(),
token_file: None,
connect_timeout: Duration::from_secs(10),
read_timeout: Duration::from_secs(60),
disable_xet: true,
}
}
}
/// How a mutable ref resolves on second and later requests.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum RefResolution {
/// First resolution wins, forever, until an operator repins.
Freeze,
/// Re-resolve upstream when reachable. Breaks guarantee 4; development
/// only, and warned about at startup.
Follow,
}
/// What rustingface is willing to fetch.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Policy {
pub ref_resolution: RefResolution,
/// `false` restricts fetching to repos already present in the bucket.
pub allow_new_repos: bool,
/// Optional glob allowlist over `repo_id`, e.g. `["Qwen/*"]`. Empty
/// allows everything.
pub allowlist: Vec<String>,
}
impl Default for Policy {
fn default() -> Self {
Self {
ref_resolution: RefResolution::Freeze,
allow_new_repos: true,
allowlist: Vec::new(),
}
}
}
/// Downstream client authentication.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AuthMode {
None,
Bearer,
}
/// Client-facing auth settings. Distinct from `upstream.token_file`, which is
/// never echoed downstream.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Auth {
pub mode: AuthMode,
/// One accepted token per line; `#` comments and blanks ignored.
pub token_file: Option<PathBuf>,
}
impl Default for Auth {
fn default() -> Self {
Self {
mode: AuthMode::None,
token_file: None,
}
}
}
/// Logging and metrics.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Observability {
/// Serve Prometheus text format at `/metrics`.
pub metrics: bool,
pub log_format: LogFormat,
}
impl Default for Observability {
fn default() -> Self {
Self {
metrics: true,
log_format: LogFormat::Json,
}
}
}
/// Log encoding.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LogFormat {
Json,
Text,
}
impl Config {
/// Reject configurations that are internally inconsistent, before the
/// listener binds and the first client depends on them.
pub fn validate(&self) -> Result<(), Error> {
if self.storage.local_path.is_none() && self.storage.bucket.is_empty() {
return Err(Error::Config(
"storage.bucket must be set unless storage.local_path is used".into(),
));
}
if self.storage.multipart_part_size < 5 * 1024 * 1024 {
return Err(Error::Config(
"storage.multipart_part_size must be at least 5MiB, the S3 minimum part size"
.into(),
));
}
if self.storage.blob_delivery == BlobDelivery::Redirect && self.storage.local_path.is_some()
{
return Err(Error::Config(
"storage.blob_delivery = \"redirect\" requires S3; a local store cannot presign"
.into(),
));
}
if self.auth.mode == AuthMode::Bearer && self.auth.token_file.is_none() {
return Err(Error::Config(
"auth.mode = \"bearer\" requires auth.token_file".into(),
));
}
if !self.upstream.disable_xet {
return Err(Error::Config(
"upstream.disable_xet = false is unsupported: rustingface implements no Xet client, \
and a Xet-backed fetch would leave the bucket unreadable without one (spec §8)"
.into(),
));
}
Ok(())
}
/// Whether the service is sealed: no upstream, serve only what is stored.
pub fn sealed(&self) -> bool {
!self.upstream.enabled
}
}
/// Parse `16MiB`, `1GB`, `524288` and friends into bytes.
pub fn parse_byte_size(s: &str) -> Result<u64, Error> {
let s = s.trim();
let split = s
.find(|c: char| !c.is_ascii_digit() && c != '_' && c != '.')
.unwrap_or(s.len());
let (digits, unit) = s.split_at(split);
let digits = digits.replace('_', "");
let value: f64 = digits
.parse()
.map_err(|_| Error::Config(format!("{s:?} is not a byte size")))?;
let multiplier: f64 = match unit.trim().to_ascii_lowercase().as_str() {
"" | "b" => 1.0,
"k" | "kb" => 1_000.0,
"ki" | "kib" => 1024.0,
"m" | "mb" => 1_000_000.0,
"mi" | "mib" => 1024.0 * 1024.0,
"g" | "gb" => 1_000_000_000.0,
"gi" | "gib" => 1024.0 * 1024.0 * 1024.0,
other => {
return Err(Error::Config(format!(
"unknown byte-size unit {other:?} in {s:?}"
)));
}
};
Ok((value * multiplier) as u64)
}
/// Accept either a bare integer or a suffixed string for a byte size, so
/// `16MiB` in TOML and `16777216` from an environment override both work.
fn de_byte_size<'de, D: Deserializer<'de>>(d: D) -> Result<u64, D::Error> {
#[derive(Deserialize)]
#[serde(untagged)]
enum Repr {
Int(u64),
Str(String),
}
match Repr::deserialize(d)? {
Repr::Int(n) => Ok(n),
Repr::Str(s) => parse_byte_size(&s).map_err(serde::de::Error::custom),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn byte_sizes_parse_binary_and_decimal_units() {
assert_eq!(parse_byte_size("16MiB").unwrap(), 16 * 1024 * 1024);
assert_eq!(parse_byte_size("16MB").unwrap(), 16_000_000);
assert_eq!(parse_byte_size("1GiB").unwrap(), 1024 * 1024 * 1024);
assert_eq!(parse_byte_size("5242880").unwrap(), 5_242_880);
assert_eq!(parse_byte_size(" 8 KiB ").unwrap(), 8192);
assert!(parse_byte_size("16 parsecs").is_err());
}
#[test]
fn the_documented_example_config_parses() {
let toml = r#"
[server]
listen = "127.0.0.1:20482"
external_url = "https://rf.internal"
request_timeout = "300s"
[storage]
endpoint = "http://caveman.kosherinata.internal:9000"
bucket = "rustingface"
region = "us-east-1"
path_style = true
blob_delivery = "proxy"
presign_ttl = "15m"
multipart_part_size = "16MiB"
[upstream]
enabled = true
endpoint = "https://huggingface.co"
connect_timeout = "10s"
read_timeout = "60s"
disable_xet = true
[policy]
ref_resolution = "freeze"
allow_new_repos = true
allowlist = ["Qwen/*"]
[auth]
mode = "none"
[observability]
metrics = true
log_format = "json"
"#;
let cfg: Config = toml::from_str(toml).unwrap();
cfg.validate().unwrap();
assert_eq!(cfg.storage.multipart_part_size, 16 * 1024 * 1024);
assert_eq!(cfg.policy.ref_resolution, RefResolution::Freeze);
assert_eq!(cfg.server.request_timeout, Duration::from_secs(300));
assert!(!cfg.sealed());
}
#[test]
fn defaults_are_valid_and_freeze_by_default() {
let cfg = Config::default();
cfg.validate().unwrap();
assert_eq!(cfg.policy.ref_resolution, RefResolution::Freeze);
assert!(cfg.upstream.disable_xet);
}
#[test]
fn a_part_size_below_the_s3_minimum_is_refused() {
let mut cfg = Config::default();
cfg.storage.multipart_part_size = 1024;
assert!(cfg.validate().is_err());
}
#[test]
fn bearer_without_a_token_file_is_refused() {
let mut cfg = Config::default();
cfg.auth.mode = AuthMode::Bearer;
assert!(cfg.validate().is_err());
}
#[test]
fn enabling_xet_is_refused() {
let mut cfg = Config::default();
cfg.upstream.disable_xet = false;
assert!(cfg.validate().is_err());
}
#[test]
fn unknown_keys_are_rejected_rather_than_silently_ignored() {
let err = toml::from_str::<Config>("[server]\nlissen = \"x\"\n").unwrap_err();
assert!(err.to_string().contains("lissen"), "{err}");
}
}
crates/rustingface-entities/src/error.rs
+201
View File
@@ -0,0 +1,201 @@
//! Error taxonomy.
//!
//! `huggingface_hub` maps the `X-Error-Code` response header onto typed
//! exceptions, so the header is part of the wire contract and not a nicety.
//! The strings below are client-internal constants upstream and should be
//! re-checked whenever the pinned client version moves (spec §6).
use std::fmt;
/// Convenience alias used throughout the workspace.
pub type Result<T, E = Error> = std::result::Result<T, E>;
/// The `X-Error-Code` values rustingface emits.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorCode {
RepoNotFound,
RevisionNotFound,
EntryNotFound,
GatedRepo,
DisabledRepo,
}
impl ErrorCode {
/// The header value, exactly as the client compares it.
pub const fn as_str(self) -> &'static str {
match self {
Self::RepoNotFound => "RepoNotFound",
Self::RevisionNotFound => "RevisionNotFound",
Self::EntryNotFound => "EntryNotFound",
Self::GatedRepo => "GatedRepo",
Self::DisabledRepo => "DisabledRepo",
}
}
}
impl fmt::Display for ErrorCode {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
/// Everything that can go wrong on the way to a response.
///
/// Failures are loud by construction: there is no variant that yields an empty
/// body or a partial stream, because a silently short file is the one outcome
/// this service exists to prevent.
#[derive(Debug, thiserror::Error)]
pub enum Error {
/// The repo is not stored and upstream is disabled or unreachable.
#[error("repo not found: {0}")]
RepoNotFound(String),
/// The revision is unknown and cannot be resolved.
#[error("revision not found: {0}")]
RevisionNotFound(String),
/// The path is not in the manifest and cannot be fetched.
#[error("entry not found: {0}")]
EntryNotFound(String),
/// Upstream returned 403 for a gated repo.
#[error("gated repo: {0}")]
GatedRepo(String),
/// Upstream returned 403 for a disabled repo.
#[error("disabled repo: {0}")]
DisabledRepo(String),
/// Upstream authentication failed (bad or missing `upstream.token_file`).
#[error("upstream authentication failed: {0}")]
UpstreamUnauthorized(String),
/// The client presented no or an unrecognised bearer token.
#[error("unauthorized")]
Unauthorized,
/// Policy (allowlist or `allow_new_repos`) refused the request.
#[error("forbidden by policy: {0}")]
PolicyDenied(String),
/// A malformed request path, range header or repo id.
#[error("bad request: {0}")]
BadRequest(String),
/// A `Range` header that cannot be satisfied against the stored length.
#[error("range not satisfiable: {0}")]
RangeNotSatisfiable(String),
/// The stored bytes did not hash to the digest upstream promised.
#[error("digest mismatch for {path}: upstream said {expected}, got {actual}")]
DigestMismatch {
path: String,
expected: String,
actual: String,
},
/// Object storage failed.
#[error("storage: {0}")]
Storage(String),
/// The upstream request failed at the transport layer, or returned a
/// status rustingface has no specific mapping for.
#[error("upstream: {0}")]
Upstream(String),
/// Configuration could not be loaded or is internally inconsistent.
#[error("configuration: {0}")]
Config(String),
/// An invariant broke. These are bugs.
#[error("internal: {0}")]
Internal(String),
}
impl Error {
/// The `X-Error-Code` header value for this error, if it has one.
pub fn code(&self) -> Option<ErrorCode> {
match self {
Self::RepoNotFound(_) => Some(ErrorCode::RepoNotFound),
Self::RevisionNotFound(_) => Some(ErrorCode::RevisionNotFound),
Self::EntryNotFound(_) => Some(ErrorCode::EntryNotFound),
Self::GatedRepo(_) => Some(ErrorCode::GatedRepo),
Self::DisabledRepo(_) => Some(ErrorCode::DisabledRepo),
_ => None,
}
}
/// The HTTP status this error is served as.
pub fn status(&self) -> u16 {
match self {
Self::RepoNotFound(_) | Self::RevisionNotFound(_) | Self::EntryNotFound(_) => 404,
Self::GatedRepo(_) | Self::DisabledRepo(_) | Self::PolicyDenied(_) => 403,
Self::UpstreamUnauthorized(_) | Self::Unauthorized => 401,
Self::BadRequest(_) => 400,
Self::RangeNotSatisfiable(_) => 416,
Self::DigestMismatch { .. } | Self::Storage(_) | Self::Upstream(_) => 502,
Self::Config(_) | Self::Internal(_) => 500,
}
}
/// A short, stable label for metrics, so cardinality stays bounded.
pub fn label(&self) -> &'static str {
match self {
Self::RepoNotFound(_) => "repo_not_found",
Self::RevisionNotFound(_) => "revision_not_found",
Self::EntryNotFound(_) => "entry_not_found",
Self::GatedRepo(_) => "gated",
Self::DisabledRepo(_) => "disabled",
Self::UpstreamUnauthorized(_) => "upstream_unauthorized",
Self::Unauthorized => "unauthorized",
Self::PolicyDenied(_) => "policy_denied",
Self::BadRequest(_) => "bad_request",
Self::RangeNotSatisfiable(_) => "range_not_satisfiable",
Self::DigestMismatch { .. } => "digest_mismatch",
Self::Storage(_) => "storage",
Self::Upstream(_) => "upstream",
Self::Config(_) => "config",
Self::Internal(_) => "internal",
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_documented_status_and_code_pairs_hold() {
let cases: Vec<(Error, u16, Option<&str>)> = vec![
(Error::RepoNotFound("a".into()), 404, Some("RepoNotFound")),
(
Error::RevisionNotFound("a".into()),
404,
Some("RevisionNotFound"),
),
(Error::EntryNotFound("a".into()), 404, Some("EntryNotFound")),
(Error::GatedRepo("a".into()), 403, Some("GatedRepo")),
(Error::UpstreamUnauthorized("a".into()), 401, None),
];
for (err, status, code) in cases {
assert_eq!(err.status(), status, "{err}");
assert_eq!(err.code().map(ErrorCode::as_str), code, "{err}");
}
}
#[test]
fn transport_and_integrity_failures_are_bad_gateway_not_not_found() {
assert_eq!(Error::Upstream("connect".into()).status(), 502);
assert_eq!(Error::Storage("timeout".into()).status(), 502);
assert_eq!(
Error::DigestMismatch {
path: "model.safetensors".into(),
expected: "sha256:aa".into(),
actual: "sha256:bb".into(),
}
.status(),
502,
"a corrupt fetch must never look like a missing file"
);
}
}
crates/rustingface-entities/src/key.rs
+184
View File
@@ -0,0 +1,184 @@
//! Bucket key layout.
//!
//! Prefixes are stable and human-inspectable, so an operator with `mc ls` can
//! answer any question about system state without this binary (spec §4).
use percent_encoding::{AsciiSet, CONTROLS, utf8_percent_encode};
use crate::repo::{RepoId, RepoType};
/// Characters escaped in a ref name before it becomes a key segment.
///
/// Real Hub refs contain slashes (`refs/pr/1`, `refs/convert/parquet`), which
/// would otherwise silently become directory levels and make the ref listing
/// for a repo non-flat.
const REF_ESCAPE: &AsciiSet = &CONTROLS.add(b'/').add(b'%').add(b'?').add(b'#').add(b' ');
/// Key of a content-addressed blob.
///
/// Two levels of hex fanout keep large repos off a single key prefix on
/// backends that shard by key range.
pub fn blob(digest_hex: &str) -> String {
debug_assert_eq!(digest_hex.len(), 64, "blob keys are sha256 hex");
format!(
"blobs/sha256/{}/{}/{}",
&digest_hex[0..2],
&digest_hex[2..4],
digest_hex
)
}
/// Extract the hex digest from a blob key, if it is one.
pub fn blob_digest(key: &str) -> Option<&str> {
let rest = key.strip_prefix("blobs/sha256/")?;
let hex = rest.rsplit('/').next()?;
(hex.len() == 64 && hex.bytes().all(|b| b.is_ascii_hexdigit())).then_some(hex)
}
/// Prefix under which every blob lives.
pub const BLOB_PREFIX: &str = "blobs/sha256/";
/// Prefix holding every manifest for a repo.
pub fn manifest_prefix(repo_type: RepoType, repo_id: &RepoId) -> String {
format!(
"manifests/{}/{}/{}/",
repo_type.as_str(),
repo_id.namespace_segment(),
repo_id.name()
)
}
/// Key of the manifest for one commit.
pub fn manifest(repo_type: RepoType, repo_id: &RepoId, commit: &str) -> String {
format!("{}{commit}.json", manifest_prefix(repo_type, repo_id))
}
/// Prefix holding every ref for a repo.
pub fn ref_prefix(repo_type: RepoType, repo_id: &RepoId) -> String {
format!(
"refs/{}/{}/{}/",
repo_type.as_str(),
repo_id.namespace_segment(),
repo_id.name()
)
}
/// Key of one ref record.
pub fn ref_key(repo_type: RepoType, repo_id: &RepoId, ref_name: &str) -> String {
format!(
"{}{}.json",
ref_prefix(repo_type, repo_id),
utf8_percent_encode(ref_name, REF_ESCAPE)
)
}
/// Recover a ref name from its key, reversing the escaping.
pub fn ref_name_from_key(key: &str) -> Option<String> {
let file = key.rsplit('/').next()?.strip_suffix(".json")?;
percent_encoding::percent_decode_str(file)
.decode_utf8()
.ok()
.map(|s| s.into_owned())
}
/// Key of a repo record.
pub fn repo(repo_type: RepoType, repo_id: &RepoId) -> String {
format!(
"repos/{}/{}/{}/repo.json",
repo_type.as_str(),
repo_id.namespace_segment(),
repo_id.name()
)
}
/// Prefix holding every repo record of one type.
pub fn repo_type_prefix(repo_type: RepoType) -> String {
format!("repos/{}/", repo_type.as_str())
}
/// Recover `(repo_type, repo_id)` from a repo record key.
pub fn repo_from_key(key: &str) -> Option<(RepoType, RepoId)> {
let rest = key.strip_prefix("repos/")?.strip_suffix("/repo.json")?;
let mut parts = rest.split('/');
let repo_type = parts.next()?.parse().ok()?;
let namespace = parts.next()?;
let name = parts.next()?;
if parts.next().is_some() {
return None;
}
Some((repo_type, RepoId::from_segments(namespace, name).ok()?))
}
#[cfg(test)]
mod tests {
use super::*;
fn qwen() -> RepoId {
"Qwen/Qwen3-32B".parse().unwrap()
}
#[test]
fn blob_key_fans_out_two_levels_and_round_trips() {
let hex = "9c1eff00112233445566778899aabbccddeeff00112233445566778899aabbcc";
let key = blob(hex);
assert_eq!(
key,
format!("blobs/sha256/9c/1e/{hex}"),
"two-level fanout on the first four hex characters"
);
assert_eq!(blob_digest(&key), Some(hex));
}
#[test]
fn blob_digest_rejects_non_blob_keys() {
assert_eq!(
blob_digest("manifests/models/Qwen/Qwen3-32B/abc.json"),
None
);
assert_eq!(blob_digest("blobs/sha256/aa/bb/short"), None);
}
#[test]
fn manifest_and_ref_keys_use_the_documented_layout() {
assert_eq!(
manifest(RepoType::Models, &qwen(), "8f2ac31"),
"manifests/models/Qwen/Qwen3-32B/8f2ac31.json"
);
assert_eq!(
ref_key(RepoType::Models, &qwen(), "main"),
"refs/models/Qwen/Qwen3-32B/main.json"
);
}
#[test]
fn slashed_ref_names_stay_one_key_segment() {
let key = ref_key(RepoType::Datasets, &qwen(), "refs/pr/1");
assert_eq!(key, "refs/datasets/Qwen/Qwen3-32B/refs%2Fpr%2F1.json");
assert_eq!(
key.strip_prefix(&ref_prefix(RepoType::Datasets, &qwen()))
.unwrap()
.matches('/')
.count(),
0,
"the ref listing for a repo stays flat"
);
assert_eq!(ref_name_from_key(&key).as_deref(), Some("refs/pr/1"));
}
#[test]
fn bare_named_repos_round_trip_through_the_repo_key() {
let id: RepoId = "gpt2".parse().unwrap();
let key = repo(RepoType::Models, &id);
assert_eq!(key, "repos/models/__/gpt2/repo.json");
let (rt, back) = repo_from_key(&key).unwrap();
assert_eq!(rt, RepoType::Models);
assert_eq!(back.to_string(), "gpt2");
}
#[test]
fn repo_from_key_rejects_malformed_keys() {
assert!(repo_from_key("repos/models/Qwen/repo.json").is_none());
assert!(repo_from_key("repos/spaces/a/b/repo.json").is_none());
assert!(repo_from_key("manifests/models/a/b/repo.json").is_none());
}
}
crates/rustingface-entities/src/lib.rs
+19
View File
@@ -0,0 +1,19 @@
//! Domain types for rustingface.
//!
//! This crate is pure: it defines the on-disk schemas (manifest, ref, repo
//! record), the bucket key layout, the configuration shape and the error
//! taxonomy that the HTTP surface replays to clients. It performs no I/O and
//! pulls in no async runtime.
pub mod config;
pub mod error;
pub mod key;
pub mod manifest;
pub mod refs;
pub mod repo;
pub use config::Config;
pub use error::{Error, ErrorCode, Result};
pub use manifest::{Entry, MANIFEST_SCHEMA, Manifest};
pub use refs::{PinnedBy, REF_SCHEMA, RefRecord};
pub use repo::{REPO_SCHEMA, RepoId, RepoRecord, RepoType};
crates/rustingface-entities/src/manifest.rs
+192
View File
@@ -0,0 +1,192 @@
//! The manifest: the immutable record of what upstream said about one commit.
//!
//! rustingface never derives these values. Every field sourced from upstream is
//! stored exactly as received and replayed byte-identically (spec §2), because
//! client cache layouts are keyed on them.
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use crate::repo::{RepoId, RepoType};
/// Schema version of `manifests/…/<commit>.json`.
pub const MANIFEST_SCHEMA: u32 = 1;
/// One recorded file within a commit.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Entry {
/// Repo-relative path, exactly as the client requested it.
pub path: String,
/// True byte length of the stored blob.
pub size: u64,
/// rustingface blob key digest, `sha256:<64 hex>`.
pub digest: String,
/// Upstream `ETag`, verbatim, quotes included.
pub etag: Option<String>,
/// Upstream `X-Linked-Etag`, verbatim. `None` if upstream omitted it.
pub linked_etag: Option<String>,
/// Upstream `X-Linked-Size`, verbatim. `None` if upstream omitted it.
pub linked_size: Option<u64>,
/// Whether upstream served this via the LFS bridge.
pub lfs: bool,
}
/// The manifest for one `(repo, commit)` pair.
///
/// `entries` is append-only and expected to be partial: clients fetch subsets
/// of a repo, and rustingface records what it has actually seen. A path absent
/// from `entries` is a miss, not a 404, while upstream is reachable.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Manifest {
pub schema: u32,
pub repo_type: RepoType,
pub repo_id: RepoId,
pub commit: String,
#[serde(with = "time::serde::rfc3339")]
pub recorded_at: OffsetDateTime,
pub upstream: Option<String>,
/// The upstream `/api/{type}/{id}/revision/{commit}` response, verbatim.
///
/// Stored whole and filtered on read (spec §13): recording the full
/// document means a field a client asks for later was never dropped.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub repo_meta: Option<Value>,
/// The upstream `/api/{type}/{id}/tree/{commit}?recursive=1&expand=1`
/// response, verbatim, when one has been recorded.
///
/// Additive to schema 1: absent on manifests written before a tree was
/// requested, and the tree handler falls back to `repo_meta.siblings`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub repo_tree: Option<Value>,
#[serde(default)]
pub entries: Vec<Entry>,
}
impl Manifest {
/// An empty manifest for a commit whose metadata has just been recorded.
pub fn new(
repo_type: RepoType,
repo_id: RepoId,
commit: String,
upstream: Option<String>,
) -> Self {
Self {
schema: MANIFEST_SCHEMA,
repo_type,
repo_id,
commit,
recorded_at: OffsetDateTime::now_utc(),
upstream,
repo_meta: None,
repo_tree: None,
entries: Vec::new(),
}
}
/// The recorded entry for `path`, if any.
pub fn entry(&self, path: &str) -> Option<&Entry> {
self.entries.iter().find(|e| e.path == path)
}
/// Insert or replace an entry, keeping `entries` sorted by path.
///
/// Entries are deterministic for a given commit, so a replace is a no-op in
/// content and last-writer-wins is correct (spec §4).
pub fn upsert(&mut self, entry: Entry) {
match self.entries.binary_search_by(|e| e.path.cmp(&entry.path)) {
Ok(i) => self.entries[i] = entry,
Err(i) => self.entries.insert(i, entry),
}
}
/// Merge `other`'s entries and metadata into `self`.
///
/// Used by the read-merge-write retry that resolves concurrent manifest
/// writes for the same commit.
pub fn merge_from(&mut self, other: &Manifest) {
for entry in &other.entries {
if self.entry(&entry.path).is_none() {
self.upsert(entry.clone());
}
}
if self.repo_meta.is_none() {
self.repo_meta.clone_from(&other.repo_meta);
}
if self.repo_tree.is_none() {
self.repo_tree.clone_from(&other.repo_tree);
}
}
/// Total stored bytes across recorded entries.
pub fn stored_bytes(&self) -> u64 {
self.entries.iter().map(|e| e.size).sum()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn entry(path: &str, size: u64) -> Entry {
Entry {
path: path.into(),
size,
digest: format!("sha256:{:064x}", size),
etag: Some("\"abc\"".into()),
linked_etag: None,
linked_size: None,
lfs: false,
}
}
fn manifest() -> Manifest {
Manifest::new(
RepoType::Models,
"Qwen/Qwen3-32B".parse().unwrap(),
"8f2ac31".into(),
Some("https://huggingface.co".into()),
)
}
#[test]
fn upsert_keeps_entries_sorted_and_replaces_in_place() {
let mut m = manifest();
m.upsert(entry("model.safetensors", 2));
m.upsert(entry("config.json", 1));
m.upsert(entry("config.json", 3));
assert_eq!(m.entries.len(), 2);
assert_eq!(m.entries[0].path, "config.json");
assert_eq!(m.entries[0].size, 3);
assert_eq!(m.entries[1].path, "model.safetensors");
}
#[test]
fn merge_takes_missing_entries_but_keeps_local_ones() {
let mut a = manifest();
a.upsert(entry("config.json", 1));
let mut b = manifest();
b.upsert(entry("config.json", 99));
b.upsert(entry("tokenizer.json", 2));
b.repo_meta = Some(serde_json::json!({"id": "Qwen/Qwen3-32B"}));
a.merge_from(&b);
assert_eq!(a.entries.len(), 2);
assert_eq!(a.entry("config.json").unwrap().size, 1, "local entry wins");
assert!(a.repo_meta.is_some(), "absent metadata is adopted");
}
#[test]
fn round_trips_through_json_with_optional_fields_absent() {
let mut m = manifest();
m.upsert(entry("config.json", 1204));
let json = serde_json::to_string(&m).unwrap();
assert!(
!json.contains("repo_meta"),
"absent optionals stay out of the document"
);
let back: Manifest = serde_json::from_str(&json).unwrap();
assert_eq!(back.commit, m.commit);
assert_eq!(back.entries, m.entries);
}
}
crates/rustingface-entities/src/refs.rs
+96
View File
@@ -0,0 +1,96 @@
//! Refs: the only mutable objects in the bucket.
//!
//! A ref record pins a mutable name (`main`, a tag, `refs/pr/1`) to the commit
//! it resolved to the first time it was seen. Guarantee 4 — reference
//! stability — is entirely this file plus a first-write-wins put.
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
/// Schema version of `refs/…/<ref>.json`.
pub const REF_SCHEMA: u32 = 1;
/// What caused a ref to hold its current value.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum PinnedBy {
/// First resolution through the resolve path, frozen automatically.
AutoFreeze,
/// An explicit `rustingface pin` or `rustingface refresh`.
Operator,
}
impl PinnedBy {
/// The string as it appears in the stored document.
pub const fn as_str(self) -> &'static str {
match self {
Self::AutoFreeze => "auto-freeze",
Self::Operator => "operator",
}
}
}
/// The `refs/<repo_type>/<namespace>/<name>/<ref>.json` record.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RefRecord {
pub schema: u32,
#[serde(rename = "ref")]
pub ref_name: String,
pub commit: String,
#[serde(with = "time::serde::rfc3339")]
pub pinned_at: OffsetDateTime,
pub pinned_by: PinnedBy,
}
impl RefRecord {
/// A pin recorded now.
pub fn new(
ref_name: impl Into<String>,
commit: impl Into<String>,
pinned_by: PinnedBy,
) -> Self {
Self {
schema: REF_SCHEMA,
ref_name: ref_name.into(),
commit: commit.into(),
pinned_at: OffsetDateTime::now_utc(),
pinned_by,
}
}
}
/// Whether `revision` is a commit SHA rather than a ref name.
///
/// The Hub uses 40-hex git SHA-1 commit ids. A revision that looks like one is
/// used directly and never written to `refs/`, which is what keeps immutable
/// revisions off the mutable-object path entirely.
pub fn is_commit_sha(revision: &str) -> bool {
revision.len() == 40 && revision.bytes().all(|b| b.is_ascii_hexdigit())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn recognises_a_commit_sha() {
assert!(is_commit_sha("8f2a1c9e0b7d4f6a3e5c2b8d9f0a1c3e5b7d9f01"));
}
#[test]
fn rejects_ref_names_and_short_or_long_hex() {
assert!(!is_commit_sha("main"));
assert!(!is_commit_sha("v1.0"));
assert!(!is_commit_sha("8f2a1c9e"));
assert!(!is_commit_sha("8f2a1c9e0b7d4f6a3e5c2b8d9f0a1c3e5b7d9f012"));
assert!(!is_commit_sha("8f2a1c9e0b7d4f6a3e5c2b8d9f0a1c3e5b7d9fzz"));
}
#[test]
fn pinned_by_serialises_kebab_case() {
let r = RefRecord::new("main", "abc", PinnedBy::AutoFreeze);
let json = serde_json::to_string(&r).unwrap();
assert!(json.contains("\"auto-freeze\""), "{json}");
assert!(json.contains("\"ref\":\"main\""), "{json}");
}
}
crates/rustingface-entities/src/repo.rs
+232
View File
@@ -0,0 +1,232 @@
//! Repository identity: type, id, and the `repos/…/repo.json` record.
use std::fmt;
use std::str::FromStr;
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use crate::error::Error;
/// Schema version of the `repos/…/repo.json` record.
pub const REPO_SCHEMA: u32 = 1;
/// The two repository namespaces rustingface serves.
///
/// Deliberately not an open enum: spaces are a non-goal (see spec §1).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum RepoType {
Models,
Datasets,
}
impl RepoType {
/// The bucket prefix segment and the `/api/{…}` path segment. These are
/// the same string upstream, and we keep them the same here.
pub const fn as_str(self) -> &'static str {
match self {
Self::Models => "models",
Self::Datasets => "datasets",
}
}
/// Every repo type rustingface knows about, for iteration in the CLI.
pub const ALL: [Self; 2] = [Self::Models, Self::Datasets];
}
impl fmt::Display for RepoType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
impl FromStr for RepoType {
type Err = Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"models" | "model" => Ok(Self::Models),
"datasets" | "dataset" => Ok(Self::Datasets),
other => Err(Error::BadRequest(format!("unknown repo type {other:?}"))),
}
}
}
/// A Hub repository id: either `namespace/name` or a bare `name`.
///
/// Bare names are legitimate on the Hub (`gpt2`, `bert-base-uncased`) and must
/// round-trip through the bucket layout, so [`RepoId::namespace_segment`]
/// substitutes a sentinel that cannot collide with a real namespace.
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct RepoId {
namespace: Option<String>,
name: String,
}
/// Key segment standing in for the namespace of a bare-named repo.
///
/// Hub namespaces must begin with an alphanumeric character, so a segment of
/// two underscores is unreachable by any real namespace and the mapping is
/// injective in both directions.
pub const BARE_NAMESPACE: &str = "__";
impl RepoId {
/// Construct from parts. Neither part may be empty or contain a `/`.
pub fn new(namespace: Option<String>, name: String) -> Result<Self, Error> {
if name.is_empty() || name.contains('/') {
return Err(Error::BadRequest(format!("invalid repo name {name:?}")));
}
if let Some(ns) = &namespace
&& (ns.is_empty() || ns.contains('/') || ns == BARE_NAMESPACE)
{
return Err(Error::BadRequest(format!("invalid namespace {ns:?}")));
}
Ok(Self { namespace, name })
}
/// The namespace, if the repo has one.
pub fn namespace(&self) -> Option<&str> {
self.namespace.as_deref()
}
/// The repo name, without namespace.
pub fn name(&self) -> &str {
&self.name
}
/// The namespace as it appears in a bucket key, substituting
/// [`BARE_NAMESPACE`] for a bare-named repo.
pub fn namespace_segment(&self) -> &str {
self.namespace.as_deref().unwrap_or(BARE_NAMESPACE)
}
/// Reconstruct from the two key segments produced by
/// [`RepoId::namespace_segment`] and [`RepoId::name`].
pub fn from_segments(namespace: &str, name: &str) -> Result<Self, Error> {
let namespace = if namespace == BARE_NAMESPACE {
None
} else {
Some(namespace.to_owned())
};
Self::new(namespace, name.to_owned())
}
/// How many path segments this id occupies in a URL (1 or 2).
pub fn segment_count(&self) -> usize {
if self.namespace.is_some() { 2 } else { 1 }
}
}
impl fmt::Display for RepoId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match &self.namespace {
Some(ns) => write!(f, "{ns}/{}", self.name),
None => f.write_str(&self.name),
}
}
}
impl FromStr for RepoId {
type Err = Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let s = s.trim_matches('/');
let mut parts = s.split('/');
let first = parts.next().unwrap_or_default();
match parts.next() {
None => Self::new(None, first.to_owned()),
Some(second) => {
if parts.next().is_some() {
return Err(Error::BadRequest(format!(
"repo id {s:?} has too many segments"
)));
}
Self::new(Some(first.to_owned()), second.to_owned())
}
}
}
}
impl Serialize for RepoId {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
s.serialize_str(&self.to_string())
}
}
impl<'de> Deserialize<'de> for RepoId {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
let s = String::deserialize(d)?;
s.parse().map_err(serde::de::Error::custom)
}
}
/// The `repos/<repo_type>/<namespace>/<name>/repo.json` record.
///
/// Its only job is to make "which repos does this bucket know about" a prefix
/// listing rather than a full scan of every manifest.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RepoRecord {
pub schema: u32,
pub repo_type: RepoType,
pub repo_id: RepoId,
#[serde(with = "time::serde::rfc3339")]
pub first_seen: OffsetDateTime,
/// Upstream this repo was first recorded from, for provenance.
pub upstream: Option<String>,
}
impl RepoRecord {
/// A record for a repo being seen for the first time, now.
pub fn new(repo_type: RepoType, repo_id: RepoId, upstream: Option<String>) -> Self {
Self {
schema: REPO_SCHEMA,
repo_type,
repo_id,
first_seen: OffsetDateTime::now_utc(),
upstream,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn namespaced_id_round_trips() {
let id: RepoId = "Qwen/Qwen3-32B".parse().unwrap();
assert_eq!(id.namespace(), Some("Qwen"));
assert_eq!(id.name(), "Qwen3-32B");
assert_eq!(id.to_string(), "Qwen/Qwen3-32B");
assert_eq!(RepoId::from_segments("Qwen", "Qwen3-32B").unwrap(), id);
}
#[test]
fn bare_id_round_trips_through_key_segments() {
let id: RepoId = "gpt2".parse().unwrap();
assert_eq!(id.namespace(), None);
assert_eq!(id.namespace_segment(), BARE_NAMESPACE);
assert_eq!(id.to_string(), "gpt2");
let back = RepoId::from_segments(id.namespace_segment(), id.name()).unwrap();
assert_eq!(back, id);
assert_eq!(back.to_string(), "gpt2");
}
#[test]
fn three_segment_id_is_rejected() {
assert!("a/b/c".parse::<RepoId>().is_err());
}
#[test]
fn sentinel_namespace_is_not_accepted_as_a_real_one() {
assert!(RepoId::new(Some(BARE_NAMESPACE.into()), "x".into()).is_err());
}
#[test]
fn repo_type_parses_singular_and_plural() {
assert_eq!("models".parse::<RepoType>().unwrap(), RepoType::Models);
assert_eq!("dataset".parse::<RepoType>().unwrap(), RepoType::Datasets);
assert!("spaces".parse::<RepoType>().is_err());
}
}
crates/rustingface/Cargo.toml
+38
View File
@@ -0,0 +1,38 @@
[package]
name = "rustingface"
description = "A sovereign, HuggingFace-compatible model registry: server and admin CLI in one binary."
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
authors.workspace = true
repository.workspace = true
[[bin]]
name = "rustingface"
path = "src/main.rs"
[dependencies]
rustingface-api.workspace = true
rustingface-core.workspace = true
rustingface-data.workspace = true
rustingface-entities.workspace = true
anyhow.workspace = true
axum.workspace = true
clap.workspace = true
figment.workspace = true
futures.workspace = true
serde_json.workspace = true
time.workspace = true
tokio.workspace = true
toml.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
[dev-dependencies]
bytes.workspace = true
futures.workspace = true
http-body-util.workspace = true
reqwest.workspace = true
tower.workspace = true
crates/rustingface/src/cmd/doctor.rs
+145
View File
@@ -0,0 +1,145 @@
//! `rustingface doctor` — prove the deployment works before a client does.
//!
//! Every check reports what it actually attempted. A doctor that only proves a
//! TCP connection opened is a doctor that passes on a bucket the service
//! cannot write to.
use anyhow::{Result, bail};
use rustingface_entities::config::{AuthMode, BlobDelivery, RefResolution};
use rustingface_entities::repo::RepoType;
use crate::cmd::human_bytes;
use crate::wire;
/// Run the checks, reporting each and failing on the first that matters.
pub async fn run(config: rustingface_entities::config::Config) -> Result<()> {
let mut problems = Vec::new();
println!("configuration");
println!(" listen {}", config.server.listen);
println!(
" external url {}",
config.server.external_url.as_deref().unwrap_or("(unset)")
);
println!(
" storage {} bucket {:?}{}",
config
.storage
.endpoint
.as_deref()
.unwrap_or("(aws default endpoint)"),
config.storage.bucket,
if config.storage.path_style {
" path-style"
} else {
""
}
);
println!(" blob delivery {:?}", config.storage.blob_delivery);
println!(
" part size {}",
human_bytes(config.storage.multipart_part_size)
);
println!(
" upstream {}",
if config.upstream.enabled {
config.upstream.endpoint.clone()
} else {
"sealed (upstream.enabled = false)".into()
}
);
println!(" ref resolution {:?}", config.policy.ref_resolution);
println!(" client auth {:?}", config.auth.mode);
println!(
" allowlist {}",
if config.policy.allowlist.is_empty() {
"(everything)".into()
} else {
config.policy.allowlist.join(", ")
}
);
if config.policy.ref_resolution == RefResolution::Follow {
problems.push(
"policy.ref_resolution = \"follow\" breaks reference stability; use \"freeze\" \
outside development"
.to_owned(),
);
}
if config.storage.blob_delivery == BlobDelivery::Redirect && config.storage.endpoint.is_none() {
problems.push(
"blob_delivery = \"redirect\" mints presigned URLs against storage.endpoint, which \
is unset; clients would be redirected to AWS"
.to_owned(),
);
}
if config.auth.mode == AuthMode::None && config.server.listen.starts_with("0.0.0.0") {
problems.push(
"auth.mode = \"none\" with a wildcard bind publishes the registry on every address \
this host carries; bind the mesh address or enable bearer auth"
.to_owned(),
);
}
println!("\nstorage");
let registry = wire::registry(&config)?;
match registry.probe_store().await {
Ok(()) => println!(" write/read/delete ok"),
Err(err) => {
println!(" write/read/delete FAILED: {err}");
problems.push(format!("object storage is not usable: {err}"));
}
}
let mut repos = 0;
for repo_type in RepoType::ALL {
match registry.list_repos(repo_type).await {
Ok(list) => {
println!(" {:<19} {} stored", repo_type.as_str(), list.len());
repos += list.len();
}
Err(err) => {
println!(" {:<19} FAILED: {err}", repo_type.as_str());
problems.push(format!("listing {} failed: {err}", repo_type.as_str()));
}
}
}
if repos == 0 && config.sealed() {
problems.push(
"this instance is sealed and the bucket holds no repositories, so every request \
will 404"
.to_owned(),
);
}
println!("\nupstream");
if config.upstream.enabled {
let probe: rustingface_entities::repo::RepoId = "gpt2".parse()?;
match registry
.upstream()
.resolve_ref(RepoType::Models, &probe, "main")
.await
{
Ok(commit) => println!(" reachable yes (gpt2@main = {commit})"),
Err(err) => {
println!(" reachable no: {err}");
problems.push(format!(
"upstream is enabled but unreachable: {err}. Cold fetches will fail; \
anything already stored still serves."
));
}
}
} else {
println!(" reachable n/a, sealed");
}
if problems.is_empty() {
println!("\nno problems found");
return Ok(());
}
println!("\n{} problem(s):", problems.len());
for problem in &problems {
println!(" - {problem}");
}
bail!("doctor found {} problem(s)", problems.len());
}
crates/rustingface/src/cmd/fetch.rs
+123
View File
@@ -0,0 +1,123 @@
//! `rustingface fetch` — warm the bucket ahead of need.
//!
//! Drives the same resolve path the server does, so what lands in the bucket
//! is byte-identical to what a client pull would have left there.
use anyhow::{Context, Result, bail};
use futures::StreamExt;
use rustingface_entities::config::Config;
use rustingface_entities::repo::RepoType;
use serde_json::Value;
use crate::cmd::{human_bytes, split_repo};
use crate::wire;
/// Fetch a revision, or a named subset of its files.
pub async fn run(config: Config, repo_type: RepoType, repo: &str, files: &[String]) -> Result<()> {
if config.sealed() {
bail!("fetch needs upstream; this instance is sealed (upstream.enabled = false)");
}
let (repo_id, revision) = split_repo(repo)?;
let revision = revision.unwrap_or_else(|| "main".into());
let registry = wire::registry(&config)?;
let commit = registry
.resolve_revision(repo_type, &repo_id, &revision)
.await
.with_context(|| format!("resolving {repo_id}@{revision}"))?;
println!("{repo_id}@{revision} -> {commit}");
let wanted: Vec<String> = if files.is_empty() {
let tree = registry
.repo_tree(repo_type, &repo_id, &commit, "", true)
.await
.with_context(|| format!("listing {repo_id}@{commit}"))?;
tree_paths(&tree)
} else {
files.to_vec()
};
if wanted.is_empty() {
println!("nothing to fetch: the revision lists no files");
return Ok(());
}
let mut fetched = 0u64;
let mut bytes = 0u64;
let mut failures = Vec::new();
for path in &wanted {
// Draining the body is what commits the bytes: the same tee the server
// runs, with this process standing in for the client.
match registry
.get_file(repo_type, &repo_id, &commit, path, None)
.await
{
Ok((file, _)) => {
let mut size = 0u64;
if let rustingface_core::Body::Stream { mut stream, .. } = file.body {
while let Some(chunk) = stream.next().await {
size += chunk.with_context(|| format!("fetching {path}"))?.len() as u64;
}
} else {
size = file.size;
}
fetched += 1;
bytes += size;
println!(" {path} {}", human_bytes(size));
}
Err(err) => {
eprintln!(" {path} FAILED: {err}");
failures.push(path.clone());
}
}
}
println!(
"{fetched}/{} files, {} stored",
wanted.len(),
human_bytes(bytes)
);
if !failures.is_empty() {
bail!(
"{} file(s) could not be fetched: {}",
failures.len(),
failures.join(", ")
);
}
Ok(())
}
/// Every file path in a tree document.
fn tree_paths(tree: &Value) -> Vec<String> {
tree.as_array()
.map(|items| {
items
.iter()
.filter(|item| item.get("type").and_then(Value::as_str) != Some("directory"))
.filter_map(|item| item.get("path").and_then(Value::as_str))
.map(str::to_owned)
.collect()
})
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn directories_are_not_fetched_as_files() {
let tree = serde_json::json!([
{"type": "file", "path": "config.json"},
{"type": "directory", "path": "onnx"},
{"type": "file", "path": "onnx/model.onnx"},
]);
assert_eq!(tree_paths(&tree), vec!["config.json", "onnx/model.onnx"]);
}
#[test]
fn a_non_array_tree_yields_nothing_rather_than_panicking() {
assert!(tree_paths(&serde_json::json!({"error": "nope"})).is_empty());
}
}
crates/rustingface/src/cmd/gc.rs
+65
View File
@@ -0,0 +1,65 @@
//! `rustingface gc` and `rustingface verify`.
use anyhow::{Result, bail};
use rustingface_entities::config::Config;
use rustingface_entities::repo::RepoType;
use crate::cmd::{human_bytes, split_repo};
use crate::wire;
/// Delete blobs no manifest references.
pub async fn run(config: Config, dry_run: bool) -> Result<()> {
let registry = wire::registry(&config)?;
let result = registry.gc(dry_run).await?;
for key in &result.deleted {
println!("{} {key}", if dry_run { "would delete" } else { "deleted" });
}
println!(
"{} blobs scanned, {} referenced, {} {} ({})",
result.scanned,
result.referenced,
result.deleted.len(),
if dry_run { "reclaimable" } else { "reclaimed" },
human_bytes(result.bytes)
);
if dry_run && !result.deleted.is_empty() {
println!("re-run without --dry-run to reclaim them");
}
Ok(())
}
/// Re-read blobs and check them against their content-addressed keys.
pub async fn verify(config: Config, repo_type: RepoType, repo: Option<&str>) -> Result<()> {
let registry = wire::registry(&config)?;
let parsed = repo.map(split_repo).transpose()?;
let target = parsed.as_ref().map(|(id, _)| (repo_type, id));
let result = registry.verify(target).await?;
for key in &result.missing {
println!("MISSING {key}");
}
for corruption in &result.corrupt {
println!(
"CORRUPT {}\n expected {}\n found {}",
corruption.key, corruption.expected, corruption.found
);
}
println!(
"{} blobs read ({}), {} missing, {} corrupt",
result.blobs,
human_bytes(result.bytes),
result.missing.len(),
result.corrupt.len()
);
if !result.corrupt.is_empty() || !result.missing.is_empty() {
bail!(
"verify found {} corrupt and {} missing blob(s)",
result.corrupt.len(),
result.missing.len()
);
}
Ok(())
}
crates/rustingface/src/cmd/list.rs
+96
View File
@@ -0,0 +1,96 @@
//! `rustingface list` and `rustingface show`.
use anyhow::{Context, Result};
use rustingface_entities::config::Config;
use rustingface_entities::repo::RepoType;
use crate::cmd::{human_bytes, split_repo};
use crate::wire;
/// List stored repos with their pins and sizes.
pub async fn run(config: Config, repo_type: Option<RepoType>) -> Result<()> {
let registry = wire::registry(&config)?;
let types: Vec<RepoType> = repo_type
.map(|t| vec![t])
.unwrap_or_else(|| RepoType::ALL.to_vec());
let mut total = 0u64;
let mut count = 0usize;
for repo_type in types {
for repo_id in registry.list_repos(repo_type).await? {
let summary = registry.summarise(repo_type, &repo_id).await?;
let pins = if summary.refs.is_empty() {
"-".to_owned()
} else {
summary
.refs
.iter()
.map(|(name, commit)| format!("{name}={}", short(commit)))
.collect::<Vec<_>>()
.join(" ")
};
println!(
"{:<9} {:<44} {:>3} rev {:>5} files {:>10} {pins}",
repo_type.as_str(),
repo_id.to_string(),
summary.commits.len(),
summary.files,
human_bytes(summary.bytes),
);
total += summary.bytes;
count += 1;
}
}
if count == 0 {
println!("no repositories stored in {}", config.storage.bucket);
} else {
println!("\n{count} repositories, {} referenced", human_bytes(total));
}
Ok(())
}
/// Print a manifest.
pub async fn show(config: Config, repo_type: RepoType, repo: &str) -> Result<()> {
let (repo_id, revision) = split_repo(repo)?;
let registry = wire::registry(&config)?;
// Resolve against the bucket only: `show` must never have the side effect
// of pulling a fresh pin from upstream.
let commit = registry
.resolve_stored_revision(
repo_type,
&repo_id,
&revision.unwrap_or_else(|| "main".into()),
)
.await
.with_context(|| format!("resolving {repo_id} in the bucket"))?;
let manifest = registry
.load_manifest(repo_type, &repo_id, &commit)
.await?
.with_context(|| format!("no manifest stored for {repo_id}@{commit}"))?;
println!("{}", serde_json::to_string_pretty(&manifest)?);
Ok(())
}
/// Abbreviate a commit for a table.
fn short(commit: &str) -> &str {
&commit[..commit.len().min(12)]
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn commits_are_abbreviated_without_panicking_on_short_input() {
assert_eq!(
short("8f2a1c9e0b7d4f6a3e5c2b8d9f0a1c3e5b7d9f01"),
"8f2a1c9e0b7d"
);
assert_eq!(short("abc"), "abc");
assert_eq!(short(""), "");
}
}
crates/rustingface/src/cmd/mod.rs
+69
View File
@@ -0,0 +1,69 @@
//! Subcommand implementations.
//!
//! Each one builds the same registry the server would and drives the same core
//! operations, so the CLI cannot drift from the service's behaviour.
pub mod doctor;
pub mod fetch;
pub mod gc;
pub mod list;
pub mod pin;
pub mod rm;
pub mod serve;
use anyhow::{Result, bail};
use rustingface_entities::repo::RepoId;
/// Split a `namespace/name[@revision]` argument.
pub fn split_repo(argument: &str) -> Result<(RepoId, Option<String>)> {
let (repo, revision) = match argument.split_once('@') {
Some((repo, revision)) if !revision.is_empty() => (repo, Some(revision.to_owned())),
Some((_, _)) => bail!("{argument:?} ends with `@` but names no revision"),
None => (argument, None),
};
Ok((repo.parse()?, revision))
}
/// Render a byte count the way an operator reads it.
pub fn human_bytes(bytes: u64) -> String {
const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
let mut value = bytes as f64;
let mut unit = 0;
while value >= 1024.0 && unit < UNITS.len() - 1 {
value /= 1024.0;
unit += 1;
}
if unit == 0 {
format!("{bytes} B")
} else {
format!("{value:.1} {}", UNITS[unit])
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_repo_argument_splits_on_the_revision_marker() {
let (repo, revision) = split_repo("Qwen/Qwen3-32B@main").unwrap();
assert_eq!(repo.to_string(), "Qwen/Qwen3-32B");
assert_eq!(revision.as_deref(), Some("main"));
let (repo, revision) = split_repo("gpt2").unwrap();
assert_eq!(repo.to_string(), "gpt2");
assert_eq!(revision, None);
}
#[test]
fn a_trailing_marker_with_no_revision_is_an_error() {
assert!(split_repo("gpt2@").is_err());
}
#[test]
fn byte_counts_read_the_way_an_operator_expects() {
assert_eq!(human_bytes(512), "512 B");
assert_eq!(human_bytes(1536), "1.5 KiB");
assert_eq!(human_bytes(4 * 1024 * 1024 * 1024), "4.0 GiB");
}
}
crates/rustingface/src/cmd/pin.rs
+88
View File
@@ -0,0 +1,88 @@
//! `rustingface pin` and `rustingface refresh`.
//!
//! Together these are the only ways a ref moves. That is the point of freeze
//! mode: upstream cannot change what a fleet resolves to, an operator can.
use anyhow::{Context, Result, bail};
use rustingface_entities::config::Config;
use rustingface_entities::refs::{PinnedBy, is_commit_sha};
use rustingface_entities::repo::RepoType;
use crate::cmd::split_repo;
use crate::wire;
/// Set a ref to a commit explicitly.
pub async fn run(
config: Config,
repo_type: RepoType,
repo: &str,
ref_name: &str,
commit: &str,
) -> Result<()> {
if !is_commit_sha(commit) {
bail!("{commit:?} is not a 40-character commit SHA");
}
let (repo_id, _) = split_repo(repo)?;
let registry = wire::registry(&config)?;
let before = registry.load_ref(repo_type, &repo_id, ref_name).await?;
registry
.write_ref(repo_type, &repo_id, ref_name, commit, PinnedBy::Operator)
.await
.with_context(|| format!("pinning {repo_id}@{ref_name}"))?;
registry.ensure_repo_record(repo_type, &repo_id).await?;
match before {
Some(before) if before.commit == commit => {
println!("{repo_id}@{ref_name} already at {commit}")
}
Some(before) => println!("{repo_id}@{ref_name}: {} -> {commit}", before.commit),
None => println!("{repo_id}@{ref_name} -> {commit}"),
}
if registry
.load_manifest(repo_type, &repo_id, commit)
.await?
.is_none()
{
println!(
"note: no manifest is stored for {commit}. \
Run `rustingface fetch {repo_id}@{commit}` to populate it before sealing."
);
}
Ok(())
}
/// Re-resolve a ref upstream and repin it.
pub async fn refresh(
config: Config,
repo_type: RepoType,
repo: &str,
ref_name: Option<&str>,
) -> Result<()> {
let (repo_id, revision) = split_repo(repo)?;
let ref_name = ref_name
.map(str::to_owned)
.or(revision)
.unwrap_or_else(|| "main".into());
let registry = wire::registry(&config)?;
let (before, after) = registry
.refresh(repo_type, &repo_id, &ref_name)
.await
.with_context(|| format!("refreshing {repo_id}@{ref_name}"))?;
match before {
Some(before) if before == after => println!("{repo_id}@{ref_name} unchanged at {after}"),
Some(before) => {
println!("{repo_id}@{ref_name}: {before} -> {after}");
println!(
"note: the previous manifest for {before} is untouched, so anything pinned to \
that commit still resolves. Run `rustingface fetch {repo_id}@{after}` to store \
the new revision."
);
}
None => println!("{repo_id}@{ref_name} -> {after}"),
}
Ok(())
}
crates/rustingface/src/cmd/rm.rs
+51
View File
@@ -0,0 +1,51 @@
//! `rustingface rm` — detach metadata.
//!
//! Blobs survive. Reclaiming them is `gc`, a second explicit step that can be
//! dry-run first, so nothing deletes bytes without an operator typing `gc`.
use anyhow::{Result, bail};
use rustingface_entities::config::Config;
use rustingface_entities::repo::RepoType;
use crate::cmd::split_repo;
use crate::wire;
/// Remove a repo, or one revision of it.
pub async fn run(config: Config, repo_type: RepoType, repo: &str, confirmed: bool) -> Result<()> {
let (repo_id, revision) = split_repo(repo)?;
let registry = wire::registry(&config)?;
if !confirmed {
let summary = registry.summarise(repo_type, &repo_id).await?;
match &revision {
Some(revision) => println!("would remove {repo_id}@{revision}"),
None => println!(
"would remove {repo_id}: {} revisions, {} files, {} refs",
summary.commits.len(),
summary.files,
summary.refs.len()
),
}
bail!("refusing to remove anything without --yes");
}
let removed = registry
.remove(repo_type, &repo_id, revision.as_deref())
.await?;
for key in &removed.manifests {
println!("removed manifest {key}");
}
for key in &removed.refs {
println!("removed ref {key}");
}
if removed.repo_record {
println!("removed repo record for {repo_id}");
}
if removed.manifests.is_empty() && removed.refs.is_empty() && !removed.repo_record {
println!("nothing to remove for {repo_id}");
return Ok(());
}
println!("blobs are untouched; run `rustingface gc --dry-run` to see what is now reclaimable");
Ok(())
}
crates/rustingface/src/cmd/serve.rs
+70
View File
@@ -0,0 +1,70 @@
//! `rustingface serve`.
use anyhow::{Context, Result};
use rustingface_api::auth::Tokens;
use rustingface_api::{AppState, router, telemetry};
use rustingface_entities::config::Config;
use tokio::net::TcpListener;
use tokio::signal::unix::{SignalKind, signal};
use tracing::info;
use crate::{config::read_optional_secret, wire};
/// Run the HTTP service until `SIGTERM` or `SIGINT`.
pub async fn run(config: Config) -> Result<()> {
let registry = wire::registry(&config)?;
let token_file =
read_optional_secret(config.auth.token_file.as_ref()).context("auth.token_file")?;
let tokens = Tokens::new(config.auth.mode, token_file.as_deref())?;
let metrics = if config.observability.metrics {
Some(telemetry::install()?)
} else {
None
};
let listener = TcpListener::bind(&config.server.listen)
.await
.with_context(|| format!("binding {}", config.server.listen))?;
let local = listener.local_addr()?;
info!(
listen = %local,
bucket = %config.storage.bucket,
upstream = if config.upstream.enabled { config.upstream.endpoint.as_str() } else { "sealed" },
delivery = ?config.storage.blob_delivery,
auth = ?config.auth.mode,
"rustingface listening"
);
let app = router(AppState::new(registry, tokens.into(), metrics));
axum::serve(listener, app)
.with_graceful_shutdown(shutdown())
.await
.context("serving")
}
/// Stop accepting on `SIGTERM`/`SIGINT` and let in-flight requests drain.
///
/// Detached transfers are not requests and are not waited on here: the process
/// exiting mid-upload leaves an orphan blob, which is safe and reclaimable,
/// while a manifest is only ever written after an upload completes.
async fn shutdown() {
let mut term = match signal(SignalKind::terminate()) {
Ok(signal) => signal,
Err(err) => {
tracing::error!(%err, "could not install the SIGTERM handler");
return;
}
};
tokio::select! {
_ = term.recv() => info!("SIGTERM received; draining"),
result = tokio::signal::ctrl_c() => {
if result.is_ok() {
info!("interrupt received; draining");
}
}
}
}
crates/rustingface/src/config.rs
+139
View File
@@ -0,0 +1,139 @@
//! Loading configuration and the secrets it points at.
//!
//! Layering is file → environment → CLI, and every credential arrives from a
//! file path rather than inline, so `pass` or systemd credentials can own it
//! and it never appears in the config, the environment, or a log line.
use std::path::{Path, PathBuf};
use anyhow::{Context, Result, bail};
use figment::Figment;
use figment::providers::{Env, Format, Toml};
use rustingface_entities::config::{Config, ENV_PREFIX};
/// Load and validate configuration.
///
/// A missing file is only acceptable when the caller did not name one: an
/// explicit `--config` that does not exist is an error, not a silent fallback
/// to defaults that would point at the wrong bucket.
pub fn load(path: Option<&Path>, explicit: bool) -> Result<Config> {
let mut figment = Figment::new();
if let Some(path) = path {
if explicit && !path.exists() {
bail!("{} does not exist", path.display());
}
if path.exists() {
figment = figment.merge(Toml::file(path));
}
}
let config: Config = figment
.merge(Env::prefixed(ENV_PREFIX).split("__"))
.extract()
.context("loading configuration")?;
config.validate()?;
Ok(config)
}
/// Read a secret from a file, trimming the trailing newline an editor adds.
///
/// systemd's `LoadCredential=` places credentials under
/// `$CREDENTIALS_DIRECTORY` with names rather than paths, so a configured path
/// that is absent is retried there by basename. That lets one config file work
/// both under systemd and when run by hand.
pub fn read_secret(path: &Path) -> Result<String> {
if path.exists() {
return read_trimmed(path);
}
if let Some(dir) = std::env::var_os("CREDENTIALS_DIRECTORY")
&& let Some(name) = path.file_name()
{
let candidate = PathBuf::from(dir).join(name);
if candidate.exists() {
return read_trimmed(&candidate);
}
}
bail!(
"{} does not exist, and no systemd credential of that name was provided",
path.display()
)
}
/// Read a secret if its path is configured at all.
pub fn read_optional_secret(path: Option<&PathBuf>) -> Result<Option<String>> {
match path {
None => Ok(None),
Some(path) => read_secret(path).map(Some).map(|s| {
// An empty credential file is almost always a deployment mistake
// — a rendered template that got no value — and it fails far more
// usefully here than as a 401 from upstream.
match s {
Some(v) if v.is_empty() => None,
other => other,
}
}),
}
}
fn read_trimmed(path: &Path) -> Result<String> {
let raw =
std::fs::read_to_string(path).with_context(|| format!("reading {}", path.display()))?;
Ok(raw.trim().to_owned())
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
fn tempdir() -> PathBuf {
let dir = std::env::temp_dir().join(format!(
"rustingface-config-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&dir).unwrap();
dir
}
#[test]
fn a_missing_explicit_config_is_an_error_not_a_silent_default() {
let missing = tempdir().join("nope.toml");
assert!(load(Some(&missing), true).is_err());
assert!(
load(Some(&missing), false).is_ok(),
"the default path may be absent; defaults then apply"
);
}
#[test]
fn a_secret_file_loses_its_trailing_newline() {
let dir = tempdir();
let path = dir.join("token");
let mut f = std::fs::File::create(&path).unwrap();
writeln!(f, "hf_secret_value").unwrap();
assert_eq!(read_secret(&path).unwrap(), "hf_secret_value");
}
#[test]
fn an_empty_secret_file_reads_as_no_secret_at_all() {
let dir = tempdir();
let path = dir.join("empty");
std::fs::write(&path, "\n").unwrap();
assert_eq!(read_optional_secret(Some(&path)).unwrap(), None);
}
#[test]
fn a_systemd_credential_stands_in_for_a_missing_path() {
let dir = tempdir();
std::fs::write(dir.join("hf-token"), "from-systemd\n").unwrap();
// SAFETY: single-threaded test process, and the variable is read back
// only by this test.
unsafe { std::env::set_var("CREDENTIALS_DIRECTORY", &dir) };
let configured = PathBuf::from("/etc/rustingface/hf-token");
assert_eq!(read_secret(&configured).unwrap(), "from-systemd");
unsafe { std::env::remove_var("CREDENTIALS_DIRECTORY") };
}
}
crates/rustingface/src/logging.rs
+45
View File
@@ -0,0 +1,45 @@
//! Log setup.
//!
//! Structured JSON when systemd is listening, pretty text when a person is.
//! Detection is by configuration first and `JOURNAL_STREAM` second, so a
//! deliberate `log_format = "text"` is not overridden by the environment.
use rustingface_entities::config::{Config, LogFormat};
use tracing_subscriber::EnvFilter;
/// A minimal subscriber for the window before configuration is loaded, so a
/// failure to parse the config is still reported through the same channel.
pub fn init_bootstrap(interactive: bool) {
let _ = tracing_subscriber::fmt()
.with_env_filter(filter(interactive))
.with_writer(std::io::stderr)
.try_init();
}
/// Install the configured subscriber.
///
/// A no-op when the bootstrap subscriber is already installed, which is the
/// normal case; the bootstrap format is close enough that swapping it out
/// mid-process would cost more than it buys.
pub fn init(config: &Config, interactive: bool) {
let json = !interactive
&& (config.observability.log_format == LogFormat::Json
|| std::env::var_os("JOURNAL_STREAM").is_some());
if json {
let _ = tracing_subscriber::fmt()
.json()
.with_env_filter(filter(interactive))
.with_writer(std::io::stderr)
.try_init();
}
}
/// The default filter, quieter for the CLI than for the daemon.
fn filter(interactive: bool) -> EnvFilter {
let default = if interactive {
"rustingface=warn,warn"
} else {
"rustingface=info,rustingface_core=info,rustingface_data=info,rustingface_api=info,warn"
};
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(default))
}
crates/rustingface/src/main.rs
+185
View File
@@ -0,0 +1,185 @@
//! rustingface: a sovereign, HuggingFace-compatible model registry.
//!
//! One binary, two jobs. `serve` runs the HTTP service; the other subcommands
//! are the operator interface to the bucket. They share the same configuration
//! and the same code paths, so what the CLI does to the bucket is exactly what
//! the server would have done.
mod cmd;
mod config;
mod logging;
mod wire;
use std::path::PathBuf;
use anyhow::Result;
use clap::{Parser, Subcommand};
use rustingface_entities::config::DEFAULT_CONFIG_PATH;
use rustingface_entities::repo::RepoType;
/// Command-line interface.
#[derive(Debug, Parser)]
#[command(name = "rustingface", version, about, long_about = None)]
struct Cli {
/// Configuration file.
#[arg(short, long, global = true, env = "RUSTINGFACE_CONFIG")]
config: Option<PathBuf>,
#[command(subcommand)]
command: Option<Command>,
}
/// The subcommands.
#[derive(Debug, Subcommand)]
enum Command {
/// Run the HTTP service. The default when no subcommand is given.
Serve,
/// Warm the bucket ahead of need.
Fetch {
/// `namespace/name[@revision]`.
repo: String,
/// Restrict to these paths. Without it, every file in the revision's
/// tree is fetched.
#[arg(long = "files", num_args = 1..)]
files: Vec<String>,
#[arg(long, value_enum, default_value_t = RepoKind::Models)]
repo_type: RepoKind,
},
/// Set a ref to a commit explicitly.
Pin {
repo: String,
r#ref: String,
commit: String,
#[arg(long, value_enum, default_value_t = RepoKind::Models)]
repo_type: RepoKind,
},
/// Re-resolve a ref upstream and repin it.
Refresh {
repo: String,
/// Defaults to `main`.
r#ref: Option<String>,
#[arg(long, value_enum, default_value_t = RepoKind::Models)]
repo_type: RepoKind,
},
/// List stored repos, revisions and sizes.
List {
#[arg(long, value_enum)]
repo_type: Option<RepoKind>,
},
/// Print a manifest.
Show {
/// `namespace/name[@revision]`.
repo: String,
#[arg(long, value_enum, default_value_t = RepoKind::Models)]
repo_type: RepoKind,
},
/// Remove refs and manifests. Blobs are left for `gc`.
Rm {
/// `namespace/name[@revision]`.
repo: String,
#[arg(long, value_enum, default_value_t = RepoKind::Models)]
repo_type: RepoKind,
/// Required, because this is the step that makes data unreachable.
#[arg(long)]
yes: bool,
},
/// Delete blobs no manifest references.
Gc {
/// Report what would be deleted without deleting it.
#[arg(long)]
dry_run: bool,
},
/// Re-read stored blobs and check them against their digests.
Verify {
/// Restrict to one repository.
repo: Option<String>,
#[arg(long, value_enum, default_value_t = RepoKind::Models)]
repo_type: RepoKind,
},
/// Check configuration, storage reachability and permissions.
Doctor,
}
/// Repo type as a CLI value.
#[derive(Debug, Clone, Copy, clap::ValueEnum)]
enum RepoKind {
Models,
Datasets,
}
impl From<RepoKind> for RepoType {
fn from(kind: RepoKind) -> Self {
match kind {
RepoKind::Models => Self::Models,
RepoKind::Datasets => Self::Datasets,
}
}
}
fn main() -> Result<()> {
let cli = Cli::parse();
let explicit = cli.config.is_some();
let path = cli
.config
.unwrap_or_else(|| PathBuf::from(DEFAULT_CONFIG_PATH));
let command = cli.command.unwrap_or(Command::Serve);
// The server logs structurally for journald; the CLI writes to a terminal
// an operator is reading, so it gets human output and its results go to
// stdout rather than into a log line.
let interactive = !matches!(command, Command::Serve);
let config = {
logging::init_bootstrap(interactive);
config::load(Some(&path), explicit)?
};
logging::init(&config, interactive);
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?;
runtime.block_on(async move {
match command {
Command::Serve => cmd::serve::run(config).await,
Command::Fetch {
repo,
files,
repo_type,
} => cmd::fetch::run(config, repo_type.into(), &repo, &files).await,
Command::Pin {
repo,
r#ref,
commit,
repo_type,
} => cmd::pin::run(config, repo_type.into(), &repo, &r#ref, &commit).await,
Command::Refresh {
repo,
r#ref,
repo_type,
} => cmd::pin::refresh(config, repo_type.into(), &repo, r#ref.as_deref()).await,
Command::List { repo_type } => cmd::list::run(config, repo_type.map(Into::into)).await,
Command::Show { repo, repo_type } => {
cmd::list::show(config, repo_type.into(), &repo).await
}
Command::Rm {
repo,
repo_type,
yes,
} => cmd::rm::run(config, repo_type.into(), &repo, yes).await,
Command::Gc { dry_run } => cmd::gc::run(config, dry_run).await,
Command::Verify { repo, repo_type } => {
cmd::gc::verify(config, repo_type.into(), repo.as_deref()).await
}
Command::Doctor => cmd::doctor::run(config).await,
}
})
}
crates/rustingface/src/wire.rs
+48
View File
@@ -0,0 +1,48 @@
//! Wiring: turn a validated [`Config`] into a running [`Registry`].
//!
//! Every binary subcommand needs the same stack, so it is assembled once here
//! rather than in each command.
use std::sync::Arc;
use anyhow::{Context, Result};
use rustingface_core::Registry;
use rustingface_core::ports::{Store, Upstream};
use rustingface_data::{HubClient, ObjectStoreAdapter, SealedUpstream};
use rustingface_entities::config::Config;
use tracing::warn;
use crate::config::read_optional_secret;
/// Build the registry described by `config`.
pub fn registry(config: &Config) -> Result<Registry> {
let access_key_id = read_optional_secret(config.storage.access_key_id_file.as_ref())
.context("storage.access_key_id_file")?;
let secret_access_key = read_optional_secret(config.storage.secret_access_key_file.as_ref())
.context("storage.secret_access_key_file")?;
if config.storage.local_path.is_some() {
warn!(
"storage.local_path is set: state lives on this host's filesystem rather than in a \
bucket, so losing the host loses the registry. This forfeits state portability and \
is intended for tests and trials only."
);
}
let store: Arc<dyn Store> = Arc::new(
ObjectStoreAdapter::new(&config.storage, access_key_id, secret_access_key)
.context("configuring object storage")?,
);
let upstream: Arc<dyn Upstream> = if config.upstream.enabled {
let token = read_optional_secret(config.upstream.token_file.as_ref())
.context("upstream.token_file")?;
Arc::new(HubClient::new(&config.upstream, token).context("configuring upstream client")?)
} else {
// Sealed: no client is constructed at all, so there is no code path
// that could reach the network even if a bug tried to.
Arc::new(SealedUpstream::new(config.upstream.endpoint.clone()))
};
Registry::new(store, upstream, config.clone()).context("assembling the registry")
}
rust-toolchain.toml
+4
View File
@@ -0,0 +1,4 @@
[toolchain]
channel = "stable"
components = ["rustfmt", "clippy"]
targets = ["x86_64-unknown-linux-musl"]
script/infra-setup.sh
Executable
+295
View File
@@ -0,0 +1,295 @@
#!/usr/bin/env bash
#
# One-time host provisioning for rustingface.
#
# Run by an operator from a workstation with full sudo (not by CI, and not as
# the scoped gitea_ci account). Idempotent: re-running with nothing to change
# is a no-op. Skips past unreachable hosts so one offline node does not block
# the rest.
#
# script/infra-setup.sh # provision everything
# script/infra-setup.sh --skip-minio # leave the bucket and credentials alone
#
# What it does:
# 1. bob — gitea_ci user, runner key, scoped sudoers, SELinux port label
# 2. caveman — the rustingface bucket and a scoped MinIO service account
# 3. hanzalova— rf.internal cert, nginx vhost, step@ renewal timer
#
# Application config is NOT shipped here. The deploy workflow renders it from
# Gitea secrets on every deploy.
set -euo pipefail
readonly APP=rustingface
readonly PORT=20482
readonly SERVICE_HOST=bob.hanzalova.internal
readonly PROXY_HOST=hanzalova.internal
readonly MINIO_HOST=caveman.kosherinata.internal
readonly MINIO_URL=http://caveman.kosherinata.internal:9000
readonly BUCKET=rustingface
readonly VHOST=rf
readonly RUNNER_PUBKEY="${RUNNER_PUBKEY:-$HOME/.ssh/id_gitea_ci.pub}"
readonly PROVISIONER_PASSWORD="${PROVISIONER_PASSWORD:-$HOME/.step/secrets/provisioner}"
readonly ROOT_CA=/etc/pki/ca-trust/source/anchors/root-internal.pem
SKIP_MINIO=0
for arg in "$@"; do
case "$arg" in
--skip-minio) SKIP_MINIO=1 ;;
-h|--help) sed -n '2,22p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
*) echo "unknown argument: $arg" >&2; exit 2 ;;
esac
done
info() { printf '\033[1;34m==>\033[0m %s\n' "$*"; }
ok() { printf ' \033[32mok\033[0m %s\n' "$*"; }
warn() { printf ' \033[33m!!\033[0m %s\n' "$*" >&2; }
die() { printf '\033[1;31m!!\033[0m %s\n' "$*" >&2; exit 1; }
# Returns non-zero (visibly) when a host cannot be reached, so the caller can
# skip it deliberately rather than by swallowing an error.
reachable() {
local host=$1
if ssh -o BatchMode=yes -o ConnectTimeout=8 -o StrictHostKeyChecking=accept-new \
"$host" true; then
return 0
fi
warn "$host is unreachable; skipping it. Re-run this script when it is back."
return 1
}
# ---------------------------------------------------------------------------
# 1. The service host: bob
# ---------------------------------------------------------------------------
provision_service_host() {
info "service host $SERVICE_HOST"
[ -f "$RUNNER_PUBKEY" ] || die "runner public key not found at $RUNNER_PUBKEY.
The keypair is maintained at ~/.ssh/id_gitea_ci on the operator workstation and is
shared by every project's deploy. Copy it from there; do not generate a new one, or
every other project's deploy breaks."
# The gitea_ci account, its key, and journal access. Creating the user is
# idempotent via `useradd` guarded on `id`.
ssh "$SERVICE_HOST" 'sudo bash -s' <<'REMOTE'
set -euo pipefail
if id gitea_ci >/dev/null 2>&1; then
echo " gitea_ci exists"
else
useradd --system --create-home --home-dir /var/lib/gitea_ci --shell /bin/bash gitea_ci
echo " created gitea_ci"
fi
install -d -o gitea_ci -g gitea_ci -m 0700 /var/lib/gitea_ci/.ssh
usermod -aG systemd-journal gitea_ci
REMOTE
ok "gitea_ci account"
# shellcheck disable=SC2029 # the key must expand locally
ssh "$SERVICE_HOST" "sudo bash -c '
umask 077
touch /var/lib/gitea_ci/.ssh/authorized_keys
grep -qxF \"$(cat "$RUNNER_PUBKEY")\" /var/lib/gitea_ci/.ssh/authorized_keys \
|| echo \"$(cat "$RUNNER_PUBKEY")\" >> /var/lib/gitea_ci/.ssh/authorized_keys
chown -R gitea_ci:gitea_ci /var/lib/gitea_ci/.ssh
chmod 600 /var/lib/gitea_ci/.ssh/authorized_keys'"
ok "runner key installed"
# Scoped sudoers: exactly the commands the deploy runs, pinned to literal
# destinations, and verified before it can lock anyone out.
ssh "$SERVICE_HOST" 'sudo bash -s' <<REMOTE
set -euo pipefail
cat > /etc/sudoers.d/.${APP}_gitea_ci.new <<'SUDOERS'
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /usr/local/bin/${APP}
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/${APP}/config.toml
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/${APP}/s3-access-key
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/${APP}/s3-secret-key
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/${APP}/hf-token
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/systemd/system/${APP}.service
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/sysusers.d/${APP}.conf
gitea_ci ALL=(root) NOPASSWD: /usr/bin/rsync * /etc/firewalld/services/${APP}.xml
gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemd-sysusers
gitea_ci ALL=(root) NOPASSWD: /usr/bin/install -d -o root -g ${APP} -m 0750 /etc/${APP}
gitea_ci ALL=(root) NOPASSWD: /usr/bin/chown -R root\:${APP} /etc/${APP}
gitea_ci ALL=(root) NOPASSWD: /usr/bin/chmod 0640 /etc/${APP}/config.toml /etc/${APP}/s3-access-key /etc/${APP}/s3-secret-key /etc/${APP}/hf-token
gitea_ci ALL=(root) NOPASSWD: /usr/bin/chmod 0755 /usr/local/bin/${APP}
gitea_ci ALL=(root) NOPASSWD: /usr/sbin/restorecon -R /usr/local/bin/${APP} /etc/${APP}
gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemctl daemon-reload
gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemctl enable ${APP}.service
gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemctl restart ${APP}.service
gitea_ci ALL=(root) NOPASSWD: /usr/bin/systemctl is-active ${APP}.service
gitea_ci ALL=(root) NOPASSWD: /usr/bin/firewall-cmd --reload
gitea_ci ALL=(root) NOPASSWD: /usr/bin/firewall-cmd --get-default-zone
gitea_ci ALL=(root) NOPASSWD: /usr/bin/firewall-cmd --zone=* --query-service=${APP}
gitea_ci ALL=(root) NOPASSWD: /usr/bin/firewall-cmd --permanent --zone=* --add-service=${APP}
gitea_ci ALL=(root) NOPASSWD: /usr/bin/firewall-cmd --zone=* --add-service=${APP}
SUDOERS
chmod 0440 /etc/sudoers.d/.${APP}_gitea_ci.new
# Verify before moving it into place: a typo in a sudoers file that is already
# installed can lock the host out of sudo entirely.
visudo -cf /etc/sudoers.d/.${APP}_gitea_ci.new
mv /etc/sudoers.d/.${APP}_gitea_ci.new /etc/sudoers.d/${APP}_gitea_ci
REMOTE
ok "scoped sudoers installed and visudo-verified"
# SELinux: the daemon cannot bind an unlabelled port. This must happen
# before the first service start, which is why it lives here and not in
# the deploy workflow.
ssh "$SERVICE_HOST" "sudo bash -s" <<REMOTE
set -euo pipefail
if semanage port -l | grep -Eq "^http_port_t .*\\b${PORT}\\b"; then
echo " tcp/${PORT} already labelled http_port_t"
else
semanage port -a -t http_port_t -p tcp ${PORT} \
|| semanage port -m -t http_port_t -p tcp ${PORT}
echo " labelled tcp/${PORT} http_port_t"
fi
REMOTE
ok "SELinux port label"
}
# ---------------------------------------------------------------------------
# 2. Object storage: caveman
# ---------------------------------------------------------------------------
provision_bucket() {
info "object storage $MINIO_HOST"
if ! ssh "$MINIO_HOST" 'command -v mc'; then
warn "mc (the MinIO client) is not installed on $MINIO_HOST.
Create the bucket and a scoped service account by hand, then re-run with --skip-minio:
mc alias set local $MINIO_URL <root-user> <root-password>
mc mb local/$BUCKET
mc admin user svcacct add local <root-user> --name rustingface
Put the resulting access key and secret in the Gitea repo secrets
S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY."
return 0
fi
warn "bucket and service-account creation needs MinIO root credentials, which
this script deliberately does not read. Run on $MINIO_HOST as an operator:
mc alias set local $MINIO_URL \"\$MINIO_ROOT_USER\" \"\$MINIO_ROOT_PASSWORD\"
mc mb --ignore-existing local/$BUCKET
mc admin user svcacct add local \"\$MINIO_ROOT_USER\" --name rustingface
then set the Gitea repo secrets S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY.
Re-run this script with --skip-minio once that is done."
}
# ---------------------------------------------------------------------------
# 3. The edge proxy: hanzalova
# ---------------------------------------------------------------------------
provision_proxy() {
info "edge proxy $PROXY_HOST"
local cert=/etc/nginx/tls/cert/${VHOST}.internal.pem
local key=/etc/nginx/tls/key/${VHOST}.internal.pem
# `step certificate verify` checks chain and expiry, not the name, which is
# enough to decide whether a fresh mint is needed.
local state
state=$(ssh "$PROXY_HOST" "if sudo test -f $cert && sudo step certificate verify $cert --roots $ROOT_CA; then echo valid; else echo missing; fi" \
| tail -1)
if [ "$state" != valid ]; then
[ -f "$PROVISIONER_PASSWORD" ] \
|| die "provisioner password not found at $PROVISIONER_PASSWORD; cannot mint ${VHOST}.internal"
rsync -az --rsync-path='sudo rsync' --chmod=0600 \
"$PROVISIONER_PASSWORD" "$PROXY_HOST:/tmp/${VHOST}-provisioner"
# Remove the credential whether or not the mint succeeds, then
# propagate the failure.
ssh "$PROXY_HOST" "
set -euo pipefail
sudo mkdir -p /etc/nginx/tls/cert /etc/nginx/tls/key
rc=0
sudo step ca certificate --force \
--provisioner lair \
--provisioner-password-file /tmp/${VHOST}-provisioner \
--ca-url https://ca.internal \
--root $ROOT_CA \
--san ${VHOST}.internal \
${VHOST}.internal $cert $key || rc=\$?
sudo rm -f /tmp/${VHOST}-provisioner
[ \$rc -eq 0 ] || { echo 'minting ${VHOST}.internal failed' >&2; exit \$rc; }
sudo chown root:root $cert $key
sudo chmod 0644 $cert
sudo chmod 0640 $key
sudo setfacl -m u:nginx:r $key"
ok "minted ${VHOST}.internal"
else
ok "${VHOST}.internal cert is present and valid"
fi
ssh "$PROXY_HOST" "sudo systemctl enable --now step@${VHOST}.timer"
ok "renewal timer enabled"
# The vhost references the cert paths, so it is only installed after the
# cert exists: nginx -t fails on a missing ssl_certificate and that blocks
# every unrelated reload on the host.
rsync -az --rsync-path='sudo rsync' --mkpath \
asset/nginx/${VHOST}.internal.conf \
"$PROXY_HOST:/etc/nginx/sites-available/${VHOST}.internal.conf"
ssh "$PROXY_HOST" "
set -euo pipefail
sudo ln -sfn ../sites-available/${VHOST}.internal.conf \
/etc/nginx/sites-enabled/${VHOST}.internal.conf
sudo nginx -t
sudo systemctl reload nginx"
ok "vhost installed and nginx reloaded"
# A reload is a request, not a guarantee: nginx keeps its old cycle if it
# cannot rebind, and then serves the certificate it loaded then. Compare
# what is actually on the wire against what is on disk, by serial.
local served disk
served=$(echo | openssl s_client -servername ${VHOST}.internal -connect ${PROXY_HOST}:443 2>/dev/null \
| openssl x509 -noout -serial || echo "serial=unavailable")
disk=$(ssh "$PROXY_HOST" "sudo openssl x509 -noout -serial -in $cert")
if [ "$served" = "$disk" ]; then
ok "the certificate on the wire matches the one on disk"
else
warn "nginx is serving a different certificate than the one on disk
wire: $served
disk: $disk
A reload was requested and reported success, but nginx kept its previous cycle.
With 24h certs this puts an expired certificate on the wire within a day.
Investigate the bind failure in the nginx error log, then restart nginx."
fi
info "DNS"
echo " Add the split-horizon record if it is not already present:"
echo " opn-cli --config ~/.opn-cli/hanzalova.yml unbound host create \\"
echo " --hostname ${VHOST} --domain internal --server 10.6.0.46 \\"
echo " --description 'rustingface via hanzalova reverse proxy (bob backend)'"
echo " opn-cli --config ~/.opn-cli/kosherinata.yml unbound host create \\"
echo " --hostname ${VHOST} --domain internal --server 10.6.0.46 \\"
echo " --description 'rustingface via hanzalova reverse proxy (bob backend)'"
echo " Never give ${VHOST}.internal a public / Cloudflare record."
}
# ---------------------------------------------------------------------------
main() {
cd "$(dirname "$0")/.."
if reachable "$SERVICE_HOST"; then
provision_service_host
fi
if [ "$SKIP_MINIO" -eq 0 ] && reachable "$MINIO_HOST"; then
provision_bucket
fi
if reachable "$PROXY_HOST"; then
provision_proxy
fi
info "done"
echo " Next: set the Gitea repo secrets (RSYNC_SSH_KEY, S3_ACCESS_KEY_ID,"
echo " S3_SECRET_ACCESS_KEY, HF_TOKEN) and push to main to deploy."
}
main "$@"
test/conformance/conformance.py
Executable
+555
View File
@@ -0,0 +1,555 @@
#!/usr/bin/env python3
"""Conformance suite: test against real clients, not against the document.
The spec is only as good as the client compatibility, so this drives an
unmodified ``huggingface_hub`` (and the ``hf`` CLI) through rustingface and
checks the behaviour a stack pinned to a model revision actually depends on.
The test that matters is the offline one. Its steps here mirror spec §11:
1. fresh bucket, ``HF_ENDPOINT`` set, pull a model, confirm bytes landed;
2. cut rustingface off from upstream;
3. wipe every client-side cache;
4. pull again -- it must succeed identically, with no upstream traffic;
5. kill rustingface, start a *new* process against the same bucket, repeat.
Step 2 differs from the spec in one honest respect. The spec null-routes
``huggingface.co`` at the site router and verifies there, which CI cannot do.
This suite instead does something strictly stronger for the sealed run --
``upstream.enabled = false`` means no upstream client is constructed at all, so
there is no code path that *could* reach the network -- and, for the
"upstream enabled but unreachable" run, points the endpoint at a closed port
and asserts that a miss fails loudly rather than hanging. Run the real
null-route test by hand before trusting a deployment.
Usage::
python3 test/conformance/conformance.py --binary target/release/rustingface
Set ``--s3-endpoint``/``--s3-bucket`` (plus ``AWS_ACCESS_KEY_ID`` and
``AWS_SECRET_ACCESS_KEY``) to run against a real bucket instead of a directory.
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from pathlib import Path
# A small, ungated, stable model. Small enough that CI can pull it repeatedly,
# real enough that it exercises the LFS bridge and a safetensors file.
REPO = os.environ.get("CONFORMANCE_REPO", "hf-internal-testing/tiny-random-gpt2")
REVISION = os.environ.get("CONFORMANCE_REVISION", "main")
# The X-Error-Code values rustingface emits. These are client-internal
# constants upstream and are re-checked against the installed client below.
EXPECTED_ERROR_CODES = {
"RepoNotFound",
"RevisionNotFound",
"EntryNotFound",
"GatedRepo",
}
# A disabled repo is the odd one out: the client dispatches DisabledRepoError
# on this exact X-Error-Message, not on an X-Error-Code, so rustingface emits
# the message verbatim. If this string moves, the mapping in
# rustingface-api/src/error.rs moves with it.
DISABLED_REPO_MESSAGE = "Access to this resource is disabled."
class Failure(Exception):
"""A conformance check that did not hold."""
@dataclass
class Results:
passed: list[str] = field(default_factory=list)
failed: list[tuple[str, str]] = field(default_factory=list)
def check(self, name: str, fn) -> None:
print(f"\n=== {name}")
try:
fn()
except Exception as err: # noqa: BLE001 - a suite reports, it does not raise
print(f" FAILED: {err}")
self.failed.append((name, str(err)))
else:
print(" ok")
self.passed.append(name)
def report(self) -> int:
print(f"\n{len(self.passed)} passed, {len(self.failed)} failed")
for name, err in self.failed:
print(f" FAILED {name}: {err}")
return 1 if self.failed else 0
def free_port() -> int:
with socket.socket() as sock:
sock.bind(("127.0.0.1", 0))
return sock.getsockname()[1]
def closed_port() -> int:
"""A port nothing is listening on, standing in for an unreachable Hub."""
return free_port()
class Instance:
"""A running rustingface, configured for one phase of the suite."""
def __init__(self, binary: Path, workdir: Path, storage: dict, upstream: dict):
self.binary = binary
self.workdir = workdir
self.port = free_port()
self.config_path = workdir / f"config-{self.port}.toml"
self.process: subprocess.Popen | None = None
storage_lines = "\n".join(f"{k} = {v}" for k, v in storage.items())
upstream_lines = "\n".join(f"{k} = {v}" for k, v in upstream.items())
self.config_path.write_text(
f"""
[server]
listen = "127.0.0.1:{self.port}"
request_timeout = "600s"
[storage]
{storage_lines}
[upstream]
{upstream_lines}
[policy]
ref_resolution = "freeze"
[observability]
metrics = true
log_format = "text"
"""
)
@property
def endpoint(self) -> str:
return f"http://127.0.0.1:{self.port}"
def __enter__(self) -> "Instance":
self.process = subprocess.Popen(
[str(self.binary), "--config", str(self.config_path), "serve"],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
)
for _ in range(100):
if self.process.poll() is not None:
raise Failure(f"rustingface exited early:\n{self.process.stdout.read()}")
try:
with urllib.request.urlopen(f"{self.endpoint}/healthz", timeout=1):
return self
except (urllib.error.URLError, ConnectionError, TimeoutError):
time.sleep(0.1)
raise Failure("rustingface did not become healthy")
def __exit__(self, *_exc) -> None:
if self.process is not None:
self.process.terminate()
try:
self.process.wait(timeout=30)
except subprocess.TimeoutExpired:
self.process.kill()
def metrics(self) -> str:
with urllib.request.urlopen(f"{self.endpoint}/metrics", timeout=5) as response:
return response.read().decode()
def cli(self, *args: str) -> str:
result = subprocess.run(
[str(self.binary), "--config", str(self.config_path), *args],
capture_output=True,
text=True,
)
if result.returncode != 0:
raise Failure(f"`rustingface {' '.join(args)}` failed:\n{result.stdout}\n{result.stderr}")
return result.stdout
def wipe_client_caches(home: Path) -> None:
"""Remove every client-side cache, so a pull cannot be served locally."""
for sub in ("hub", "assets", "xet"):
shutil.rmtree(home / sub, ignore_errors=True)
shutil.rmtree(home.parent / ".cache" / "huggingface", ignore_errors=True)
def client_env(endpoint: str, home: Path) -> dict:
env = dict(os.environ)
env.update(
{
"HF_ENDPOINT": endpoint,
"HF_HOME": str(home),
"HF_HUB_CACHE": str(home / "hub"),
# rustingface implements no Xet client and never advertises the
# capability; setting this makes the client's side explicit too.
"HF_HUB_DISABLE_XET": "1",
"HF_HUB_DISABLE_TELEMETRY": "1",
"HF_HUB_DISABLE_PROGRESS_BARS": "1",
}
)
env.pop("HF_TOKEN", None)
return env
def run_client(script: str, endpoint: str, home: Path) -> str:
"""Run a snippet in a fresh interpreter, so no client state leaks between phases."""
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
env=client_env(endpoint, home),
)
if result.returncode != 0:
raise Failure(f"client failed:\n{result.stdout}\n{result.stderr}")
return result.stdout
DOWNLOAD_ONE = """
import hashlib, json, sys
from huggingface_hub import hf_hub_download
path = hf_hub_download({repo!r}, "config.json", revision={revision!r})
data = open(path, "rb").read()
print(json.dumps({{"sha256": hashlib.sha256(data).hexdigest(), "len": len(data)}}))
"""
SNAPSHOT = """
import hashlib, json, os
from huggingface_hub import snapshot_download
root = snapshot_download({repo!r}, revision={revision!r})
files = {{}}
for base, _dirs, names in os.walk(root):
for name in names:
full = os.path.join(base, name)
if os.path.islink(full) and not os.path.exists(full):
continue
rel = os.path.relpath(full, root)
files[rel] = hashlib.sha256(open(full, "rb").read()).hexdigest()
print(json.dumps(files, sort_keys=True))
"""
MODEL_INFO = """
import json
from huggingface_hub import HfApi
info = HfApi().model_info({repo!r}, revision={revision!r})
print(json.dumps({{"sha": info.sha, "siblings": sorted(s.rfilename for s in info.siblings)}}))
"""
FILE_METADATA = """
import json
from huggingface_hub import get_hf_file_metadata, hf_hub_url
url = hf_hub_url({repo!r}, "config.json", revision={revision!r})
meta = get_hf_file_metadata(url)
print(json.dumps({{"commit": meta.commit_hash, "etag": meta.etag, "size": meta.size}}))
"""
MISSING_FILE = """
import json
from huggingface_hub import hf_hub_download
from huggingface_hub.errors import EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError
try:
hf_hub_download({repo!r}, "no-such-file.bin", revision={revision!r})
except EntryNotFoundError:
print(json.dumps({{"raised": "EntryNotFoundError"}}))
except (RepositoryNotFoundError, RevisionNotFoundError) as err:
print(json.dumps({{"raised": type(err).__name__}}))
else:
print(json.dumps({{"raised": None}}))
"""
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--binary", type=Path, default=Path("target/release/rustingface"))
parser.add_argument("--s3-endpoint")
parser.add_argument("--s3-bucket")
parser.add_argument("--keep", action="store_true", help="keep the work directory")
args = parser.parse_args()
if not args.binary.exists():
print(f"{args.binary} does not exist; build it with `cargo build --release`", file=sys.stderr)
return 2
results = Results()
workdir = Path(tempfile.mkdtemp(prefix="rustingface-conformance-"))
print(f"work directory: {workdir}")
bucket_dir = workdir / "bucket"
bucket_dir.mkdir()
hf_home = workdir / "hf"
if args.s3_endpoint:
storage = {
"endpoint": json.dumps(args.s3_endpoint),
"bucket": json.dumps(args.s3_bucket or "rustingface-conformance"),
"region": '"us-east-1"',
"path_style": "true",
"multipart_part_size": '"5MiB"',
}
for var, key in (("AWS_ACCESS_KEY_ID", "access_key_id_file"),
("AWS_SECRET_ACCESS_KEY", "secret_access_key_file")):
value = os.environ.get(var)
if not value:
print(f"{var} must be set to run against S3", file=sys.stderr)
return 2
secret = workdir / key
secret.write_text(value)
storage[key] = json.dumps(str(secret))
else:
storage = {"local_path": json.dumps(str(bucket_dir)), "multipart_part_size": '"5MiB"'}
open_upstream = {"enabled": "true", "endpoint": '"https://huggingface.co"'}
unreachable_upstream = {"enabled": "true", "endpoint": f'"http://127.0.0.1:{closed_port()}"',
"connect_timeout": '"2s"'}
sealed_upstream = {"enabled": "false"}
# ------------------------------------------------------------------
# The client's own constants, re-checked rather than assumed.
# ------------------------------------------------------------------
def check_error_codes() -> None:
import huggingface_hub
print(f" huggingface_hub {huggingface_hub.__version__}")
source = Path(huggingface_hub.__file__).parent
found = set()
for path in source.rglob("*.py"):
text = path.read_text(errors="ignore")
for code in EXPECTED_ERROR_CODES:
if f'"{code}"' in text or f"'{code}'" in text:
found.add(code)
missing = EXPECTED_ERROR_CODES - found
if missing:
raise Failure(
f"this client does not mention {sorted(missing)}; the X-Error-Code strings "
"rustingface emits are version-coupled and need review"
)
http = (source / "utils" / "_http.py").read_text()
if DISABLED_REPO_MESSAGE not in http:
raise Failure(
f"this client no longer dispatches DisabledRepoError on "
f"{DISABLED_REPO_MESSAGE!r}; rustingface-api/src/error.rs needs review"
)
print(f" {len(found)} codes plus the disabled-repo message all match")
results.check("X-Error-Code strings match the pinned client", check_error_codes)
baseline: dict = {}
snapshot_baseline: dict = {}
# ------------------------------------------------------------------
# 1. Cold pull through rustingface.
# ------------------------------------------------------------------
def cold_pull() -> None:
nonlocal baseline, snapshot_baseline
wipe_client_caches(hf_home)
with Instance(args.binary, workdir, storage, open_upstream) as rf:
baseline = json.loads(
run_client(DOWNLOAD_ONE.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
snapshot_baseline = json.loads(
run_client(SNAPSHOT.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
print(f" {len(snapshot_baseline)} files in the snapshot")
listing = rf.cli("list")
if REPO not in listing:
raise Failure(f"`rustingface list` does not mention {REPO}:\n{listing}")
if "rustingface_bytes_served_total" not in rf.metrics():
raise Failure("/metrics did not expose the served-bytes counter")
results.check("a cold pull succeeds and lands in the bucket", cold_pull)
# ------------------------------------------------------------------
# 2-4. Cut off upstream, wipe the client, pull again.
# ------------------------------------------------------------------
def sealed_pull() -> None:
wipe_client_caches(hf_home)
with Instance(args.binary, workdir, storage, sealed_upstream) as rf:
again = json.loads(
run_client(DOWNLOAD_ONE.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if again != baseline:
raise Failure(f"sealed pull differs from the first: {again} != {baseline}")
snapshot_again = json.loads(
run_client(SNAPSHOT.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if snapshot_again != snapshot_baseline:
missing = set(snapshot_baseline) - set(snapshot_again)
raise Failure(f"sealed snapshot differs; missing {sorted(missing)}")
results.check("a sealed instance serves the identical snapshot", sealed_pull)
# ------------------------------------------------------------------
# 5. A brand-new process against the same bucket.
# ------------------------------------------------------------------
def portable_state() -> None:
wipe_client_caches(hf_home)
# Nothing carries over: a fresh process, a fresh client cache, only the
# bucket in common. This is the step that tests sovereignty rather
# than a cache.
with Instance(args.binary, workdir, storage, sealed_upstream) as rf:
again = json.loads(
run_client(DOWNLOAD_ONE.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if again != baseline:
raise Failure(f"reinstalled instance differs: {again} != {baseline}")
results.check("a reinstalled instance serves from the bucket alone", portable_state)
# ------------------------------------------------------------------
# Metadata replay.
# ------------------------------------------------------------------
def metadata_replay() -> None:
with Instance(args.binary, workdir, storage, open_upstream) as rf:
live = json.loads(
run_client(FILE_METADATA.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
with Instance(args.binary, workdir, storage, sealed_upstream) as rf:
replayed = json.loads(
run_client(FILE_METADATA.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if live != replayed:
raise Failure(f"file metadata changed once sealed: {live} != {replayed}")
print(f" commit {live['commit']} etag {live['etag']} size {live['size']}")
results.check("file metadata replays byte-identically when sealed", metadata_replay)
def model_info() -> None:
with Instance(args.binary, workdir, storage, open_upstream) as rf:
info = json.loads(
run_client(MODEL_INFO.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if not info["sha"]:
raise Failure("model_info returned no commit sha")
print(f" {info['sha']} with {len(info['siblings'])} siblings")
results.check("model_info answers through rustingface", model_info)
# ------------------------------------------------------------------
# Failures are loud.
# ------------------------------------------------------------------
def missing_file_raises() -> None:
with Instance(args.binary, workdir, storage, sealed_upstream) as rf:
result = json.loads(
run_client(MISSING_FILE.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if result["raised"] is None:
raise Failure("a missing file returned successfully instead of raising")
print(f" client raised {result['raised']}")
results.check("a missing file raises a typed client exception", missing_file_raises)
def unreachable_upstream_fails_loudly() -> None:
wipe_client_caches(hf_home)
with Instance(args.binary, workdir, storage, unreachable_upstream) as rf:
# Something already stored still serves.
served = json.loads(
run_client(DOWNLOAD_ONE.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if served != baseline:
raise Failure("a stored file stopped serving when upstream became unreachable")
# Something not stored fails, rather than returning an empty file.
result = json.loads(
run_client(MISSING_FILE.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
)
if result["raised"] is None:
raise Failure("an unfetchable file returned successfully")
results.check(
"an unreachable upstream serves what is stored and refuses what is not",
unreachable_upstream_fails_loudly,
)
# ------------------------------------------------------------------
# The hf CLI.
# ------------------------------------------------------------------
def hf_cli() -> None:
with Instance(args.binary, workdir, storage, sealed_upstream) as rf:
env = client_env(rf.endpoint, hf_home)
whoami = subprocess.run(
["hf", "auth", "whoami"], capture_output=True, text=True, env=env
)
# `hf auth whoami` without a token is allowed to report "not logged
# in"; what must not happen is a crash against /api/whoami-v2.
if "Traceback" in whoami.stderr:
raise Failure(f"hf auth whoami crashed:\n{whoami.stderr}")
download = subprocess.run(
["hf", "download", REPO, "config.json", "--revision", REVISION],
capture_output=True,
text=True,
env=env,
)
if download.returncode != 0:
raise Failure(f"hf download failed:\n{download.stdout}\n{download.stderr}")
if shutil.which("hf"):
results.check("the hf CLI downloads through rustingface", hf_cli)
else:
print("\n=== the hf CLI downloads through rustingface\n skipped: hf not on PATH")
# ------------------------------------------------------------------
# safetensors, the path an inference stack actually takes.
# ------------------------------------------------------------------
def safetensors_load() -> None:
script = """
import glob, json, os
from huggingface_hub import snapshot_download
from safetensors import safe_open
root = snapshot_download({repo!r}, revision={revision!r})
files = glob.glob(os.path.join(root, "**", "*.safetensors"), recursive=True)
if not files:
print(json.dumps({{"tensors": 0, "note": "repo has no safetensors"}}))
else:
total = 0
for path in files:
with safe_open(path, framework="np") as handle:
total += len(handle.keys())
print(json.dumps({{"tensors": total}}))
""".format(repo=REPO, revision=REVISION)
with Instance(args.binary, workdir, storage, sealed_upstream) as rf:
out = json.loads(run_client(script, rf.endpoint, hf_home))
print(f" {out}")
results.check("safetensors files load from a sealed snapshot", safetensors_load)
# ------------------------------------------------------------------
# Admin operations against the same bucket.
# ------------------------------------------------------------------
def admin_operations() -> None:
with Instance(args.binary, workdir, storage, sealed_upstream) as rf:
show = rf.cli("show", REPO)
manifest = json.loads(show)
if not manifest["entries"]:
raise Failure("the manifest records no entries")
verify = rf.cli("verify", REPO)
if "0 corrupt" not in verify:
raise Failure(f"verify reported problems:\n{verify}")
dry = rf.cli("gc", "--dry-run")
if "0 reclaimable" not in dry:
raise Failure(f"gc found reclaimable blobs while everything is referenced:\n{dry}")
results.check("show, verify and gc agree with what was stored", admin_operations)
if not args.keep:
shutil.rmtree(workdir, ignore_errors=True)
else:
print(f"\nwork directory kept at {workdir}")
return results.report()
if __name__ == "__main__":
sys.exit(main())
test/conformance/requirements.txt
+9
View File
@@ -0,0 +1,9 @@
# Pinned deliberately. A client upgrade is a spec-review trigger (spec §11):
# the X-Error-Code strings rustingface emits are client-internal constants, not
# a stable API, and `conformance.py` re-checks them against whatever is
# installed here on every run.
huggingface_hub==0.36.0
safetensors==0.6.2
# safetensors needs a framework to open a file; numpy is the lightest.
numpy==2.3.4
requests==2.32.5