diff --git a/cmd/spinifex/cmd/admin_ochre.go b/cmd/spinifex/cmd/admin_ochre.go index 00512a812..22e468edc 100644 --- a/cmd/spinifex/cmd/admin_ochre.go +++ b/cmd/spinifex/cmd/admin_ochre.go @@ -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///)") + 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 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", diff --git a/cmd/spinifex/cmd/admin_ochre_pull_test.go b/cmd/spinifex/cmd/admin_ochre_pull_test.go new file mode 100644 index 000000000..21b6fae11 --- /dev/null +++ b/cmd/spinifex/cmd/admin_ochre_pull_test.go @@ -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//. +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) +} diff --git a/spinifex/gateway/bedrock/hfhub/hfhub.go b/spinifex/gateway/bedrock/hfhub/hfhub.go new file mode 100644 index 000000000..c35e23a6f --- /dev/null +++ b/spinifex/gateway/bedrock/hfhub/hfhub.go @@ -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)) + } +} diff --git a/spinifex/gateway/bedrock/hfhub/hfhub_test.go b/spinifex/gateway/bedrock/hfhub/hfhub_test.go new file mode 100644 index 000000000..5de3720b3 --- /dev/null +++ b/spinifex/gateway/bedrock/hfhub/hfhub_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/weights.go b/spinifex/gateway/bedrock/weights.go index 075418e8a..9ac3abc10 100644 --- a/spinifex/gateway/bedrock/weights.go +++ b/spinifex/gateway/bedrock/weights.go @@ -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 diff --git a/spinifex/gateway/bedrock/weights_test.go b/spinifex/gateway/bedrock/weights_test.go index 692485ba1..a790d405e 100644 --- a/spinifex/gateway/bedrock/weights_test.go +++ b/spinifex/gateway/bedrock/weights_test.go @@ -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) {