feat(wanda): add Artifact struct to spec for extraction

Add support for defining artifacts to extract from built container images. The Artifact struct specifies:
- src: path inside container (supports variable expansion)
- dst: destination path on host (relative to artifacts dir)
- optional: if true, extraction failure warns instead of failing build

Topic: wanda-artifact-spec
Labels: draft

Signed-off-by: andrew <andrew@anyscale.com>

Author: andrew-anyscale

wanda/spec.go

@@ -8,6 +8,23 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+// Artifact defines a file or directory to extract from the built image.
+type Artifact struct {
+	// Src is the path inside the container to extract.
+	// Can be a file or directory. Supports variable expansion.
+	Src string `yaml:"src"`
+
+	// Dst is the destination path on the host filesystem.
+	// If it ends with "/", src is copied into the directory.
+	// Otherwise, src is copied as the file/directory named dst.
+	// Relative paths are relative to ArtifactsDir.
+	Dst string `yaml:"dst"`
+
+	// Optional marks this artifact as best-effort.
+	// If true, extraction failure will be logged but won't fail the build.
+	Optional bool `yaml:"optional,omitempty"`
+}
+
 // Spec is a specification for a container image.
 type Spec struct {
 	Name string `yaml:"name,omitempty"`
@@ -28,6 +45,9 @@ type Spec struct {
 
 	// DisableCaching disables use of caching.
 	DisableCaching bool `yaml:"disable_caching,omitempty"`
+
+	// Artifacts defines files and directories to extract from the built image.
+	Artifacts []*Artifact `yaml:"artifacts,omitempty"`
 }
 
 func parseSpecFile(f string) (*Spec, error) {
@@ -116,6 +136,21 @@ func stringsExpandVar(slice []string, lookup lookupFunc) []string {
 	return result
 }
 
+func artifactsExpandVar(artifacts []*Artifact, lookup lookupFunc) []*Artifact {
+	if len(artifacts) == 0 {
+		return nil
+	}
+	result := make([]*Artifact, len(artifacts))
+	for i, a := range artifacts {
+		result[i] = &Artifact{
+			Src:      expandVar(a.Src, lookup),
+			Dst:      expandVar(a.Dst, lookup),
+			Optional: a.Optional,
+		}
+	}
+	return result
+}
+
 func (s *Spec) expandVar(lookup lookupFunc) *Spec {
 	result := new(Spec)
 
@@ -127,6 +162,7 @@ func (s *Spec) expandVar(lookup lookupFunc) *Spec {
 	result.BuildArgs = stringsExpandVar(s.BuildArgs, lookup)
 	result.BuildHintArgs = stringsExpandVar(s.BuildHintArgs, lookup)
 	result.DisableCaching = s.DisableCaching
+	result.Artifacts = artifactsExpandVar(s.Artifacts, lookup)
 
 	return result
 }

wanda/spec_test.go

@@ -178,3 +178,120 @@ func TestParseSpecFile(t *testing.T) {
 		t.Errorf("got %+v, want %+v", got, want)
 	}
 }
+
+func TestParseSpecFileWithArtifacts(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	specFile := filepath.Join(tmpDir, "spec.yaml")
+	spec := strings.Join([]string{
+		"name: wheel-builder",
+		"froms: [python:3.11]",
+		"dockerfile: Dockerfile",
+		"artifacts:",
+		"  - src: /build/dist/*.whl",
+		"    dst: ./wheels/",
+		"  - src: /build/docs/",
+		"    dst: ./docs-output/",
+		"  - src: /app/bin/myapp",
+		"    dst: ./bin/myapp",
+	}, "\n") + "\n"
+
+	if err := os.WriteFile(specFile, []byte(spec), 0644); err != nil {
+		t.Fatalf("write spec file: %v", err)
+	}
+
+	got, err := parseSpecFile(specFile)
+	if err != nil {
+		t.Fatalf("parse spec file: %v", err)
+	}
+
+	want := &Spec{
+		Name:       "wheel-builder",
+		Froms:      []string{"python:3.11"},
+		Dockerfile: "Dockerfile",
+		Artifacts: []*Artifact{
+			{Src: "/build/dist/*.whl", Dst: "./wheels/"},
+			{Src: "/build/docs/", Dst: "./docs-output/"},
+			{Src: "/app/bin/myapp", Dst: "./bin/myapp"},
+		},
+	}
+
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("got %+v, want %+v", got, want)
+	}
+}
+
+func TestSpecExpandWithArtifacts(t *testing.T) {
+	spec := &Spec{
+		Name:       "wheel-builder",
+		Froms:      []string{"python:$PYTHON_VERSION"},
+		Dockerfile: "Dockerfile",
+		Artifacts: []*Artifact{
+			{Src: "/build/$PROJECT/dist/*.whl", Dst: "$OUTPUT_DIR/wheels/"},
+			{Src: "/build/docs/", Dst: "./docs/"},
+		},
+	}
+
+	envs := map[string]string{
+		"PYTHON_VERSION": "3.11",
+		"PROJECT":        "myproject",
+		"OUTPUT_DIR":     "/tmp/artifacts",
+	}
+
+	expanded := spec.expandVar(func(k string) (string, bool) {
+		v, ok := envs[k]
+		return v, ok
+	})
+
+	want := &Spec{
+		Name:       "wheel-builder",
+		Froms:      []string{"python:3.11"},
+		Dockerfile: "Dockerfile",
+		Artifacts: []*Artifact{
+			{Src: "/build/myproject/dist/*.whl", Dst: "/tmp/artifacts/wheels/"},
+			{Src: "/build/docs/", Dst: "./docs/"},
+		},
+	}
+
+	if !reflect.DeepEqual(expanded, want) {
+		t.Errorf("got %+v, want %+v", expanded, want)
+	}
+}
+
+func TestSpecMarshalLoopbackWithArtifacts(t *testing.T) {
+	spec := &Spec{
+		Name:       "test",
+		Froms:      []string{"ubuntu:22.04"},
+		Dockerfile: "Dockerfile",
+		Artifacts: []*Artifact{
+			{Src: "/app/output.bin", Dst: "./output/"},
+			{Src: "/app/data.txt", Dst: "./data/"},
+		},
+	}
+
+	bs, err := yaml.Marshal(spec)
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+
+	loopback := new(Spec)
+	if err := yaml.Unmarshal(bs, loopback); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+
+	if spec.Name != loopback.Name {
+		t.Errorf("Name: got %q, want %q", loopback.Name, spec.Name)
+	}
+	if spec.Dockerfile != loopback.Dockerfile {
+		t.Errorf("Dockerfile: got %q, want %q", loopback.Dockerfile, spec.Dockerfile)
+	}
+	if len(spec.Artifacts) != len(loopback.Artifacts) {
+		t.Fatalf("Artifacts length: got %d, want %d", len(loopback.Artifacts), len(spec.Artifacts))
+	}
+	for i, a := range spec.Artifacts {
+		lb := loopback.Artifacts[i]
+		if a.Src != lb.Src || a.Dst != lb.Dst {
+			t.Errorf("Artifacts[%d]: got %+v, want %+v", i, lb, a)
+		}
+	}
+}