From b0f35adb10b60db6c244785f1623fc541a940803 Mon Sep 17 00:00:00 2001 From: chokevin Date: Fri, 7 Aug 2026 19:33:25 -0700 Subject: [PATCH 1/3] cli: retrieve acknowledged artifact bundles Replace temporary PVC-reader pods with a generation-bound bundle contract and direct Azure Blob retrieval. Have tau-core stamp trusted non-secret transport metadata and fail closed across staged publication, metrics, and checkpoints. --- .../tau-core-controller/kustomize/rbac.yaml | 9 + .../tau-core-controller/templates/rbac.yaml | 9 + cli/README.md | 5 +- cli/internal/artifactbundle/bundle.go | 305 ++++++++++++++++ cli/internal/artifactbundle/bundle_test.go | 293 +++++++++++++++ cli/internal/artifactbundle/publish.go | 143 ++++++++ cli/internal/artifactbundle/retrieve.go | 340 ++++++++++++++++++ cli/internal/artifactindex/finalize.go | 9 +- cli/internal/artifactindex/finalize_test.go | 4 + cli/internal/cli/pvc_helpers.go | 6 +- cli/internal/cli/run_blob_store.go | 231 ++++++++++++ cli/internal/cli/run_blob_store_test.go | 134 +++++++ cli/internal/cli/run_bundle.go | 58 +++ cli/internal/cli/run_bundle_test.go | 93 +++++ cli/internal/cli/run_get.go | 173 ++++++++- cli/internal/cli/run_get_test.go | 10 +- cli/internal/cli/run_job.go | 29 ++ cli/internal/cli/run_ray.go | 19 + cli/internal/jobrender/render.go | 12 + cli/internal/jobrender/render_test.go | 46 +++ cli/internal/rayjobrender/render.go | 9 + cli/internal/rayjobrender/render_test.go | 29 ++ controllers/tau-core/cmd/controller/main.go | 4 + .../controller/artifact_store_controller.go | 224 ++++++++++++ .../artifact_store_controller_test.go | 95 +++++ .../internal/labelkeys/contract_test.go | 8 +- .../tau-core/internal/labelkeys/labelkeys.go | 14 +- core/workloadmeta/metadata.go | 2 + site/content/en/docs/reference/cli.md | 22 +- site/content/en/docs/reference/run-config.md | 18 + .../en/docs/tasks/researcher/first-run.md | 20 +- 31 files changed, 2347 insertions(+), 26 deletions(-) create mode 100644 cli/internal/artifactbundle/bundle.go create mode 100644 cli/internal/artifactbundle/bundle_test.go create mode 100644 cli/internal/artifactbundle/publish.go create mode 100644 cli/internal/artifactbundle/retrieve.go create mode 100644 cli/internal/cli/run_blob_store.go create mode 100644 cli/internal/cli/run_blob_store_test.go create mode 100644 cli/internal/cli/run_bundle.go create mode 100644 cli/internal/cli/run_bundle_test.go create mode 100644 controllers/tau-core/internal/controller/artifact_store_controller.go create mode 100644 controllers/tau-core/internal/controller/artifact_store_controller_test.go diff --git a/charts/tau-core-controller/kustomize/rbac.yaml b/charts/tau-core-controller/kustomize/rbac.yaml index 55792a3a..4623c5d6 100644 --- a/charts/tau-core-controller/kustomize/rbac.yaml +++ b/charts/tau-core-controller/kustomize/rbac.yaml @@ -80,6 +80,15 @@ rules: - apiGroups: [""] resources: ["events"] verbs: ["create", "patch"] + - apiGroups: [""] + resources: ["persistentvolumeclaims", "persistentvolumes"] + verbs: ["get", "list", "watch"] + - apiGroups: ["batch"] + resources: ["jobs"] + verbs: ["get", "list", "watch", "patch", "update"] + - apiGroups: ["ray.io"] + resources: ["rayjobs"] + verbs: ["get", "list", "watch", "patch", "update"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRoleBinding diff --git a/charts/tau-core-controller/templates/rbac.yaml b/charts/tau-core-controller/templates/rbac.yaml index 81290b8b..6cf70c28 100644 --- a/charts/tau-core-controller/templates/rbac.yaml +++ b/charts/tau-core-controller/templates/rbac.yaml @@ -78,6 +78,15 @@ rules: - apiGroups: [""] resources: ["events"] verbs: ["create", "patch"] + - apiGroups: [""] + resources: ["persistentvolumeclaims", "persistentvolumes"] + verbs: ["get", "list", "watch"] + - apiGroups: ["batch"] + resources: ["jobs"] + verbs: ["get", "list", "watch", "patch", "update"] + - apiGroups: ["ray.io"] + resources: ["rayjobs"] + verbs: ["get", "list", "watch", "patch", "update"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRoleBinding diff --git a/cli/README.md b/cli/README.md index 54c8b86f..0a5fb08f 100644 --- a/cli/README.md +++ b/cli/README.md @@ -182,7 +182,8 @@ Use the run lifecycle after submission: ```bash tau run status --watch tau run logs -tau run get # Job-backed run with persisted results +tau run get +tau run get --destination ./artifacts/ # complete acknowledged Job/RayJob bundle tau run resume --config tau/train.yaml tau run cancel ``` @@ -227,7 +228,7 @@ Repositories that require catalog routing should set | --- | --- | | `tau cluster` | Install/uninstall the TauGrid Helm distribution and validate an existing cluster. Tau does not provision AKS infrastructure. | | `tau workspace` | Inspect workspace connections and readiness, request quota, and scaffold repositories. | -| `tau run` | Resolve a project target; submit, inspect, resume, or cancel the run, and fetch persisted results for Job-backed runs. | +| `tau run` | Resolve a project target; submit, inspect, resume, or cancel the run, and fetch persisted Job/RayJob result bundles. | | `tau serve` | Deploy and operate online model endpoints. | | `tau data` | Manage dataset and model registries. | | `tau python` | Invoke Python SDK helpers while keeping Go Tau as the executor. | diff --git a/cli/internal/artifactbundle/bundle.go b/cli/internal/artifactbundle/bundle.go new file mode 100644 index 00000000..8ad5bbfa --- /dev/null +++ b/cli/internal/artifactbundle/bundle.go @@ -0,0 +1,305 @@ +// Package artifactbundle owns Tau's durable run-bundle completion contract. +// +// A PVC directory is storage, not a retrieval API: result publication, metrics +// offload, and checkpoint indexing finish at different times and in different +// trees. This package writes one final acknowledgement only after those nested +// producers return successfully, then describes every durable tree or glob that +// belongs to the run. Readers can therefore fail closed without mounting the PVC. +package artifactbundle + +import ( + "encoding/json" + "fmt" + "path" + "strings" +) + +const ( + SchemaVersion = "tau.run.artifact-bundle.v1" + BundleDir = ".tau/bundles" + CurrentManifest = ".tau/bundle.json" + CurrentCompletion = ".tau/bundle.complete" +) + +type Publication struct { + Mode string `json:"mode"` + ID string `json:"id"` + Root string `json:"root"` + Completion string `json:"completion"` +} + +type Metrics struct { + SessionID string `json:"session_id,omitempty"` + History []string `json:"history,omitempty"` + OffloadRoot string `json:"offload_root,omitempty"` + Acknowledged bool `json:"acknowledged"` +} + +type Checkpoint struct { + Artifact string `json:"artifact,omitempty"` + Root string `json:"root"` + Index string `json:"index"` +} + +type References struct { + Artifacts string `json:"artifacts"` + Checkpoint string `json:"checkpoint,omitempty"` + Logs string `json:"logs"` +} + +type PathSpec struct { + Name string `json:"name"` + Path string `json:"path"` + Kind string `json:"kind"` + Optional bool `json:"optional,omitempty"` +} + +type Manifest struct { + SchemaVersion string `json:"schema_version"` + BundleID string `json:"bundle_id"` + Run string `json:"run"` + Namespace string `json:"namespace"` + ResultPVC string `json:"result_pvc"` + ResultRoot string `json:"result_root"` + Publication *Publication `json:"publication,omitempty"` + Metrics *Metrics `json:"metrics,omitempty"` + Checkpoint *Checkpoint `json:"checkpoint,omitempty"` + References References `json:"references"` + Paths []PathSpec `json:"paths"` +} + +type Runtime struct { + BundleID string + Run string + Namespace string + ResultPVC string + OutputDir string + PublicationMode string + PublicationID string + PublicationRoot string + PublicationMarker string + MetricsSessionID string + MetricsHistory []string + MetricsOffloadDir string + MetricsEnabled bool + CheckpointArtifact string + CheckpointRoot string + CheckpointIndex string +} + +func (r Runtime) Enabled() bool { + return strings.TrimSpace(r.BundleID) != "" +} + +func (r Runtime) Validate() error { + if !r.Enabled() { + return nil + } + for name, value := range map[string]string{ + "bundle ID": r.BundleID, + "run": r.Run, + "namespace": r.Namespace, + "result PVC": r.ResultPVC, + "output": r.OutputDir, + } { + if strings.TrimSpace(value) == "" { + return fmt.Errorf("artifact bundle %s is required", name) + } + } + id := strings.TrimSpace(r.BundleID) + if id == "." || id == ".." || strings.ContainsAny(id, `/\`) { + return fmt.Errorf("artifact bundle ID must be a single path segment") + } + if err := validateDurablePath("output", r.OutputDir); err != nil { + return err + } + for label, value := range map[string]string{ + "publication root": r.PublicationRoot, + "publication marker": r.PublicationMarker, + "metrics offload": r.MetricsOffloadDir, + "checkpoint root": r.CheckpointRoot, + "checkpoint index": r.CheckpointIndex, + } { + if strings.TrimSpace(value) != "" { + if err := validateDurablePath(label, value); err != nil { + return err + } + } + } + for _, history := range r.MetricsHistory { + if err := validateDurablePath("metrics history", history); err != nil { + return err + } + } + if r.PublicationMode != "" { + if strings.TrimSpace(r.PublicationID) == "" || + strings.TrimSpace(r.PublicationRoot) == "" || + strings.TrimSpace(r.PublicationMarker) == "" { + return fmt.Errorf("artifact bundle publication requires ID, root, and completion marker") + } + } + return nil +} + +func validateDurablePath(label, value string) error { + clean := path.Clean(strings.TrimSpace(value)) + if clean == "." || (clean != "/data" && !strings.HasPrefix(clean, "/data/")) { + return fmt.Errorf("artifact bundle %s must be under /data", label) + } + return nil +} + +func (r Runtime) Manifest() (Manifest, error) { + if err := r.Validate(); err != nil { + return Manifest{}, err + } + m := Manifest{ + SchemaVersion: SchemaVersion, + BundleID: strings.TrimSpace(r.BundleID), + Run: strings.TrimSpace(r.Run), + Namespace: strings.TrimSpace(r.Namespace), + ResultPVC: strings.TrimSpace(r.ResultPVC), + ResultRoot: path.Clean(r.OutputDir), + References: References{ + Artifacts: path.Clean(r.OutputDir), + Logs: fmt.Sprintf("tau run logs %s -n %s", strings.TrimSpace(r.Run), strings.TrimSpace(r.Namespace)), + }, + Paths: []PathSpec{{ + Name: "results", + Path: path.Clean(r.OutputDir), + Kind: "tree", + }, { + Name: "bundle-manifest", + Path: GenerationManifestPath(r.OutputDir, r.BundleID), + Kind: "file", + }, { + Name: "bundle-acknowledgement", + Path: GenerationCompletionPath(r.OutputDir, r.BundleID), + Kind: "file", + }}, + } + if r.PublicationMode != "" { + m.Publication = &Publication{ + Mode: strings.TrimSpace(r.PublicationMode), + ID: strings.TrimSpace(r.PublicationID), + Root: path.Clean(r.PublicationRoot), + Completion: path.Clean(r.PublicationMarker), + } + m.Paths[0].Path = path.Clean(r.PublicationRoot) + } + if r.MetricsEnabled { + m.Metrics = &Metrics{ + SessionID: strings.TrimSpace(r.MetricsSessionID), + History: append([]string(nil), r.MetricsHistory...), + OffloadRoot: cleanOptionalPath(r.MetricsOffloadDir), + Acknowledged: true, + } + for i, history := range r.MetricsHistory { + m.Paths = append(m.Paths, PathSpec{ + Name: fmt.Sprintf("metrics-history-%d", i+1), + Path: path.Clean(history), + Kind: "glob", + Optional: true, + }) + } + if strings.TrimSpace(r.MetricsOffloadDir) != "" { + m.Paths = append(m.Paths, PathSpec{ + Name: "metrics-offload", + Path: path.Clean(r.MetricsOffloadDir), + Kind: "tree", + }) + } + } + if strings.TrimSpace(r.CheckpointRoot) != "" { + m.Checkpoint = &Checkpoint{ + Artifact: strings.TrimSpace(r.CheckpointArtifact), + Root: path.Clean(r.CheckpointRoot), + Index: path.Clean(r.CheckpointIndex), + } + m.References.Checkpoint = path.Clean(r.CheckpointIndex) + m.Paths = append(m.Paths, PathSpec{ + Name: "checkpoint-index", + Path: path.Clean(r.CheckpointIndex), + Kind: "file", + }, PathSpec{ + Name: "checkpoints", + Path: path.Clean(r.CheckpointRoot), + Kind: "tree", + }) + } + return m, nil +} + +func cleanOptionalPath(value string) string { + if strings.TrimSpace(value) == "" { + return "" + } + return path.Clean(value) +} + +func GenerationManifestPath(outputDir, bundleID string) string { + return path.Join(path.Clean(outputDir), BundleDir, strings.TrimSpace(bundleID)+".json") +} + +func GenerationCompletionPath(outputDir, bundleID string) string { + return path.Join(path.Clean(outputDir), BundleDir, strings.TrimSpace(bundleID)+".complete") +} + +func CurrentManifestPath(outputDir string) string { + return path.Join(path.Clean(outputDir), CurrentManifest) +} + +func CurrentCompletionPath(outputDir string) string { + return path.Join(path.Clean(outputDir), CurrentCompletion) +} + +func Marshal(m Manifest) ([]byte, error) { + raw, err := json.MarshalIndent(m, "", " ") + if err != nil { + return nil, err + } + return append(raw, '\n'), nil +} + +func Parse(raw []byte) (Manifest, error) { + var m Manifest + decoder := json.NewDecoder(strings.NewReader(string(raw))) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&m); err != nil { + return Manifest{}, fmt.Errorf("decode artifact bundle manifest: %w", err) + } + if m.SchemaVersion != SchemaVersion { + return Manifest{}, fmt.Errorf("artifact bundle schema %q is unsupported; expected %q", m.SchemaVersion, SchemaVersion) + } + if strings.TrimSpace(m.BundleID) == "" || strings.TrimSpace(m.Run) == "" || + strings.TrimSpace(m.Namespace) == "" || + strings.TrimSpace(m.ResultPVC) == "" || strings.TrimSpace(m.ResultRoot) == "" { + return Manifest{}, fmt.Errorf("artifact bundle manifest is missing required identity fields") + } + if err := validateDurablePath("result root", m.ResultRoot); err != nil { + return Manifest{}, err + } + if m.Metrics != nil && !m.Metrics.Acknowledged { + return Manifest{}, fmt.Errorf("artifact bundle metrics are not acknowledged") + } + if m.Publication != nil { + if strings.TrimSpace(m.Publication.ID) == "" || strings.TrimSpace(m.Publication.Completion) == "" { + return Manifest{}, fmt.Errorf("artifact bundle publication is missing required fields") + } + if err := validateDurablePath("publication completion", m.Publication.Completion); err != nil { + return Manifest{}, err + } + } + if len(m.Paths) == 0 { + return Manifest{}, fmt.Errorf("artifact bundle manifest has no durable paths") + } + for _, spec := range m.Paths { + if spec.Kind != "tree" && spec.Kind != "glob" && spec.Kind != "file" { + return Manifest{}, fmt.Errorf("artifact bundle path %q has unsupported kind %q", spec.Name, spec.Kind) + } + if err := validateDurablePath("path "+spec.Name, spec.Path); err != nil { + return Manifest{}, err + } + } + return m, nil +} diff --git a/cli/internal/artifactbundle/bundle_test.go b/cli/internal/artifactbundle/bundle_test.go new file mode 100644 index 00000000..b9da76cb --- /dev/null +++ b/cli/internal/artifactbundle/bundle_test.go @@ -0,0 +1,293 @@ +package artifactbundle + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" +) + +type memoryStore map[string][]byte + +func (s memoryStore) Read(_ context.Context, name string) ([]byte, error) { + raw, ok := s[name] + if !ok { + return nil, os.ErrNotExist + } + return append([]byte(nil), raw...), nil +} + +func (s memoryStore) List(_ context.Context, prefix string) ([]Object, error) { + var out []Object + for name, raw := range s { + if strings.HasPrefix(name, prefix) { + out = append(out, Object{Name: name, Size: int64(len(raw))}) + } + } + return out, nil +} + +func (s memoryStore) Download(_ context.Context, name string, out io.Writer) error { + raw, ok := s[name] + if !ok { + return os.ErrNotExist + } + _, err := out.Write(raw) + return err +} + +type listErrorStore struct { + memoryStore + err error +} + +func (s listErrorStore) List(context.Context, string) ([]Object, error) { + return nil, s.err +} + +func testRuntime() Runtime { + return Runtime{ + BundleID: "bundle-1", + Run: "training-1", + Namespace: "research", + ResultPVC: "blob-training", + OutputDir: "/data/runs/training-1", + PublicationMode: "staged", + PublicationID: "publication-1", + PublicationRoot: "/data/runs/training-1/.tau-artifacts/publication-1", + PublicationMarker: "/data/runs/training-1/.tau-artifacts/publication-1/.tau-artifacts-complete", + MetricsSessionID: "metrics-1", + MetricsHistory: []string{"/data/runs/training-1/metrics/*.jsonl"}, + MetricsOffloadDir: "/data/runs/training-1/.tau/metrics/metrics-1/offload", + MetricsEnabled: true, + CheckpointArtifact: "last.safetensors", + CheckpointRoot: "/data/checkpoints/finetunes/training-1", + CheckpointIndex: "/data/checkpoints/finetunes/training-1/artifacts.json", + } +} + +func TestWrapperCommitsOnlyAfterNestedAcknowledgements(t *testing.T) { + runtime := testRuntime() + root := t.TempDir() + replace := func(script string) string { + return strings.ReplaceAll(script, "/data", filepath.Join(root, "data")) + } + publicationMarker := replace(runtime.PublicationMarker) + if err := os.MkdirAll(filepath.Dir(publicationMarker), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(publicationMarker, []byte("complete publication-1\n"), 0o644); err != nil { + t.Fatal(err) + } + checkpointIndex := replace(runtime.CheckpointIndex) + if err := os.MkdirAll(filepath.Dir(checkpointIndex), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(checkpointIndex, []byte(`{"bundle_id":"bundle-1","artifacts":[{"status":"ready"}]}`), 0o644); err != nil { + t.Fatal(err) + } + script, err := WrapShellScript("true", runtime) + if err != nil { + t.Fatal(err) + } + if out, err := exec.Command("bash", "-c", replace(script)).CombinedOutput(); err != nil { + t.Fatalf("bundle wrapper failed: %v\n%s", err, out) + } + manifestPath := replace(GenerationManifestPath(runtime.OutputDir, runtime.BundleID)) + if _, err := os.Stat(manifestPath); err != nil { + t.Fatalf("generation manifest not committed: %v", err) + } + markerPath := replace(GenerationCompletionPath(runtime.OutputDir, runtime.BundleID)) + if raw, err := os.ReadFile(markerPath); err != nil || string(raw) != "complete bundle-1\n" { + t.Fatalf("bundle acknowledgement = %q, %v", raw, err) + } +} + +func TestWrapperFailsClosedWithoutPublicationAcknowledgement(t *testing.T) { + runtime := testRuntime() + root := t.TempDir() + script, err := WrapShellScript("true", runtime) + if err != nil { + t.Fatal(err) + } + script = strings.ReplaceAll(script, "/data", filepath.Join(root, "data")) + out, err := exec.Command("bash", "-c", script).CombinedOutput() + if err == nil || !strings.Contains(string(out), "publication acknowledgement") { + t.Fatalf("missing publication marker = %v\n%s", err, out) + } +} + +func TestWrapperFailsClosedWithoutDeclaredCheckpointIndex(t *testing.T) { + runtime := testRuntime() + root := t.TempDir() + replace := func(script string) string { + return strings.ReplaceAll(script, "/data", filepath.Join(root, "data")) + } + publicationMarker := replace(runtime.PublicationMarker) + if err := os.MkdirAll(filepath.Dir(publicationMarker), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(publicationMarker, []byte("complete publication-1\n"), 0o644); err != nil { + t.Fatal(err) + } + script, err := WrapShellScript("true", runtime) + if err != nil { + t.Fatal(err) + } + out, err := exec.Command("bash", "-c", replace(script)).CombinedOutput() + if err == nil || !strings.Contains(string(out), "checkpoint index is missing or belongs to another bundle") { + t.Fatalf("missing checkpoint index = %v\n%s", err, out) + } +} + +func TestWrapperFailsClosedWithCheckpointIndexFromAnotherBundle(t *testing.T) { + runtime := testRuntime() + root := t.TempDir() + replace := func(script string) string { + return strings.ReplaceAll(script, "/data", filepath.Join(root, "data")) + } + publicationMarker := replace(runtime.PublicationMarker) + if err := os.MkdirAll(filepath.Dir(publicationMarker), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(publicationMarker, []byte("complete publication-1\n"), 0o644); err != nil { + t.Fatal(err) + } + checkpointIndex := replace(runtime.CheckpointIndex) + if err := os.MkdirAll(filepath.Dir(checkpointIndex), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(checkpointIndex, []byte(`{"bundle_id":"old-bundle"}`), 0o644); err != nil { + t.Fatal(err) + } + script, err := WrapShellScript("true", runtime) + if err != nil { + t.Fatal(err) + } + out, err := exec.Command("bash", "-c", replace(script)).CombinedOutput() + if err == nil || !strings.Contains(string(out), "belongs to another bundle") { + t.Fatalf("stale checkpoint index = %v\n%s", err, out) + } +} + +func TestCompleteBundleLocalFixtureEnumeratesAndDownloadsWithoutKubernetes(t *testing.T) { + runtime := testRuntime() + manifest, err := runtime.Manifest() + if err != nil { + t.Fatal(err) + } + rawManifest, err := Marshal(manifest) + if err != nil { + t.Fatal(err) + } + store := memoryStore{ + "runs/training-1/.tau/bundles/bundle-1.json": rawManifest, + "runs/training-1/.tau/bundles/bundle-1.complete": []byte("complete bundle-1\n"), + "runs/training-1/.tau-artifacts/publication-1/.tau-artifacts-complete": []byte("complete publication-1\n"), + "runs/training-1/.tau-artifacts/publication-1/result.json": []byte(`{"score":0.9}`), + "runs/training-1/metrics/epoch-0001.jsonl": []byte("{\"epoch\":1}\n"), + "runs/training-1/.tau/metrics/metrics-1/offload/metrics-status/run-status.jsonl": []byte("{\"state\":\"succeeded\"}\n"), + "checkpoints/finetunes/training-1/artifacts.json": []byte(`{"bundle_id":"bundle-1","artifacts":[{"status":"ready"}]}`), + "checkpoints/finetunes/training-1/artifacts/last.safetensors": []byte("weights"), + } + loaded, err := Load(context.Background(), store, runtime.OutputDir, runtime.BundleID) + if err != nil { + t.Fatal(err) + } + objects, err := Enumerate(context.Background(), store, loaded) + if err != nil { + t.Fatal(err) + } + if len(objects) != len(store) { + t.Fatalf("enumerated %d objects, want all %d: %+v", len(objects), len(store), objects) + } + destination := filepath.Join(t.TempDir(), "bundle") + files, err := Download(context.Background(), store, loaded, objects, destination) + if err != nil { + t.Fatal(err) + } + if len(files) != len(store) { + t.Fatalf("downloaded %d files, want %d", len(files), len(store)) + } + for name, want := range store { + got, err := os.ReadFile(filepath.Join(destination, filepath.FromSlash(name))) + if err != nil { + t.Fatalf("read downloaded %s: %v", name, err) + } + if !bytes.Equal(got, want) { + t.Fatalf("downloaded %s = %q, want %q", name, got, want) + } + } +} + +func TestLoadRejectsMissingBundleAcknowledgement(t *testing.T) { + manifest, err := testRuntime().Manifest() + if err != nil { + t.Fatal(err) + } + raw, err := Marshal(manifest) + if err != nil { + t.Fatal(err) + } + store := memoryStore{"runs/training-1/.tau/bundles/bundle-1.json": raw} + _, err = Load(context.Background(), store, manifest.ResultRoot, manifest.BundleID) + if err == nil || !errors.Is(err, os.ErrNotExist) { + t.Fatalf("missing acknowledgement error = %v", err) + } +} + +func TestEnumerateDoesNotSwallowOptionalPathErrors(t *testing.T) { + manifest, err := testRuntime().Manifest() + if err != nil { + t.Fatal(err) + } + want := errors.New("storage unavailable") + _, err = Enumerate(context.Background(), listErrorStore{memoryStore: memoryStore{}, err: want}, manifest) + if err == nil || !errors.Is(err, want) { + t.Fatalf("optional path listing error = %v", err) + } +} + +func TestDownloadRejectsExistingDestinationWithoutReplacingIt(t *testing.T) { + store := memoryStore{"runs/training-1/result.json": []byte("new")} + root := t.TempDir() + target := filepath.Join(root, "runs", "training-1", "result.json") + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(target, []byte("existing"), 0o600); err != nil { + t.Fatal(err) + } + _, err := Download(context.Background(), store, Manifest{}, []Object{{ + Name: "runs/training-1/result.json", + Size: 3, + }}, root) + if err == nil || !strings.Contains(err.Error(), "already exists") { + t.Fatalf("existing destination error = %v", err) + } + raw, readErr := os.ReadFile(target) + if readErr != nil || string(raw) != "existing" { + t.Fatalf("existing destination changed to %q, err=%v", raw, readErr) + } +} + +func TestDownloadChecksSizeBeforePublishingDestination(t *testing.T) { + store := memoryStore{"runs/training-1/result.json": []byte("actual")} + root := t.TempDir() + target := filepath.Join(root, "runs", "training-1", "result.json") + _, err := Download(context.Background(), store, Manifest{}, []Object{{ + Name: "runs/training-1/result.json", + Size: 999, + }}, root) + if err == nil || !strings.Contains(err.Error(), "expected 999") { + t.Fatalf("size mismatch error = %v", err) + } + if _, statErr := os.Stat(target); !os.IsNotExist(statErr) { + t.Fatalf("size mismatch published destination: %v", statErr) + } +} diff --git a/cli/internal/artifactbundle/publish.go b/cli/internal/artifactbundle/publish.go new file mode 100644 index 00000000..e1eeae1d --- /dev/null +++ b/cli/internal/artifactbundle/publish.go @@ -0,0 +1,143 @@ +package artifactbundle + +import ( + "fmt" + "path" + "strings" +) + +func WrapCommand(command []string, runtime Runtime) ([]string, error) { + if !runtime.Enabled() { + return command, nil + } + if len(command) == 0 { + return nil, fmt.Errorf("artifact bundle completion requires a Tau-wrappable command") + } + script, err := wrapperScript(runtime, `"$@" &`) + if err != nil { + return nil, err + } + return append([]string{"bash", "-c", script, "tau-bundle-entrypoint"}, command...), nil +} + +func WrapShellScript(command string, runtime Runtime) (string, error) { + if !runtime.Enabled() { + return command, nil + } + if strings.TrimSpace(command) == "" { + return "", fmt.Errorf("artifact bundle completion requires a non-empty entrypoint") + } + return wrapperScript(runtime, "(\n"+command+"\n) &") +} + +func wrapperScript(runtime Runtime, launch string) (string, error) { + manifest, err := runtime.Manifest() + if err != nil { + return "", err + } + raw, err := Marshal(manifest) + if err != nil { + return "", err + } + generationManifest := GenerationManifestPath(runtime.OutputDir, runtime.BundleID) + generationCompletion := GenerationCompletionPath(runtime.OutputDir, runtime.BundleID) + currentManifest := CurrentManifestPath(runtime.OutputDir) + currentCompletion := CurrentCompletionPath(runtime.OutputDir) + publicationCheck := "" + if manifest.Publication != nil { + publicationCheck = fmt.Sprintf(`if [ ! -f %s ] || [ "$(cat %s)" != %s ]; then + echo "artifact bundle refused: staged publication acknowledgement is missing or invalid" >&2 + exit 126 +fi +`, shellQuote(manifest.Publication.Completion), shellQuote(manifest.Publication.Completion), + shellQuote("complete "+manifest.Publication.ID)) + } + checkpointCheck := "" + if strings.TrimSpace(runtime.CheckpointIndex) != "" { + checkpointCheck = fmt.Sprintf(`if [ ! -f %s ] || + ! python3 - %s %s <<'TAU_BUNDLE_CHECKPOINT_EOF' +import json, pathlib, sys +index = json.loads(pathlib.Path(sys.argv[1]).read_text()) +if index.get("bundle_id") != sys.argv[2]: + raise SystemExit(1) +TAU_BUNDLE_CHECKPOINT_EOF +then + echo "artifact bundle refused: declared checkpoint index is missing or belongs to another bundle" >&2 + exit 126 +fi +`, shellQuote(runtime.CheckpointIndex), shellQuote(runtime.CheckpointIndex), shellQuote(runtime.BundleID)) + } + return fmt.Sprintf(`tau_bundle_child="" +tau_bundle_forward_signal() { + if [ -n "${tau_bundle_child:-}" ]; then + kill -TERM "$tau_bundle_child" 2>/dev/null || true + fi +} +trap tau_bundle_forward_signal TERM INT +mkdir -p %s +rm -f %s %s +%s +tau_bundle_child=$! +while :; do + wait "$tau_bundle_child" + tau_bundle_status=$? + if ! kill -0 "$tau_bundle_child" 2>/dev/null; then + break + fi +done +trap - TERM INT +if [ "$tau_bundle_status" -ne 0 ]; then + exit "$tau_bundle_status" +fi +%s%stau_bundle_tmp=%s +if ! printf '%%s' %s > "$tau_bundle_tmp"; then + rm -f "$tau_bundle_tmp" + echo "artifact bundle refused: could not write completion manifest" >&2 + exit 126 +fi +if [ -e %s ]; then + if ! cmp -s "$tau_bundle_tmp" %s; then + rm -f "$tau_bundle_tmp" + echo "artifact bundle refused: immutable generation manifest already differs" >&2 + exit 126 + fi + rm -f "$tau_bundle_tmp" +else + if ! mv -n "$tau_bundle_tmp" %s; then + rm -f "$tau_bundle_tmp" + echo "artifact bundle refused: could not commit immutable generation manifest" >&2 + exit 126 + fi +fi +tau_bundle_current_tmp=%s +if ! cp %s "$tau_bundle_current_tmp" || + ! mv -f "$tau_bundle_current_tmp" %s; then + rm -f "$tau_bundle_current_tmp" + echo "artifact bundle refused: could not publish current manifest" >&2 + exit 126 +fi +tau_bundle_marker_tmp=%s +if ! printf 'complete %%s\n' %s > "$tau_bundle_marker_tmp" || + ! mv -f "$tau_bundle_marker_tmp" %s; then + rm -f "$tau_bundle_marker_tmp" + echo "artifact bundle refused: could not commit generation acknowledgement" >&2 + exit 126 +fi +tau_bundle_current_marker_tmp=%s +if ! printf 'complete %%s\n' %s > "$tau_bundle_current_marker_tmp" || + ! mv -f "$tau_bundle_current_marker_tmp" %s; then + rm -f "$tau_bundle_current_marker_tmp" + echo "artifact bundle refused: could not commit current acknowledgement" >&2 + exit 126 +fi +`, shellQuote(path.Dir(generationManifest)), shellQuote(generationCompletion), shellQuote(currentCompletion), + launch, publicationCheck, checkpointCheck, shellQuote(generationManifest+".tmp.$$"), shellQuote(string(raw)), + shellQuote(generationManifest), shellQuote(generationManifest), shellQuote(generationManifest), + shellQuote(currentManifest+".tmp.$$"), shellQuote(generationManifest), shellQuote(currentManifest), + shellQuote(generationCompletion+".tmp.$$"), shellQuote(runtime.BundleID), shellQuote(generationCompletion), + shellQuote(currentCompletion+".tmp.$$"), shellQuote(runtime.BundleID), shellQuote(currentCompletion)), nil +} + +func shellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", `'"'"'`) + "'" +} diff --git a/cli/internal/artifactbundle/retrieve.go b/cli/internal/artifactbundle/retrieve.go new file mode 100644 index 00000000..d77fd143 --- /dev/null +++ b/cli/internal/artifactbundle/retrieve.go @@ -0,0 +1,340 @@ +package artifactbundle + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "os" + "path" + "path/filepath" + "sort" + "strings" +) + +type Object struct { + Name string `json:"name"` + Size int64 `json:"size_bytes"` +} + +type Store interface { + Read(context.Context, string) ([]byte, error) + List(context.Context, string) ([]Object, error) + Download(context.Context, string, io.Writer) error +} + +type DownloadedFile struct { + Path string `json:"path"` + Size int64 `json:"size_bytes"` + SHA256 string `json:"sha256"` +} + +func PVCRelativePath(absolute string) (string, error) { + clean := path.Clean(strings.TrimSpace(absolute)) + if clean == "/data" { + return "", nil + } + if !strings.HasPrefix(clean, "/data/") { + return "", fmt.Errorf("durable artifact path %q is not under the Blob CSI mount /data", absolute) + } + return strings.TrimPrefix(clean, "/data/"), nil +} + +func Load(ctx context.Context, store Store, resultRoot, bundleID string) (Manifest, error) { + manifestPath := CurrentManifestPath(resultRoot) + completionPath := CurrentCompletionPath(resultRoot) + if strings.TrimSpace(bundleID) != "" { + manifestPath = GenerationManifestPath(resultRoot, bundleID) + completionPath = GenerationCompletionPath(resultRoot, bundleID) + } + manifestKey, err := PVCRelativePath(manifestPath) + if err != nil { + return Manifest{}, err + } + raw, err := store.Read(ctx, manifestKey) + if err != nil { + return Manifest{}, fmt.Errorf("read artifact bundle manifest %s: %w", manifestPath, err) + } + manifest, err := Parse(raw) + if err != nil { + return Manifest{}, err + } + if path.Clean(manifest.ResultRoot) != path.Clean(resultRoot) { + return Manifest{}, fmt.Errorf("artifact bundle result root %q does not match requested root %q", manifest.ResultRoot, resultRoot) + } + if strings.TrimSpace(bundleID) != "" && manifest.BundleID != strings.TrimSpace(bundleID) { + return Manifest{}, fmt.Errorf("artifact bundle manifest ID %q does not match workload ID %q", manifest.BundleID, bundleID) + } + completionKey, err := PVCRelativePath(completionPath) + if err != nil { + return Manifest{}, err + } + completion, err := store.Read(ctx, completionKey) + if err != nil { + return Manifest{}, fmt.Errorf("artifact bundle is not acknowledged at %s: %w", completionPath, err) + } + if strings.TrimSpace(string(completion)) != "complete "+manifest.BundleID { + return Manifest{}, fmt.Errorf("artifact bundle acknowledgement %s is invalid", completionPath) + } + if manifest.Publication != nil { + key, err := PVCRelativePath(manifest.Publication.Completion) + if err != nil { + return Manifest{}, err + } + raw, err := store.Read(ctx, key) + if err != nil { + return Manifest{}, fmt.Errorf("staged artifacts are not completely published: %w", err) + } + if strings.TrimSpace(string(raw)) != "complete "+manifest.Publication.ID { + return Manifest{}, fmt.Errorf("staged artifact publication marker %s is invalid", manifest.Publication.Completion) + } + } + return manifest, nil +} + +func Enumerate(ctx context.Context, store Store, manifest Manifest) ([]Object, error) { + objects := map[string]Object{} + for _, spec := range manifest.Paths { + prefixPath := spec.Path + if spec.Kind == "glob" { + prefixPath = globPrefix(spec.Path) + } + prefix, err := PVCRelativePath(prefixPath) + if err != nil { + return nil, err + } + if prefix != "" && !strings.HasSuffix(prefix, "/") && spec.Kind == "tree" { + prefix += "/" + } + listed, err := store.List(ctx, prefix) + if err != nil { + return nil, fmt.Errorf("enumerate artifact bundle path %s: %w", spec.Path, err) + } + matches := 0 + for _, object := range listed { + if spec.Kind == "glob" { + absolute := "/data/" + strings.TrimPrefix(object.Name, "/") + ok, matchErr := path.Match(spec.Path, absolute) + if matchErr != nil { + return nil, fmt.Errorf("match artifact bundle glob %q: %w", spec.Path, matchErr) + } + if !ok { + continue + } + } else if spec.Kind == "file" && object.Name != prefix { + continue + } + objects[object.Name] = object + matches++ + } + if matches == 0 && !spec.Optional { + return nil, fmt.Errorf("artifact bundle required path %s is empty", spec.Path) + } + } + out := make([]Object, 0, len(objects)) + for _, object := range objects { + out = append(out, object) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out, nil +} + +func globPrefix(pattern string) string { + index := strings.IndexAny(pattern, "*?[") + if index < 0 { + return pattern + } + prefix := pattern[:index] + if slash := strings.LastIndex(prefix, "/"); slash >= 0 { + return prefix[:slash+1] + } + return "/data/" +} + +func Download(ctx context.Context, store Store, manifest Manifest, objects []Object, destination string) ([]DownloadedFile, error) { + if strings.TrimSpace(destination) == "" { + return nil, fmt.Errorf("artifact bundle destination is required") + } + root, err := filepath.Abs(destination) + if err != nil { + return nil, fmt.Errorf("resolve artifact bundle destination: %w", err) + } + if err := os.MkdirAll(root, 0o755); err != nil { + return nil, fmt.Errorf("create artifact bundle destination: %w", err) + } + if info, err := os.Lstat(root); err != nil { + return nil, err + } else if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { + return nil, fmt.Errorf("artifact bundle destination %s must be a real directory", root) + } + targets := make([]string, len(objects)) + seenTargets := make(map[string]string, len(objects)) + for i, object := range objects { + target, err := downloadTarget(root, object.Name) + if err != nil { + return nil, err + } + if previous, exists := seenTargets[target]; exists { + return nil, fmt.Errorf("artifact bundle objects %q and %q resolve to the same destination", previous, object.Name) + } + seenTargets[target] = object.Name + if _, err := os.Lstat(target); err == nil { + return nil, fmt.Errorf("artifact bundle destination file already exists: %s", target) + } else if !os.IsNotExist(err) { + return nil, err + } + targets[i] = target + } + metadataDir := filepath.Join(root, ".tau-bundle") + for _, target := range []string{ + filepath.Join(metadataDir, "manifest.json"), + filepath.Join(metadataDir, "files.json"), + } { + if _, err := os.Lstat(target); err == nil { + return nil, fmt.Errorf("artifact bundle metadata file already exists: %s", target) + } else if !os.IsNotExist(err) { + return nil, err + } + } + files := make([]DownloadedFile, 0, len(objects)) + for i, object := range objects { + target := targets[i] + if err := ensureSafeDirectory(root, filepath.Dir(target)); err != nil { + return nil, err + } + file, err := os.CreateTemp(filepath.Dir(target), ".tau-download-*") + if err != nil { + return nil, err + } + tmp := file.Name() + if err := file.Chmod(0o600); err != nil { + _ = file.Close() + _ = os.Remove(tmp) + return nil, err + } + hash := sha256.New() + writer := io.MultiWriter(file, hash) + downloadErr := store.Download(ctx, object.Name, writer) + closeErr := file.Close() + if downloadErr != nil { + _ = os.Remove(tmp) + return nil, fmt.Errorf("download artifact bundle object %s: %w", object.Name, downloadErr) + } + if closeErr != nil { + _ = os.Remove(tmp) + return nil, closeErr + } + info, err := os.Stat(tmp) + if err != nil { + _ = os.Remove(tmp) + return nil, err + } + if object.Size >= 0 && info.Size() != object.Size { + _ = os.Remove(tmp) + return nil, fmt.Errorf("downloaded artifact %s has %d bytes, expected %d", object.Name, info.Size(), object.Size) + } + if err := os.Link(tmp, target); err != nil { + _ = os.Remove(tmp) + return nil, err + } + if err := os.Remove(tmp); err != nil { + _ = os.Remove(target) + return nil, err + } + files = append(files, DownloadedFile{ + Path: object.Name, + Size: info.Size(), + SHA256: hex.EncodeToString(hash.Sum(nil)), + }) + } + if err := ensureSafeDirectory(root, metadataDir); err != nil { + return nil, err + } + if err := writeJSONAtomic(filepath.Join(metadataDir, "manifest.json"), manifest); err != nil { + return nil, err + } + if err := writeJSONAtomic(filepath.Join(metadataDir, "files.json"), files); err != nil { + return nil, err + } + return files, nil +} + +func downloadTarget(root, objectName string) (string, error) { + if objectName == "" || strings.HasPrefix(objectName, "/") || path.Clean(objectName) != objectName { + return "", fmt.Errorf("artifact bundle object name %q is not a canonical relative path", objectName) + } + target := filepath.Join(root, filepath.FromSlash(objectName)) + if !pathWithin(root, target) { + return "", fmt.Errorf("artifact bundle object %q escapes destination", objectName) + } + return target, nil +} + +func ensureSafeDirectory(root, directory string) error { + rel, err := filepath.Rel(root, directory) + if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return fmt.Errorf("artifact bundle directory %q escapes destination", directory) + } + current := root + for _, part := range strings.Split(rel, string(filepath.Separator)) { + if part == "" || part == "." { + continue + } + current = filepath.Join(current, part) + info, err := os.Lstat(current) + switch { + case os.IsNotExist(err): + if err := os.Mkdir(current, 0o755); err != nil { + return err + } + case err != nil: + return err + case info.Mode()&os.ModeSymlink != 0 || !info.IsDir(): + return fmt.Errorf("artifact bundle destination parent %s must be a real directory", current) + } + } + return nil +} + +func pathWithin(root, target string) bool { + rel, err := filepath.Rel(root, target) + return err == nil && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) +} + +func writeJSONAtomic(target string, value any) error { + raw, err := json.MarshalIndent(value, "", " ") + if err != nil { + return err + } + raw = append(raw, '\n') + file, err := os.CreateTemp(filepath.Dir(target), ".tau-metadata-*") + if err != nil { + return err + } + tmp := file.Name() + if err := file.Chmod(0o600); err != nil { + _ = file.Close() + _ = os.Remove(tmp) + return err + } + if _, err := file.Write(raw); err != nil { + _ = file.Close() + _ = os.Remove(tmp) + return err + } + if err := file.Close(); err != nil { + _ = os.Remove(tmp) + return err + } + if err := os.Link(tmp, target); err != nil { + _ = os.Remove(tmp) + return err + } + if err := os.Remove(tmp); err != nil { + _ = os.Remove(target) + return err + } + return nil +} diff --git a/cli/internal/artifactindex/finalize.go b/cli/internal/artifactindex/finalize.go index e81a9c70..ab1e5734 100644 --- a/cli/internal/artifactindex/finalize.go +++ b/cli/internal/artifactindex/finalize.go @@ -33,6 +33,9 @@ type Config struct { // ResourceName and Namespace are recorded in the index for provenance. ResourceName string Namespace string + // BundleID binds the emitted index to the final bundle acknowledgement. + // Empty preserves the legacy index schema for callers without bundles. + BundleID string } // Script returns a POSIX shell snippet that finalizes the declared checkpoint @@ -59,7 +62,8 @@ func Script(cfg Config) string { b.WriteString("TAU_ARTIFACT_RUN=" + shellQuote(cfg.Run) + "\n") b.WriteString("TAU_ARTIFACT_RESOURCE=" + shellQuote(cfg.ResourceName) + "\n") b.WriteString("TAU_ARTIFACT_NAMESPACE=" + shellQuote(cfg.Namespace) + "\n") - b.WriteString("export TAU_ARTIFACT_CHECKPOINT TAU_ARTIFACT_RUN TAU_ARTIFACT_RESOURCE TAU_ARTIFACT_NAMESPACE\n") + b.WriteString("TAU_ARTIFACT_BUNDLE_ID=" + shellQuote(cfg.BundleID) + "\n") + b.WriteString("export TAU_ARTIFACT_CHECKPOINT TAU_ARTIFACT_RUN TAU_ARTIFACT_RESOURCE TAU_ARTIFACT_NAMESPACE TAU_ARTIFACT_BUNDLE_ID\n") b.WriteString(finalizeScript) return b.String() } @@ -107,6 +111,7 @@ import datetime, json, os, pathlib, shutil, sys artifact = os.environ.get("TAU_ARTIFACT_CHECKPOINT", "").strip() run = os.environ.get("TAU_ARTIFACT_RUN", "").strip() +bundle_id = os.environ.get("TAU_ARTIFACT_BUNDLE_ID", "").strip() if not artifact or not run: sys.exit(0) @@ -204,6 +209,8 @@ index = { "durable_root": durable.as_posix(), "artifacts": [record], } +if bundle_id: + index["bundle_id"] = bundle_id index_path = run_dir / "artifacts.json" tmp = index_path.with_suffix(".json.tmp") diff --git a/cli/internal/artifactindex/finalize_test.go b/cli/internal/artifactindex/finalize_test.go index 1eda889c..d4f7e292 100644 --- a/cli/internal/artifactindex/finalize_test.go +++ b/cli/internal/artifactindex/finalize_test.go @@ -123,6 +123,7 @@ func TestScriptWritesIndexMatchingReaderSchema(t *testing.T) { Run: run, ResourceName: "demo-run-rayjob", Namespace: "research", + BundleID: "bundle-1", }) if script == "" { t.Fatal("Script returned empty for a fully-specified config") @@ -154,6 +155,9 @@ func TestScriptWritesIndexMatchingReaderSchema(t *testing.T) { // required: the reader's optional fields (e.g. storage_probe) are written // by other producers. assertSubset(t, index, readerStructTags(t, "managedWorkflowArtifactIndex"), "index") + if index["bundle_id"] != "bundle-1" { + t.Fatalf("bundle_id = %#v, want bundle-1", index["bundle_id"]) + } artifacts, ok := index["artifacts"].([]any) if !ok || len(artifacts) != 1 { diff --git a/cli/internal/cli/pvc_helpers.go b/cli/internal/cli/pvc_helpers.go index 6ed0d2b8..f641a0b2 100644 --- a/cli/internal/cli/pvc_helpers.go +++ b/cli/internal/cli/pvc_helpers.go @@ -34,6 +34,7 @@ const ( type managedWorkflowArtifactIndex struct { SchemaVersion int `json:"schema_version"` Run string `json:"run"` + BundleID string `json:"bundle_id,omitempty"` Namespace string `json:"namespace,omitempty"` ResourceName string `json:"resource_name,omitempty"` CreatedAt string `json:"created_at,omitempty"` @@ -234,11 +235,6 @@ func fetchPVCList(ctx context.Context, kubeContext, namespace, runName, pvcName, return fetchPVCListWithMode(ctx, kubeContext, namespace, runName, pvcName, dirPath, false) } -// fetchPVCListRecursive returns every descendant as a relative path. -func fetchPVCListRecursive(ctx context.Context, kubeContext, namespace, runName, pvcName, dirPath string) ([]string, error) { - return fetchPVCListWithMode(ctx, kubeContext, namespace, runName, pvcName, dirPath, true) -} - func fetchPVCListWithMode(ctx context.Context, kubeContext, namespace, runName, pvcName, dirPath string, recursive bool) ([]string, error) { if pvcName == "" { pvcName = defaultTauPVCName diff --git a/cli/internal/cli/run_blob_store.go b/cli/internal/cli/run_blob_store.go new file mode 100644 index 00000000..ee1840e9 --- /dev/null +++ b/cli/internal/cli/run_blob_store.go @@ -0,0 +1,231 @@ +package cli + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/url" + "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore/to" + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container" + + "github.com/Azure/taugrid/cli/internal/artifactbundle" +) + +const azureBlobCSIDriver = "blob.csi.azure.com" + +const runBlobVolumeSchema = "tau.run.blob-volume.v1" + +type runBlobVolume struct { + SchemaVersion string `json:"schema_version"` + AccountURL string `json:"account_url"` + Container string `json:"container"` +} + +func resolveRunBlobVolume(ctx context.Context, reader runResultReader, namespace, pvcName string) (runBlobVolume, error) { + rawPVC, err := reader.Raw(ctx, []string{"get", "pvc", pvcName, "-n", namespace, "-o", "json"}, nil) + if err != nil { + return runBlobVolume{}, fmt.Errorf("resolve artifact transport from PVC %s/%s: %w", namespace, pvcName, err) + } + var pvc struct { + Spec struct { + VolumeName string `json:"volumeName"` + } `json:"spec"` + } + if err := json.Unmarshal([]byte(rawPVC), &pvc); err != nil { + return runBlobVolume{}, fmt.Errorf("decode PVC %s/%s: %w", namespace, pvcName, err) + } + if strings.TrimSpace(pvc.Spec.VolumeName) == "" { + return runBlobVolume{}, fmt.Errorf("PVC %s/%s is not bound to a PersistentVolume", namespace, pvcName) + } + rawPV, err := reader.Raw(ctx, []string{"get", "pv", pvc.Spec.VolumeName, "-o", "json"}, nil) + if err != nil { + return runBlobVolume{}, fmt.Errorf("resolve artifact transport from PV %s: %w", pvc.Spec.VolumeName, err) + } + var pv struct { + Spec struct { + CSI struct { + Driver string `json:"driver"` + VolumeHandle string `json:"volumeHandle"` + VolumeAttributes map[string]string `json:"volumeAttributes"` + } `json:"csi"` + } `json:"spec"` + } + if err := json.Unmarshal([]byte(rawPV), &pv); err != nil { + return runBlobVolume{}, fmt.Errorf("decode PV %s: %w", pvc.Spec.VolumeName, err) + } + if !strings.EqualFold(strings.TrimSpace(pv.Spec.CSI.Driver), azureBlobCSIDriver) { + return runBlobVolume{}, fmt.Errorf( + "PVC %s/%s uses CSI driver %q; complete Tau bundle retrieval currently requires %s and will not create a PVC-reader pod", + namespace, pvcName, pv.Spec.CSI.Driver, azureBlobCSIDriver, + ) + } + attributes, err := foldVolumeAttributes(pv.Spec.CSI.VolumeAttributes) + if err != nil { + return runBlobVolume{}, fmt.Errorf("PV %s: %w", pvc.Spec.VolumeName, err) + } + account := strings.TrimSpace(attributes["storageaccount"]) + containerName := strings.TrimSpace(attributes["containername"]) + parts := strings.Split(pv.Spec.CSI.VolumeHandle, "#") + if account == "" && len(parts) > 1 { + account = strings.TrimSpace(parts[1]) + } + if containerName == "" && len(parts) > 2 { + containerName = strings.TrimSpace(parts[2]) + } + if account == "" || containerName == "" { + return runBlobVolume{}, fmt.Errorf( + "PV %s does not expose a Blob CSI storageAccount/containerName identity; Tau will not read Secret credentials", + pvc.Spec.VolumeName, + ) + } + server := strings.TrimSpace(attributes["server"]) + if server == "" { + suffix := firstNonEmpty(attributes["storageendpointsuffix"], "core.windows.net") + server = account + ".blob." + suffix + } + if !strings.Contains(server, "://") { + server = "https://" + server + } + parsed, err := url.Parse(server) + if err != nil || parsed.Scheme != "https" || parsed.Host == "" || parsed.User != nil { + return runBlobVolume{}, fmt.Errorf("PV %s has invalid Blob server %q", pvc.Spec.VolumeName, server) + } + return runBlobVolume{ + SchemaVersion: runBlobVolumeSchema, + AccountURL: strings.TrimSuffix(parsed.String(), "/"), + Container: containerName, + }, nil +} + +func parseRunBlobVolume(raw string) (runBlobVolume, error) { + var volume runBlobVolume + decoder := json.NewDecoder(strings.NewReader(raw)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&volume); err != nil { + return runBlobVolume{}, fmt.Errorf("decode Blob artifact transport annotation: %w", err) + } + if err := validateRunBlobVolume(volume); err != nil { + return runBlobVolume{}, err + } + if !trustedAzureBlobHost(volume.AccountURL) { + return runBlobVolume{}, fmt.Errorf("Blob artifact transport account URL %q is not a trusted Azure Blob endpoint", volume.AccountURL) + } + return volume, nil +} + +func validateRunBlobVolume(volume runBlobVolume) error { + if volume.SchemaVersion != runBlobVolumeSchema { + return fmt.Errorf("Blob artifact transport schema %q is unsupported", volume.SchemaVersion) + } + parsed, err := url.Parse(strings.TrimSpace(volume.AccountURL)) + if err != nil || parsed.Scheme != "https" || parsed.Host == "" || parsed.User != nil || + parsed.RawQuery != "" || parsed.Fragment != "" || (parsed.Path != "" && parsed.Path != "/") { + return fmt.Errorf("Blob artifact transport account URL %q is invalid", volume.AccountURL) + } + containerName := strings.TrimSpace(volume.Container) + if containerName == "" || strings.ContainsAny(containerName, `/\`) { + return fmt.Errorf("Blob artifact transport container %q is invalid", volume.Container) + } + return nil +} + +func trustedAzureBlobHost(accountURL string) bool { + parsed, err := url.Parse(strings.TrimSpace(accountURL)) + if err != nil { + return false + } + host := strings.ToLower(strings.TrimSuffix(parsed.Hostname(), ".")) + for _, suffix := range []string{ + ".blob.core.windows.net", + ".blob.core.usgovcloudapi.net", + ".blob.core.chinacloudapi.cn", + ".blob.core.cloudapi.de", + } { + if strings.HasSuffix(host, suffix) && len(host) > len(suffix) { + return true + } + } + return false +} + +func foldVolumeAttributes(attributes map[string]string) (map[string]string, error) { + out := make(map[string]string, len(attributes)) + for key, value := range attributes { + folded := strings.ToLower(strings.TrimSpace(key)) + if previous, exists := out[folded]; exists && previous != value { + return nil, fmt.Errorf("volumeAttributes contains conflicting case variants for %q", key) + } + out[folded] = value + } + return out, nil +} + +type azureRunArtifactStore struct { + container *container.Client +} + +func newAzureRunArtifactStore(volume runBlobVolume) (*azureRunArtifactStore, error) { + if err := validateRunBlobVolume(volume); err != nil { + return nil, err + } + if !trustedAzureBlobHost(volume.AccountURL) { + return nil, fmt.Errorf("Blob artifact transport account URL %q is not a trusted Azure Blob endpoint", volume.AccountURL) + } + credential, err := azidentity.NewDefaultAzureCredential(nil) + if err != nil { + return nil, fmt.Errorf("create Azure credential for artifact retrieval: %w", err) + } + client, err := container.NewClient(volume.AccountURL+"/"+url.PathEscape(volume.Container), credential, nil) + if err != nil { + return nil, fmt.Errorf("create Azure Blob artifact client: %w", err) + } + return &azureRunArtifactStore{container: client}, nil +} + +func (s *azureRunArtifactStore) Read(ctx context.Context, name string) ([]byte, error) { + var out bytes.Buffer + if err := s.Download(ctx, name, &out); err != nil { + return nil, err + } + return out.Bytes(), nil +} + +func (s *azureRunArtifactStore) List(ctx context.Context, prefix string) ([]artifactbundle.Object, error) { + pager := s.container.NewListBlobsFlatPager(&container.ListBlobsFlatOptions{Prefix: to.Ptr(prefix)}) + var objects []artifactbundle.Object + for pager.More() { + page, err := pager.NextPage(ctx) + if err != nil { + return nil, err + } + if page.Segment == nil { + continue + } + for _, item := range page.Segment.BlobItems { + if item == nil || item.Name == nil { + continue + } + size := int64(-1) + if item.Properties != nil && item.Properties.ContentLength != nil { + size = *item.Properties.ContentLength + } + objects = append(objects, artifactbundle.Object{Name: *item.Name, Size: size}) + } + } + return objects, nil +} + +func (s *azureRunArtifactStore) Download(ctx context.Context, name string, out io.Writer) error { + response, err := s.container.NewBlockBlobClient(name).DownloadStream(ctx, nil) + if err != nil { + return err + } + defer response.Body.Close() + _, err = io.Copy(out, response.Body) + return err +} diff --git a/cli/internal/cli/run_blob_store_test.go b/cli/internal/cli/run_blob_store_test.go new file mode 100644 index 00000000..be446bc5 --- /dev/null +++ b/cli/internal/cli/run_blob_store_test.go @@ -0,0 +1,134 @@ +package cli + +import ( + "context" + "errors" + "strings" + "testing" +) + +type fakeBlobVolumeReader struct { + responses map[string]string + errors map[string]error + calls []string +} + +func TestRunBlobVolumeAnnotationParsesTrustedEndpoint(t *testing.T) { + want := runBlobVolume{ + SchemaVersion: runBlobVolumeSchema, + AccountURL: "https://trainingacct.blob.core.windows.net", + Container: "results", + } + got, err := parseRunBlobVolume(`{"schema_version":"tau.run.blob-volume.v1","account_url":"https://trainingacct.blob.core.windows.net","container":"results"}`) + if err != nil { + t.Fatal(err) + } + if got != want { + t.Fatalf("round trip = %+v, want %+v", got, want) + } +} + +func TestRunBlobVolumeAnnotationRejectsCredentialBearingURL(t *testing.T) { + _, err := parseRunBlobVolume(`{"schema_version":"tau.run.blob-volume.v1","account_url":"https://user:secret@example.test","container":"results"}`) + if err == nil || !strings.Contains(err.Error(), "invalid") { + t.Fatalf("credential-bearing URL error = %v", err) + } +} + +func TestRunBlobVolumeAnnotationRejectsUntrustedHost(t *testing.T) { + _, err := parseRunBlobVolume(`{"schema_version":"tau.run.blob-volume.v1","account_url":"https://attacker.example","container":"results"}`) + if err == nil || !strings.Contains(err.Error(), "trusted Azure Blob endpoint") { + t.Fatalf("untrusted host error = %v", err) + } +} + +func TestAzureRunArtifactStoreRejectsUntrustedHostBeforeCredentialUse(t *testing.T) { + _, err := newAzureRunArtifactStore(runBlobVolume{ + SchemaVersion: runBlobVolumeSchema, + AccountURL: "https://attacker.example", + Container: "results", + }) + if err == nil || !strings.Contains(err.Error(), "trusted Azure Blob endpoint") { + t.Fatalf("untrusted store error = %v", err) + } +} + +func (f *fakeBlobVolumeReader) Raw(_ context.Context, args []string, _ []byte) (string, error) { + key := strings.Join(args, " ") + f.calls = append(f.calls, key) + if err := f.errors[key]; err != nil { + return "", err + } + return f.responses[key], nil +} + +func TestResolveRunBlobVolumeFromDynamicCSIHandle(t *testing.T) { + reader := &fakeBlobVolumeReader{responses: map[string]string{ + "get pvc blob-training -n research -o json": `{"spec":{"volumeName":"pvc-123"}}`, + "get pv pvc-123 -o json": `{"spec":{"csi":{"driver":"blob.csi.azure.com","volumeHandle":"rg#trainingacct#results#uuid#research#subscription","volumeAttributes":{"storageEndpointSuffix":"core.windows.net"}}}}`, + }} + volume, err := resolveRunBlobVolume(context.Background(), reader, "research", "blob-training") + if err != nil { + t.Fatal(err) + } + if volume.AccountURL != "https://trainingacct.blob.core.windows.net" || volume.Container != "results" { + t.Fatalf("resolved volume = %+v", volume) + } + if len(reader.calls) != 2 { + t.Fatalf("calls = %v", reader.calls) + } +} + +func TestResolveRunBlobVolumePrefersStaticAttributesAndPrivateServer(t *testing.T) { + reader := &fakeBlobVolumeReader{responses: map[string]string{ + "get pvc blob-training -n research -o json": `{"spec":{"volumeName":"static-pv"}}`, + "get pv static-pv -o json": `{"spec":{"csi":{"driver":"blob.csi.azure.com","volumeHandle":"opaque","volumeAttributes":{"storageAccount":"trainingacct","containerName":"artifacts","server":"trainingacct.privatelink.blob.core.windows.net"}}}}`, + }} + volume, err := resolveRunBlobVolume(context.Background(), reader, "research", "blob-training") + if err != nil { + t.Fatal(err) + } + if volume.AccountURL != "https://trainingacct.privatelink.blob.core.windows.net" || volume.Container != "artifacts" { + t.Fatalf("resolved volume = %+v", volume) + } +} + +func TestResolveRunBlobVolumeRejectsUnsupportedStorageWithoutCreatingReader(t *testing.T) { + reader := &fakeBlobVolumeReader{responses: map[string]string{ + "get pvc shared -n research -o json": `{"spec":{"volumeName":"nfs-pv"}}`, + "get pv nfs-pv -o json": `{"spec":{"csi":{"driver":"file.csi.azure.com","volumeHandle":"opaque"}}}`, + }} + _, err := resolveRunBlobVolume(context.Background(), reader, "research", "shared") + if err == nil || !strings.Contains(err.Error(), "will not create a PVC-reader pod") { + t.Fatalf("unsupported storage error = %v", err) + } +} + +func TestResolveRunBlobVolumeDoesNotReadSecretBackedIdentity(t *testing.T) { + reader := &fakeBlobVolumeReader{responses: map[string]string{ + "get pvc shared -n research -o json": `{"spec":{"volumeName":"secret-pv"}}`, + "get pv secret-pv -o json": `{"spec":{"csi":{"driver":"blob.csi.azure.com","volumeHandle":"opaque","volumeAttributes":{"secretName":"storage-key"}}}}`, + }} + _, err := resolveRunBlobVolume(context.Background(), reader, "research", "shared") + if err == nil || !strings.Contains(err.Error(), "will not read Secret credentials") { + t.Fatalf("secret-backed identity error = %v", err) + } + for _, call := range reader.calls { + if strings.Contains(call, " secret ") { + t.Fatalf("resolver attempted to read a Secret: %v", reader.calls) + } + } +} + +func TestResolveRunBlobVolumeSurfacesMetadataAuthorizationFailure(t *testing.T) { + reader := &fakeBlobVolumeReader{ + responses: map[string]string{}, + errors: map[string]error{ + "get pvc blob-training -n research -o json": errors.New("forbidden"), + }, + } + _, err := resolveRunBlobVolume(context.Background(), reader, "research", "blob-training") + if err == nil || !strings.Contains(err.Error(), "forbidden") { + t.Fatalf("authorization error = %v", err) + } +} diff --git a/cli/internal/cli/run_bundle.go b/cli/internal/cli/run_bundle.go new file mode 100644 index 00000000..61b1e951 --- /dev/null +++ b/cli/internal/cli/run_bundle.go @@ -0,0 +1,58 @@ +package cli + +import ( + "path" + "strings" + + "github.com/Azure/taugrid/cli/internal/artifactbundle" + "github.com/Azure/taugrid/cli/internal/artifactpublish" + "github.com/Azure/taugrid/cli/internal/metricsoffload" + "github.com/Azure/taugrid/cli/internal/storage" +) + +func resolveArtifactBundle( + run, namespace, submissionID, outputDir, resultPVC string, + outputWritable bool, + publication artifactpublish.Runtime, + metrics metricsoffload.Runtime, + metricsSessionID string, + checkpointArtifact string, +) (artifactbundle.Runtime, error) { + outputDir = path.Clean(strings.TrimSpace(outputDir)) + if outputDir == "." || strings.TrimSpace(resultPVC) == "" || !outputWritable { + return artifactbundle.Runtime{}, nil + } + if outputDir != "/data" && !strings.HasPrefix(outputDir, "/data/") { + return artifactbundle.Runtime{}, nil + } + bundleID := firstNonEmpty(publication.PublicationID, submissionID) + if bundleID == "" { + return artifactbundle.Runtime{}, nil + } + runtime := artifactbundle.Runtime{ + BundleID: bundleID, + Run: run, + Namespace: namespace, + ResultPVC: resultPVC, + OutputDir: outputDir, + PublicationMode: publication.Mode, + PublicationID: publication.PublicationID, + MetricsSessionID: strings.TrimSpace(metricsSessionID), + MetricsHistory: append([]string(nil), metrics.History...), + MetricsOffloadDir: metrics.Out, + MetricsEnabled: metrics.Enabled(), + CheckpointArtifact: strings.TrimSpace(checkpointArtifact), + } + if publication.Enabled() { + runtime.PublicationRoot = publication.PublishedDir() + runtime.PublicationMarker = path.Join(publication.PublishedDir(), artifactpublish.CompletionMarker) + } + if runtime.CheckpointArtifact != "" { + runtime.CheckpointRoot = storage.DurableFinetuneDir(run) + runtime.CheckpointIndex = storage.DurableFinetuneArtifactsFile(run) + } + if err := runtime.Validate(); err != nil { + return artifactbundle.Runtime{}, err + } + return runtime, nil +} diff --git a/cli/internal/cli/run_bundle_test.go b/cli/internal/cli/run_bundle_test.go new file mode 100644 index 00000000..73530a50 --- /dev/null +++ b/cli/internal/cli/run_bundle_test.go @@ -0,0 +1,93 @@ +package cli + +import ( + "path" + "testing" + "time" + + "github.com/Azure/taugrid/cli/internal/artifactbundle" + "github.com/Azure/taugrid/cli/internal/artifactpublish" + "github.com/Azure/taugrid/cli/internal/metricsoffload" +) + +func TestResolveArtifactBundleOwnsAllProducerPaths(t *testing.T) { + publication := artifactpublish.Runtime{ + Mode: artifactpublish.ModeStaged, + OutputDir: "/data/runs/training-1", + StagingDir: "/mnt/tau-output/training-1", + PublicationID: "publication-1", + } + metrics := metricsoffload.Runtime{ + Image: "registry.example/tau:v1", + RunID: "training-1", + Project: "project", + Experiment: "experiment", + Group: "group", + Store: "/var/run/tau/metrics/session/expstore", + Out: "/data/runs/training-1/.tau/metrics/session/offload", + History: []string{"/data/runs/training-1/metrics/*.jsonl"}, + CompletionFile: "/var/run/tau/metrics-completion.json", + RemoteWriteEndpoint: "https://metrics.example/receive", + Interval: time.Second, + DoneFile: "/var/run/tau/metrics-done", + } + runtime, err := resolveArtifactBundle( + "training-1", + "research", + "submission-1", + "/data/runs/training-1", + "blob-training", + true, + publication, + metrics, + "session-1", + "last.safetensors", + ) + if err != nil { + t.Fatal(err) + } + if runtime.BundleID != publication.PublicationID { + t.Fatalf("bundle ID = %q, want publication ID", runtime.BundleID) + } + manifest, err := runtime.Manifest() + if err != nil { + t.Fatal(err) + } + if manifest.Metrics == nil || !manifest.Metrics.Acknowledged || manifest.Metrics.SessionID != "session-1" { + t.Fatalf("metrics contract = %+v", manifest.Metrics) + } + if manifest.Checkpoint == nil || manifest.Checkpoint.Index != "/data/checkpoints/finetunes/training-1/artifacts.json" { + t.Fatalf("checkpoint contract = %+v", manifest.Checkpoint) + } + if manifest.Publication == nil || + manifest.Publication.Completion != path.Join(publication.PublishedDir(), artifactpublish.CompletionMarker) { + t.Fatalf("publication contract = %+v", manifest.Publication) + } + if got := artifactbundle.GenerationManifestPath(runtime.OutputDir, runtime.BundleID); got != "/data/runs/training-1/.tau/bundles/publication-1.json" { + t.Fatalf("generation manifest path = %q", got) + } +} + +func TestResolveArtifactBundleSkipsReadOnlyAndEphemeralResults(t *testing.T) { + for _, test := range []struct { + output string + pvc string + writable bool + }{ + {output: "", pvc: "blob-training", writable: true}, + {output: "/data/runs/training-1", pvc: "", writable: true}, + {output: "/data/runs/training-1", pvc: "blob-training", writable: false}, + {output: "/data-nfs/runs/training-1", pvc: "shared-nfs", writable: true}, + } { + runtime, err := resolveArtifactBundle( + "training-1", "research", "submission-1", test.output, test.pvc, test.writable, + artifactpublish.Runtime{}, metricsoffload.Runtime{}, "", "", + ) + if err != nil { + t.Fatal(err) + } + if runtime.Enabled() { + t.Fatalf("runtime enabled for %+v", test) + } + } +} diff --git a/cli/internal/cli/run_get.go b/cli/internal/cli/run_get.go index 8847335e..9bbe0847 100644 --- a/cli/internal/cli/run_get.go +++ b/cli/internal/cli/run_get.go @@ -6,10 +6,12 @@ import ( "fmt" "path" "path/filepath" + "sort" "strings" "github.com/spf13/cobra" + "github.com/Azure/taugrid/cli/internal/artifactbundle" "github.com/Azure/taugrid/cli/internal/artifactpublish" "github.com/Azure/taugrid/core/kube" "github.com/Azure/taugrid/core/workloadmeta" @@ -20,6 +22,8 @@ type runResultRef struct { PVC string Publication string PublicationID string + BundleID string + ArtifactStore string // CheckpointArtifact is the storage.checkpoint value the run declared, // empty when it declared none. Presence is what lets an empty result // directory be reported as a missing promised artifact rather than as an @@ -45,24 +49,34 @@ func newRunGetCmd() *cobra.Command { artifact string pathOverride string pvcOverride string + destination string output string ) cmd := &cobra.Command{ Use: "get NAME", Short: "Fetch the result file or directory recorded by storage.output", - Long: `Fetch artifacts a run Job wrote to its configured storage.output path. + Long: `Fetch artifacts a run Job or RayJob wrote to its configured storage.output path. Reads the ` + workloadmeta.AnnotationResultPath + ` and ` + workloadmeta.AnnotationResultPVC + ` annotations the run recorded on the Job. If the path is a file, it's catted directly. If the path is a directory, its recursive listing is printed; pass --artifact NAME to -fetch one file from it. Object-backed mounts are allowed to settle before Tau -accepts a zero-entry listing, so a populated directory is not reported as empty -during the BlobFuse mount-time list suppression window. +fetch one file from it. + +For runs submitted by a current Tau version, --destination downloads the complete +acknowledged bundle: staged terminal artifacts, immutable metrics histories and +offload metadata, plus the durable checkpoint tree when declared. Tau records the +non-secret Blob CSI account/container identity on new workloads (and falls back +to bound-PV discovery for legacy workloads), then reads through Azure RBAC with +DefaultAzureCredential. It never reads storage Secrets, account keys, or SAS +tokens and does not create a PVC-reader pod. The command fails closed if either +the staged publication marker or final bundle acknowledgement is absent, and it +never replaces existing destination files. Override the recorded path/pvc with --path/--pvc. Examples: tau run get swordfish-bench-001 -n ray tau run get swordfish-bench-001 -n ray --artifact profile/rank-0.summary.md + tau run get swordfish-bench-001 -n ray --destination ./results tau run get my-job --path /data/my-job/results -o json`, Args: cobra.ExactArgs(1), RunE: func(cmd *cobra.Command, args []string) error { @@ -86,10 +100,16 @@ Examples: } } if pathOverride != "" { + if ref.Path != "" && path.Clean(pathOverride) != path.Clean(ref.Path) { + ref.BundleID = "" + ref.Publication = "" + ref.PublicationID = "" + } ref.Path = pathOverride } if pvcOverride != "" { ref.PVC = pvcOverride + ref.ArtifactStore = "" } if ref.Path == "" { return fmt.Errorf("run workload %s/%s has no %s; resubmit with storage.output, or pass --path", ns, name, workloadmeta.AnnotationResultPath) @@ -97,6 +117,60 @@ Examples: if ref.PVC == "" { return fmt.Errorf("run workload %s/%s has no %s; pass --pvc to override", ns, name, workloadmeta.AnnotationResultPVC) } + if destination != "" && artifact != "" { + return fmt.Errorf("--destination and --artifact cannot be combined") + } + var blobVolume runBlobVolume + if strings.TrimSpace(ref.ArtifactStore) != "" { + blobVolume, err = parseRunBlobVolume(ref.ArtifactStore) + } else { + blobVolume, err = resolveRunBlobVolume(cmd.Context(), kube.New(resolvedContext), ns, ref.PVC) + } + if err != nil { + return err + } + store, err := newAzureRunArtifactStore(blobVolume) + if err != nil { + return err + } + if ref.BundleID != "" || destination != "" { + manifest, loadErr := artifactbundle.Load(cmd.Context(), store, ref.Path, ref.BundleID) + if loadErr != nil { + if destination != "" { + return fmt.Errorf( + "complete artifact bundle is unavailable: %w; this run may predate Tau's final bundle acknowledgement", + loadErr, + ) + } + return loadErr + } + if manifest.ResultPVC != ref.PVC || path.Clean(manifest.ResultRoot) != path.Clean(ref.Path) { + return fmt.Errorf("artifact bundle identity does not match workload result metadata") + } + if destination != "" { + objects, err := artifactbundle.Enumerate(cmd.Context(), store, manifest) + if err != nil { + return err + } + files, err := artifactbundle.Download(cmd.Context(), store, manifest, objects, destination) + if err != nil { + return err + } + return writeRunBundleDownload(cmd, output, manifest, destination, files) + } + if artifact == "" { + objects, err := artifactbundle.Enumerate(cmd.Context(), store, manifest) + if err != nil { + return err + } + entries := make([]string, 0, len(objects)) + for _, object := range objects { + entries = append(entries, object.Name) + } + return writeRunGet(cmd, output, nil, entries, manifest.ResultRoot, manifest.ResultPVC, ref.CheckpointArtifact) + } + } + resultPath := ref.Path if ref.Publication == artifactpublish.ModeStaged { if strings.TrimSpace(ref.PublicationID) == "" { @@ -104,7 +178,7 @@ Examples: } resultPath = path.Join(ref.Path, artifactpublish.GenerationsDir, ref.PublicationID) marker := path.Join(resultPath, artifactpublish.CompletionMarker) - raw, err := fetchPVCFile(cmd.Context(), resolvedContext, ns, name, ref.PVC, marker) + raw, err := readRunBlobPath(cmd.Context(), store, marker) if err != nil { return fmt.Errorf("staged artifacts are not completely published: %w", err) } @@ -121,20 +195,20 @@ Examples: return err } file := path.Join(resultPath, cleanArtifact) - raw, err := fetchPVCFile(cmd.Context(), resolvedContext, ns, name, ref.PVC, file) + raw, err := readRunBlobPath(cmd.Context(), store, file) if err != nil { return err } return writeRunGet(cmd, output, raw, nil, file, ref.PVC, "") } if !isDir { - raw, err := fetchPVCFile(cmd.Context(), resolvedContext, ns, name, ref.PVC, resultPath) + raw, err := readRunBlobPath(cmd.Context(), store, resultPath) if err != nil { return err } return writeRunGet(cmd, output, raw, nil, resultPath, ref.PVC, "") } - entries, err := fetchPVCListRecursive(cmd.Context(), resolvedContext, ns, name, ref.PVC, resultPath) + entries, err := listRunBlobPath(cmd.Context(), store, resultPath) if err != nil { return err } @@ -145,6 +219,7 @@ Examples: cmd.Flags().StringVar(&artifact, "artifact", "", "fetch this filename under the result directory") cmd.Flags().StringVar(&pathOverride, "path", "", "override the recorded result path") cmd.Flags().StringVar(&pvcOverride, "pvc", "", "override the recorded result PVC") + cmd.Flags().StringVarP(&destination, "destination", "d", "", "download the complete acknowledged bundle into this directory") cmd.Flags().StringVarP(&output, "output", "o", "table", "table|json|raw") return cmd } @@ -212,10 +287,92 @@ func parseRunResultRef(raw []byte, resource string) (runResultRef, error) { PVC: obj.Metadata.Annotations[workloadmeta.AnnotationResultPVC], Publication: obj.Metadata.Annotations[workloadmeta.AnnotationArtifactPublication], PublicationID: obj.Metadata.Annotations[workloadmeta.AnnotationArtifactPublicationID], + BundleID: obj.Metadata.Annotations[workloadmeta.AnnotationArtifactBundleID], + ArtifactStore: obj.Metadata.Annotations[workloadmeta.AnnotationArtifactStore], CheckpointArtifact: obj.Metadata.Annotations[workloadmeta.AnnotationCheckpointArtifact], }, nil } +func readRunBlobPath(ctx context.Context, store artifactbundle.Store, absolutePath string) ([]byte, error) { + key, err := artifactbundle.PVCRelativePath(absolutePath) + if err != nil { + return nil, err + } + raw, err := store.Read(ctx, key) + if err != nil { + return nil, fmt.Errorf("read durable artifact %s: %w", absolutePath, err) + } + return raw, nil +} + +func listRunBlobPath(ctx context.Context, store artifactbundle.Store, absolutePath string) ([]string, error) { + prefix, err := artifactbundle.PVCRelativePath(absolutePath) + if err != nil { + return nil, err + } + if prefix != "" && !strings.HasSuffix(prefix, "/") { + prefix += "/" + } + objects, err := store.List(ctx, prefix) + if err != nil { + return nil, fmt.Errorf("enumerate durable artifact directory %s: %w", absolutePath, err) + } + entries := make([]string, 0, len(objects)) + for _, object := range objects { + entries = append(entries, strings.TrimPrefix(object.Name, prefix)) + } + sort.Strings(entries) + return entries, nil +} + +func writeRunBundleDownload( + cmd *cobra.Command, + output string, + manifest artifactbundle.Manifest, + destination string, + files []artifactbundle.DownloadedFile, +) error { + var total int64 + for _, file := range files { + total += file.Size + } + switch output { + case "json": + raw, err := json.MarshalIndent(map[string]any{ + "bundle_id": manifest.BundleID, + "destination": destination, + "file_count": len(files), + "size_bytes": total, + "files": files, + "references": manifest.References, + }, "", " ") + if err != nil { + return err + } + raw = append(raw, '\n') + _, err = cmd.OutOrStdout().Write(raw) + return err + case "raw": + for _, file := range files { + if _, err := fmt.Fprintln(cmd.OutOrStdout(), file.Path); err != nil { + return err + } + } + return nil + default: + _, err := fmt.Fprintf( + cmd.OutOrStdout(), + "Downloaded bundle %s: %d files, %d bytes\nDestination: %s\nLogs: %s\n", + manifest.BundleID, + len(files), + total, + destination, + manifest.References.Logs, + ) + return err + } +} + // looksLikeDirectory is a heuristic: a path with no extension on its final // segment is treated as a directory. Researchers can always force file or dir // semantics with --path or --artifact. diff --git a/cli/internal/cli/run_get_test.go b/cli/internal/cli/run_get_test.go index 0773bfa8..2e1d71d4 100644 --- a/cli/internal/cli/run_get_test.go +++ b/cli/internal/cli/run_get_test.go @@ -5,6 +5,7 @@ import ( "context" "encoding/base64" "errors" + "strconv" "strings" "testing" @@ -118,10 +119,11 @@ func TestRunGetJSONRepresentsSettledEmptyDirectory(t *testing.T) { } func TestRunResultRefFallsBackFromJobToRayJob(t *testing.T) { + artifactStore := `{"schema_version":"tau.run.blob-volume.v1","account_url":"https://trainingacct.blob.core.windows.net","container":"results"}` reader := &fakeRunResultReader{ responses: map[string]string{ "job": "", - "rayjob.ray.io": `{"metadata":{"annotations":{"` + workloadmeta.AnnotationResultPath + `":"/data/research-workspace/runs/modernbert-ray","` + workloadmeta.AnnotationResultPVC + `":"research-workspace","` + workloadmeta.AnnotationArtifactPublication + `":"staged","` + workloadmeta.AnnotationArtifactPublicationID + `":"publication-1"}}}`, + "rayjob.ray.io": `{"metadata":{"annotations":{"` + workloadmeta.AnnotationResultPath + `":"/data/research-workspace/runs/modernbert-ray","` + workloadmeta.AnnotationResultPVC + `":"research-workspace","` + workloadmeta.AnnotationArtifactPublication + `":"staged","` + workloadmeta.AnnotationArtifactPublicationID + `":"publication-1","` + workloadmeta.AnnotationArtifactBundleID + `":"bundle-1","` + workloadmeta.AnnotationArtifactStore + `":` + strconv.Quote(artifactStore) + `}}}`, }, } ref, err := runResultRefWithReader(context.Background(), reader, "research-workspace", "modernbert-ray") @@ -137,6 +139,12 @@ func TestRunResultRefFallsBackFromJobToRayJob(t *testing.T) { if ref.PublicationID != "publication-1" { t.Fatalf("RayJob publication ID = %q", ref.PublicationID) } + if ref.BundleID != "bundle-1" { + t.Fatalf("RayJob bundle ID = %q", ref.BundleID) + } + if ref.ArtifactStore != artifactStore { + t.Fatalf("RayJob artifact store = %q", ref.ArtifactStore) + } if strings.Join(reader.calls, ",") != "job,rayjob.ray.io" { t.Fatalf("lookup order = %v", reader.calls) } diff --git a/cli/internal/cli/run_job.go b/cli/internal/cli/run_job.go index 995ed5f0..5d91f1a7 100644 --- a/cli/internal/cli/run_job.go +++ b/cli/internal/cli/run_job.go @@ -11,6 +11,7 @@ import ( "strings" "time" + "github.com/Azure/taugrid/cli/internal/artifactbundle" "github.com/Azure/taugrid/cli/internal/artifactpublish" "github.com/Azure/taugrid/cli/internal/jobrender" "github.com/Azure/taugrid/cli/internal/metricsoffload" @@ -285,6 +286,34 @@ func executeRunJob(ctx context.Context, stdout, stderr io.Writer, request *runJo } opts.Annotations[experiment.AnnotationExperimentSource] = "stellar" } + artifactBundle, err := resolveArtifactBundle( + request.Name, + ns, + o.submissionID, + outputDir, + annotations[workloadmeta.AnnotationResultPVC], + outputWritable, + artifactPublication, + opts.MetricsOffload, + o.metricsSessionID, + o.checkpointArtifact, + ) + if err != nil { + return err + } + opts.ArtifactBundle = artifactBundle + if artifactBundle.Enabled() && opts.Nodes > 1 { + warnings = append(warnings, + "tau: complete bundle acknowledgement is unavailable for multi-node Indexed Jobs; no shared completion marker will be emitted") + opts.ArtifactBundle = artifactbundle.Runtime{} + artifactBundle = artifactbundle.Runtime{} + } + if artifactBundle.Enabled() { + if opts.Annotations == nil { + opts.Annotations = map[string]string{} + } + opts.Annotations[workloadmeta.AnnotationArtifactBundleID] = artifactBundle.BundleID + } if o.dryRun == "" && !opts.DisableDefaultPriorities { disabled, warning := autoDisableMissingDefaultPriorities(ctx, runner) if disabled { diff --git a/cli/internal/cli/run_ray.go b/cli/internal/cli/run_ray.go index 7fa5ce65..84fae322 100644 --- a/cli/internal/cli/run_ray.go +++ b/cli/internal/cli/run_ray.go @@ -211,6 +211,24 @@ func executeRunRay(ctx context.Context, stdout, stderr io.Writer, request *runRa annotations[experiment.AnnotationExperimentSource] = "stellar" annotations[workloadmeta.AnnotationMetricsSession] = o.metricsSessionID } + artifactBundle, err := resolveArtifactBundle( + name, + namespace, + o.submissionID, + outputDir, + annotations[workloadmeta.AnnotationResultPVC], + outputWritable, + artifactPublication, + metricsRuntime, + o.metricsSessionID, + o.checkpointArtifact, + ) + if err != nil { + return err + } + if artifactBundle.Enabled() { + annotations[workloadmeta.AnnotationArtifactBundleID] = artifactBundle.BundleID + } projectArchive, scriptName, err := buildProjectArchive(o) if err != nil { @@ -243,6 +261,7 @@ func executeRunRay(ctx context.Context, stdout, stderr io.Writer, request *runRa Annotations: annotations, OutputDir: outputDir, ArtifactPublish: artifactPublication, + ArtifactBundle: artifactBundle, CheckpointArtifact: o.checkpointArtifact, MetricsOffload: metricsRuntime, Resources: rayjobrender.Resources{ diff --git a/cli/internal/jobrender/render.go b/cli/internal/jobrender/render.go index 6f908dc0..7d74afdb 100644 --- a/cli/internal/jobrender/render.go +++ b/cli/internal/jobrender/render.go @@ -31,6 +31,7 @@ import ( "gopkg.in/yaml.v3" + "github.com/Azure/taugrid/cli/internal/artifactbundle" "github.com/Azure/taugrid/cli/internal/artifactindex" "github.com/Azure/taugrid/cli/internal/artifactpublish" "github.com/Azure/taugrid/cli/internal/metricsoffload" @@ -188,6 +189,7 @@ type Options struct { // ArtifactPublish optionally wraps the workload so closed artifacts staged // on local /mnt are copied and renamed into durable OutputDir after success. ArtifactPublish artifactpublish.Runtime + ArtifactBundle artifactbundle.Runtime // NodeSelector is a run-time placement override merged after profile and // topology selectors. ClearNodeSelector drops profile selectors before the @@ -350,6 +352,7 @@ func Render(p profile.Profile, o Options) ([]byte, error) { Run: o.Name, ResourceName: o.Name, Namespace: o.Namespace, + BundleID: o.ArtifactBundle.BundleID, }) } if o.MetricsOffload.Enabled() { @@ -370,6 +373,15 @@ func Render(p profile.Profile, o Options) ([]byte, error) { return nil, err } } + if o.ArtifactBundle.Enabled() { + if o.Nodes > 1 { + return nil, fmt.Errorf("artifact bundle completion requires a single Job pod") + } + cmd, err = artifactbundle.WrapCommand(cmd, o.ArtifactBundle) + if err != nil { + return nil, err + } + } gpuPlan, err := profile.BuildGPUSchedulingPlan(p) if err != nil { diff --git a/cli/internal/jobrender/render_test.go b/cli/internal/jobrender/render_test.go index 474e308f..48251f93 100644 --- a/cli/internal/jobrender/render_test.go +++ b/cli/internal/jobrender/render_test.go @@ -16,6 +16,7 @@ import ( "gopkg.in/yaml.v3" + "github.com/Azure/taugrid/cli/internal/artifactbundle" "github.com/Azure/taugrid/cli/internal/artifactpublish" "github.com/Azure/taugrid/cli/internal/metricsoffload" "github.com/Azure/taugrid/core/envspec" @@ -2322,6 +2323,17 @@ func TestRender_DirectJobMetricsOffloadContract(t *testing.T) { ScriptPath: script, PVCMount: "research-workspace", MetricsOffload: runtime, + ArtifactBundle: artifactbundle.Runtime{ + BundleID: "bundle-1", + Run: "modernbert-bounded", + Namespace: "research-workspace", + ResultPVC: "research-workspace", + OutputDir: "/data/research-workspace/modernbert-bounded", + MetricsSessionID: "metrics-1", + MetricsHistory: runtime.History, + MetricsOffloadDir: runtime.Out, + MetricsEnabled: true, + }, Annotations: map[string]string{ workloadmeta.AnnotationExperimentSource: "stellar", }, @@ -2372,6 +2384,17 @@ func TestRender_DirectJobMetricsOffloadContract(t *testing.T) { t.Fatalf("main lifecycle wrapper missing %q:\n%s", want, mainCommand) } } + command := main["command"].([]any) + if len(command) < 5 || command[0] != "bash" || command[1] != "-c" || + command[3] != "tau-bundle-entrypoint" || command[4] != "bash" { + t.Fatalf("artifact bundle must wrap metrics lifecycle command: %v", command) + } + bundleScript := command[2].(string) + for _, want := range []string{"tau_bundle_child", ".tau/bundle.complete", "bundle-1"} { + if !strings.Contains(bundleScript, want) { + t.Fatalf("bundle lifecycle wrapper missing %q:\n%s", want, bundleScript) + } + } volumes := fmt.Sprint(pod["volumes"]) if !strings.Contains(volumes, "tau-metrics-runtime") || !strings.Contains(volumes, "emptyDir") { t.Fatalf("pod-local metrics runtime volume missing: %s", volumes) @@ -2390,6 +2413,29 @@ func TestRender_DirectJobMetricsOffloadContract(t *testing.T) { } } +func TestRender_ArtifactBundleRejectsMultiNodeIndexedJob(t *testing.T) { + script := torchrunScript(t) + _, err := Render(trainProfile(), Options{ + Name: "multi-node", + Namespace: "research", + ScriptPath: script, + Launcher: "torchrun", + Nodes: 2, + PVCMount: "blob-training", + OutputDir: "/data/runs/multi-node", + ArtifactBundle: artifactbundle.Runtime{ + BundleID: "bundle-1", + Run: "multi-node", + Namespace: "research", + ResultPVC: "blob-training", + OutputDir: "/data/runs/multi-node", + }, + }) + if err == nil || !strings.Contains(err.Error(), "single Job pod") { + t.Fatalf("multi-node bundle error = %v", err) + } +} + func TestRender_DirectJobMetricsOffloadSafety(t *testing.T) { script := torchrunScript(t) readOnlyProfile := trainProfile() diff --git a/cli/internal/rayjobrender/render.go b/cli/internal/rayjobrender/render.go index e649b1ac..538f2ef5 100644 --- a/cli/internal/rayjobrender/render.go +++ b/cli/internal/rayjobrender/render.go @@ -15,6 +15,7 @@ import ( "gopkg.in/yaml.v3" + "github.com/Azure/taugrid/cli/internal/artifactbundle" "github.com/Azure/taugrid/cli/internal/artifactindex" "github.com/Azure/taugrid/cli/internal/artifactpublish" "github.com/Azure/taugrid/cli/internal/jsonutil" @@ -120,6 +121,7 @@ type Options struct { Resources Resources OutputDir string ArtifactPublish artifactpublish.Runtime + ArtifactBundle artifactbundle.Runtime // CheckpointArtifact is storage.checkpoint: the file or directory, // relative to the run checkpoint dir, that this run produces as its @@ -746,6 +748,7 @@ func entrypoint(o Options) (string, error) { Run: o.Name, ResourceName: o.Name, Namespace: o.Namespace, + BundleID: o.ArtifactBundle.BundleID, }) + "\n" } if o.ArtifactPublish.Enabled() { @@ -760,6 +763,12 @@ func entrypoint(o Options) (string, error) { return "", err } } + if o.ArtifactBundle.Enabled() { + script, err = artifactbundle.WrapShellScript(script, o.ArtifactBundle) + if err != nil { + return "", err + } + } return raylogoffload.WrapShellScript(script), nil } diff --git a/cli/internal/rayjobrender/render_test.go b/cli/internal/rayjobrender/render_test.go index e1abf167..327d21f2 100644 --- a/cli/internal/rayjobrender/render_test.go +++ b/cli/internal/rayjobrender/render_test.go @@ -13,6 +13,7 @@ import ( "gopkg.in/yaml.v3" + "github.com/Azure/taugrid/cli/internal/artifactbundle" "github.com/Azure/taugrid/cli/internal/artifactpublish" "github.com/Azure/taugrid/cli/internal/metricsoffload" "github.com/Azure/taugrid/cli/internal/payload" @@ -518,6 +519,21 @@ func TestRenderRayJobWithManagedMetricsAndStagedArtifacts(t *testing.T) { PublicationID: "publication-1", }, MetricsOffload: runtime, + ArtifactBundle: artifactbundle.Runtime{ + BundleID: "publication-1", + Run: "modernbert-ray", + Namespace: "research-workspace", + ResultPVC: "research-workspace", + OutputDir: "/data/research-workspace/runs/modernbert-ray", + PublicationMode: artifactpublish.ModeStaged, + PublicationID: "publication-1", + PublicationRoot: "/data/research-workspace/runs/modernbert-ray/.tau-artifacts/publication-1", + PublicationMarker: "/data/research-workspace/runs/modernbert-ray/.tau-artifacts/publication-1/.tau-artifacts-complete", + MetricsSessionID: "session", + MetricsHistory: runtime.History, + MetricsOffloadDir: runtime.Out, + MetricsEnabled: true, + }, Annotations: map[string]string{ workloadmeta.AnnotationResultPath: "/data/research-workspace/runs/modernbert-ray", workloadmeta.AnnotationResultPVC: "research-workspace", @@ -560,6 +576,19 @@ func TestRenderRayJobWithManagedMetricsAndStagedArtifacts(t *testing.T) { if got := containerNames(t, pod["containers"].([]any)); !strings.Contains(got, "metrics-offload") { t.Fatalf("head containers = %s", got) } + entrypoint := spec["entrypoint"].(string) + driverIndex := strings.Index(entrypoint, "tau_driver_child") + bundleIndex := strings.Index(entrypoint, "tau_bundle_child") + metricsIndex := strings.Index(entrypoint, "tau_metrics_child") + if driverIndex < 0 || bundleIndex < 0 || metricsIndex < 0 || + !(driverIndex < bundleIndex && bundleIndex < metricsIndex) { + t.Fatalf("lifecycle wrapper order must be logs -> bundle -> metrics:\n%s", entrypoint) + } + for _, want := range []string{".tau/bundle.complete", "publication-1", ".tau-artifacts-complete"} { + if !strings.Contains(entrypoint, want) { + t.Fatalf("bundle lifecycle wrapper missing %q:\n%s", want, entrypoint) + } + } } func TestRenderCustomOlderRayImageUsesOldProbeEndpoint(t *testing.T) { diff --git a/controllers/tau-core/cmd/controller/main.go b/controllers/tau-core/cmd/controller/main.go index f8c86185..875f528c 100644 --- a/controllers/tau-core/cmd/controller/main.go +++ b/controllers/tau-core/cmd/controller/main.go @@ -76,6 +76,10 @@ func main() { setupLog.Error(err, "unable to create controller", "controller", "TauQuotaRequest") os.Exit(1) } + if err := corecontroller.SetupArtifactStoreControllers(mgr); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "ArtifactStore") + os.Exit(1) + } if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil { setupLog.Error(err, "unable to set up health check") diff --git a/controllers/tau-core/internal/controller/artifact_store_controller.go b/controllers/tau-core/internal/controller/artifact_store_controller.go new file mode 100644 index 00000000..a84b3668 --- /dev/null +++ b/controllers/tau-core/internal/controller/artifact_store_controller.go @@ -0,0 +1,224 @@ +package controller + +import ( + "context" + "encoding/json" + "fmt" + "net/url" + "strings" + "time" + + "github.com/Azure/taugrid/controllers/tau-core/internal/labelkeys" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/meta" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/builder" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/event" + "sigs.k8s.io/controller-runtime/pkg/predicate" +) + +const ( + artifactStoreSchema = "tau.run.blob-volume.v1" + azureBlobCSIDriver = "blob.csi.azure.com" +) + +type artifactStoreReconciler struct { + client.Client + newObject func() client.Object +} + +type artifactStoreDocument struct { + SchemaVersion string `json:"schema_version"` + AccountURL string `json:"account_url"` + Container string `json:"container"` +} + +// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;patch;update +// +kubebuilder:rbac:groups=ray.io,resources=rayjobs,verbs=get;list;watch;patch;update +// +kubebuilder:rbac:groups="",resources=persistentvolumeclaims,verbs=get;list;watch +// +kubebuilder:rbac:groups="",resources=persistentvolumes,verbs=get;list;watch + +func (r *artifactStoreReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + workload := r.newObject() + if err := r.Get(ctx, req.NamespacedName, workload); err != nil { + return ctrl.Result{}, client.IgnoreNotFound(err) + } + if !artifactStoreCandidate(workload) { + return ctrl.Result{}, nil + } + + pvcName := strings.TrimSpace(workload.GetAnnotations()[labelkeys.AnnotationResultPVC]) + var pvc corev1.PersistentVolumeClaim + if err := r.Get(ctx, types.NamespacedName{Namespace: req.Namespace, Name: pvcName}, &pvc); err != nil { + return ctrl.Result{}, err + } + if strings.TrimSpace(pvc.Spec.VolumeName) == "" { + return ctrl.Result{RequeueAfter: 10 * time.Second}, nil + } + var pv corev1.PersistentVolume + if err := r.Get(ctx, types.NamespacedName{Name: pvc.Spec.VolumeName}, &pv); err != nil { + return ctrl.Result{}, err + } + document, supported, err := artifactStoreFromPV(&pv) + if err != nil { + return ctrl.Result{}, err + } + if !supported { + return ctrl.Result{}, nil + } + raw, err := json.Marshal(document) + if err != nil { + return ctrl.Result{}, fmt.Errorf("encode artifact store metadata: %w", err) + } + if workload.GetAnnotations()[labelkeys.AnnotationArtifactStore] == string(raw) { + return ctrl.Result{}, nil + } + before := workload.DeepCopyObject().(client.Object) + annotations := workload.GetAnnotations() + if annotations == nil { + annotations = map[string]string{} + } + annotations[labelkeys.AnnotationArtifactStore] = string(raw) + workload.SetAnnotations(annotations) + if err := r.Patch(ctx, workload, client.MergeFrom(before)); err != nil { + return ctrl.Result{}, fmt.Errorf("stamp artifact store metadata on %s/%s: %w", req.Namespace, req.Name, err) + } + return ctrl.Result{}, nil +} + +func artifactStoreCandidate(object client.Object) bool { + return object.GetLabels()[labelkeys.LabelManagedBy] == "tau" && + strings.TrimSpace(object.GetAnnotations()[labelkeys.AnnotationArtifactBundleID]) != "" && + strings.TrimSpace(object.GetAnnotations()[labelkeys.AnnotationResultPVC]) != "" +} + +func artifactStoreFromPV(pv *corev1.PersistentVolume) (artifactStoreDocument, bool, error) { + if pv.Spec.CSI == nil || !strings.EqualFold(strings.TrimSpace(pv.Spec.CSI.Driver), azureBlobCSIDriver) { + return artifactStoreDocument{}, false, nil + } + attributes, err := foldArtifactStoreAttributes(pv.Spec.CSI.VolumeAttributes) + if err != nil { + return artifactStoreDocument{}, false, fmt.Errorf("PV %s: %w", pv.Name, err) + } + account := strings.TrimSpace(attributes["storageaccount"]) + containerName := strings.TrimSpace(attributes["containername"]) + parts := strings.Split(pv.Spec.CSI.VolumeHandle, "#") + if account == "" && len(parts) > 1 { + account = strings.TrimSpace(parts[1]) + } + if containerName == "" && len(parts) > 2 { + containerName = strings.TrimSpace(parts[2]) + } + if account == "" || containerName == "" || strings.ContainsAny(containerName, `/\`) { + return artifactStoreDocument{}, false, nil + } + server := strings.TrimSpace(attributes["server"]) + if server == "" { + server = account + ".blob." + firstArtifactStoreValue(attributes["storageendpointsuffix"], "core.windows.net") + } + if !strings.Contains(server, "://") { + server = "https://" + server + } + parsed, err := url.Parse(server) + if err != nil || parsed.Scheme != "https" || parsed.Host == "" || parsed.User != nil || + parsed.RawQuery != "" || parsed.Fragment != "" || (parsed.Path != "" && parsed.Path != "/") { + return artifactStoreDocument{}, false, fmt.Errorf("PV %s has invalid Blob server %q", pv.Name, server) + } + accountURL := strings.TrimSuffix(parsed.String(), "/") + if !trustedArtifactStoreHost(accountURL) { + return artifactStoreDocument{}, false, nil + } + return artifactStoreDocument{ + SchemaVersion: artifactStoreSchema, + AccountURL: accountURL, + Container: containerName, + }, true, nil +} + +func foldArtifactStoreAttributes(attributes map[string]string) (map[string]string, error) { + out := make(map[string]string, len(attributes)) + for key, value := range attributes { + folded := strings.ToLower(strings.TrimSpace(key)) + if previous, exists := out[folded]; exists && previous != value { + return nil, fmt.Errorf("volumeAttributes contains conflicting case variants for %q", key) + } + out[folded] = value + } + return out, nil +} + +func trustedArtifactStoreHost(accountURL string) bool { + parsed, err := url.Parse(strings.TrimSpace(accountURL)) + if err != nil { + return false + } + host := strings.ToLower(strings.TrimSuffix(parsed.Hostname(), ".")) + for _, suffix := range []string{ + ".blob.core.windows.net", + ".blob.core.usgovcloudapi.net", + ".blob.core.chinacloudapi.cn", + ".blob.core.cloudapi.de", + } { + if strings.HasSuffix(host, suffix) && len(host) > len(suffix) { + return true + } + } + return false +} + +func firstArtifactStoreValue(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} + +func artifactStorePredicate() predicate.Predicate { + return predicate.Funcs{ + CreateFunc: func(e event.CreateEvent) bool { return artifactStoreCandidate(e.Object) }, + UpdateFunc: func(e event.UpdateEvent) bool { return artifactStoreCandidate(e.ObjectNew) }, + DeleteFunc: func(event.DeleteEvent) bool { return false }, + GenericFunc: func(e event.GenericEvent) bool { return artifactStoreCandidate(e.Object) }, + } +} + +func SetupArtifactStoreControllers(mgr ctrl.Manager) error { + job := &artifactStoreReconciler{ + Client: mgr.GetClient(), + newObject: func() client.Object { return &batchv1.Job{} }, + } + if err := ctrl.NewControllerManagedBy(mgr). + Named("job-artifact-store"). + For(&batchv1.Job{}, builder.WithPredicates(artifactStorePredicate())). + Complete(job); err != nil { + return err + } + + rayJob := &unstructured.Unstructured{} + rayJob.SetGroupVersionKind(schema.GroupVersionKind{Group: "ray.io", Version: "v1", Kind: "RayJob"}) + if _, err := mgr.GetRESTMapper().RESTMapping(rayJob.GroupVersionKind().GroupKind(), rayJob.GroupVersionKind().Version); err != nil { + if meta.IsNoMatchError(err) { + return nil + } + return fmt.Errorf("discover RayJob API for artifact store controller: %w", err) + } + ray := &artifactStoreReconciler{ + Client: mgr.GetClient(), + newObject: func() client.Object { + object := &unstructured.Unstructured{} + object.SetGroupVersionKind(rayJob.GroupVersionKind()) + return object + }, + } + return ctrl.NewControllerManagedBy(mgr). + Named("rayjob-artifact-store"). + For(rayJob, builder.WithPredicates(artifactStorePredicate())). + Complete(ray) +} diff --git a/controllers/tau-core/internal/controller/artifact_store_controller_test.go b/controllers/tau-core/internal/controller/artifact_store_controller_test.go new file mode 100644 index 00000000..68688780 --- /dev/null +++ b/controllers/tau-core/internal/controller/artifact_store_controller_test.go @@ -0,0 +1,95 @@ +package controller + +import ( + "context" + "encoding/json" + "testing" + + "github.com/Azure/taugrid/controllers/tau-core/internal/labelkeys" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" +) + +func TestArtifactStoreReconcilerStampsBlobIdentity(t *testing.T) { + pvc := &corev1.PersistentVolumeClaim{ + ObjectMeta: metav1.ObjectMeta{Name: "blob-training", Namespace: "research"}, + Spec: corev1.PersistentVolumeClaimSpec{VolumeName: "pvc-123"}, + } + pv := &corev1.PersistentVolume{ + ObjectMeta: metav1.ObjectMeta{Name: "pvc-123"}, + Spec: corev1.PersistentVolumeSpec{ + PersistentVolumeSource: corev1.PersistentVolumeSource{ + CSI: &corev1.CSIPersistentVolumeSource{ + Driver: azureBlobCSIDriver, + VolumeHandle: "rg#trainingacct#results#uuid#research#subscription", + VolumeAttributes: map[string]string{"storageEndpointSuffix": "core.windows.net"}, + }, + }, + }, + } + job := &batchv1.Job{ObjectMeta: metav1.ObjectMeta{ + Name: "training-1", + Namespace: "research", + Labels: map[string]string{labelkeys.LabelManagedBy: "tau"}, + Annotations: map[string]string{ + labelkeys.AnnotationArtifactBundleID: "bundle-1", + labelkeys.AnnotationResultPVC: "blob-training", + labelkeys.AnnotationArtifactStore: `{"schema_version":"tau.run.blob-volume.v1","account_url":"https://attacker.example","container":"stolen"}`, + }, + }} + c := fake.NewClientBuilder().WithScheme(testScheme(t)).WithObjects(pvc, pv, job).Build() + reconciler := &artifactStoreReconciler{ + Client: c, + newObject: func() client.Object { return &batchv1.Job{} }, + } + if _, err := reconciler.Reconcile(context.Background(), ctrl.Request{ + NamespacedName: types.NamespacedName{Namespace: "research", Name: "training-1"}, + }); err != nil { + t.Fatal(err) + } + + var got batchv1.Job + if err := c.Get(context.Background(), types.NamespacedName{Namespace: "research", Name: "training-1"}, &got); err != nil { + t.Fatal(err) + } + raw := got.Annotations[labelkeys.AnnotationArtifactStore] + var document artifactStoreDocument + if err := json.Unmarshal([]byte(raw), &document); err != nil { + t.Fatalf("decode stamped annotation %q: %v", raw, err) + } + if document.SchemaVersion != artifactStoreSchema || + document.AccountURL != "https://trainingacct.blob.core.windows.net" || + document.Container != "results" { + t.Fatalf("artifact store document = %+v", document) + } +} + +func TestArtifactStoreFromPVRejectsUntrustedEndpoint(t *testing.T) { + pv := &corev1.PersistentVolume{ + ObjectMeta: metav1.ObjectMeta{Name: "hostile"}, + Spec: corev1.PersistentVolumeSpec{ + PersistentVolumeSource: corev1.PersistentVolumeSource{ + CSI: &corev1.CSIPersistentVolumeSource{ + Driver: azureBlobCSIDriver, + VolumeAttributes: map[string]string{ + "storageAccount": "trainingacct", + "containerName": "results", + "server": "attacker.example", + }, + }, + }, + }, + } + _, supported, err := artifactStoreFromPV(pv) + if err != nil { + t.Fatal(err) + } + if supported { + t.Fatal("untrusted endpoint was marked supported") + } +} diff --git a/controllers/tau-core/internal/labelkeys/contract_test.go b/controllers/tau-core/internal/labelkeys/contract_test.go index f8937054..12099f0f 100644 --- a/controllers/tau-core/internal/labelkeys/contract_test.go +++ b/controllers/tau-core/internal/labelkeys/contract_test.go @@ -25,8 +25,12 @@ var tauTokenRe = regexp.MustCompile(regexp.QuoteMeta(Domain) + `[A-Za-z0-9._-]*` // sharedKeys contains only keys that cross the module boundary. Controller-only // namespace and quota keys do not belong in the CLI's canonical package. var sharedKeys = map[string]string{ - "LabelWorkspace": "LabelWorkspace", - "AnnotationResultScope": "AnnotationResultScope", + "LabelWorkspace": "LabelWorkspace", + "LabelManagedBy": "LabelManagedBy", + "AnnotationResultScope": "AnnotationResultScope", + "AnnotationResultPVC": "AnnotationResultPVC", + "AnnotationArtifactBundleID": "AnnotationArtifactBundleID", + "AnnotationArtifactStore": "AnnotationArtifactStore", } func TestSharedLabelKeysAgreeWithTauCLI(t *testing.T) { diff --git a/controllers/tau-core/internal/labelkeys/labelkeys.go b/controllers/tau-core/internal/labelkeys/labelkeys.go index c0673e32..07fc6d80 100644 --- a/controllers/tau-core/internal/labelkeys/labelkeys.go +++ b/controllers/tau-core/internal/labelkeys/labelkeys.go @@ -14,12 +14,16 @@ const ( LabelWorkspace = "tau.azure.com/workspace" LabelLocalQueue = "tau.azure.com/local-queue" LabelGPUClass = "tau.azure.com/gpu-class" + LabelManagedBy = "tau.azure.com/managed-by" - AnnotationApproved = "tau.azure.com/approved" - AnnotationRejected = "tau.azure.com/rejected" - AnnotationReviewedBy = "tau.azure.com/reviewed-by" - AnnotationResultScope = "tau.azure.com/result-scope" - AnnotationV0Primary = "tau.azure.com/v0-primary-workspace" + AnnotationApproved = "tau.azure.com/approved" + AnnotationRejected = "tau.azure.com/rejected" + AnnotationReviewedBy = "tau.azure.com/reviewed-by" + AnnotationResultScope = "tau.azure.com/result-scope" + AnnotationV0Primary = "tau.azure.com/v0-primary-workspace" + AnnotationResultPVC = "tau.azure.com/result-pvc" + AnnotationArtifactBundleID = "tau.azure.com/artifact-bundle-id" + AnnotationArtifactStore = "tau.azure.com/artifact-store" FinalizerWorkspaceCleanup = "tau.azure.com/workspace-cleanup" ) diff --git a/core/workloadmeta/metadata.go b/core/workloadmeta/metadata.go index 26f6377a..4130f22f 100644 --- a/core/workloadmeta/metadata.go +++ b/core/workloadmeta/metadata.go @@ -175,6 +175,8 @@ const ( AnnotationCheckpointURI = "tau.azure.com/checkpoint-uri" AnnotationArtifactPublication = "tau.azure.com/artifact-publication" AnnotationArtifactPublicationID = "tau.azure.com/artifact-publication-id" + AnnotationArtifactBundleID = "tau.azure.com/artifact-bundle-id" + AnnotationArtifactStore = "tau.azure.com/artifact-store" // AnnotationCheckpointArtifact records the storage.checkpoint value a run // declared, so a later command can tell "this run produced no artifacts" diff --git a/site/content/en/docs/reference/cli.md b/site/content/en/docs/reference/cli.md index 442563b1..829a4c82 100644 --- a/site/content/en/docs/reference/cli.md +++ b/site/content/en/docs/reference/cli.md @@ -139,7 +139,7 @@ subcommand** — `tau run train --dry-run=client` runs the `run` root with | `list` | List Tau-managed Jobs/RayJobs in a namespace | | `status [job-name]` | Show lifecycle state and startup phases; `--watch` to poll | | `logs ` | Stream Ray driver logs or the batch Job pod logs | -| `get ` | Fetch durable run results and artifacts | +| `get ` | List or fetch durable results; `--destination DIR` downloads the complete acknowledged artifact bundle | | `cancel ` | Delete the underlying Job/RayJob and free its Kueue quota | | `resume --config tau.yaml` | Manually restart a failed run from its checkpoint | @@ -148,6 +148,26 @@ by the `resilience.*` fields in your run config. See [recovery](../../operations/recovery/) for the full retry and resume contract. +`tau run get --destination DIR` is the supported replacement for mounting +the result PVC in a temporary reader Pod. Current Tau Jobs and RayJobs commit a +durable bundle acknowledgement only after staged artifacts, checkpoint indexing, +and enabled metrics offload have completed successfully. Retrieval validates +that acknowledgement and the staged-publication marker. New workloads record +their non-secret Blob CSI account/container identity in Tau-owned metadata, so +workspace users do not need cluster-scoped PV read access; legacy workloads fall +back to read-only PVC/PV discovery. The tau-core controller resolves and repairs +that metadata from the bound PV; the researcher never receives PV-wide access. +Tau downloads with `DefaultAzureCredential` and does not read CSI Secrets or +accept storage account keys/SAS tokens. It also refuses to replace existing +destination files. Unsupported CSI drivers, incomplete bundles, and pre-contract +bundles fail explicitly. + +Complete bundle acknowledgement currently applies to single-pod batch Jobs and +RayJobs whose result PVC is mounted at Tau's `/data` durable-storage root. +Multi-node Indexed Jobs do not emit a shared bundle marker because individual +indexes cannot safely acknowledge Job-level completion. Custom PVC mount roots +retain their existing run behavior but cannot claim this Blob bundle contract. + ## `tau serve` `deploy [name]`, `status [name]`, `scale [name]`, and `delete [name]` render diff --git a/site/content/en/docs/reference/run-config.md b/site/content/en/docs/reference/run-config.md index 39aa4188..3bc3824d 100644 --- a/site/content/en/docs/reference/run-config.md +++ b/site/content/en/docs/reference/run-config.md @@ -38,6 +38,24 @@ namespace. Tau references and mounts that claim; it does not provision or own the PVC, StorageClass, CSI configuration, or backing storage. The platform owner chooses the backend and manages its lifecycle. +Set `storage.publish: staged` when terminal artifacts must become visible only +after the workload succeeds and Tau verifies their copies. The application +writes closed regular files to `TAU_OUTPUT_STAGING_DIR`; immutable metric chunks +may continue to use their declared `metrics.history` paths under `/data`. After +staged publication, checkpoint indexing, and enabled metrics offload all +acknowledge completion, Tau atomically commits a bundle manifest under +`storage.output/.tau/`. `tau run get --destination DIR` requires that +final acknowledgement and downloads the result tree plus any declared +checkpoint tree directly from the Blob CSI container. It does not create a +reader Pod or bypass publication acknowledgement. The tau-core controller stamps +the non-secret Blob transport identity on new workloads; this metadata contains +no account key, SAS token, or other repository credential. Multi-node Indexed +Jobs intentionally do not emit a bundle acknowledgement until Tau has a +Job-level coordinator; a single index is not allowed to claim completion for the +whole Job. Bundle retrieval also requires the result PVC at Tau's `/data` +durable-storage root; custom mount roots continue to run without a bundle +acknowledgement. + Main field groups: | Group | Purpose | diff --git a/site/content/en/docs/tasks/researcher/first-run.md b/site/content/en/docs/tasks/researcher/first-run.md index 151f483d..e7f059c4 100644 --- a/site/content/en/docs/tasks/researcher/first-run.md +++ b/site/content/en/docs/tasks/researcher/first-run.md @@ -43,6 +43,7 @@ rest of this walkthrough. tau run status --watch tau run logs tau run get +tau run get --destination ./artifacts/ ``` `status --watch` renders the startup phase tree (Submitted, Kueue admission, @@ -51,7 +52,24 @@ scheduling, image pull, readiness) until the interrupt it. For RayJobs, `logs` streams the Ray driver's execution output rather than head-pod logs; for batch Jobs it streams the Job pod logs. `get` fetches durable run results and artifacts once the run -has produced them. +has produced them. With `--destination`, it downloads the complete acknowledged +bundle: staged terminal artifacts, immutable metric chunks and offload metadata, +and the durable checkpoint tree when one was declared. The downloaded +`.tau-bundle/manifest.json` records artifact, checkpoint, and log references; +`.tau-bundle/files.json` records the downloaded byte count and SHA-256 for every +file. + +Complete-bundle retrieval is read-only and creates no Kubernetes Pods. Tau reads +the non-secret Blob CSI volume identity that the tau-core controller records on +current workloads and uses your Azure RBAC identity through +`DefaultAzureCredential`; it never reads Kubernetes storage Secrets or asks for +an account key or SAS token. Legacy workloads may require permission to read +their bound PV for transport discovery. +A missing staged-publication marker, metrics acknowledgement, final bundle +marker, or storage listing is an error rather than a partial success. Downloads +refuse to replace existing destination files. Runs created before the bundle +acknowledgement contract can still list or fetch known artifacts where practical, +but cannot claim a complete bundle. If you need to stop a run before it finishes — for example, you spot a bad hyperparameter mid-training — cancel it instead of leaving it to fail on its From ddaca0cdf541f0ae3f35f7bb5b185159f27b91b1 Mon Sep 17 00:00:00 2001 From: chokevin Date: Fri, 7 Aug 2026 21:07:54 -0700 Subject: [PATCH 2/3] cli: harden artifact bundle retrieval --- cli/internal/artifactbundle/bundle_test.go | 43 ++++++- cli/internal/artifactbundle/retrieve.go | 44 ++++--- cli/internal/cli/pvc_helpers.go | 5 + cli/internal/cli/run_bundle.go | 3 + cli/internal/cli/run_bundle_test.go | 1 + cli/internal/cli/run_get.go | 107 +++++------------- cli/internal/cli/run_job.go | 5 + cli/internal/jobrender/render.go | 8 +- cli/internal/jobrender/render_test.go | 23 ++++ site/content/en/docs/reference/run-config.md | 8 +- .../en/docs/tasks/researcher/first-run.md | 3 +- 11 files changed, 139 insertions(+), 111 deletions(-) diff --git a/cli/internal/artifactbundle/bundle_test.go b/cli/internal/artifactbundle/bundle_test.go index b9da76cb..abe2a842 100644 --- a/cli/internal/artifactbundle/bundle_test.go +++ b/cli/internal/artifactbundle/bundle_test.go @@ -50,6 +50,18 @@ func (s listErrorStore) List(context.Context, string) ([]Object, error) { return nil, s.err } +type downloadErrorStore struct { + memoryStore + fail string +} + +func (s downloadErrorStore) Download(ctx context.Context, name string, out io.Writer) error { + if name == s.fail { + return errors.New("download interrupted") + } + return s.memoryStore.Download(ctx, name, out) +} + func testRuntime() Runtime { return Runtime{ BundleID: "bundle-1", @@ -278,7 +290,7 @@ func TestDownloadRejectsExistingDestinationWithoutReplacingIt(t *testing.T) { func TestDownloadChecksSizeBeforePublishingDestination(t *testing.T) { store := memoryStore{"runs/training-1/result.json": []byte("actual")} - root := t.TempDir() + root := filepath.Join(t.TempDir(), "bundle") target := filepath.Join(root, "runs", "training-1", "result.json") _, err := Download(context.Background(), store, Manifest{}, []Object{{ Name: "runs/training-1/result.json", @@ -291,3 +303,32 @@ func TestDownloadChecksSizeBeforePublishingDestination(t *testing.T) { t.Fatalf("size mismatch published destination: %v", statErr) } } + +func TestDownloadPublishesNothingWhenLaterObjectFails(t *testing.T) { + parent := t.TempDir() + root := filepath.Join(parent, "bundle") + store := downloadErrorStore{ + memoryStore: memoryStore{ + "runs/training-1/first.json": []byte("first"), + "runs/training-1/second.json": []byte("second"), + }, + fail: "runs/training-1/second.json", + } + _, err := Download(context.Background(), store, Manifest{}, []Object{ + {Name: "runs/training-1/first.json", Size: 5}, + {Name: "runs/training-1/second.json", Size: 6}, + }, root) + if err == nil || !strings.Contains(err.Error(), "download interrupted") { + t.Fatalf("interrupted download error = %v", err) + } + if _, statErr := os.Stat(root); !os.IsNotExist(statErr) { + t.Fatalf("interrupted download published destination: %v", statErr) + } + stages, globErr := filepath.Glob(filepath.Join(parent, ".bundle.tau-download-*")) + if globErr != nil { + t.Fatal(globErr) + } + if len(stages) != 0 { + t.Fatalf("interrupted download left staging directories: %v", stages) + } +} diff --git a/cli/internal/artifactbundle/retrieve.go b/cli/internal/artifactbundle/retrieve.go index d77fd143..0cb19140 100644 --- a/cli/internal/artifactbundle/retrieve.go +++ b/cli/internal/artifactbundle/retrieve.go @@ -161,18 +161,24 @@ func Download(ctx context.Context, store Store, manifest Manifest, objects []Obj if err != nil { return nil, fmt.Errorf("resolve artifact bundle destination: %w", err) } - if err := os.MkdirAll(root, 0o755); err != nil { - return nil, fmt.Errorf("create artifact bundle destination: %w", err) - } - if info, err := os.Lstat(root); err != nil { + if _, err := os.Lstat(root); err == nil { + return nil, fmt.Errorf("artifact bundle destination already exists: %s", root) + } else if !os.IsNotExist(err) { return nil, err - } else if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { - return nil, fmt.Errorf("artifact bundle destination %s must be a real directory", root) } + parent := filepath.Dir(root) + if err := os.MkdirAll(parent, 0o755); err != nil { + return nil, fmt.Errorf("create artifact bundle destination parent: %w", err) + } + stage, err := os.MkdirTemp(parent, "."+filepath.Base(root)+".tau-download-*") + if err != nil { + return nil, fmt.Errorf("create artifact bundle staging directory: %w", err) + } + defer os.RemoveAll(stage) targets := make([]string, len(objects)) seenTargets := make(map[string]string, len(objects)) for i, object := range objects { - target, err := downloadTarget(root, object.Name) + target, err := downloadTarget(stage, object.Name) if err != nil { return nil, err } @@ -180,28 +186,13 @@ func Download(ctx context.Context, store Store, manifest Manifest, objects []Obj return nil, fmt.Errorf("artifact bundle objects %q and %q resolve to the same destination", previous, object.Name) } seenTargets[target] = object.Name - if _, err := os.Lstat(target); err == nil { - return nil, fmt.Errorf("artifact bundle destination file already exists: %s", target) - } else if !os.IsNotExist(err) { - return nil, err - } targets[i] = target } - metadataDir := filepath.Join(root, ".tau-bundle") - for _, target := range []string{ - filepath.Join(metadataDir, "manifest.json"), - filepath.Join(metadataDir, "files.json"), - } { - if _, err := os.Lstat(target); err == nil { - return nil, fmt.Errorf("artifact bundle metadata file already exists: %s", target) - } else if !os.IsNotExist(err) { - return nil, err - } - } + metadataDir := filepath.Join(stage, ".tau-bundle") files := make([]DownloadedFile, 0, len(objects)) for i, object := range objects { target := targets[i] - if err := ensureSafeDirectory(root, filepath.Dir(target)); err != nil { + if err := ensureSafeDirectory(stage, filepath.Dir(target)); err != nil { return nil, err } file, err := os.CreateTemp(filepath.Dir(target), ".tau-download-*") @@ -249,7 +240,7 @@ func Download(ctx context.Context, store Store, manifest Manifest, objects []Obj SHA256: hex.EncodeToString(hash.Sum(nil)), }) } - if err := ensureSafeDirectory(root, metadataDir); err != nil { + if err := ensureSafeDirectory(stage, metadataDir); err != nil { return nil, err } if err := writeJSONAtomic(filepath.Join(metadataDir, "manifest.json"), manifest); err != nil { @@ -258,6 +249,9 @@ func Download(ctx context.Context, store Store, manifest Manifest, objects []Obj if err := writeJSONAtomic(filepath.Join(metadataDir, "files.json"), files); err != nil { return nil, err } + if err := os.Rename(stage, root); err != nil { + return nil, fmt.Errorf("publish artifact bundle destination: %w", err) + } return files, nil } diff --git a/cli/internal/cli/pvc_helpers.go b/cli/internal/cli/pvc_helpers.go index f641a0b2..0e4bb7a6 100644 --- a/cli/internal/cli/pvc_helpers.go +++ b/cli/internal/cli/pvc_helpers.go @@ -235,6 +235,11 @@ func fetchPVCList(ctx context.Context, kubeContext, namespace, runName, pvcName, return fetchPVCListWithMode(ctx, kubeContext, namespace, runName, pvcName, dirPath, false) } +// fetchPVCListRecursive returns every descendant as a relative path. +func fetchPVCListRecursive(ctx context.Context, kubeContext, namespace, runName, pvcName, dirPath string) ([]string, error) { + return fetchPVCListWithMode(ctx, kubeContext, namespace, runName, pvcName, dirPath, true) +} + func fetchPVCListWithMode(ctx context.Context, kubeContext, namespace, runName, pvcName, dirPath string, recursive bool) ([]string, error) { if pvcName == "" { pvcName = defaultTauPVCName diff --git a/cli/internal/cli/run_bundle.go b/cli/internal/cli/run_bundle.go index 61b1e951..1f103dee 100644 --- a/cli/internal/cli/run_bundle.go +++ b/cli/internal/cli/run_bundle.go @@ -25,6 +25,9 @@ func resolveArtifactBundle( if outputDir != "/data" && !strings.HasPrefix(outputDir, "/data/") { return artifactbundle.Runtime{}, nil } + if !looksLikeDirectory(outputDir) { + return artifactbundle.Runtime{}, nil + } bundleID := firstNonEmpty(publication.PublicationID, submissionID) if bundleID == "" { return artifactbundle.Runtime{}, nil diff --git a/cli/internal/cli/run_bundle_test.go b/cli/internal/cli/run_bundle_test.go index 73530a50..730fcdf3 100644 --- a/cli/internal/cli/run_bundle_test.go +++ b/cli/internal/cli/run_bundle_test.go @@ -78,6 +78,7 @@ func TestResolveArtifactBundleSkipsReadOnlyAndEphemeralResults(t *testing.T) { {output: "/data/runs/training-1", pvc: "", writable: true}, {output: "/data/runs/training-1", pvc: "blob-training", writable: false}, {output: "/data-nfs/runs/training-1", pvc: "shared-nfs", writable: true}, + {output: "/data/runs/training-1/result.json", pvc: "blob-training", writable: true}, } { runtime, err := resolveArtifactBundle( "training-1", "research", "submission-1", test.output, test.pvc, test.writable, diff --git a/cli/internal/cli/run_get.go b/cli/internal/cli/run_get.go index 9bbe0847..c97e7348 100644 --- a/cli/internal/cli/run_get.go +++ b/cli/internal/cli/run_get.go @@ -6,7 +6,6 @@ import ( "fmt" "path" "path/filepath" - "sort" "strings" "github.com/spf13/cobra" @@ -120,55 +119,39 @@ Examples: if destination != "" && artifact != "" { return fmt.Errorf("--destination and --artifact cannot be combined") } - var blobVolume runBlobVolume - if strings.TrimSpace(ref.ArtifactStore) != "" { - blobVolume, err = parseRunBlobVolume(ref.ArtifactStore) - } else { - blobVolume, err = resolveRunBlobVolume(cmd.Context(), kube.New(resolvedContext), ns, ref.PVC) - } - if err != nil { - return err - } - store, err := newAzureRunArtifactStore(blobVolume) - if err != nil { - return err - } - if ref.BundleID != "" || destination != "" { + if destination != "" { + var blobVolume runBlobVolume + if strings.TrimSpace(ref.ArtifactStore) != "" { + blobVolume, err = parseRunBlobVolume(ref.ArtifactStore) + } else { + blobVolume, err = resolveRunBlobVolume(cmd.Context(), kube.New(resolvedContext), ns, ref.PVC) + } + if err != nil { + return err + } + store, err := newAzureRunArtifactStore(blobVolume) + if err != nil { + return err + } manifest, loadErr := artifactbundle.Load(cmd.Context(), store, ref.Path, ref.BundleID) if loadErr != nil { - if destination != "" { - return fmt.Errorf( - "complete artifact bundle is unavailable: %w; this run may predate Tau's final bundle acknowledgement", - loadErr, - ) - } - return loadErr + return fmt.Errorf( + "complete artifact bundle is unavailable: %w; this run may predate Tau's final bundle acknowledgement", + loadErr, + ) } if manifest.ResultPVC != ref.PVC || path.Clean(manifest.ResultRoot) != path.Clean(ref.Path) { return fmt.Errorf("artifact bundle identity does not match workload result metadata") } - if destination != "" { - objects, err := artifactbundle.Enumerate(cmd.Context(), store, manifest) - if err != nil { - return err - } - files, err := artifactbundle.Download(cmd.Context(), store, manifest, objects, destination) - if err != nil { - return err - } - return writeRunBundleDownload(cmd, output, manifest, destination, files) + objects, err := artifactbundle.Enumerate(cmd.Context(), store, manifest) + if err != nil { + return err } - if artifact == "" { - objects, err := artifactbundle.Enumerate(cmd.Context(), store, manifest) - if err != nil { - return err - } - entries := make([]string, 0, len(objects)) - for _, object := range objects { - entries = append(entries, object.Name) - } - return writeRunGet(cmd, output, nil, entries, manifest.ResultRoot, manifest.ResultPVC, ref.CheckpointArtifact) + files, err := artifactbundle.Download(cmd.Context(), store, manifest, objects, destination) + if err != nil { + return err } + return writeRunBundleDownload(cmd, output, manifest, destination, files) } resultPath := ref.Path @@ -178,7 +161,7 @@ Examples: } resultPath = path.Join(ref.Path, artifactpublish.GenerationsDir, ref.PublicationID) marker := path.Join(resultPath, artifactpublish.CompletionMarker) - raw, err := readRunBlobPath(cmd.Context(), store, marker) + raw, err := fetchPVCFile(cmd.Context(), resolvedContext, ns, name, ref.PVC, marker) if err != nil { return fmt.Errorf("staged artifacts are not completely published: %w", err) } @@ -195,20 +178,20 @@ Examples: return err } file := path.Join(resultPath, cleanArtifact) - raw, err := readRunBlobPath(cmd.Context(), store, file) + raw, err := fetchPVCFile(cmd.Context(), resolvedContext, ns, name, ref.PVC, file) if err != nil { return err } return writeRunGet(cmd, output, raw, nil, file, ref.PVC, "") } if !isDir { - raw, err := readRunBlobPath(cmd.Context(), store, resultPath) + raw, err := fetchPVCFile(cmd.Context(), resolvedContext, ns, name, ref.PVC, resultPath) if err != nil { return err } return writeRunGet(cmd, output, raw, nil, resultPath, ref.PVC, "") } - entries, err := listRunBlobPath(cmd.Context(), store, resultPath) + entries, err := fetchPVCListRecursive(cmd.Context(), resolvedContext, ns, name, ref.PVC, resultPath) if err != nil { return err } @@ -293,38 +276,6 @@ func parseRunResultRef(raw []byte, resource string) (runResultRef, error) { }, nil } -func readRunBlobPath(ctx context.Context, store artifactbundle.Store, absolutePath string) ([]byte, error) { - key, err := artifactbundle.PVCRelativePath(absolutePath) - if err != nil { - return nil, err - } - raw, err := store.Read(ctx, key) - if err != nil { - return nil, fmt.Errorf("read durable artifact %s: %w", absolutePath, err) - } - return raw, nil -} - -func listRunBlobPath(ctx context.Context, store artifactbundle.Store, absolutePath string) ([]string, error) { - prefix, err := artifactbundle.PVCRelativePath(absolutePath) - if err != nil { - return nil, err - } - if prefix != "" && !strings.HasSuffix(prefix, "/") { - prefix += "/" - } - objects, err := store.List(ctx, prefix) - if err != nil { - return nil, fmt.Errorf("enumerate durable artifact directory %s: %w", absolutePath, err) - } - entries := make([]string, 0, len(objects)) - for _, object := range objects { - entries = append(entries, strings.TrimPrefix(object.Name, prefix)) - } - sort.Strings(entries) - return entries, nil -} - func writeRunBundleDownload( cmd *cobra.Command, output string, diff --git a/cli/internal/cli/run_job.go b/cli/internal/cli/run_job.go index 5d91f1a7..18a61c88 100644 --- a/cli/internal/cli/run_job.go +++ b/cli/internal/cli/run_job.go @@ -301,6 +301,11 @@ func executeRunJob(ctx context.Context, stdout, stderr io.Writer, request *runJo if err != nil { return err } + if artifactBundle.Enabled() && strings.TrimSpace(o.script) == "" { + warnings = append(warnings, + "tau: complete bundle acknowledgement is unavailable for Jobs that use the image ENTRYPOINT/CMD") + artifactBundle = artifactbundle.Runtime{} + } opts.ArtifactBundle = artifactBundle if artifactBundle.Enabled() && opts.Nodes > 1 { warnings = append(warnings, diff --git a/cli/internal/jobrender/render.go b/cli/internal/jobrender/render.go index 7d74afdb..c5c07c74 100644 --- a/cli/internal/jobrender/render.go +++ b/cli/internal/jobrender/render.go @@ -377,9 +377,11 @@ func Render(p profile.Profile, o Options) ([]byte, error) { if o.Nodes > 1 { return nil, fmt.Errorf("artifact bundle completion requires a single Job pod") } - cmd, err = artifactbundle.WrapCommand(cmd, o.ArtifactBundle) - if err != nil { - return nil, err + if len(cmd) > 0 { + cmd, err = artifactbundle.WrapCommand(cmd, o.ArtifactBundle) + if err != nil { + return nil, err + } } } diff --git a/cli/internal/jobrender/render_test.go b/cli/internal/jobrender/render_test.go index 48251f93..069e0774 100644 --- a/cli/internal/jobrender/render_test.go +++ b/cli/internal/jobrender/render_test.go @@ -2436,6 +2436,29 @@ func TestRender_ArtifactBundleRejectsMultiNodeIndexedJob(t *testing.T) { } } +func TestRender_ArtifactBundlePreservesImageEntrypointWithoutExplicitCommand(t *testing.T) { + out, err := Render(trainProfile(), Options{ + Name: "image-entrypoint", + Namespace: "research", + ArtifactBundle: artifactbundle.Runtime{ + BundleID: "bundle-1", + Run: "image-entrypoint", + Namespace: "research", + ResultPVC: "blob-training", + OutputDir: "/data/runs/image-entrypoint", + }, + }) + if err != nil { + t.Fatalf("render image entrypoint Job: %v", err) + } + job := parseYAML(t, out) + pod := job["spec"].(map[string]any)["template"].(map[string]any)["spec"].(map[string]any) + container := pod["containers"].([]any)[0].(map[string]any) + if command, ok := container["command"]; ok { + t.Fatalf("image entrypoint Job unexpectedly overrides command: %v", command) + } +} + func TestRender_DirectJobMetricsOffloadSafety(t *testing.T) { script := torchrunScript(t) readOnlyProfile := trainProfile() diff --git a/site/content/en/docs/reference/run-config.md b/site/content/en/docs/reference/run-config.md index 3bc3824d..8589a21d 100644 --- a/site/content/en/docs/reference/run-config.md +++ b/site/content/en/docs/reference/run-config.md @@ -52,9 +52,11 @@ the non-secret Blob transport identity on new workloads; this metadata contains no account key, SAS token, or other repository credential. Multi-node Indexed Jobs intentionally do not emit a bundle acknowledgement until Tau has a Job-level coordinator; a single index is not allowed to claim completion for the -whole Job. Bundle retrieval also requires the result PVC at Tau's `/data` -durable-storage root; custom mount roots continue to run without a bundle -acknowledgement. +whole Job. Jobs that rely entirely on the image's `ENTRYPOINT`/`CMD`, and +file-valued `storage.output` paths, also continue without a bundle +acknowledgement because Tau cannot safely wrap them as a result directory. +Bundle retrieval requires the result PVC at Tau's `/data` durable-storage root; +custom mount roots continue to run without a bundle acknowledgement. Main field groups: diff --git a/site/content/en/docs/tasks/researcher/first-run.md b/site/content/en/docs/tasks/researcher/first-run.md index e7f059c4..13cd603e 100644 --- a/site/content/en/docs/tasks/researcher/first-run.md +++ b/site/content/en/docs/tasks/researcher/first-run.md @@ -67,7 +67,8 @@ an account key or SAS token. Legacy workloads may require permission to read their bound PV for transport discovery. A missing staged-publication marker, metrics acknowledgement, final bundle marker, or storage listing is an error rather than a partial success. Downloads -refuse to replace existing destination files. Runs created before the bundle +are staged beside the requested destination and published only after every file +and metadata record succeeds; an existing destination is never replaced. Runs created before the bundle acknowledgement contract can still list or fetch known artifacts where practical, but cannot claim a complete bundle. From 6503bd20f2dc554ef3da6e3d769d629221e01158 Mon Sep 17 00:00:00 2001 From: Pengfei Ni Date: Sat, 8 Aug 2026 13:24:02 +0800 Subject: [PATCH 3/3] cli: make shell quote escaping explicit --- cli/internal/artifactbundle/bundle_test.go | 11 +++++++++++ cli/internal/artifactbundle/publish.go | 7 ++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/cli/internal/artifactbundle/bundle_test.go b/cli/internal/artifactbundle/bundle_test.go index abe2a842..cc7ea623 100644 --- a/cli/internal/artifactbundle/bundle_test.go +++ b/cli/internal/artifactbundle/bundle_test.go @@ -83,6 +83,17 @@ func testRuntime() Runtime { } } +func TestShellQuotePreservesJSONQuotes(t *testing.T) { + raw := "{\"message\":\"it's $(printf injected)\"}\n" + out, err := exec.Command("bash", "-c", "printf '%s' "+shellQuote(raw)).CombinedOutput() + if err != nil { + t.Fatalf("evaluate quoted JSON: %v\n%s", err, out) + } + if string(out) != raw { + t.Fatalf("quoted JSON = %q, want %q", out, raw) + } +} + func TestWrapperCommitsOnlyAfterNestedAcknowledgements(t *testing.T) { runtime := testRuntime() root := t.TempDir() diff --git a/cli/internal/artifactbundle/publish.go b/cli/internal/artifactbundle/publish.go index e1eeae1d..009d06a1 100644 --- a/cli/internal/artifactbundle/publish.go +++ b/cli/internal/artifactbundle/publish.go @@ -138,6 +138,11 @@ fi shellQuote(currentCompletion+".tmp.$$"), shellQuote(runtime.BundleID), shellQuote(currentCompletion)), nil } +var shellQuoteReplacer = strings.NewReplacer( + "'", `'"'"'`, + `"`, `'"\""'`, +) + func shellQuote(value string) string { - return "'" + strings.ReplaceAll(value, "'", `'"'"'`) + "'" + return "'" + shellQuoteReplacer.Replace(value) + "'" }