mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 17:11:30 +00:00
feat: pull Ochre model weights from a hub into predastore (#798)
* feat: pull Ochre model weights from a hub into predastore Add `spx admin ochre weights pull`: fetch a self-host model's weights from Hugging Face directly into predastore, producing a prefix `weights stage` consumes. Streams safetensors-only artefacts (plus config/tokenizer and the shard index), refusing pickle/.bin. Pins an immutable commit SHA, writes an ochre-pull.json manifest that stage records as SourceRevision, and aborts clean (removing any written object) on a gated or mid-stream failure. Token resolves from --hf-token, HF_TOKEN, or a stored platform credential; never logged. Also add `spx admin ochre credentials set` to store an encrypted provider credential from stdin, and SourceRevision on the weights KV record + list. * fix: stream hub downloads without a whole-request timeout http.Client.Timeout bounds the entire request including the body read, so a multi-GB safetensors shard was guillotined mid-download at 30s. Bound dial and response-header latency via the transport instead, leaving the body free to stream for as long as the server keeps sending.
This commit is contained in:
Josh Siviour
@@ -1,7 +1,9 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
@@ -18,9 +20,12 @@ import (
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
"github.com/mulgadc/spinifex/spinifex/ebsprovider"
|
||||
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
|
||||
"github.com/mulgadc/spinifex/spinifex/gateway/bedrock/hfhub"
|
||||
handlers_ec2_snapshot "github.com/mulgadc/spinifex/spinifex/handlers/ec2/snapshot"
|
||||
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
|
||||
"github.com/mulgadc/spinifex/spinifex/objectstore"
|
||||
"github.com/mulgadc/spinifex/spinifex/utils"
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/nats-io/nats.go/jetstream"
|
||||
"github.com/pterm/pterm"
|
||||
"github.com/spf13/cobra"
|
||||
@@ -73,12 +78,51 @@ explicit act.`,
|
||||
Run: runOchreWeightsRemove,
|
||||
}
|
||||
|
||||
var ochreWeightsPullCmd = &cobra.Command{
|
||||
Use: "pull",
|
||||
Short: "Fetch a self-host model's weights from Hugging Face into predastore",
|
||||
Long: `pull resolves --revision to an immutable commit SHA, lists the repo's file
|
||||
tree, downloads only *.safetensors shards plus the fixed config/tokenizer
|
||||
set, and streams each into predastore under --s3-uri (or a keyed default
|
||||
when omitted). It never touches the offline 'stage' path or the weights KV
|
||||
store itself; it prints the resulting s3:// URI for 'weights stage --s3-uri'.
|
||||
|
||||
Refuses before any object lands if --model-id is not a self-host catalog
|
||||
entry, the repo has no safetensors, or the hub returns 401/403 (a gated repo
|
||||
with no usable token). A failure partway through removes every object
|
||||
already written so 'stage' never sees a half-model.`,
|
||||
Run: runOchreWeightsPull,
|
||||
}
|
||||
|
||||
// ochreCredentialsCmd groups vendor credential management for pull's
|
||||
// optional stored-token fallback (D2); it holds no Run of its own.
|
||||
var ochreCredentialsCmd = &cobra.Command{
|
||||
Use: "credentials",
|
||||
Short: "Manage stored provider credentials for self-host model pulls",
|
||||
}
|
||||
|
||||
var ochreCredentialsSetCmd = &cobra.Command{
|
||||
Use: "set",
|
||||
Short: "Store a vendor API token/licence credential",
|
||||
Long: `set stores --vendor's token, encrypted at rest under the cluster's IAM master
|
||||
key, in the bedrock-credentials KV bucket. --account defaults to the
|
||||
platform account, since self-host models are platform-level, not per-tenant.
|
||||
|
||||
The token itself is never a command-line argument: --token must be '-', and
|
||||
the secret is read from stdin, so it never appears in shell history or a
|
||||
process listing.`,
|
||||
Run: runOchreCredentialsSet,
|
||||
}
|
||||
|
||||
func init() {
|
||||
adminCmd.AddCommand(ochreCmd)
|
||||
ochreCmd.AddCommand(ochreWeightsCmd)
|
||||
ochreCmd.AddCommand(ochreCredentialsCmd)
|
||||
ochreWeightsCmd.AddCommand(ochreWeightsStageCmd)
|
||||
ochreWeightsCmd.AddCommand(ochreWeightsListCmd)
|
||||
ochreWeightsCmd.AddCommand(ochreWeightsRemoveCmd)
|
||||
ochreWeightsCmd.AddCommand(ochreWeightsPullCmd)
|
||||
ochreCredentialsCmd.AddCommand(ochreCredentialsSetCmd)
|
||||
|
||||
ochreWeightsStageCmd.Flags().String("model-id", "", "Catalog model ID to stage weights for (required)")
|
||||
ochreWeightsStageCmd.Flags().String("s3-uri", "", "predastore S3 URI holding the model's Hugging Face files, e.g. s3://bucket/prefix/ (required)")
|
||||
@@ -88,6 +132,20 @@ func init() {
|
||||
|
||||
ochreWeightsRemoveCmd.Flags().String("model-id", "", "Model ID to remove from staged weights (required)")
|
||||
_ = ochreWeightsRemoveCmd.MarkFlagRequired("model-id")
|
||||
|
||||
ochreWeightsPullCmd.Flags().String("model-id", "", "Self-host catalog model ID this pull is for (required)")
|
||||
ochreWeightsPullCmd.Flags().String("hf-repo", "", "Hugging Face repo, e.g. meta-llama/Llama-3.2-1B-Instruct (required)")
|
||||
ochreWeightsPullCmd.Flags().String("revision", "main", "Hugging Face branch, tag or commit SHA to resolve and pin")
|
||||
ochreWeightsPullCmd.Flags().String("s3-uri", "", "predastore destination prefix, e.g. s3://bucket/prefix/ (default: s3://ochre-weights/<repo>/<sha>/)")
|
||||
ochreWeightsPullCmd.Flags().String("hf-token", "", "Hugging Face token for a gated repo (falls back to HF_TOKEN, then a stored platform credential)")
|
||||
_ = ochreWeightsPullCmd.MarkFlagRequired("model-id")
|
||||
_ = ochreWeightsPullCmd.MarkFlagRequired("hf-repo")
|
||||
|
||||
ochreCredentialsSetCmd.Flags().String("vendor", "", "Vendor to store a credential for, e.g. huggingface (required)")
|
||||
ochreCredentialsSetCmd.Flags().String("account", "", "Account ID to store the credential under (default: platform account)")
|
||||
ochreCredentialsSetCmd.Flags().String("token", "", "Must be '-': reads the secret from stdin (required)")
|
||||
_ = ochreCredentialsSetCmd.MarkFlagRequired("vendor")
|
||||
_ = ochreCredentialsSetCmd.MarkFlagRequired("token")
|
||||
}
|
||||
|
||||
// requiredWeightsFiles are the fixed-name Hugging Face artefacts stage
|
||||
@@ -105,6 +163,233 @@ var requiredWeightsFiles = []string{
|
||||
// tokenizer.model under original/, not at the prefix root. Either is enough.
|
||||
var tokenizerFileNames = []string{"tokenizer.json", "tokenizer.model"}
|
||||
|
||||
// vendorHuggingFace is the CredentialStore vendor key 'ochre credentials
|
||||
// set' and pull's stored-credential fallback share for the Hugging Face
|
||||
// licence token (D2).
|
||||
const vendorHuggingFace = "huggingface"
|
||||
|
||||
// ochrePullManifestFile is the fixed name 'weights pull' writes into the
|
||||
// destination prefix and 'weights stage' reads back (D3), recording which
|
||||
// upstream commit a staged model's weights came from.
|
||||
const ochrePullManifestFile = "ochre-pull.json"
|
||||
|
||||
// ochrePullManifest is the JSON body of ochrePullManifestFile.
|
||||
type ochrePullManifest struct {
|
||||
HFRepo string `json:"hf_repo"`
|
||||
RevisionSHA string `json:"revision_sha"`
|
||||
PulledAt time.Time `json:"pulled_at"`
|
||||
}
|
||||
|
||||
// allowedPullFileNames are the fixed-name Hugging Face artefacts pull fetches
|
||||
// alongside safetensors shards -- config and tokenizer files needed to serve,
|
||||
// never weights themselves. Matched by basename so a nested path (e.g.
|
||||
// Llama's original/tokenizer.model) still qualifies.
|
||||
var allowedPullFileNames = map[string]bool{
|
||||
"config.json": true,
|
||||
"tokenizer_config.json": true,
|
||||
"tokenizer.json": true,
|
||||
"tokenizer.model": true,
|
||||
"generation_config.json": true,
|
||||
"special_tokens_map.json": true,
|
||||
"model.safetensors.index.json": true,
|
||||
}
|
||||
|
||||
// selectPullFiles filters a Hugging Face tree listing down to the
|
||||
// safetensors-only set pull downloads (D4): every *.safetensors file plus
|
||||
// the fixed config/tokenizer set. Anything else -- notably .bin/.pt pickle
|
||||
// checkpoints, which can execute code on load -- is dropped silently rather
|
||||
// than aborting the whole pull.
|
||||
func selectPullFiles(entries []hfhub.TreeEntry) []hfhub.TreeEntry {
|
||||
var selected []hfhub.TreeEntry
|
||||
for _, e := range entries {
|
||||
if e.Type != "file" {
|
||||
continue
|
||||
}
|
||||
if strings.HasSuffix(e.Path, ".safetensors") || allowedPullFileNames[path.Base(e.Path)] {
|
||||
selected = append(selected, e)
|
||||
}
|
||||
}
|
||||
return selected
|
||||
}
|
||||
|
||||
// anySafetensors reports whether entries contains at least one *.safetensors
|
||||
// file. selectPullFiles's config/tokenizer set alone is never enough --
|
||||
// D4 aborts a pull with no actual weights to serve.
|
||||
func anySafetensors(entries []hfhub.TreeEntry) bool {
|
||||
for _, e := range entries {
|
||||
if strings.HasSuffix(e.Path, ".safetensors") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// defaultPullPrefix derives D7's fallback destination when --s3-uri is
|
||||
// omitted: keyed by the exact repo and resolved commit SHA, so re-pulling
|
||||
// the same commit lands at the same, self-deduping URI.
|
||||
func defaultPullPrefix(hfRepo, sha string) string {
|
||||
return fmt.Sprintf("s3://ochre-weights/%s/%s/", strings.Trim(hfRepo, "/"), sha)
|
||||
}
|
||||
|
||||
// putPullManifest writes ochrePullManifestFile into bucket/prefix, the last
|
||||
// step of a successful pull (D3).
|
||||
func putPullManifest(ctx context.Context, store objectstore.ObjectStore, bucket, prefix string, manifest ochrePullManifest) error {
|
||||
data, err := json.Marshal(manifest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode pull manifest: %w", err)
|
||||
}
|
||||
key := prefix + ochrePullManifestFile
|
||||
if _, err := store.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), Key: aws.String(key), Body: bytes.NewReader(data)}); err != nil {
|
||||
return fmt.Errorf("put pull manifest s3://%s/%s: %w", bucket, key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// readPullManifest reads and decodes bucket/prefix's ochre-pull.json, if
|
||||
// present. A missing manifest is not an error: stage's offline path (an
|
||||
// operator-supplied prefix with no pull manifest) must keep working
|
||||
// unchanged, so ok is simply false.
|
||||
func readPullManifest(ctx context.Context, store objectstore.ObjectStore, bucket, prefix string) (ochrePullManifest, bool, error) {
|
||||
key := prefix + ochrePullManifestFile
|
||||
out, err := store.GetObject(ctx, &s3.GetObjectInput{Bucket: aws.String(bucket), Key: aws.String(key)})
|
||||
if err != nil {
|
||||
if objectstore.IsNoSuchKeyError(err) {
|
||||
return ochrePullManifest{}, false, nil
|
||||
}
|
||||
return ochrePullManifest{}, false, fmt.Errorf("get pull manifest s3://%s/%s: %w", bucket, key, err)
|
||||
}
|
||||
defer out.Body.Close()
|
||||
|
||||
var manifest ochrePullManifest
|
||||
if err := json.NewDecoder(out.Body).Decode(&manifest); err != nil {
|
||||
return ochrePullManifest{}, false, fmt.Errorf("decode pull manifest s3://%s/%s: %w", bucket, key, err)
|
||||
}
|
||||
return manifest, true, nil
|
||||
}
|
||||
|
||||
// cleanupPulledObjects best-effort deletes every key already written by a
|
||||
// pull that failed partway through (D5): a half-uploaded prefix must never
|
||||
// be left for 'stage' to mistake for a complete model. A delete failure is
|
||||
// logged, not returned -- the original pull error is what the operator needs.
|
||||
func cleanupPulledObjects(ctx context.Context, store objectstore.ObjectStore, bucket string, keys []string) {
|
||||
for _, key := range keys {
|
||||
if _, err := store.DeleteObject(ctx, &s3.DeleteObjectInput{Bucket: aws.String(bucket), Key: aws.String(key)}); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "warning: could not clean up s3://%s/%s after a failed pull: %v\n", bucket, key, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pullOneFile downloads hfPath from the hub to a local temp file and PUTs it
|
||||
// to bucket/key. A temp file is required, not a direct HTTP-to-S3 pipe:
|
||||
// predastore's PutObject signs the request over an io.ReadSeeker, which an
|
||||
// HTTP response body is not. tmpDir is cleaned up by the caller.
|
||||
func pullOneFile(ctx context.Context, hf *hfhub.Client, store objectstore.ObjectStore, hfRepo, sha, hfPath, bucket, key, tmpDir string) error {
|
||||
body, _, err := hf.DownloadFile(ctx, hfRepo, sha, hfPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer body.Close()
|
||||
|
||||
tmpFile, err := os.CreateTemp(tmpDir, "part-*")
|
||||
if err != nil {
|
||||
return fmt.Errorf("create temp file: %w", err)
|
||||
}
|
||||
defer os.Remove(tmpFile.Name())
|
||||
defer tmpFile.Close()
|
||||
|
||||
if _, err := io.Copy(tmpFile, body); err != nil {
|
||||
return fmt.Errorf("download to temp file: %w", err)
|
||||
}
|
||||
if _, err := tmpFile.Seek(0, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("seek temp file: %w", err)
|
||||
}
|
||||
|
||||
if _, err := store.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), Key: aws.String(key), Body: tmpFile}); err != nil {
|
||||
return fmt.Errorf("put s3://%s/%s: %w", bucket, key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// pullFilesToPrefix downloads and uploads every selected file, one at a time
|
||||
// so at most one file is ever held on local disk (D6). Keys are flattened to
|
||||
// each file's basename to match the flat layout 'stage' validates -- e.g.
|
||||
// Llama's original/tokenizer.model lands at <prefix>tokenizer.model. It
|
||||
// returns the keys written so far even on error, so the caller can clean up
|
||||
// a partial prefix (D5).
|
||||
func pullFilesToPrefix(ctx context.Context, hf *hfhub.Client, store objectstore.ObjectStore, hfRepo, sha, bucket, prefix string, files []hfhub.TreeEntry) ([]string, error) {
|
||||
tmpDir, err := os.MkdirTemp("", "spinifex-ochre-pull-*")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create temp dir: %w", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
var uploaded []string
|
||||
for _, f := range files {
|
||||
key := prefix + path.Base(f.Path)
|
||||
if err := pullOneFile(ctx, hf, store, hfRepo, sha, f.Path, bucket, key, tmpDir); err != nil {
|
||||
return uploaded, fmt.Errorf("pull %s: %w", f.Path, err)
|
||||
}
|
||||
uploaded = append(uploaded, key)
|
||||
}
|
||||
return uploaded, nil
|
||||
}
|
||||
|
||||
// runPullWeights holds 'ochre weights pull' decision and side-effect logic:
|
||||
// catalog validation, ref resolution to an immutable commit SHA (D3), tree
|
||||
// listing and safetensors-only filtering (D4), streaming each selected file
|
||||
// into predastore (D6), and finally the ochre-pull.json manifest. It never
|
||||
// touches the weights KV store; 'stage' consumes the printed s3:// URI
|
||||
// separately. A failure removes every object already written (D5).
|
||||
func runPullWeights(ctx context.Context, hf *hfhub.Client, store objectstore.ObjectStore, modelID, hfRepo, revision, s3URIFlag string) (string, error) {
|
||||
if _, found, selfHost := gateway_bedrock.LookupServingSpec(modelID); !found || !selfHost {
|
||||
if !found {
|
||||
return "", fmt.Errorf("unknown model ID %q: not present in the Ochre catalog", modelID)
|
||||
}
|
||||
return "", fmt.Errorf("%q is a provider-served model, not self-host; weights pull does not apply", modelID)
|
||||
}
|
||||
|
||||
fmt.Printf("Resolving %s@%s ...\n", hfRepo, revision)
|
||||
sha, err := hf.ResolveRevision(ctx, hfRepo, revision)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
fmt.Printf("Listing files at %s@%s ...\n", hfRepo, sha)
|
||||
tree, err := hf.ListTree(ctx, hfRepo, sha)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
selected := selectPullFiles(tree)
|
||||
if !anySafetensors(selected) {
|
||||
return "", fmt.Errorf("%s@%s has no *.safetensors files; refusing to pull pickle-format (.bin) weights", hfRepo, sha)
|
||||
}
|
||||
|
||||
s3URI := s3URIFlag
|
||||
if s3URI == "" {
|
||||
s3URI = defaultPullPrefix(hfRepo, sha)
|
||||
}
|
||||
bucket, prefix, err := parseWeightsS3URI(s3URI)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
fmt.Printf("Pulling %d file(s) into s3://%s/%s ...\n", len(selected), bucket, prefix)
|
||||
uploaded, err := pullFilesToPrefix(ctx, hf, store, hfRepo, sha, bucket, prefix, selected)
|
||||
if err != nil {
|
||||
cleanupPulledObjects(ctx, store, bucket, uploaded)
|
||||
return "", err
|
||||
}
|
||||
|
||||
manifest := ochrePullManifest{HFRepo: hfRepo, RevisionSHA: sha, PulledAt: time.Now().UTC()}
|
||||
if err := putPullManifest(ctx, store, bucket, prefix, manifest); err != nil {
|
||||
cleanupPulledObjects(ctx, store, bucket, uploaded)
|
||||
return "", err
|
||||
}
|
||||
|
||||
return fmt.Sprintf("s3://%s/%s", bucket, prefix), nil
|
||||
}
|
||||
|
||||
// parseWeightsS3URI splits an s3://bucket/prefix URI into its bucket and
|
||||
// prefix. The prefix is normalised to end with '/' so downstream listing and
|
||||
// validation scope to the directory rather than any key merely sharing the
|
||||
@@ -386,6 +671,16 @@ func runStageWeights(ctx context.Context, store objectstore.ObjectStore, weights
|
||||
return "", err
|
||||
}
|
||||
|
||||
// A pull manifest is optional: absent means an operator-supplied prefix
|
||||
// from the offline path, which stage has always supported and must keep
|
||||
// supporting unchanged.
|
||||
var sourceRevision string
|
||||
if manifest, ok, err := readPullManifest(ctx, store, bucket, prefix); err != nil {
|
||||
return "", err
|
||||
} else if ok {
|
||||
sourceRevision = manifest.RevisionSHA
|
||||
}
|
||||
|
||||
tmpDir, err := os.MkdirTemp(tmpDirFlag, "spinifex-weights-tmp-*")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("could not create temp dir: %w", err)
|
||||
@@ -408,7 +703,7 @@ func runStageWeights(ctx context.Context, store objectstore.ObjectStore, weights
|
||||
return "", err
|
||||
}
|
||||
|
||||
if err := weightsStore.PutWeights(ctx, modelID, sourceURI, snapshotID); err != nil {
|
||||
if err := weightsStore.PutWeightsWithRevision(ctx, modelID, sourceURI, snapshotID, sourceRevision); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
@@ -563,7 +858,7 @@ func listWeightsOutput(ctx context.Context, weightsStore *gateway_bedrock.Weight
|
||||
return "No models staged.", nil
|
||||
}
|
||||
|
||||
tableData := pterm.TableData{{"MODEL ID", "SOURCE URI", "SNAPSHOT ID", "STATUS"}}
|
||||
tableData := pterm.TableData{{"MODEL ID", "SOURCE URI", "REVISION", "SNAPSHOT ID", "STATUS"}}
|
||||
for _, e := range entries {
|
||||
live, err := checkSnapshotLive(ctx, e.SnapshotID)
|
||||
if err != nil {
|
||||
@@ -573,7 +868,11 @@ func listWeightsOutput(ctx context.Context, weightsStore *gateway_bedrock.Weight
|
||||
if !live {
|
||||
status = weightsStatusDangling
|
||||
}
|
||||
tableData = append(tableData, []string{e.ModelID, e.SourceURI, e.SnapshotID, status})
|
||||
revision := e.SourceRevision
|
||||
if revision == "" {
|
||||
revision = "-"
|
||||
}
|
||||
tableData = append(tableData, []string{e.ModelID, e.SourceURI, revision, e.SnapshotID, status})
|
||||
}
|
||||
return pterm.DefaultTable.WithHasHeader().WithLeftAlignment().WithData(tableData).Srender()
|
||||
}
|
||||
@@ -656,6 +955,141 @@ func runOchreWeightsRemove(cmd *cobra.Command, _ []string) {
|
||||
fmt.Printf("✅ Removed staged-weights entry for %s (snapshot %s and source objects untouched).\n", modelID, entry.SnapshotID)
|
||||
}
|
||||
|
||||
// resolveHFToken applies pull's token resolution order (D2): the --hf-token
|
||||
// flag, then HF_TOKEN env (the must-have path, self-sufficient on its own),
|
||||
// then a stored platform credential for vendor "huggingface". The stored
|
||||
// credential step is a best-effort fallback: a cluster with no IAM master
|
||||
// key provisioned yet simply yields no token rather than failing the pull.
|
||||
func resolveHFToken(ctx context.Context, cmd *cobra.Command, cfg *config.ClusterConfig, nc *nats.Conn) string {
|
||||
if token, _ := cmd.Flags().GetString("hf-token"); token != "" {
|
||||
return token
|
||||
}
|
||||
if token := os.Getenv("HF_TOKEN"); token != "" {
|
||||
return token
|
||||
}
|
||||
|
||||
masterKeyPath := filepath.Join(cfg.NodeBaseDir(), "config", "master.key")
|
||||
masterKey, err := handlers_iam.LoadMasterKey(masterKeyPath)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
js, err := jetstream.New(nc)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
credStore := gateway_bedrock.NewCredentialStore(js, masterKey, len(cfg.Nodes), nil)
|
||||
token, ok, err := credStore.Resolve(ctx, utils.GlobalAccountID, vendorHuggingFace)
|
||||
if err != nil || !ok {
|
||||
return ""
|
||||
}
|
||||
return token
|
||||
}
|
||||
|
||||
func runOchreWeightsPull(cmd *cobra.Command, _ []string) {
|
||||
modelID, _ := cmd.Flags().GetString("model-id")
|
||||
hfRepo, _ := cmd.Flags().GetString("hf-repo")
|
||||
revision, _ := cmd.Flags().GetString("revision")
|
||||
s3URI, _ := cmd.Flags().GetString("s3-uri")
|
||||
|
||||
appConfig, nc, err := loadConfigAndConnectFn()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
defer nc.Close()
|
||||
node := appConfig.Nodes[appConfig.Node]
|
||||
|
||||
ctx := context.Background()
|
||||
token := resolveHFToken(ctx, cmd, appConfig, nc)
|
||||
|
||||
hf := hfhub.NewClient(token)
|
||||
store := objectstore.NewS3ObjectStoreFromConfig(node.Predastore.Host, node.Predastore.Region, node.Predastore.AccessKey, node.Predastore.SecretKey)
|
||||
|
||||
finalURI, err := runPullWeights(ctx, hf, store, modelID, hfRepo, revision, s3URI)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "%v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
fmt.Printf("✅ Pulled %s@%s into %s\n", hfRepo, revision, finalURI)
|
||||
fmt.Println(finalURI)
|
||||
}
|
||||
|
||||
// readTokenFromStdin reads a credential from r, trimming exactly one
|
||||
// trailing newline (matching 'printf "%s" "$TOKEN" | cmd --token -' or a
|
||||
// piped 'echo'). The token is never logged.
|
||||
func readTokenFromStdin(r io.Reader) (string, error) {
|
||||
data, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read token from stdin: %w", err)
|
||||
}
|
||||
token := strings.TrimSuffix(strings.TrimSuffix(string(data), "\n"), "\r")
|
||||
if token == "" {
|
||||
return "", fmt.Errorf("no token read from stdin")
|
||||
}
|
||||
return token, nil
|
||||
}
|
||||
|
||||
// ochreCredentialsStore connects to the cluster, loads the node's IAM master
|
||||
// key from disk -- the same on-disk key initIAMServiceFromConfig uses -- and
|
||||
// returns a CredentialStore ready for PutCredential/Resolve.
|
||||
func ochreCredentialsStore() (*gateway_bedrock.CredentialStore, func(), error) {
|
||||
cfg, nc, err := loadConfigAndConnectFn()
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("connect to cluster: %w", err)
|
||||
}
|
||||
masterKeyPath := filepath.Join(cfg.NodeBaseDir(), "config", "master.key")
|
||||
masterKey, err := handlers_iam.LoadMasterKey(masterKeyPath)
|
||||
if err != nil {
|
||||
nc.Close()
|
||||
return nil, nil, fmt.Errorf("load master key: %w", err)
|
||||
}
|
||||
js, err := jetstream.New(nc)
|
||||
if err != nil {
|
||||
nc.Close()
|
||||
return nil, nil, fmt.Errorf("jetstream context: %w", err)
|
||||
}
|
||||
return gateway_bedrock.NewCredentialStore(js, masterKey, len(cfg.Nodes), nil), func() { nc.Close() }, nil
|
||||
}
|
||||
|
||||
func runOchreCredentialsSet(cmd *cobra.Command, _ []string) {
|
||||
vendor, _ := cmd.Flags().GetString("vendor")
|
||||
accountID, _ := cmd.Flags().GetString("account")
|
||||
tokenFlag, _ := cmd.Flags().GetString("token")
|
||||
|
||||
if tokenFlag != "-" {
|
||||
fmt.Fprintln(os.Stderr, "Error: --token must be '-' (the secret is read from stdin, never a command-line argument)")
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
if accountID == "" {
|
||||
accountID = utils.GlobalAccountID
|
||||
}
|
||||
|
||||
token, err := readTokenFromStdin(os.Stdin)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
|
||||
store, cleanup, err := ochreCredentialsStore()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
if err := store.PutCredential(context.Background(), accountID, vendor, token); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
fmt.Printf("✅ Stored %s credential for account %s\n", vendor, accountID)
|
||||
}
|
||||
|
||||
var adminOchreAccessCmd = &cobra.Command{
|
||||
Use: "access",
|
||||
Short: "Manage per-account model access grants",
|
||||
|
||||
@@ -0,0 +1,578 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go/service/s3"
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
|
||||
"github.com/mulgadc/spinifex/spinifex/gateway/bedrock/hfhub"
|
||||
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
|
||||
"github.com/mulgadc/spinifex/spinifex/objectstore"
|
||||
"github.com/mulgadc/spinifex/spinifex/testutil"
|
||||
"github.com/mulgadc/spinifex/spinifex/utils"
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// --- selectPullFiles / anySafetensors (D4) ---
|
||||
|
||||
// TestSelectPullFiles_SafetensorsAndConfigOnly covers D4's filter: every
|
||||
// *.safetensors file, the fixed config/tokenizer set, and the shard index are
|
||||
// kept (matched by basename, so a nested path like Llama's
|
||||
// original/tokenizer.model still qualifies); a pickle .bin checkpoint and
|
||||
// directory entries are dropped.
|
||||
func TestSelectPullFiles_SafetensorsAndConfigOnly(t *testing.T) {
|
||||
entries := []hfhub.TreeEntry{
|
||||
{Type: "file", Path: "config.json"},
|
||||
{Type: "file", Path: "model-00001-of-00002.safetensors"},
|
||||
{Type: "file", Path: "model-00002-of-00002.safetensors"},
|
||||
{Type: "file", Path: "tokenizer_config.json"},
|
||||
{Type: "file", Path: "tokenizer.json"},
|
||||
{Type: "file", Path: "pytorch_model.bin"},
|
||||
{Type: "file", Path: "model.safetensors.index.json"},
|
||||
{Type: "directory", Path: "original"},
|
||||
{Type: "file", Path: "original/tokenizer.model"},
|
||||
}
|
||||
|
||||
selected := selectPullFiles(entries)
|
||||
|
||||
var paths []string
|
||||
for _, e := range selected {
|
||||
paths = append(paths, e.Path)
|
||||
}
|
||||
assert.ElementsMatch(t, []string{
|
||||
"config.json",
|
||||
"model-00001-of-00002.safetensors",
|
||||
"model-00002-of-00002.safetensors",
|
||||
"tokenizer_config.json",
|
||||
"tokenizer.json",
|
||||
"model.safetensors.index.json",
|
||||
"original/tokenizer.model",
|
||||
}, paths)
|
||||
}
|
||||
|
||||
// TestAnySafetensors_TrueWhenPresent covers the pass case: at least one
|
||||
// selected file is a safetensors shard.
|
||||
func TestAnySafetensors_TrueWhenPresent(t *testing.T) {
|
||||
assert.True(t, anySafetensors([]hfhub.TreeEntry{{Type: "file", Path: "model.safetensors"}}))
|
||||
}
|
||||
|
||||
// TestAnySafetensors_FalseWhenOnlyPickleAndConfig covers D4's abort trigger:
|
||||
// a repo whose only weights are pickle .bin, even with config/tokenizer
|
||||
// files also selected, has nothing pull will actually download as weights.
|
||||
func TestAnySafetensors_FalseWhenOnlyPickleAndConfig(t *testing.T) {
|
||||
selected := selectPullFiles([]hfhub.TreeEntry{
|
||||
{Type: "file", Path: "config.json"},
|
||||
{Type: "file", Path: "pytorch_model.bin"},
|
||||
})
|
||||
assert.False(t, anySafetensors(selected))
|
||||
}
|
||||
|
||||
// --- defaultPullPrefix (D7) ---
|
||||
|
||||
// TestDefaultPullPrefix_KeyedByRepoAndSHA covers the fallback destination
|
||||
// when --s3-uri is omitted: immutable and self-deduping, keyed by the exact
|
||||
// repo and resolved commit SHA.
|
||||
func TestDefaultPullPrefix_KeyedByRepoAndSHA(t *testing.T) {
|
||||
got := defaultPullPrefix("meta-llama/Llama-3.2-1B-Instruct", "abc123def456")
|
||||
assert.Equal(t, "s3://ochre-weights/meta-llama/Llama-3.2-1B-Instruct/abc123def456/", got)
|
||||
}
|
||||
|
||||
// --- ochre-pull.json manifest (D3) ---
|
||||
|
||||
// TestPullManifest_RoundTrips covers put then read of the manifest 'pull'
|
||||
// writes and 'stage' reads back.
|
||||
func TestPullManifest_RoundTrips(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
ctx := context.Background()
|
||||
want := ochrePullManifest{HFRepo: "meta-llama/Llama-3.2-1B-Instruct", RevisionSHA: "abc123", PulledAt: time.Now().UTC().Truncate(time.Second)}
|
||||
|
||||
require.NoError(t, putPullManifest(ctx, store, "ochre-weights", "meta-llama/Llama-3.2-1B-Instruct/abc123/", want))
|
||||
|
||||
got, ok, err := readPullManifest(ctx, store, "ochre-weights", "meta-llama/Llama-3.2-1B-Instruct/abc123/")
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, want.HFRepo, got.HFRepo)
|
||||
assert.Equal(t, want.RevisionSHA, got.RevisionSHA)
|
||||
assert.WithinDuration(t, want.PulledAt, got.PulledAt, time.Second)
|
||||
}
|
||||
|
||||
// TestReadPullManifest_MissingIsNotError covers stage's offline path: a
|
||||
// prefix with no manifest must report ok=false, not an error.
|
||||
func TestReadPullManifest_MissingIsNotError(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
_, ok, err := readPullManifest(context.Background(), store, "ochre-weights", "some/prefix/")
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok)
|
||||
}
|
||||
|
||||
// --- runStageWeights x manifest integration ---
|
||||
|
||||
// TestRunStageWeights_RecordsSourceRevisionFromManifest covers D3's stage
|
||||
// side: when the source prefix carries a pull manifest, stage must record
|
||||
// its commit SHA against the model.
|
||||
func TestRunStageWeights_RecordsSourceRevisionFromManifest(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
seedCompleteWeightsPrefix(t, store, "models", "llama-3.2-1b/")
|
||||
require.NoError(t, putPullManifest(ctx, store, "models", "llama-3.2-1b/", ochrePullManifest{
|
||||
HFRepo: "meta-llama/Llama-3.2-1B-Instruct", RevisionSHA: "abc123def456", PulledAt: time.Now().UTC(),
|
||||
}))
|
||||
|
||||
materialize := func(context.Context, string, int64) (string, error) { return "snap-0001", nil }
|
||||
_, err := runStageWeights(ctx, store, weightsStore, t.TempDir(), selfHostModelID, "s3://models/llama-3.2-1b/", materialize, explodingSnapshotChecker(t))
|
||||
require.NoError(t, err)
|
||||
|
||||
entry, ok, err := weightsStore.GetWeights(ctx, selfHostModelID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "abc123def456", entry.SourceRevision)
|
||||
}
|
||||
|
||||
// TestRunStageWeights_NoManifestLeavesSourceRevisionEmpty covers the offline
|
||||
// path explicitly: an operator-supplied prefix with no pull manifest stages
|
||||
// fine and records an empty SourceRevision, not an error or a fabricated one.
|
||||
func TestRunStageWeights_NoManifestLeavesSourceRevisionEmpty(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
seedCompleteWeightsPrefix(t, store, "models", "llama-3.2-1b/")
|
||||
|
||||
materialize := func(context.Context, string, int64) (string, error) { return "snap-0001", nil }
|
||||
_, err := runStageWeights(ctx, store, weightsStore, t.TempDir(), selfHostModelID, "s3://models/llama-3.2-1b/", materialize, explodingSnapshotChecker(t))
|
||||
require.NoError(t, err)
|
||||
|
||||
entry, ok, err := weightsStore.GetWeights(ctx, selfHostModelID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Empty(t, entry.SourceRevision)
|
||||
}
|
||||
|
||||
// --- fake Hugging Face server ---
|
||||
|
||||
type fakeHFFile struct {
|
||||
path string
|
||||
content string
|
||||
}
|
||||
|
||||
// newFakeHFServer serves a Hugging Face-shaped API for repo@revision
|
||||
// resolving to sha: revision resolution, a recursive tree listing built from
|
||||
// files, and each file's content at /resolve/<sha>/<path>.
|
||||
func newFakeHFServer(t *testing.T, repo, revision, sha string, files []fakeHFFile) *httptest.Server {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc(fmt.Sprintf("/api/models/%s/revision/%s", repo, revision), func(w http.ResponseWriter, _ *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode(map[string]string{"sha": sha})
|
||||
})
|
||||
mux.HandleFunc(fmt.Sprintf("/api/models/%s/tree/%s", repo, sha), func(w http.ResponseWriter, _ *http.Request) {
|
||||
entries := make([]hfhub.TreeEntry, len(files))
|
||||
for i, f := range files {
|
||||
entries[i] = hfhub.TreeEntry{Type: "file", Path: f.path, Size: int64(len(f.content))}
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(entries)
|
||||
})
|
||||
for _, f := range files {
|
||||
mux.HandleFunc(fmt.Sprintf("/%s/resolve/%s/%s", repo, sha, f.path), func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = io.WriteString(w, f.content)
|
||||
})
|
||||
}
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
// gatedHFServer answers every request with status, simulating a gated repo
|
||||
// with no usable token (D5).
|
||||
func gatedHFServer(t *testing.T, status int) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(status)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
// deadHFClient is a Client pointed at an address nothing listens on, so a
|
||||
// path that must never touch the network fails fast instead of reaching the
|
||||
// real Hugging Face hub.
|
||||
func deadHFClient() *hfhub.Client {
|
||||
return &hfhub.Client{BaseURL: "http://127.0.0.1:1", HTTP: &http.Client{Timeout: time.Second}}
|
||||
}
|
||||
|
||||
const (
|
||||
testHFRepo = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
testHFSHA = "abc123def456"
|
||||
)
|
||||
|
||||
// --- runPullWeights (core logic) ---
|
||||
|
||||
// TestRunPullWeights_UnknownModelRefused covers the catalog gate: an
|
||||
// unrecognised model ID must refuse before any Hugging Face or S3 work.
|
||||
func TestRunPullWeights_UnknownModelRefused(t *testing.T) {
|
||||
_, err := runPullWeights(context.Background(), deadHFClient(), explodingObjectStore{t: t}, "not-a-real-model", testHFRepo, "main", "")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "not present in the Ochre catalog")
|
||||
}
|
||||
|
||||
// TestRunPullWeights_ProviderModelRefused covers the self-host gate: a
|
||||
// provider-served model must refuse before any network work, same as stage.
|
||||
func TestRunPullWeights_ProviderModelRefused(t *testing.T) {
|
||||
_, err := runPullWeights(context.Background(), deadHFClient(), explodingObjectStore{t: t}, providerModelID, testHFRepo, "main", "")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "not self-host")
|
||||
}
|
||||
|
||||
// TestRunPullWeights_ResolvesUploadsAndWritesManifest is the happy path: it
|
||||
// resolves main to a pinned SHA, uploads only the safetensors + config set
|
||||
// flattened under the default prefix (D7), skips the pickle checkpoint, and
|
||||
// writes the ochre-pull.json manifest (D3).
|
||||
func TestRunPullWeights_ResolvesUploadsAndWritesManifest(t *testing.T) {
|
||||
files := []fakeHFFile{
|
||||
{"config.json", "{}"},
|
||||
{"tokenizer_config.json", "{}"},
|
||||
{"tokenizer.json", "{}"},
|
||||
{"model.safetensors", "fake-weights-bytes"},
|
||||
{"pytorch_model.bin", "must-not-be-uploaded"},
|
||||
}
|
||||
srv := newFakeHFServer(t, testHFRepo, "main", testHFSHA, files)
|
||||
hf := &hfhub.Client{BaseURL: srv.URL, HTTP: srv.Client()}
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
|
||||
finalURI, err := runPullWeights(context.Background(), hf, store, selfHostModelID, testHFRepo, "main", "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, defaultPullPrefix(testHFRepo, testHFSHA), finalURI)
|
||||
|
||||
bucket, prefix, err := parseWeightsS3URI(finalURI)
|
||||
require.NoError(t, err)
|
||||
|
||||
out, err := store.GetObject(context.Background(), &s3.GetObjectInput{Bucket: strPtr(bucket), Key: strPtr(prefix + "model.safetensors")})
|
||||
require.NoError(t, err)
|
||||
got, _ := io.ReadAll(out.Body)
|
||||
assert.Equal(t, "fake-weights-bytes", string(got))
|
||||
|
||||
_, err = store.GetObject(context.Background(), &s3.GetObjectInput{Bucket: strPtr(bucket), Key: strPtr(prefix + "pytorch_model.bin")})
|
||||
assert.True(t, objectstore.IsNoSuchKeyError(err), "pickle checkpoint must never be uploaded")
|
||||
|
||||
manifest, ok, err := readPullManifest(context.Background(), store, bucket, prefix)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, testHFRepo, manifest.HFRepo)
|
||||
assert.Equal(t, testHFSHA, manifest.RevisionSHA)
|
||||
}
|
||||
|
||||
// TestRunPullWeights_ExplicitS3URIOverridesDefault covers D7's primary path:
|
||||
// an operator-named --s3-uri is used verbatim (normalised) instead of the
|
||||
// keyed default.
|
||||
func TestRunPullWeights_ExplicitS3URIOverridesDefault(t *testing.T) {
|
||||
files := []fakeHFFile{{"config.json", "{}"}, {"tokenizer_config.json", "{}"}, {"tokenizer.json", "{}"}, {"model.safetensors", "x"}}
|
||||
srv := newFakeHFServer(t, testHFRepo, "main", testHFSHA, files)
|
||||
hf := &hfhub.Client{BaseURL: srv.URL, HTTP: srv.Client()}
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
|
||||
finalURI, err := runPullWeights(context.Background(), hf, store, selfHostModelID, testHFRepo, "main", "s3://custom-bucket/my-prefix")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "s3://custom-bucket/my-prefix/", finalURI)
|
||||
}
|
||||
|
||||
// TestRunPullWeights_NoSafetensorsAborts covers D4's abort case: a repo with
|
||||
// only pickle checkpoints must be refused before any object lands.
|
||||
func TestRunPullWeights_NoSafetensorsAborts(t *testing.T) {
|
||||
files := []fakeHFFile{{"config.json", "{}"}, {"pytorch_model.bin", "x"}}
|
||||
srv := newFakeHFServer(t, testHFRepo, "main", testHFSHA, files)
|
||||
hf := &hfhub.Client{BaseURL: srv.URL, HTTP: srv.Client()}
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
|
||||
_, err := runPullWeights(context.Background(), hf, store, selfHostModelID, testHFRepo, "main", "")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "no *.safetensors")
|
||||
assert.Equal(t, 0, store.Count(), "a refused pull must land nothing")
|
||||
}
|
||||
|
||||
// TestRunPullWeights_GatedRepoAbortsCleanBeforeAnyUpload covers D5: a
|
||||
// 401/403 from the hub must fail before any object lands, with a clear
|
||||
// awserrors-coded licence/credential error.
|
||||
func TestRunPullWeights_GatedRepoAbortsCleanBeforeAnyUpload(t *testing.T) {
|
||||
for _, status := range []int{http.StatusUnauthorized, http.StatusForbidden} {
|
||||
srv := gatedHFServer(t, status)
|
||||
hf := &hfhub.Client{BaseURL: srv.URL, HTTP: srv.Client()}
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
|
||||
_, err := runPullWeights(context.Background(), hf, store, selfHostModelID, testHFRepo, "main", "")
|
||||
require.Error(t, err)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorAccessDeniedException), "status %d: got %v", status, err)
|
||||
assert.Equal(t, 0, store.Count(), "a gated pull must land nothing")
|
||||
}
|
||||
}
|
||||
|
||||
// failAfterNPutsStore wraps a MemoryObjectStore, succeeding the first n
|
||||
// PutObject calls and failing every one after, so a mid-pull failure and its
|
||||
// cleanup (D5) are exercised deterministically.
|
||||
type failAfterNPutsStore struct {
|
||||
*objectstore.MemoryObjectStore
|
||||
|
||||
n int
|
||||
calls int
|
||||
}
|
||||
|
||||
func (f *failAfterNPutsStore) PutObject(ctx context.Context, input *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
|
||||
f.calls++
|
||||
if f.calls > f.n {
|
||||
return nil, errors.New("simulated upload failure")
|
||||
}
|
||||
return f.MemoryObjectStore.PutObject(ctx, input)
|
||||
}
|
||||
|
||||
var _ objectstore.ObjectStore = (*failAfterNPutsStore)(nil)
|
||||
|
||||
// TestRunPullWeights_MidStreamFailureCleansUpPartialPrefix covers D5's core
|
||||
// promise: a failure partway through the upload loop must remove every
|
||||
// object already written, so 'stage' never sees a half-model.
|
||||
func TestRunPullWeights_MidStreamFailureCleansUpPartialPrefix(t *testing.T) {
|
||||
files := []fakeHFFile{
|
||||
{"config.json", "{}"},
|
||||
{"tokenizer_config.json", "{}"},
|
||||
{"tokenizer.json", "{}"},
|
||||
{"model.safetensors", "weights"},
|
||||
}
|
||||
srv := newFakeHFServer(t, testHFRepo, "main", testHFSHA, files)
|
||||
hf := &hfhub.Client{BaseURL: srv.URL, HTTP: srv.Client()}
|
||||
store := &failAfterNPutsStore{MemoryObjectStore: objectstore.NewMemoryObjectStore(), n: 2}
|
||||
|
||||
_, err := runPullWeights(context.Background(), hf, store, selfHostModelID, testHFRepo, "main", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, 0, store.Count(), "objects uploaded before the failure must be cleaned up")
|
||||
}
|
||||
|
||||
// --- resolveHFToken (D2) ---
|
||||
|
||||
func newOchreWeightsPullTestCmd(t *testing.T, modelID, hfRepo, revision, s3URI, hfToken string) *cobra.Command {
|
||||
t.Helper()
|
||||
require.NoError(t, ochreWeightsPullCmd.Flags().Set("model-id", modelID))
|
||||
require.NoError(t, ochreWeightsPullCmd.Flags().Set("hf-repo", hfRepo))
|
||||
require.NoError(t, ochreWeightsPullCmd.Flags().Set("revision", revision))
|
||||
require.NoError(t, ochreWeightsPullCmd.Flags().Set("s3-uri", s3URI))
|
||||
require.NoError(t, ochreWeightsPullCmd.Flags().Set("hf-token", hfToken))
|
||||
return ochreWeightsPullCmd
|
||||
}
|
||||
|
||||
// TestResolveHFToken_PrefersFlagOverEnvAndStored covers the top of D2's
|
||||
// resolution order.
|
||||
func TestResolveHFToken_PrefersFlagOverEnvAndStored(t *testing.T) {
|
||||
t.Setenv("HF_TOKEN", "hf_env_token")
|
||||
cmd := newOchreWeightsPullTestCmd(t, selfHostModelID, testHFRepo, "main", "", "hf_flag_token")
|
||||
cfg := &config.ClusterConfig{Node: "node1", Nodes: map[string]config.Config{"node1": {BaseDir: t.TempDir()}}}
|
||||
|
||||
got := resolveHFToken(context.Background(), cmd, cfg, nil)
|
||||
assert.Equal(t, "hf_flag_token", got)
|
||||
}
|
||||
|
||||
// TestResolveHFToken_FallsBackToEnv covers the second link: no flag, HF_TOKEN set.
|
||||
func TestResolveHFToken_FallsBackToEnv(t *testing.T) {
|
||||
t.Setenv("HF_TOKEN", "hf_env_token")
|
||||
cmd := newOchreWeightsPullTestCmd(t, selfHostModelID, testHFRepo, "main", "", "")
|
||||
cfg := &config.ClusterConfig{Node: "node1", Nodes: map[string]config.Config{"node1": {BaseDir: t.TempDir()}}}
|
||||
|
||||
got := resolveHFToken(context.Background(), cmd, cfg, nil)
|
||||
assert.Equal(t, "hf_env_token", got)
|
||||
}
|
||||
|
||||
// TestResolveHFToken_NoTokenAnywhereReturnsEmpty covers the terminal case:
|
||||
// no flag, no env, and no master key on disk to even attempt the stored
|
||||
// credential -- must return empty, not error or panic, so a public repo
|
||||
// pull still works standalone.
|
||||
func TestResolveHFToken_NoTokenAnywhereReturnsEmpty(t *testing.T) {
|
||||
t.Setenv("HF_TOKEN", "")
|
||||
cmd := newOchreWeightsPullTestCmd(t, selfHostModelID, testHFRepo, "main", "", "")
|
||||
cfg := &config.ClusterConfig{Node: "node1", Nodes: map[string]config.Config{"node1": {BaseDir: t.TempDir()}}}
|
||||
|
||||
got := resolveHFToken(context.Background(), cmd, cfg, nil)
|
||||
assert.Empty(t, got)
|
||||
}
|
||||
|
||||
// TestResolveHFToken_FallsBackToStoredPlatformCredential covers D2's
|
||||
// third link: with no flag or env, a credential stored via 'ochre
|
||||
// credentials set' (platform account, vendor "huggingface") is used.
|
||||
func TestResolveHFToken_FallsBackToStoredPlatformCredential(t *testing.T) {
|
||||
t.Setenv("HF_TOKEN", "")
|
||||
_, nc, js := testutil.StartTestJetStream(t)
|
||||
|
||||
baseDir := t.TempDir()
|
||||
masterKey := bytes32Key()
|
||||
require.NoError(t, os.MkdirAll(filepath.Join(baseDir, "config"), 0750))
|
||||
require.NoError(t, handlers_iam.SaveMasterKey(filepath.Join(baseDir, "config", "master.key"), masterKey))
|
||||
|
||||
credStore := gateway_bedrock.NewCredentialStore(js, masterKey, 1, nil)
|
||||
require.NoError(t, credStore.PutCredential(context.Background(), utils.GlobalAccountID, vendorHuggingFace, "hf_stored_token"))
|
||||
|
||||
cmd := newOchreWeightsPullTestCmd(t, selfHostModelID, testHFRepo, "main", "", "")
|
||||
cfg := &config.ClusterConfig{Node: "node1", Nodes: map[string]config.Config{"node1": {BaseDir: baseDir}}}
|
||||
|
||||
got := resolveHFToken(context.Background(), cmd, cfg, nc)
|
||||
assert.Equal(t, "hf_stored_token", got)
|
||||
}
|
||||
|
||||
func bytes32Key() []byte {
|
||||
key := make([]byte, 32)
|
||||
for i := range key {
|
||||
key[i] = byte(i + 1)
|
||||
}
|
||||
return key
|
||||
}
|
||||
|
||||
func strPtr(s string) *string { return &s }
|
||||
|
||||
// --- runOchreWeightsPull wrapper ---
|
||||
|
||||
// TestRunOchreWeightsPull_ConnectFailureExits1 mirrors stage's own coverage:
|
||||
// pull cannot even reach the catalog check without a cluster connection.
|
||||
func TestRunOchreWeightsPull_ConnectFailureExits1(t *testing.T) {
|
||||
stubConnect(t, nil, nil, errors.New("dial failed"))
|
||||
|
||||
cmd := newOchreWeightsPullTestCmd(t, selfHostModelID, testHFRepo, "main", "", "")
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsPull(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsPull_UnknownModelExits1 drives the real cobra Run
|
||||
// function end to end, confirming the wrapper propagates runPullWeights's
|
||||
// catalog refusal into a process exit.
|
||||
func TestRunOchreWeightsPull_UnknownModelExits1(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
stubConnect(t, fakeClusterConfig(), nc, nil)
|
||||
|
||||
cmd := newOchreWeightsPullTestCmd(t, "not-a-real-model", testHFRepo, "main", "", "")
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsPull(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// --- ochre credentials set ---
|
||||
|
||||
// withStdin swaps os.Stdin for a pipe fed with content, restoring the real
|
||||
// os.Stdin afterwards, mirroring captureStdout's os.Stdout swap.
|
||||
func withStdin(t *testing.T, content string, fn func()) {
|
||||
t.Helper()
|
||||
r, w, err := os.Pipe()
|
||||
require.NoError(t, err)
|
||||
orig := os.Stdin
|
||||
os.Stdin = r //nolint:reassign // test-local stdin swap, restored below
|
||||
t.Cleanup(func() { os.Stdin = orig }) //nolint:reassign // restoring the real os.Stdin captured above
|
||||
|
||||
go func() {
|
||||
_, _ = io.WriteString(w, content)
|
||||
_ = w.Close()
|
||||
}()
|
||||
|
||||
fn()
|
||||
}
|
||||
|
||||
func newOchreCredentialsSetTestCmd(t *testing.T, vendor, account, token string) *cobra.Command {
|
||||
t.Helper()
|
||||
require.NoError(t, ochreCredentialsSetCmd.Flags().Set("vendor", vendor))
|
||||
require.NoError(t, ochreCredentialsSetCmd.Flags().Set("account", account))
|
||||
require.NoError(t, ochreCredentialsSetCmd.Flags().Set("token", token))
|
||||
return ochreCredentialsSetCmd
|
||||
}
|
||||
|
||||
// TestReadTokenFromStdin_TrimsTrailingNewline covers the common 'echo token
|
||||
// | cmd --token -' shape.
|
||||
func TestReadTokenFromStdin_TrimsTrailingNewline(t *testing.T) {
|
||||
token, err := readTokenFromStdin(strings.NewReader("hf_mytoken\n"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "hf_mytoken", token)
|
||||
}
|
||||
|
||||
// TestReadTokenFromStdin_EmptyIsError covers a piped-in nothing (e.g. an
|
||||
// empty file) refusing rather than silently storing an empty credential.
|
||||
func TestReadTokenFromStdin_EmptyIsError(t *testing.T) {
|
||||
_, err := readTokenFromStdin(strings.NewReader(""))
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// TestRunOchreCredentialsSet_RejectsNonDashToken covers the no-secrets-as-args
|
||||
// guard: --token must be literally '-', never the secret itself.
|
||||
func TestRunOchreCredentialsSet_RejectsNonDashToken(t *testing.T) {
|
||||
cmd := newOchreCredentialsSetTestCmd(t, vendorHuggingFace, "", "hf_literal_secret")
|
||||
code := withOchreExitCapture(t, func() { runOchreCredentialsSet(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreCredentialsSet_ConnectFailureExits1 covers the connect
|
||||
// failure path, same shape as every other ochre wrapper.
|
||||
func TestRunOchreCredentialsSet_ConnectFailureExits1(t *testing.T) {
|
||||
stubConnect(t, nil, nil, errors.New("dial failed"))
|
||||
|
||||
cmd := newOchreCredentialsSetTestCmd(t, vendorHuggingFace, "", "-")
|
||||
withStdin(t, "hf_mytoken\n", func() {
|
||||
code := withOchreExitCapture(t, func() { runOchreCredentialsSet(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunOchreCredentialsSet_MasterKeyLoadFailureExits1 covers the store's
|
||||
// own setup failing (no master.key provisioned yet) before PutCredential is
|
||||
// ever reached.
|
||||
func TestRunOchreCredentialsSet_MasterKeyLoadFailureExits1(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
cfg := &config.ClusterConfig{Node: "node1", Nodes: map[string]config.Config{"node1": {BaseDir: t.TempDir()}}}
|
||||
stubConnect(t, cfg, nc, nil)
|
||||
|
||||
cmd := newOchreCredentialsSetTestCmd(t, vendorHuggingFace, "", "-")
|
||||
withStdin(t, "hf_mytoken\n", func() {
|
||||
code := withOchreExitCapture(t, func() { runOchreCredentialsSet(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
})
|
||||
}
|
||||
|
||||
// TestRunOchreCredentialsSet_RoundTripsAndDefaultsToPlatformAccount drives
|
||||
// the real Run function end to end: stores a credential read from stdin
|
||||
// under the platform account (D2's default), then confirms a separately
|
||||
// constructed CredentialStore resolves it, and that the token never appears
|
||||
// on stdout.
|
||||
func TestRunOchreCredentialsSet_RoundTripsAndDefaultsToPlatformAccount(t *testing.T) {
|
||||
ns, nc, _ := testutil.StartTestJetStream(t)
|
||||
baseDir := t.TempDir()
|
||||
masterKey := bytes32Key()
|
||||
require.NoError(t, os.MkdirAll(filepath.Join(baseDir, "config"), 0750))
|
||||
require.NoError(t, handlers_iam.SaveMasterKey(filepath.Join(baseDir, "config", "master.key"), masterKey))
|
||||
cfg := &config.ClusterConfig{Node: "node1", Nodes: map[string]config.Config{"node1": {BaseDir: baseDir}}}
|
||||
stubConnect(t, cfg, nc, nil)
|
||||
|
||||
cmd := newOchreCredentialsSetTestCmd(t, vendorHuggingFace, "", "-")
|
||||
var out string
|
||||
withStdin(t, "hf_supersecret\n", func() {
|
||||
out = captureStdout(t, func() { runOchreCredentialsSet(cmd, nil) })
|
||||
})
|
||||
|
||||
assert.NotContains(t, out, "hf_supersecret", "the token must never be printed")
|
||||
|
||||
// runOchreCredentialsSet closes its own NATS connection on return, so
|
||||
// verification dials a second connection to the same embedded server.
|
||||
verifyConn, err := nats.Connect(ns.ClientURL())
|
||||
require.NoError(t, err)
|
||||
defer verifyConn.Close()
|
||||
verifyJS := testutil.NewJetStream(t, verifyConn)
|
||||
|
||||
credStore := gateway_bedrock.NewCredentialStore(verifyJS, masterKey, 1, nil)
|
||||
token, ok, err := credStore.Resolve(context.Background(), utils.GlobalAccountID, vendorHuggingFace)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "hf_supersecret", token)
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
// Package hfhub is a minimal Hugging Face Hub client: resolve a repo
|
||||
// revision to an immutable commit SHA, list its file tree, and stream
|
||||
// individual files. It is the network half of 'ochre weights pull'; it never
|
||||
// writes to predastore or touches the weights KV store itself.
|
||||
package hfhub
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
)
|
||||
|
||||
// DefaultBaseURL is the public Hugging Face Hub API and file host.
|
||||
const DefaultBaseURL = "https://huggingface.co"
|
||||
|
||||
// headerTimeout bounds dial + time-to-first-byte, NOT the body stream.
|
||||
// A whole-request http.Client.Timeout would guillotine a multi-GB shard
|
||||
// mid-download; a header timeout still fails fast on a dead/hung server.
|
||||
const headerTimeout = 30 * time.Second
|
||||
|
||||
// Client resolves a Hugging Face repo revision to an immutable commit SHA,
|
||||
// lists its file tree, and streams individual files. BaseURL is overridable
|
||||
// so tests run against an httptest fake instead of the real hub.
|
||||
type Client struct {
|
||||
BaseURL string
|
||||
Token string
|
||||
HTTP *http.Client
|
||||
}
|
||||
|
||||
// NewClient builds a Client against the public Hugging Face hub. token may
|
||||
// be empty for public repos; a gated or private repo without one fails with
|
||||
// a licence/credential error from ResolveRevision or ListTree.
|
||||
func NewClient(token string) *Client {
|
||||
return &Client{
|
||||
BaseURL: DefaultBaseURL,
|
||||
Token: token,
|
||||
HTTP: defaultHTTPClient(),
|
||||
}
|
||||
}
|
||||
|
||||
// defaultHTTPClient streams downloads without a whole-request timeout,
|
||||
// bounding only dial and response-header latency so large shards can read
|
||||
// for as long as the server keeps sending.
|
||||
func defaultHTTPClient() *http.Client {
|
||||
return &http.Client{
|
||||
Transport: &http.Transport{
|
||||
DialContext: (&net.Dialer{Timeout: headerTimeout}).DialContext,
|
||||
ResponseHeaderTimeout: headerTimeout,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// TreeEntry is one file or directory in a Hugging Face repo's tree listing.
|
||||
type TreeEntry struct {
|
||||
Type string `json:"type"` // "file" or "directory"
|
||||
Path string `json:"path"`
|
||||
Size int64 `json:"size"`
|
||||
OID string `json:"oid"`
|
||||
}
|
||||
|
||||
// revisionInfo is the subset of the Hugging Face model-info response pull
|
||||
// needs: the immutable commit SHA a mutable ref (branch/tag) resolves to.
|
||||
type revisionInfo struct {
|
||||
SHA string `json:"sha"`
|
||||
}
|
||||
|
||||
// ResolveRevision resolves repo's revision (branch, tag, or SHA; e.g. "main")
|
||||
// to its immutable commit SHA, so a pull is pinned to weights that cannot
|
||||
// silently change under an already-served model ID.
|
||||
func (c *Client) ResolveRevision(ctx context.Context, repo, revision string) (string, error) {
|
||||
url := fmt.Sprintf("%s/api/models/%s/revision/%s", c.baseURL(), repo, revision)
|
||||
body, err := c.getJSON(ctx, repo, url)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer body.Close()
|
||||
|
||||
var info revisionInfo
|
||||
if err := json.NewDecoder(body).Decode(&info); err != nil {
|
||||
return "", fmt.Errorf("decode revision info for %s@%s: %w", repo, revision, err)
|
||||
}
|
||||
if info.SHA == "" {
|
||||
return "", fmt.Errorf("hugging face returned no commit sha for %s@%s", repo, revision)
|
||||
}
|
||||
return info.SHA, nil
|
||||
}
|
||||
|
||||
// ListTree lists every file and directory in repo at commit sha, recursively.
|
||||
func (c *Client) ListTree(ctx context.Context, repo, sha string) ([]TreeEntry, error) {
|
||||
url := fmt.Sprintf("%s/api/models/%s/tree/%s?recursive=true", c.baseURL(), repo, sha)
|
||||
body, err := c.getJSON(ctx, repo, url)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer body.Close()
|
||||
|
||||
var entries []TreeEntry
|
||||
if err := json.NewDecoder(body).Decode(&entries); err != nil {
|
||||
return nil, fmt.Errorf("decode file tree for %s@%s: %w", repo, sha, err)
|
||||
}
|
||||
return entries, nil
|
||||
}
|
||||
|
||||
// DownloadFile streams path from repo at commit sha. The caller must close
|
||||
// the returned reader. contentLength is -1 when the server does not report
|
||||
// it (http.Response.ContentLength's own convention).
|
||||
func (c *Client) DownloadFile(ctx context.Context, repo, sha, path string) (io.ReadCloser, int64, error) {
|
||||
url := fmt.Sprintf("%s/%s/resolve/%s/%s", c.baseURL(), repo, sha, path)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("build request for %s: %w", path, err)
|
||||
}
|
||||
c.authorize(req)
|
||||
|
||||
resp, err := c.httpClient().Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("download %s from %s: %w", path, repo, err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode > 299 {
|
||||
defer resp.Body.Close()
|
||||
return nil, 0, mapStatusError(repo, resp)
|
||||
}
|
||||
return resp.Body, resp.ContentLength, nil
|
||||
}
|
||||
|
||||
// getJSON issues an authenticated GET against url and returns the response
|
||||
// body on 2xx; the caller decodes and closes it. Non-2xx is mapped to a
|
||||
// clear licence/not-found/generic error naming repo.
|
||||
func (c *Client) getJSON(ctx context.Context, repo, url string) (io.ReadCloser, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build request: %w", err)
|
||||
}
|
||||
c.authorize(req)
|
||||
|
||||
resp, err := c.httpClient().Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request %s: %w", repo, err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode > 299 {
|
||||
defer resp.Body.Close()
|
||||
return nil, mapStatusError(repo, resp)
|
||||
}
|
||||
return resp.Body, nil
|
||||
}
|
||||
|
||||
func (c *Client) authorize(req *http.Request) {
|
||||
if c.Token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+c.Token)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) baseURL() string {
|
||||
if c.BaseURL != "" {
|
||||
return c.BaseURL
|
||||
}
|
||||
return DefaultBaseURL
|
||||
}
|
||||
|
||||
func (c *Client) httpClient() *http.Client {
|
||||
if c.HTTP != nil {
|
||||
return c.HTTP
|
||||
}
|
||||
return defaultHTTPClient()
|
||||
}
|
||||
|
||||
// mapStatusError maps a non-2xx Hugging Face response to a clear
|
||||
// awserrors-coded error naming repo: 401/403 as a licence/credential
|
||||
// failure (D5 -- pull must abort before any object lands), 404 as
|
||||
// not-found, anything else as a generic upstream error.
|
||||
func mapStatusError(repo string, resp *http.Response) error {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
switch resp.StatusCode {
|
||||
case http.StatusUnauthorized, http.StatusForbidden:
|
||||
return awserrors.Errorf(awserrors.ErrorAccessDeniedException,
|
||||
"hugging face repo %q requires a licence/credential (HTTP %d): %s", repo, resp.StatusCode, string(body))
|
||||
case http.StatusNotFound:
|
||||
return awserrors.Errorf(awserrors.ErrorResourceNotFoundException,
|
||||
"hugging face repo %q not found (HTTP %d): %s", repo, resp.StatusCode, string(body))
|
||||
default:
|
||||
return fmt.Errorf("hugging face request for %q failed (HTTP %d): %s", repo, resp.StatusCode, string(body))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
package hfhub
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// newTestClient builds a Client pointed at srv, with token attached to every
|
||||
// request so gating tests can assert it was actually sent.
|
||||
func newTestClient(srv *httptest.Server, token string) *Client {
|
||||
return &Client{BaseURL: srv.URL, Token: token, HTTP: srv.Client()}
|
||||
}
|
||||
|
||||
// TestResolveRevision_PinsToSHA covers the ref-to-immutable-SHA pin: a
|
||||
// mutable ref like "main" must resolve to the fixed commit sha the fake
|
||||
// server reports, which is what a pull then downloads from.
|
||||
func TestResolveRevision_PinsToSHA(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "/api/models/meta-llama/Llama-3.2-1B-Instruct/revision/main", r.URL.Path)
|
||||
_ = json.NewEncoder(w).Encode(revisionInfo{SHA: "abc123def456"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := newTestClient(srv, "")
|
||||
sha, err := c.ResolveRevision(context.Background(), "meta-llama/Llama-3.2-1B-Instruct", "main")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "abc123def456", sha)
|
||||
}
|
||||
|
||||
// TestResolveRevision_SendsBearerToken confirms a non-empty token is sent as
|
||||
// an Authorization header, so a gated repo with a valid token is served.
|
||||
func TestResolveRevision_SendsBearerToken(t *testing.T) {
|
||||
var gotAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
_ = json.NewEncoder(w).Encode(revisionInfo{SHA: "sha1"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := newTestClient(srv, "hf_secrettoken")
|
||||
_, err := c.ResolveRevision(context.Background(), "meta-llama/Llama-3.2-1B-Instruct", "main")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Bearer hf_secrettoken", gotAuth)
|
||||
}
|
||||
|
||||
// TestResolveRevision_GatedRepoReturnsAccessDeniedError covers D5: a 401/403
|
||||
// from the hub must surface as a clean, awserrors-coded licence/credential
|
||||
// error naming the repo, before any download is attempted.
|
||||
func TestResolveRevision_GatedRepoReturnsAccessDeniedError(t *testing.T) {
|
||||
for _, status := range []int{http.StatusUnauthorized, http.StatusForbidden} {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(status)
|
||||
_, _ = w.Write([]byte(`{"error":"gated"}`))
|
||||
}))
|
||||
|
||||
c := newTestClient(srv, "")
|
||||
_, err := c.ResolveRevision(context.Background(), "meta-llama/Llama-3.2-1B-Instruct", "main")
|
||||
require.Error(t, err)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorAccessDeniedException), "status %d: got %v", status, err)
|
||||
assert.Contains(t, err.Error(), "meta-llama/Llama-3.2-1B-Instruct")
|
||||
srv.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveRevision_UnknownRepoReturnsNotFoundError covers a 404 mapping
|
||||
// to a not-found error distinguishable from the gated/licence case.
|
||||
func TestResolveRevision_UnknownRepoReturnsNotFoundError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := newTestClient(srv, "")
|
||||
_, err := c.ResolveRevision(context.Background(), "nope/does-not-exist", "main")
|
||||
require.Error(t, err)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorResourceNotFoundException))
|
||||
assert.False(t, awserrors.IsErrorCode(err, awserrors.ErrorAccessDeniedException))
|
||||
}
|
||||
|
||||
// TestListTree_ReturnsEntries covers the recursive tree listing pull filters
|
||||
// down to the safetensors-only set.
|
||||
func TestListTree_ReturnsEntries(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "/api/models/meta-llama/Llama-3.2-1B-Instruct/tree/abc123", r.URL.Path)
|
||||
assert.Equal(t, "true", r.URL.Query().Get("recursive"))
|
||||
_ = json.NewEncoder(w).Encode([]TreeEntry{
|
||||
{Type: "file", Path: "config.json", Size: 10, OID: "oid1"},
|
||||
{Type: "file", Path: "model.safetensors", Size: 1000, OID: "oid2"},
|
||||
{Type: "file", Path: "pytorch_model.bin", Size: 1000, OID: "oid3"},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := newTestClient(srv, "")
|
||||
entries, err := c.ListTree(context.Background(), "meta-llama/Llama-3.2-1B-Instruct", "abc123")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, entries, 3)
|
||||
assert.Equal(t, "model.safetensors", entries[1].Path)
|
||||
}
|
||||
|
||||
// TestDownloadFile_StreamsBody covers the resolve URL shape and that the
|
||||
// full body is readable from the returned reader.
|
||||
func TestDownloadFile_StreamsBody(t *testing.T) {
|
||||
const content = "fake safetensors bytes"
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "/meta-llama/Llama-3.2-1B-Instruct/resolve/abc123/model.safetensors", r.URL.Path)
|
||||
_, _ = io.WriteString(w, content)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := newTestClient(srv, "")
|
||||
body, _, err := c.DownloadFile(context.Background(), "meta-llama/Llama-3.2-1B-Instruct", "abc123", "model.safetensors")
|
||||
require.NoError(t, err)
|
||||
defer body.Close()
|
||||
|
||||
got, err := io.ReadAll(body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, content, string(got))
|
||||
}
|
||||
|
||||
// TestDownloadFile_GatedReturnsAccessDeniedError covers a mid-stream-request
|
||||
// gate (e.g. a token that expired between resolve and download) still
|
||||
// mapping cleanly rather than surfacing a raw HTTP error.
|
||||
func TestDownloadFile_GatedReturnsAccessDeniedError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := newTestClient(srv, "")
|
||||
body, _, err := c.DownloadFile(context.Background(), "meta-llama/Llama-3.2-1B-Instruct", "abc123", "model.safetensors")
|
||||
require.Error(t, err)
|
||||
assert.Nil(t, body)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorAccessDeniedException))
|
||||
}
|
||||
|
||||
// TestNewClient_DefaultsBaseURLAndToken covers the constructor's defaults so
|
||||
// a caller who passes an empty token still gets a usable client against the
|
||||
// real hub.
|
||||
func TestNewClient_DefaultsBaseURLAndToken(t *testing.T) {
|
||||
c := NewClient("")
|
||||
assert.Equal(t, DefaultBaseURL, c.BaseURL)
|
||||
assert.Empty(t, c.Token)
|
||||
assert.NotNil(t, c.HTTP)
|
||||
}
|
||||
@@ -35,6 +35,10 @@ func weightsKey(modelID string) string {
|
||||
type weightsRecord struct {
|
||||
SnapshotID string `json:"snapshot_id"`
|
||||
SourceURI string `json:"source_uri"`
|
||||
// SourceRevision is the upstream hub commit SHA the weights were pulled
|
||||
// at, if known. Empty for weights staged from an operator-supplied S3
|
||||
// prefix with no 'ochre-pull.json' manifest -- the offline stage path.
|
||||
SourceRevision string `json:"source_revision,omitempty"`
|
||||
}
|
||||
|
||||
// WeightsResolver resolves a self-hosted model's deployment-specific weights
|
||||
@@ -119,17 +123,26 @@ func (s *WeightsStore) GetWeights(ctx context.Context, modelID string) (WeightsE
|
||||
if err != nil || !ok {
|
||||
return WeightsEntry{}, ok, err
|
||||
}
|
||||
return WeightsEntry{ModelID: modelID, SourceURI: rec.SourceURI, SnapshotID: rec.SnapshotID}, true, nil
|
||||
return WeightsEntry{ModelID: modelID, SourceURI: rec.SourceURI, SnapshotID: rec.SnapshotID, SourceRevision: rec.SourceRevision}, true, nil
|
||||
}
|
||||
|
||||
// PutWeights records modelID's staged weights artifact: the snapshot ID
|
||||
// endpoints COW-clone from, and the source S3 URI it was materialised from.
|
||||
// It leaves SourceRevision empty; use PutWeightsWithRevision when 'stage'
|
||||
// found a pull manifest at the source prefix.
|
||||
func (s *WeightsStore) PutWeights(ctx context.Context, modelID, sourceURI, snapshotID string) error {
|
||||
return s.PutWeightsWithRevision(ctx, modelID, sourceURI, snapshotID, "")
|
||||
}
|
||||
|
||||
// PutWeightsWithRevision records modelID's staged weights artifact along with
|
||||
// the upstream hub commit SHA it was pulled at, if known (empty when 'stage'
|
||||
// found no 'ochre-pull.json' manifest at the source prefix).
|
||||
func (s *WeightsStore) PutWeightsWithRevision(ctx context.Context, modelID, sourceURI, snapshotID, sourceRevision string) error {
|
||||
kv, err := s.bucket(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
value, err := json.Marshal(weightsRecord{SnapshotID: snapshotID, SourceURI: sourceURI})
|
||||
value, err := json.Marshal(weightsRecord{SnapshotID: snapshotID, SourceURI: sourceURI, SourceRevision: sourceRevision})
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode weights record for %s: %w", modelID, err)
|
||||
}
|
||||
@@ -142,9 +155,10 @@ func (s *WeightsStore) PutWeights(ctx context.Context, modelID, sourceURI, snaps
|
||||
// WeightsEntry is one staged model's KV record, decoded back from
|
||||
// weightsKey's base64url-encoded modelID for operator-facing listing.
|
||||
type WeightsEntry struct {
|
||||
ModelID string
|
||||
SourceURI string
|
||||
SnapshotID string
|
||||
ModelID string
|
||||
SourceURI string
|
||||
SnapshotID string
|
||||
SourceRevision string
|
||||
}
|
||||
|
||||
// ListWeights returns every staged model and its record, sorted by model
|
||||
@@ -179,7 +193,7 @@ func (s *WeightsStore) ListWeights(ctx context.Context) ([]WeightsEntry, error)
|
||||
if err := json.Unmarshal(entry.Value(), &rec); err != nil {
|
||||
continue
|
||||
}
|
||||
entries = append(entries, WeightsEntry{ModelID: string(modelID), SourceURI: rec.SourceURI, SnapshotID: rec.SnapshotID})
|
||||
entries = append(entries, WeightsEntry{ModelID: string(modelID), SourceURI: rec.SourceURI, SnapshotID: rec.SnapshotID, SourceRevision: rec.SourceRevision})
|
||||
}
|
||||
sort.Slice(entries, func(i, j int) bool { return entries[i].ModelID < entries[j].ModelID })
|
||||
return entries, nil
|
||||
|
||||
@@ -95,6 +95,43 @@ func TestWeightsStore_GetWeights_ReturnsSourceURIAndSnapshot(t *testing.T) {
|
||||
}, entry)
|
||||
}
|
||||
|
||||
// TestWeightsStore_PutWeightsWithRevision_RoundTrips covers D3: a pull-staged
|
||||
// model's upstream commit SHA survives the KV round trip and comes back on
|
||||
// both GetWeights and ListWeights.
|
||||
func TestWeightsStore_PutWeightsWithRevision_RoundTrips(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
|
||||
|
||||
ctx := context.Background()
|
||||
require.NoError(t, store.PutWeightsWithRevision(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://ochre-weights/meta-llama/Llama-3.2-1B-Instruct/abc123/", "snap-0001", "abc123def456"))
|
||||
|
||||
entry, ok, err := store.GetWeights(ctx, "meta.llama3-2-1b-instruct-v1:0")
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "abc123def456", entry.SourceRevision)
|
||||
|
||||
entries, err := store.ListWeights(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, entries, 1)
|
||||
assert.Equal(t, "abc123def456", entries[0].SourceRevision)
|
||||
}
|
||||
|
||||
// TestWeightsStore_PutWeights_LeavesSourceRevisionEmpty covers the offline
|
||||
// stage path: weights staged with no pull manifest must record an empty
|
||||
// SourceRevision, not fail or fabricate one.
|
||||
func TestWeightsStore_PutWeights_LeavesSourceRevisionEmpty(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
|
||||
|
||||
ctx := context.Background()
|
||||
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
|
||||
entry, ok, err := store.GetWeights(ctx, "meta.llama3-2-1b-instruct-v1:0")
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Empty(t, entry.SourceRevision)
|
||||
}
|
||||
|
||||
// TestWeightsStore_ListWeights_Empty covers the no-keys-yet case: an empty
|
||||
// bucket must report an empty list, not an error (jetstream.ErrNoKeysFound).
|
||||
func TestWeightsStore_ListWeights_Empty(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user