feat(wanda): implement artifact copying from built images

Extract artifacts from built images using `docker cp`. This approach is cross-platform reliable on Windows.

- Creates container from image without starting it
- Copies each artifact using `docker cp`
- Cleans up container when done
- Optional artifacts log warning instead of failing
- Extraction runs for root spec only, even on cache hit
- Logs extraction duration for performance monitoring

Topic: wanda-artifact-copy
Relative: wanda-artifact-spec
Labels: draft

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

Author: andrew-anyscale

wanda/docker_cmd.go

@@ -126,6 +126,28 @@ func (c *dockerCmd) tag(src, asTag string) error {
 	return c.run("tag", src, asTag)
 }
 
+// createContainer creates a container from an image without starting it.
+// Returns the container ID.
+func (c *dockerCmd) createContainer(image string) (string, error) {
+	cmd := c.cmd("create", image)
+	buf := new(bytes.Buffer)
+	cmd.Stdout = buf
+	if err := cmd.Run(); err != nil {
+		return "", err
+	}
+	return strings.TrimSpace(buf.String()), nil
+}
+
+// copyFromContainer copies a file or directory from a container to the host.
+func (c *dockerCmd) copyFromContainer(containerID, src, dst string) error {
+	return c.run("cp", containerID+":"+src, dst)
+}
+
+// removeContainer removes a container.
+func (c *dockerCmd) removeContainer(containerID string) error {
+	return c.run("rm", containerID)
+}
+
 func (c *dockerCmd) build(in *buildInput, core *buildInputCore, hints *buildInputHints) error {
 	if hints == nil {
 		hints = newBuildInputHints(nil, nil)

wanda/docker_cmd_test.go

@@ -1,9 +1,9 @@
 package wanda
 
 import (
-	"testing"
-
+	"os"
 	"strings"
+	"testing"
 
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/v1/daemon"
@@ -135,3 +135,80 @@ func TestDockerCmdBuild_withHints(t *testing.T) {
 		t.Errorf("MESSAGE env got %q, want `MESSAGE=hint message`", messageEnv)
 	}
 }
+
+func TestDockerCmdCopyFromContainer(t *testing.T) {
+	cmd := newDockerCmd(&dockerCmdConfig{})
+
+	const testImage = "alpine:latest"
+
+	if err := cmd.run("pull", testImage); err != nil {
+		t.Fatalf("pull image: %v", err)
+	}
+
+	containerID, err := cmd.createContainer(testImage)
+	if err != nil {
+		t.Fatalf("createContainer: %v", err)
+	}
+	defer cmd.removeContainer(containerID)
+
+	tmpDir := t.TempDir()
+
+	// Copy a known file from the container
+	if err := cmd.copyFromContainer(containerID, "/etc/alpine-release", tmpDir+"/alpine-release"); err != nil {
+		t.Fatalf("copyFromContainer: %v", err)
+	}
+
+	if _, err := os.Stat(tmpDir + "/alpine-release"); os.IsNotExist(err) {
+		t.Error("alpine-release was not copied")
+	}
+}
+
+func TestDockerCmdCopyFromContainer_directory(t *testing.T) {
+	cmd := newDockerCmd(&dockerCmdConfig{})
+
+	const testImage = "alpine:latest"
+
+	if err := cmd.run("pull", testImage); err != nil {
+		t.Fatalf("pull image: %v", err)
+	}
+
+	containerID, err := cmd.createContainer(testImage)
+	if err != nil {
+		t.Fatalf("createContainer: %v", err)
+	}
+	defer cmd.removeContainer(containerID)
+
+	tmpDir := t.TempDir()
+
+	// Copy a directory from the container
+	if err := cmd.copyFromContainer(containerID, "/etc", tmpDir+"/etc"); err != nil {
+		t.Fatalf("copyFromContainer: %v", err)
+	}
+
+	if _, err := os.Stat(tmpDir + "/etc/alpine-release"); os.IsNotExist(err) {
+		t.Error("alpine-release was not copied from /etc directory")
+	}
+}
+
+func TestDockerCmdCopyFromContainer_notFound(t *testing.T) {
+	cmd := newDockerCmd(&dockerCmdConfig{})
+
+	const testImage = "alpine:latest"
+
+	if err := cmd.run("pull", testImage); err != nil {
+		t.Fatalf("pull image: %v", err)
+	}
+
+	containerID, err := cmd.createContainer(testImage)
+	if err != nil {
+		t.Fatalf("createContainer: %v", err)
+	}
+	defer cmd.removeContainer(containerID)
+
+	tmpDir := t.TempDir()
+
+	// Copying a non-existent file should fail
+	if err := cmd.copyFromContainer(containerID, "/nonexistent/file", tmpDir+"/file"); err == nil {
+		t.Error("copyFromContainer should fail for non-existent file")
+	}
+}

wanda/forge.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"time"
 
 	"github.com/google/go-containerregistry/pkg/authn"
 	cranename "github.com/google/go-containerregistry/pkg/name"
@@ -69,6 +70,17 @@ func Build(specFile string, config *ForgeConfig) error {
 		}
 	}
 
+	// Extract artifacts only for the root spec.
+	if config.ArtifactsDir != "" {
+		rootSpec := graph.Specs[graph.Root].Spec
+		if len(rootSpec.Artifacts) > 0 {
+			rootTag := forge.workTag(rootSpec.Name)
+			if err := forge.ExtractArtifacts(rootSpec, rootTag); err != nil {
+				return fmt.Errorf("extract artifacts: %w", err)
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -182,7 +194,68 @@ func (f *Forge) resolveBases(froms []string) (map[string]*imageSource, error) {
 	return m, nil
 }
 
+// ExtractArtifacts extracts artifacts from a built image to the artifacts directory.
+// NOTE(andrew-anyscale): We use individual `docker cp` commands rather than
+// `docker export | tar` or any other potentially faster method currently because
+// docker cp takes care of many cross-platform headache inducing issues.
+// If this becomes a bottleneck, we can consider using a different approach.
+func (f *Forge) ExtractArtifacts(spec *Spec, imageTag string) error {
+	d := f.newDockerCmd()
+	artifactsDir := f.config.ArtifactsDir
+
+	if err := os.MkdirAll(artifactsDir, 0755); err != nil {
+		return fmt.Errorf("create artifacts dir: %w", err)
+	}
+
+	log.Printf("extracting %d artifact(s) from %s", len(spec.Artifacts), imageTag)
+	extractStart := time.Now()
+
+	containerID, err := d.createContainer(imageTag)
+	if err != nil {
+		return fmt.Errorf("create container: %w", err)
+	}
+	defer func() {
+		if err := d.removeContainer(containerID); err != nil {
+			log.Printf("warning: failed to remove container %s: %v", containerID, err)
+		}
+	}()
+
+	for _, a := range spec.Artifacts {
+		dst := resolveArtifactDst(a.Dst, artifactsDir)
+		log.Printf("  %s -> %s", a.Src, dst)
+
+		// Create parent directory for the destination.
+		dstDir := filepath.Dir(dst)
+		if strings.HasSuffix(a.Dst, "/") {
+			dstDir = dst
+		}
+		if err := os.MkdirAll(dstDir, 0755); err != nil {
+			return fmt.Errorf("create dir for artifact %s: %w", a.Dst, err)
+		}
+
+		if err := d.copyFromContainer(containerID, a.Src, dst); err != nil {
+			if a.Optional {
+				log.Printf("warning: optional artifact not found: %s", a.Src)
+				continue
+			}
+			return fmt.Errorf("copy artifact %s: %w", a.Src, err)
+		}
+	}
+
+	log.Printf("extraction completed in %v", time.Since(extractStart))
+	return nil
+}
+
+// resolveArtifactDst resolves the destination path for an artifact.
+func resolveArtifactDst(dst, artifactsDir string) string {
+	if filepath.IsAbs(dst) {
+		return dst
+	}
+	return filepath.Join(artifactsDir, dst)
+}
+
 // Build builds a container image from the given specification.
+// If the spec has artifacts defined, they will be extracted after build.
 func (f *Forge) Build(spec *Spec) error {
 	// Prepare the tar stream.
 	ts := newTarStream()
@@ -282,6 +355,7 @@ func (f *Forge) Build(spec *Spec) error {
 						}
 					}
 				}
+
 				return nil // and we are done.
 			}
 		}

wanda/forge_config.go

@@ -15,6 +15,7 @@ type ForgeConfig struct {
 	Epoch          string
 	WandaSpecsFile string
 	EnvFile        string
+	ArtifactsDir   string
 
 	RayCI   bool
 	Rebuild bool

wanda/forge_test.go

@@ -846,3 +846,137 @@ func TestBuild_EnvfileCacheInvalidation(t *testing.T) {
 		t.Errorf("expected cache miss on changed envfile, got %d hits", forge2.cacheHit())
 	}
 }
+
+func TestForgeConfigArtifactsDir(t *testing.T) {
+	config := &ForgeConfig{ArtifactsDir: "/custom/artifacts", WorkDir: "/work"}
+	if got := config.ArtifactsDir; got != "/custom/artifacts" {
+		t.Errorf("ArtifactsDir = %q, want %q", got, "/custom/artifacts")
+	}
+}
+
+func TestBuild_WithArtifacts_exact(t *testing.T) {
+	tmpDir := t.TempDir()
+	artifactsDir := filepath.Join(tmpDir, "artifacts")
+
+	config := &ForgeConfig{
+		WorkDir:      "testdata",
+		NamePrefix:   "cr.ray.io/rayproject/",
+		ArtifactsDir: artifactsDir,
+	}
+
+	if err := Build("testdata/artifact-exact.wanda.yaml", config); err != nil {
+		t.Fatalf("build: %v", err)
+	}
+
+	extractedFile := filepath.Join(artifactsDir, "bin/myapp")
+	content, err := os.ReadFile(extractedFile)
+	if err != nil {
+		t.Fatalf("read extracted file: %v", err)
+	}
+
+	want := "binary-content\n"
+	if got := string(content); got != want {
+		t.Errorf("extracted content = %q, want %q", got, want)
+	}
+}
+
+func TestBuild_WithArtifacts_optional(t *testing.T) {
+	tmpDir := t.TempDir()
+	artifactsDir := filepath.Join(tmpDir, "artifacts")
+
+	config := &ForgeConfig{
+		WorkDir:      "testdata",
+		NamePrefix:   "cr.ray.io/rayproject/",
+		ArtifactsDir: artifactsDir,
+	}
+
+	// Build should succeed even though the optional artifact doesn't exist
+	if err := Build("testdata/artifact-optional.wanda.yaml", config); err != nil {
+		t.Fatalf("build: %v", err)
+	}
+
+	// Required artifact should be extracted
+	extractedFile := filepath.Join(artifactsDir, "bin/myapp")
+	content, err := os.ReadFile(extractedFile)
+	if err != nil {
+		t.Fatalf("read extracted file: %v", err)
+	}
+
+	want := "binary-content\n"
+	if got := string(content); got != want {
+		t.Errorf("extracted content = %q, want %q", got, want)
+	}
+
+	// Optional artifact should not exist (since source doesn't exist)
+	optionalFile := filepath.Join(artifactsDir, "optional.txt")
+	if _, err := os.Stat(optionalFile); !os.IsNotExist(err) {
+		t.Errorf("optional file should not exist, but got err: %v", err)
+	}
+}
+
+func TestBuild_WithArtifacts_rootOnly(t *testing.T) {
+	wandaSpecs := filepath.Join(t.TempDir(), ".wandaspecs")
+	absTestdata, err := filepath.Abs("testdata")
+	if err != nil {
+		t.Fatalf("abs testdata: %v", err)
+	}
+	if err := os.WriteFile(wandaSpecs, []byte(absTestdata), 0644); err != nil {
+		t.Fatalf("write wandaspecs: %v", err)
+	}
+
+	tmpDir := t.TempDir()
+	artifactsDir := filepath.Join(tmpDir, "artifacts")
+
+	config := &ForgeConfig{
+		WorkDir:        "testdata",
+		NamePrefix:     "cr.ray.io/rayproject/",
+		WandaSpecsFile: wandaSpecs,
+		ArtifactsDir:   artifactsDir,
+	}
+
+	if err := Build("testdata/artifact-dep-top.wanda.yaml", config); err != nil {
+		t.Fatalf("build with deps: %v", err)
+	}
+
+	topFile := filepath.Join(artifactsDir, "top.txt")
+	if _, err := os.Stat(topFile); os.IsNotExist(err) {
+		t.Error("root spec artifact should have been extracted")
+	}
+
+	depDocsFile := filepath.Join(artifactsDir, "docs/readme.md")
+	if _, err := os.Stat(depDocsFile); !os.IsNotExist(err) {
+		t.Error("dependency artifact should NOT have been extracted")
+	}
+}
+
+func TestBuild_WithArtifacts_cacheHit(t *testing.T) {
+	tmpDir := t.TempDir()
+	artifactsDir1 := filepath.Join(tmpDir, "artifacts1")
+	artifactsDir2 := filepath.Join(tmpDir, "artifacts2")
+
+	config := &ForgeConfig{
+		WorkDir:      "testdata",
+		NamePrefix:   "cr.ray.io/rayproject/",
+		ArtifactsDir: artifactsDir1,
+	}
+
+	if err := Build("testdata/artifact-exact.wanda.yaml", config); err != nil {
+		t.Fatalf("first build: %v", err)
+	}
+
+	extractedFile1 := filepath.Join(artifactsDir1, "bin/myapp")
+	if _, err := os.Stat(extractedFile1); os.IsNotExist(err) {
+		t.Fatal("first build should have extracted artifact")
+	}
+
+	config.ArtifactsDir = artifactsDir2
+
+	if err := Build("testdata/artifact-exact.wanda.yaml", config); err != nil {
+		t.Fatalf("second build: %v", err)
+	}
+
+	extractedFile2 := filepath.Join(artifactsDir2, "bin/myapp")
+	if _, err := os.Stat(extractedFile2); os.IsNotExist(err) {
+		t.Error("artifact should be extracted even on cache hit")
+	}
+}

wanda/testdata/Dockerfile.artifact

@@ -0,0 +1,8 @@
+FROM alpine:latest
+
+RUN mkdir -p /build/dist /build/docs /app/bin \
+    && echo "wheel-content" > /build/dist/mypackage-1.0.0.whl \
+    && echo "wheel-content2" > /build/dist/mypackage-1.0.1.whl \
+    && echo "docs-content" > /build/docs/README.md \
+    && echo "binary-content" > /app/bin/myapp \
+    && chmod +x /app/bin/myapp

wanda/testdata/Dockerfile.artifact-top

@@ -0,0 +1,4 @@
+ARG BASE
+FROM ${BASE}
+
+RUN echo "top-layer-file" > /top-layer.txt

wanda/testdata/artifact-dep-base.wanda.yaml

@@ -0,0 +1,5 @@
+name: artifact-dep-base
+dockerfile: Dockerfile.artifact
+artifacts:
+  - src: /build/docs/README.md
+    dst: docs/readme.md