diff --git a/cmd/ecs-agent/agent.go b/cmd/ecs-agent/agent.go index 982df46db..c28d51732 100644 --- a/cmd/ecs-agent/agent.go +++ b/cmd/ecs-agent/agent.go @@ -43,6 +43,11 @@ type Agent struct { // cred serves task IAM role credentials at 169.254.170.2 (nil in unit tests). cred *credEndpoint + // gpu is the local pinning ledger over the host's discovered GPU UUIDs. + // Always non-nil (newAgent seeds an empty one); New() reseeds it from + // nvidia-smi discovery. + gpu *gpuLedger + closers []func() error } @@ -61,6 +66,7 @@ func newAgent(cfg config, id identity, cp controlPlane, puller ctrruntime.ImageP reg: newRegistrar(cp, id), hb: newHeartbeater(cp, id, cfg.Heartbeat), netns: newTaskNetns(execNetRunner{}), + gpu: newGPULedger(id.Capacity.GPUIDs), } } @@ -84,7 +90,7 @@ func New(cfg config) (*Agent, error) { InstanceID: meta.InstanceID, AZ: meta.AZ, Hostname: host, - Capacity: detectCapacity(), + Capacity: detectCapacity(discoverNvidiaGPUs(execCommandRunner)), AgentVersion: version, } @@ -153,11 +159,15 @@ func (a *Agent) Stop() error { return firstErr } -// detectCapacity reports the host's total CPU units (1 vCPU = 1024) and memory. -func detectCapacity() bus.InstanceCapacity { +// detectCapacity reports the host's total CPU units (1 vCPU = 1024), memory, +// and GPU count/UUIDs from the caller's nvidia-smi discovery (nil on a +// non-GPU host). +func detectCapacity(gpuUUIDs []string) bus.InstanceCapacity { return bus.InstanceCapacity{ CPU: goruntime.NumCPU() * 1024, MemoryMiB: memTotalMiB(), + GPU: len(gpuUUIDs), + GPUIDs: gpuUUIDs, } } diff --git a/cmd/ecs-agent/agent_test.go b/cmd/ecs-agent/agent_test.go index d076608a9..1c035dddd 100644 --- a/cmd/ecs-agent/agent_test.go +++ b/cmd/ecs-agent/agent_test.go @@ -42,7 +42,7 @@ func TestAgent_RunRegistersThenStopsOnContext(t *testing.T) { } func TestDetectCapacity_Positive(t *testing.T) { - c := detectCapacity() + c := detectCapacity(nil) if c.CPU <= 0 { t.Errorf("CPU = %d, want > 0", c.CPU) } @@ -50,4 +50,17 @@ func TestDetectCapacity_Positive(t *testing.T) { if c.MemoryMiB < 0 { t.Errorf("MemoryMiB = %d, want >= 0", c.MemoryMiB) } + if c.GPU != 0 || c.GPUIDs != nil { + t.Errorf("no GPU UUIDs given: want GPU=0/GPUIDs=nil, got %d/%v", c.GPU, c.GPUIDs) + } +} + +func TestDetectCapacity_WithGPUUUIDs(t *testing.T) { + c := detectCapacity([]string{"GPU-aaa", "GPU-bbb"}) + if c.GPU != 2 { + t.Errorf("GPU = %d, want 2", c.GPU) + } + if len(c.GPUIDs) != 2 { + t.Errorf("GPUIDs = %v, want 2 entries", c.GPUIDs) + } } diff --git a/cmd/ecs-agent/controlplane.go b/cmd/ecs-agent/controlplane.go index c72009649..d31cb2324 100644 --- a/cmd/ecs-agent/controlplane.go +++ b/cmd/ecs-agent/controlplane.go @@ -29,6 +29,9 @@ type controlPlane interface { // assign taskIDs and stop taskIDs accepted on the previous poll. Returns the // still-pending assignments and stop directives. PollAssignments(cluster, instance string, ackAssigns, ackStops []string) ([]bus.Assign, []bus.StopDirective, error) + // ReportTaskGPU reports the local ledger's pinned device UUIDs for a task's + // containers, so the control plane can populate DescribeTasks gpuIds. + ReportTaskGPU(cluster, task string, containers []handlers_ecs.ContainerGPUReport) error } // gatewayControlPlane implements controlPlane over the SigV4 ECS gateway client. @@ -74,6 +77,11 @@ func (g *gatewayControlPlane) Register(id identity) error { }, VersionInfo: &ecs.VersionInfo{AgentVersion: aws.String(id.AgentVersion)}, } + if len(id.Capacity.GPUIDs) > 0 { + in.TotalResources = append(in.TotalResources, &ecs.Resource{ + Name: aws.String("GPU"), Type: aws.String("STRINGSET"), StringSetValue: aws.StringSlice(id.Capacity.GPUIDs), + }) + } body, err := json.Marshal(in) if err != nil { return fmt.Errorf("marshal register: %w", err) @@ -133,3 +141,15 @@ func (g *gatewayControlPlane) PollAssignments(cluster, instance string, ackAssig } return out.Assignments, out.Stops, nil } + +func (g *gatewayControlPlane) ReportTaskGPU(cluster, task string, containers []handlers_ecs.ContainerGPUReport) error { + in := &handlers_ecs.ReportTaskGPUInput{Cluster: cluster, Task: task, Containers: containers} + body, err := json.Marshal(in) + if err != nil { + return fmt.Errorf("marshal report-task-gpu: %w", err) + } + if _, err := g.client.Call("ReportTaskGPU", body); err != nil { + return fmt.Errorf("report task gpu: %w", err) + } + return nil +} diff --git a/cmd/ecs-agent/controlplane_test.go b/cmd/ecs-agent/controlplane_test.go new file mode 100644 index 000000000..8620d5bc4 --- /dev/null +++ b/cmd/ecs-agent/controlplane_test.go @@ -0,0 +1,136 @@ +package main + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/aws/aws-sdk-go/service/ecs" + "github.com/mulgadc/spinifex/cmd/ecs-agent/credentials" + handlers_ecs "github.com/mulgadc/spinifex/spinifex/handlers/ecs" +) + +// fixedCredsProvider is a static credentials.CredentialsProvider for tests +// that only need a valid SigV4 signature, not real IMDS rotation. +type fixedCredsProvider struct{} + +func (fixedCredsProvider) Retrieve(context.Context) (credentials.Credentials, error) { + return credentials.Credentials{AccessKeyID: "AKIATEST", SecretAccessKey: "secret"}, nil +} + +func newTestGatewayControlPlane(t *testing.T, srv *httptest.Server) *gatewayControlPlane { + t.Helper() + cp, err := newGatewayControlPlane(config{GatewayURL: srv.URL, Region: "us-east-1"}, fixedCredsProvider{}) + if err != nil { + t.Fatalf("newGatewayControlPlane: %v", err) + } + return cp +} + +// Register carries the discovered GPU UUIDs as a STRINGSET resource, matching +// the control plane's RegisterContainerInstance GPU parsing (Epic C3). +func TestGatewayControlPlane_Register_SendsGPUStringSet(t *testing.T) { + var body []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ = io.ReadAll(r.Body) + _, _ = w.Write([]byte(`{}`)) + })) + defer srv.Close() + + cp := newTestGatewayControlPlane(t, srv) + id := identity{ + ClusterName: "default", InstanceID: "i-1", + Capacity: detectCapacity([]string{"GPU-aaa", "GPU-bbb"}), + } + if err := cp.Register(id); err != nil { + t.Fatalf("Register: %v", err) + } + + var in ecs.RegisterContainerInstanceInput + if err := json.Unmarshal(body, &in); err != nil { + t.Fatalf("decode body: %v", err) + } + var gpu *ecs.Resource + for _, r := range in.TotalResources { + if r.Name != nil && *r.Name == "GPU" { + gpu = r + } + } + if gpu == nil { + t.Fatal("no GPU resource in TotalResources") + } + if *gpu.Type != "STRINGSET" { + t.Errorf("GPU resource type = %q, want STRINGSET", *gpu.Type) + } + got := make([]string, len(gpu.StringSetValue)) + for i, v := range gpu.StringSetValue { + got[i] = *v + } + if len(got) != 2 || got[0] != "GPU-aaa" || got[1] != "GPU-bbb" { + t.Errorf("GPU UUIDs = %v, want [GPU-aaa GPU-bbb]", got) + } +} + +// A non-GPU host (no discovered UUIDs) registers with no GPU resource at all. +func TestGatewayControlPlane_Register_NoGPUOmitsResource(t *testing.T) { + var body []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ = io.ReadAll(r.Body) + _, _ = w.Write([]byte(`{}`)) + })) + defer srv.Close() + + cp := newTestGatewayControlPlane(t, srv) + id := identity{ClusterName: "default", InstanceID: "i-1", Capacity: detectCapacity(nil)} + if err := cp.Register(id); err != nil { + t.Fatalf("Register: %v", err) + } + + var in ecs.RegisterContainerInstanceInput + if err := json.Unmarshal(body, &in); err != nil { + t.Fatalf("decode body: %v", err) + } + for _, r := range in.TotalResources { + if r.Name != nil && *r.Name == "GPU" { + t.Fatalf("unexpected GPU resource on a non-GPU host: %+v", r) + } + } +} + +// ReportTaskGPU posts the per-container device UUIDs under the ReportTaskGPU +// action, decodable as the handlers_ecs internal shape. +func TestGatewayControlPlane_ReportTaskGPU(t *testing.T) { + var gotTarget string + var body []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotTarget = r.Header.Get("X-Amz-Target") + body, _ = io.ReadAll(r.Body) + _, _ = w.Write([]byte(`{}`)) + })) + defer srv.Close() + + cp := newTestGatewayControlPlane(t, srv) + err := cp.ReportTaskGPU("web", "t-1", []handlers_ecs.ContainerGPUReport{ + {Name: "trainer", GPUIDs: []string{"GPU-aaa"}}, + }) + if err != nil { + t.Fatalf("ReportTaskGPU: %v", err) + } + if gotTarget == "" { + t.Fatal("no X-Amz-Target header sent") + } + + var in handlers_ecs.ReportTaskGPUInput + if err := json.Unmarshal(body, &in); err != nil { + t.Fatalf("decode body: %v", err) + } + if in.Cluster != "web" || in.Task != "t-1" { + t.Errorf("cluster/task = %q/%q, want web/t-1", in.Cluster, in.Task) + } + if len(in.Containers) != 1 || in.Containers[0].Name != "trainer" || len(in.Containers[0].GPUIDs) != 1 { + t.Errorf("containers = %+v, want one trainer entry with one UUID", in.Containers) + } +} diff --git a/cmd/ecs-agent/gpu_discovery.go b/cmd/ecs-agent/gpu_discovery.go new file mode 100644 index 000000000..947e52742 --- /dev/null +++ b/cmd/ecs-agent/gpu_discovery.go @@ -0,0 +1,42 @@ +package main + +import ( + "os/exec" + "strings" +) + +// commandRunner abstracts exec.Command so discovery is unit-testable without +// shelling out. Production uses execCommandRunner; tests inject a fake. +type commandRunner func(name string, args ...string) ([]byte, error) + +// execCommandRunner runs a real host command and returns its stdout. +func execCommandRunner(name string, args ...string) ([]byte, error) { + return exec.Command(name, args...).Output() +} + +// discoverNvidiaGPUs enumerates local NVIDIA GPU device UUIDs via nvidia-smi. +// A missing binary or any run failure (no GPU host) yields an empty list, not +// an error — the agent still registers with zero GPU capacity. +func discoverNvidiaGPUs(run commandRunner) []string { + out, err := run("nvidia-smi", "--query-gpu=uuid", "--format=csv,noheader") + if err != nil { + return nil + } + return parseNvidiaSMIUUIDs(string(out)) +} + +// parseNvidiaSMIUUIDs splits nvidia-smi's newline-delimited UUID output, +// dropping blank lines. Exported as a pure function so fixture strings can be +// tested without exec. +func parseNvidiaSMIUUIDs(out string) []string { + lines := strings.Split(strings.TrimSpace(out), "\n") + uuids := make([]string, 0, len(lines)) + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" { + continue + } + uuids = append(uuids, line) + } + return uuids +} diff --git a/cmd/ecs-agent/gpu_discovery_test.go b/cmd/ecs-agent/gpu_discovery_test.go new file mode 100644 index 000000000..d5f5651c0 --- /dev/null +++ b/cmd/ecs-agent/gpu_discovery_test.go @@ -0,0 +1,71 @@ +package main + +import ( + "errors" + "reflect" + "testing" +) + +func TestParseNvidiaSMIUUIDs(t *testing.T) { + cases := []struct { + name string + in string + want []string + }{ + {"empty", "", nil}, + {"whitespace only", " \n\n ", nil}, + {"single", "GPU-11111111-1111-1111-1111-111111111111\n", []string{"GPU-11111111-1111-1111-1111-111111111111"}}, + {"multi", "GPU-aaa\nGPU-bbb\nGPU-ccc\n", []string{"GPU-aaa", "GPU-bbb", "GPU-ccc"}}, + {"blank lines between", "GPU-aaa\n\nGPU-bbb\n", []string{"GPU-aaa", "GPU-bbb"}}, + {"trailing whitespace per line", "GPU-aaa \r\nGPU-bbb\r\n", []string{"GPU-aaa", "GPU-bbb"}}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got := parseNvidiaSMIUUIDs(c.in) + if !reflect.DeepEqual(got, c.want) && (len(got) != 0 || len(c.want) != 0) { + t.Errorf("parseNvidiaSMIUUIDs(%q) = %v, want %v", c.in, got, c.want) + } + }) + } +} + +func TestDiscoverNvidiaGPUs_AbsentBinaryReturnsEmptyNoError(t *testing.T) { + run := func(name string, args ...string) ([]byte, error) { + return nil, errors.New(`exec: "nvidia-smi": executable file not found in $PATH`) + } + got := discoverNvidiaGPUs(run) + if len(got) != 0 { + t.Errorf("want empty list on absent nvidia-smi, got %v", got) + } +} + +func TestDiscoverNvidiaGPUs_ParsesFixtureOutput(t *testing.T) { + var gotName string + var gotArgs []string + run := func(name string, args ...string) ([]byte, error) { + gotName = name + gotArgs = args + return []byte("GPU-aaa\nGPU-bbb\n"), nil + } + got := discoverNvidiaGPUs(run) + want := []string{"GPU-aaa", "GPU-bbb"} + if !reflect.DeepEqual(got, want) { + t.Errorf("got %v, want %v", got, want) + } + if gotName != "nvidia-smi" { + t.Errorf("command = %q, want nvidia-smi", gotName) + } + if len(gotArgs) != 2 || gotArgs[0] != "--query-gpu=uuid" || gotArgs[1] != "--format=csv,noheader" { + t.Errorf("args = %v, want the uuid/csv,noheader query", gotArgs) + } +} + +func TestDiscoverNvidiaGPUs_NoDevicesReturnsEmpty(t *testing.T) { + run := func(name string, args ...string) ([]byte, error) { + return []byte(""), nil + } + got := discoverNvidiaGPUs(run) + if len(got) != 0 { + t.Errorf("want empty list, got %v", got) + } +} diff --git a/cmd/ecs-agent/gpu_ledger.go b/cmd/ecs-agent/gpu_ledger.go new file mode 100644 index 000000000..b1e55e87e --- /dev/null +++ b/cmd/ecs-agent/gpu_ledger.go @@ -0,0 +1,62 @@ +package main + +import ( + "fmt" + "strings" + "sync" +) + +// gpuLedger is the agent's local source of truth for which of the host's +// discovered GPU UUIDs are free vs pinned to a running task's container. C4 +// reads a task's pinned UUIDs to build the CDI device list for containerd. +type gpuLedger struct { + mu sync.Mutex + free []string + pinned map[string][]string // key ("taskID/containerName") -> uuids +} + +// newGPULedger seeds the ledger with the host's discovered device UUIDs. A nil +// or empty slice is a valid non-GPU host: every Pin then fails, which callers +// treat as best-effort (no device, no CDI injection later). +func newGPULedger(uuids []string) *gpuLedger { + free := append([]string(nil), uuids...) + return &gpuLedger{free: free, pinned: make(map[string][]string)} +} + +// gpuKey builds the ledger key for one container's pin/release. +func gpuKey(taskID, containerName string) string { + return taskID + "/" + containerName +} + +// Pin reserves n free UUIDs under key, returning them. An error means fewer +// than n devices were free; the key is left unpinned so a retry (or a +// best-effort caller) can proceed without a partial reservation. +func (l *gpuLedger) Pin(key string, n int) ([]string, error) { + if n <= 0 { + return nil, nil + } + l.mu.Lock() + defer l.mu.Unlock() + if len(l.free) < n { + return nil, fmt.Errorf("gpu ledger: requested %d GPUs, %d free", n, len(l.free)) + } + uuids := append([]string(nil), l.free[:n]...) + l.free = l.free[n:] + l.pinned[key] = uuids + return uuids, nil +} + +// ReleaseTask frees every UUID pinned under any container key belonging to +// taskID, called once the task's terminal STOPPED state is reported. +func (l *gpuLedger) ReleaseTask(taskID string) { + prefix := taskID + "/" + l.mu.Lock() + defer l.mu.Unlock() + for key, uuids := range l.pinned { + if !strings.HasPrefix(key, prefix) { + continue + } + l.free = append(l.free, uuids...) + delete(l.pinned, key) + } +} diff --git a/cmd/ecs-agent/gpu_ledger_test.go b/cmd/ecs-agent/gpu_ledger_test.go new file mode 100644 index 000000000..92a064b8b --- /dev/null +++ b/cmd/ecs-agent/gpu_ledger_test.go @@ -0,0 +1,155 @@ +package main + +import ( + "sync" + "testing" +) + +func TestGPULedger_PinReturnsDistinctUUIDs(t *testing.T) { + l := newGPULedger([]string{"GPU-a", "GPU-b", "GPU-c"}) + got, err := l.Pin("t-1/web", 2) + if err != nil { + t.Fatalf("Pin: %v", err) + } + if len(got) != 2 { + t.Fatalf("want 2 UUIDs, got %v", got) + } + if got[0] == got[1] { + t.Errorf("pinned UUIDs must be distinct, got %v", got) + } + if len(l.free) != 1 { + t.Errorf("free = %d, want 1 remaining", len(l.free)) + } +} + +func TestGPULedger_PinInsufficientFreeErrors(t *testing.T) { + l := newGPULedger([]string{"GPU-a"}) + _, err := l.Pin("t-1/web", 2) + if err == nil { + t.Fatal("want error when requesting more GPUs than free") + } + if len(l.free) != 1 { + t.Errorf("failed pin must not consume any device, free = %d", len(l.free)) + } +} + +func TestGPULedger_PinZeroIsNoop(t *testing.T) { + l := newGPULedger([]string{"GPU-a"}) + got, err := l.Pin("t-1/web", 0) + if err != nil || got != nil { + t.Errorf("Pin(0) = %v, %v; want nil, nil", got, err) + } +} + +func TestGPULedger_ReleaseTaskFreesOnlyThatTasksDevices(t *testing.T) { + l := newGPULedger([]string{"GPU-a", "GPU-b", "GPU-c", "GPU-d"}) + if _, err := l.Pin(gpuKey("t-1", "web"), 2); err != nil { + t.Fatalf("pin t-1: %v", err) + } + if _, err := l.Pin(gpuKey("t-2", "worker"), 1); err != nil { + t.Fatalf("pin t-2: %v", err) + } + if len(l.free) != 1 { + t.Fatalf("free = %d, want 1 before release", len(l.free)) + } + + l.ReleaseTask("t-1") + if len(l.free) != 3 { + t.Errorf("free = %d, want 3 after releasing t-1's two devices", len(l.free)) + } + if _, ok := l.pinned[gpuKey("t-2", "worker")]; !ok { + t.Error("t-2's pin should survive releasing t-1") + } + if _, ok := l.pinned[gpuKey("t-1", "web")]; ok { + t.Error("t-1's pin should be gone after release") + } +} + +// Multi-device pinning across concurrent task assigns must hand out disjoint +// UUID sets — no two tasks may ever share a device. +func TestGPULedger_ConcurrentPinsHandOutDistinctUUIDs(t *testing.T) { + const nDevices = 20 + uuids := make([]string, nDevices) + for i := range uuids { + uuids[i] = "GPU-" + string(rune('a'+i)) + } + l := newGPULedger(uuids) + + var wg sync.WaitGroup + var mu sync.Mutex + seen := make(map[string]bool) + errs := 0 + for i := range nDevices { + wg.Add(1) + go func(n int) { + defer wg.Done() + key := gpuKey("t-concurrent", string(rune('A'+n))) + got, err := l.Pin(key, 1) + mu.Lock() + defer mu.Unlock() + if err != nil { + errs++ + return + } + for _, u := range got { + if seen[u] { + t.Errorf("UUID %s handed out to more than one concurrent pin", u) + } + seen[u] = true + } + }(i) + } + wg.Wait() + if errs != 0 { + t.Errorf("want all %d pins to succeed (exactly enough devices), got %d errors", nDevices, errs) + } + if len(seen) != nDevices { + t.Errorf("distinct UUIDs handed out = %d, want %d", len(seen), nDevices) + } + if len(l.free) != 0 { + t.Errorf("free = %d, want 0 after all devices pinned", len(l.free)) + } +} + +// Releasing a task's devices concurrently with new pins never double-hands a +// UUID: a device is only ever free or pinned to exactly one key at a time. +func TestGPULedger_ConcurrentPinAndReleaseNoDoubleAssignment(t *testing.T) { + l := newGPULedger([]string{"GPU-a", "GPU-b"}) + var wg sync.WaitGroup + for round := range 50 { + wg.Add(2) + taskID := "t-" + string(rune('a'+round%26)) + go func(id string) { + defer wg.Done() + _, _ = l.Pin(gpuKey(id, "web"), 1) + }(taskID) + go func(id string) { + defer wg.Done() + l.ReleaseTask(id) + }(taskID) + } + wg.Wait() + // Invariant: every currently-pinned UUID appears in exactly one pinned entry + // and never simultaneously in free. + l.mu.Lock() + defer l.mu.Unlock() + freeSet := make(map[string]bool, len(l.free)) + for _, u := range l.free { + if freeSet[u] { + t.Errorf("duplicate UUID %s in free list", u) + } + freeSet[u] = true + } + pinnedSet := make(map[string]bool) + for _, uuids := range l.pinned { + for _, u := range uuids { + if pinnedSet[u] { + t.Errorf("UUID %s pinned under more than one key", u) + } + if freeSet[u] { + t.Errorf("UUID %s is both free and pinned", u) + } + pinnedSet[u] = true + } + } +} diff --git a/cmd/ecs-agent/registration_test.go b/cmd/ecs-agent/registration_test.go index 03bc03956..0999c140b 100644 --- a/cmd/ecs-agent/registration_test.go +++ b/cmd/ecs-agent/registration_test.go @@ -7,16 +7,19 @@ import ( "testing" "time" + handlers_ecs "github.com/mulgadc/spinifex/spinifex/handlers/ecs" "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) -// fakeCP is a test double for the gateway control-plane. It records registers and -// task-state reports, and replays a scripted queue of poll responses. +// fakeCP is a test double for the gateway control-plane. It records registers, +// task-state reports and GPU reports, and replays a scripted queue of poll +// responses. type fakeCP struct { mu sync.Mutex registers int states []bus.TaskState + gpuReports []gpuReport pollAcks [][]string pollStopAcks [][]string pollReplies [][]bus.Assign @@ -27,6 +30,13 @@ type fakeCP struct { submitErr bool } +// gpuReport is one recorded ReportTaskGPU call. +type gpuReport struct { + cluster string + task string + containers []handlers_ecs.ContainerGPUReport +} + var _ controlPlane = (*fakeCP)(nil) func (f *fakeCP) Register(_ identity) error { @@ -66,6 +76,25 @@ func (f *fakeCP) PollAssignments(_, _ string, ackAssigns, ackStops []string) ([] return out, stops, nil } +func (f *fakeCP) ReportTaskGPU(cluster, task string, containers []handlers_ecs.ContainerGPUReport) error { + f.mu.Lock() + defer f.mu.Unlock() + f.gpuReports = append(f.gpuReports, gpuReport{cluster: cluster, task: task, containers: containers}) + return nil +} + +func (f *fakeCP) gpuReportsFor(taskID string) []gpuReport { + f.mu.Lock() + defer f.mu.Unlock() + var out []gpuReport + for _, r := range f.gpuReports { + if r.task == taskID { + out = append(out, r) + } + } + return out +} + func (f *fakeCP) registerCount() int { f.mu.Lock() defer f.mu.Unlock() @@ -102,6 +131,15 @@ func testIdentity() identity { } } +// testIdentityWithGPUs is testIdentity seeded with discovered GPU UUIDs, so +// newAgent's ledger has real devices to pin from. +func testIdentityWithGPUs(uuids ...string) identity { + id := testIdentity() + id.Capacity.GPU = len(uuids) + id.Capacity.GPUIDs = uuids + return id +} + func TestRegistrar_Register(t *testing.T) { cp := &fakeCP{} if err := newRegistrar(cp, testIdentity()).Register(); err != nil { diff --git a/cmd/ecs-agent/runtime/containerd_run.go b/cmd/ecs-agent/runtime/containerd_run.go index ccec62579..79cfb46ae 100644 --- a/cmd/ecs-agent/runtime/containerd_run.go +++ b/cmd/ecs-agent/runtime/containerd_run.go @@ -6,12 +6,41 @@ import ( "sort" containerd "github.com/containerd/containerd/v2/client" + "github.com/containerd/containerd/v2/pkg/cdi" "github.com/containerd/containerd/v2/pkg/cio" "github.com/containerd/containerd/v2/pkg/namespaces" "github.com/containerd/containerd/v2/pkg/oci" specs "github.com/opencontainers/runtime-spec/specs-go" ) +// cdiVendorGPU is the NVIDIA CDI vendor/class prefix for whole-GPU devices, per +// the NVIDIA CDI naming convention (nvidia.com/gpu=). +const cdiVendorGPU = "nvidia.com/gpu" + +// cdiDeviceNames maps a container's pinned GPU device UUIDs to CDI device +// names. Empty input (non-GPU container) returns nil so no CDI opt is added. +func cdiDeviceNames(uuids []string) []string { + if len(uuids) == 0 { + return nil + } + names := make([]string, len(uuids)) + for i, id := range uuids { + names[i] = cdiVendorGPU + "=" + id + } + return names +} + +// cdiSpecOpts returns the oci.SpecOpts needed to inject gpuIDs (via +// containerd's CDI machinery) as devices into a container's OCI spec. It +// returns nil for a non-GPU container so Run adds nothing to the opts list. +func cdiSpecOpts(gpuIDs []string) []oci.SpecOpts { + devices := cdiDeviceNames(gpuIDs) + if len(devices) == 0 { + return nil + } + return []oci.SpecOpts{cdi.WithCDIDevices(devices...)} +} + var _ Runner = (*containerdPuller)(nil) // Run creates a container from an already-pulled image and starts its task. The @@ -51,6 +80,7 @@ func (p *containerdPuller) Run(ctx context.Context, id string, spec RunSpec) (st if len(spec.Env) > 0 { specOpts = append(specOpts, oci.WithEnv(envSlice(spec.Env))) } + specOpts = append(specOpts, cdiSpecOpts(spec.GPUIDs)...) container, err := p.client.NewContainer(ctx, id, containerd.WithNewSnapshot(id+"-snapshot", image), diff --git a/cmd/ecs-agent/runtime/containerd_run_test.go b/cmd/ecs-agent/runtime/containerd_run_test.go new file mode 100644 index 000000000..147dd577e --- /dev/null +++ b/cmd/ecs-agent/runtime/containerd_run_test.go @@ -0,0 +1,47 @@ +package runtime + +import ( + "reflect" + "testing" +) + +// cdiDeviceNames maps a container's pinned GPU UUIDs to CDI device names +// (nvidia.com/gpu=); Run wires the result into oci.WithCDIDevices when +// non-empty, so this is the injection logic's unit-testable core. +func TestCDIDeviceNames(t *testing.T) { + cases := []struct { + name string + uuids []string + want []string + }{ + {name: "empty", uuids: nil, want: nil}, + {name: "single", uuids: []string{"GPU-aaaaaaaa-1111-2222-3333-444444444444"}, + want: []string{"nvidia.com/gpu=GPU-aaaaaaaa-1111-2222-3333-444444444444"}}, + {name: "multiple preserves order", uuids: []string{"GPU-aaa", "GPU-bbb"}, + want: []string{"nvidia.com/gpu=GPU-aaa", "nvidia.com/gpu=GPU-bbb"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := cdiDeviceNames(tc.uuids) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("cdiDeviceNames(%v) = %v, want %v", tc.uuids, got, tc.want) + } + }) + } +} + +// TestCDISpecOpts_Wiring confirms Run's specOpts list picks up exactly one CDI +// opt when GPUIDs is non-empty, and none for a non-GPU container — the wiring +// Run relies on (a live oci.Spec is needed to unit test the opt's effect). +func TestCDISpecOpts_Wiring(t *testing.T) { + if opts := cdiSpecOpts(nil); len(opts) != 0 { + t.Errorf("non-GPU container: want no spec opts, got %d", len(opts)) + } + opts := cdiSpecOpts([]string{"GPU-aaa", "GPU-bbb"}) + if len(opts) != 1 { + t.Fatalf("GPU container: want exactly one CDI spec opt, got %d", len(opts)) + } + if opts[0] == nil { + t.Error("GPU container: spec opt is nil") + } +} diff --git a/cmd/ecs-agent/runtime/runtime.go b/cmd/ecs-agent/runtime/runtime.go index 84c5562bd..c7707fc7b 100644 --- a/cmd/ecs-agent/runtime/runtime.go +++ b/cmd/ecs-agent/runtime/runtime.go @@ -48,6 +48,12 @@ type RunSpec struct { // (awsvpc task ENI). Empty keeps the container in the host (VM) netns — // bridge/host mode behaviour. NetnsPath string + // GPU is the whole-GPU count requested for this container. + GPU int + // GPUIDs are the ledger-pinned device UUIDs assigned to this container, one + // per GPU requested. Empty for non-GPU containers or if pinning fell short. + // The runner injects these as CDI devices into the OCI spec. + GPUIDs []string } // RunStatus is a finished container's outcome. diff --git a/cmd/ecs-agent/taskrun.go b/cmd/ecs-agent/taskrun.go index 198f3d9b4..c229c052d 100644 --- a/cmd/ecs-agent/taskrun.go +++ b/cmd/ecs-agent/taskrun.go @@ -8,6 +8,7 @@ import ( "time" ctrruntime "github.com/mulgadc/spinifex/cmd/ecs-agent/runtime" + handlers_ecs "github.com/mulgadc/spinifex/spinifex/handlers/ecs" "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" ) @@ -124,12 +125,15 @@ func (a *Agent) runTask(ctx context.Context, as *bus.Assign) { } cid := containerID(as.TaskID, c.Name) + gpuIDs := a.pinContainerGPUs(as.TaskID, c.Name, c.GPU) spec := ctrruntime.RunSpec{ Image: c.Image, Command: c.Command, Env: withCredEnv(c.Environment, credID), Labels: taskLabels(as, c.Name), NetnsPath: netnsPath, + GPU: c.GPU, + GPUIDs: gpuIDs, } id, err := a.runner.Run(ctx, cid, spec) if err != nil { @@ -138,7 +142,9 @@ func (a *Agent) runTask(ctx context.Context, as *bus.Assign) { a.reportTaskState(as, bus.TaskStatusStopped, "container start failed: "+err.Error(), statuses) return } - statuses = append(statuses, bus.ContainerStatus{Name: c.Name, Status: bus.TaskStatusRunning, ContainerID: id}) + statuses = append(statuses, bus.ContainerStatus{ + Name: c.Name, Status: bus.TaskStatusRunning, ContainerID: id, GPUIDs: gpuIDs, + }) go a.waitContainer(ctx, as, c.Name, id) } @@ -234,8 +240,26 @@ func withCredEnv(env map[string]string, credID string) map[string]string { return out } +// pinContainerGPUs reserves n device UUIDs for a container from the local +// ledger, before the container is started so the runner can inject them as +// CDI devices. A short ledger (fewer free devices than requested, e.g. +// discovery found none) is logged and treated as best-effort: the container +// still runs without pinned UUIDs (and without GPU access) rather than +// failing the task. +func (a *Agent) pinContainerGPUs(taskID, containerName string, n int) []string { + if n <= 0 || a.gpu == nil { + return nil + } + uuids, err := a.gpu.Pin(gpuKey(taskID, containerName), n) + if err != nil { + slog.Warn("ecs-agent: gpu pin short", "task", taskID, "container", containerName, "err", err) + } + return uuids +} + // reportTaskState reports a task transition through the gateway's -// SubmitTaskStateChange action. +// SubmitTaskStateChange action, then (RUNNING with pinned GPUs) reports the +// per-container device UUIDs, and (STOPPED) releases them back to the ledger. func (a *Agent) reportTaskState(as *bus.Assign, status, reason string, containers []bus.ContainerStatus) { if a.cp == nil { return @@ -253,6 +277,28 @@ func (a *Agent) reportTaskState(as *bus.Assign, status, reason string, container if err := a.cp.SubmitTaskState(msg); err != nil { slog.Error("ecs-agent: report task-state failed", "task", as.TaskID, "err", err) } + a.reportTaskGPUs(as.TaskID, containers) + if status == bus.TaskStatusStopped && a.gpu != nil { + a.gpu.ReleaseTask(as.TaskID) + } +} + +// reportTaskGPUs sends the control plane the pinned device UUIDs for any +// container in the report that has them, so DescribeTasks can surface gpuIds. +// A no-op when nothing was pinned (the common non-GPU-task case). +func (a *Agent) reportTaskGPUs(taskID string, containers []bus.ContainerStatus) { + reports := make([]handlers_ecs.ContainerGPUReport, 0, len(containers)) + for _, c := range containers { + if len(c.GPUIDs) > 0 { + reports = append(reports, handlers_ecs.ContainerGPUReport{Name: c.Name, GPUIDs: c.GPUIDs}) + } + } + if len(reports) == 0 { + return + } + if err := a.cp.ReportTaskGPU(a.id.ClusterName, taskID, reports); err != nil { + slog.Error("ecs-agent: report task-gpu failed", "task", taskID, "err", err) + } } // containerID composes a containerd-valid container ID from a task + container. diff --git a/cmd/ecs-agent/taskrun_test.go b/cmd/ecs-agent/taskrun_test.go index d264846cb..d51e8f50f 100644 --- a/cmd/ecs-agent/taskrun_test.go +++ b/cmd/ecs-agent/taskrun_test.go @@ -70,6 +70,160 @@ func TestRunTask_PullRunReportsRunning(t *testing.T) { } } +// TestRunTask_GPUCarriedToRunSpec covers mulga-11opz (Epic C task C1): the +// AssignContainer's GPU count reaches the runtime RunSpec unchanged; a +// container with no GPU request keeps RunSpec.GPU at zero (regression). +func TestRunTask_GPUCarriedToRunSpec(t *testing.T) { + cp := &fakeCP{} + rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} + a := newRunAgent(cp, rt) + as := testAssign() + as.Containers = append(as.Containers, bus.AssignContainer{ + Name: "trainer", Image: "registry/trainer:1", GPU: 2, + }) + a.runTask(context.Background(), as) + + if len(rt.Runs) != 2 { + t.Fatalf("expected two runs, got %+v", rt.Runs) + } + if rt.Runs[0].GPU != 0 { + t.Errorf("web container: want GPU=0, got %d", rt.Runs[0].GPU) + } + if rt.Runs[1].GPU != 2 { + t.Errorf("trainer container: want GPU=2, got %d", rt.Runs[1].GPU) + } +} + +// TestRunTask_GPUPinsAndReportsUUIDs covers C3: a container's GPU request is +// pinned from the agent's local ledger, carried on the RUNNING report's +// ContainerStatus, and forwarded to the control plane via ReportTaskGPU. +func TestRunTask_GPUPinsAndReportsUUIDs(t *testing.T) { + cp := &fakeCP{} + rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} + a := newAgent(config{}, testIdentityWithGPUs("GPU-aaa", "GPU-bbb"), cp, rt, rt, nil) + as := testAssign() + as.Containers = append(as.Containers, bus.AssignContainer{ + Name: "trainer", Image: "registry/trainer:1", GPU: 1, + }) + a.runTask(context.Background(), as) + + st := cp.taskStates() + if len(st) == 0 || st[len(st)-1].LastStatus != bus.TaskStatusRunning { + t.Fatalf("want RUNNING final state, got %+v", st) + } + last := st[len(st)-1] + var trainer *bus.ContainerStatus + for i := range last.Containers { + if last.Containers[i].Name == "trainer" { + trainer = &last.Containers[i] + } + } + if trainer == nil || len(trainer.GPUIDs) != 1 { + t.Fatalf("trainer container missing pinned GPU UUID: %+v", last.Containers) + } + if trainer.GPUIDs[0] != "GPU-aaa" && trainer.GPUIDs[0] != "GPU-bbb" { + t.Errorf("unexpected pinned UUID %q", trainer.GPUIDs[0]) + } + + reports := cp.gpuReportsFor(as.TaskID) + if len(reports) != 1 { + t.Fatalf("want one ReportTaskGPU call, got %d", len(reports)) + } + if reports[0].cluster != as.ClusterName { + t.Errorf("report cluster = %q, want %q", reports[0].cluster, as.ClusterName) + } + if len(reports[0].containers) != 1 || reports[0].containers[0].Name != "trainer" { + t.Errorf("report containers = %+v, want one trainer entry", reports[0].containers) + } +} + +// TestRunTask_GPUIDsReachRunSpec covers C4: the ledger's pinned UUIDs for a +// container reach RunSpec.GPUIDs (not just the RunSpec.GPU count), so the +// runner can inject them as CDI devices; a non-GPU container's RunSpec +// carries no GPUIDs. +func TestRunTask_GPUIDsReachRunSpec(t *testing.T) { + cp := &fakeCP{} + rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} + a := newAgent(config{}, testIdentityWithGPUs("GPU-aaa", "GPU-bbb"), cp, rt, rt, nil) + as := testAssign() + as.Containers = append(as.Containers, bus.AssignContainer{ + Name: "trainer", Image: "registry/trainer:1", GPU: 1, + }) + a.runTask(context.Background(), as) + + if len(rt.Runs) != 2 { + t.Fatalf("expected two runs, got %+v", rt.Runs) + } + if len(rt.Runs[0].GPUIDs) != 0 { + t.Errorf("web container: want no GPUIDs, got %+v", rt.Runs[0].GPUIDs) + } + trainer := rt.Runs[1] + if len(trainer.GPUIDs) != 1 { + t.Fatalf("trainer container: want one pinned GPUID, got %+v", trainer.GPUIDs) + } + if trainer.GPUIDs[0] != "GPU-aaa" && trainer.GPUIDs[0] != "GPU-bbb" { + t.Errorf("unexpected pinned UUID %q", trainer.GPUIDs[0]) + } +} + +// A container with no GPU request never triggers a ReportTaskGPU call. +func TestRunTask_NoGPUSkipsReport(t *testing.T) { + cp := &fakeCP{} + rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} + a := newRunAgent(cp, rt) + as := testAssign() + a.runTask(context.Background(), as) + + if reports := cp.gpuReportsFor(as.TaskID); len(reports) != 0 { + t.Errorf("want no GPU report for a non-GPU task, got %+v", reports) + } +} + +// A GPU request exceeding the ledger's free devices (discovery found none, or +// fewer than requested) is best-effort: the container still runs, it just +// carries no pinned UUIDs and so no CDI devices are injected. +func TestRunTask_GPUShortLedgerStillRuns(t *testing.T) { + cp := &fakeCP{} + rt := &ctrruntime.FakePuller{WaitErr: errors.New("blocked")} + a := newRunAgent(cp, rt) // testIdentity() has no discovered GPU UUIDs + as := testAssign() + as.Containers = append(as.Containers, bus.AssignContainer{ + Name: "trainer", Image: "registry/trainer:1", GPU: 2, + }) + a.runTask(context.Background(), as) + + if len(rt.Runs) != 2 { + t.Fatalf("expected both containers to run despite the short ledger, got %+v", rt.Runs) + } + if reports := cp.gpuReportsFor(as.TaskID); len(reports) != 0 { + t.Errorf("want no GPU report when nothing was pinned, got %+v", reports) + } +} + +// stopTask releases the task's pinned GPU UUIDs back to the ledger so a later +// task on the same instance can reuse them. +func TestStopTask_ReleasesGPUsBackToLedger(t *testing.T) { + cp := &fakeCP{} + rt := &ctrruntime.FakePuller{Containers: []ctrruntime.Container{ + {ID: "t-001-trainer", Running: true, Labels: map[string]string{ + "mulga.ecs.taskID": "t-001", "mulga.ecs.containerName": "trainer", + }}, + }} + a := newAgent(config{}, testIdentityWithGPUs("GPU-aaa"), cp, rt, rt, nil) + if _, err := a.gpu.Pin(gpuKey("t-001", "trainer"), 1); err != nil { + t.Fatalf("seed pin: %v", err) + } + if len(a.gpu.free) != 0 { + t.Fatalf("setup: want ledger fully pinned, free = %v", a.gpu.free) + } + + a.stopTask(context.Background(), bus.StopDirective{TaskID: "t-001", Reason: "bye"}) + + if len(a.gpu.free) != 1 { + t.Errorf("want the GPU released back to the ledger, free = %v", a.gpu.free) + } +} + // A pull failure reports STOPPED and never starts the container. func TestRunTask_PullFailureReportsStopped(t *testing.T) { cp := &fakeCP{} diff --git a/go.mod b/go.mod index 2ea3ef507..58646e62e 100644 --- a/go.mod +++ b/go.mod @@ -143,6 +143,7 @@ require ( github.com/minio/crc64nvme v1.1.1 // indirect github.com/minio/highwayhash v1.0.4 // indirect github.com/moby/locker v1.0.1 // indirect + github.com/moby/sys/capability v0.4.0 // indirect github.com/moby/sys/mountinfo v0.7.2 // indirect github.com/moby/sys/sequential v0.6.0 // indirect github.com/moby/sys/signal v0.7.1 // indirect @@ -157,6 +158,7 @@ require ( github.com/nats-io/nuid v1.0.1 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect github.com/opencontainers/image-spec v1.1.1 // indirect + github.com/opencontainers/runtime-tools v0.9.1-0.20251114084447-edf4cb3d2116 // indirect github.com/pierrec/lz4/v4 v4.1.18 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect @@ -193,6 +195,9 @@ require ( google.golang.org/grpc v1.81.1 // indirect google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/yaml.v3 v3.0.1 // indirect + sigs.k8s.io/yaml v1.6.0 // indirect + tags.cncf.io/container-device-interface v1.1.0 // indirect + tags.cncf.io/container-device-interface/specs-go v1.1.0 // indirect ) // go.opencensus.io -> grpc@v1.33.2 drags in the pre-split monolithic genproto, diff --git a/go.sum b/go.sum index dcc72f28a..43859c8c7 100644 --- a/go.sum +++ b/go.sum @@ -402,6 +402,8 @@ codeberg.org/go-latex/latex v0.0.1/go.mod h1:AiC91vVG2uURZRd4ZN1j3mAac0XBrLsxK6+ codeberg.org/go-latex/latex v0.1.0/go.mod h1:LA0q/AyWIYrqVd+A9Upkgsb+IqPcmSTKc9Dny04MHMw= codeberg.org/go-pdf/fpdf v0.10.0/go.mod h1:Y0DGRAdZ0OmnZPvjbMp/1bYxmIPxm0ws4tfoPOc4LjU= contrib.go.opencensus.io/exporter/stackdriver v0.13.15-0.20230702191903-2de6d2748484/go.mod h1:uxw+4/0SiKbbVSD/F2tk5pJTdVcfIBBcsQ8gwcu4X+E= +cyphar.com/go-pathrs v0.2.1 h1:9nx1vOgwVvX1mNBWDu93+vaceedpbsDqo+XuBGL40b8= +cyphar.com/go-pathrs v0.2.1/go.mod h1:y8f1EMG7r+hCuFf/rXsKqMJrJAUoADZGNh5/vZPKcGc= dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= dmitri.shuralyov.com/gpu/mtl v0.0.0-20201218220906-28db891af037/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= dmitri.shuralyov.com/gpu/mtl v0.0.0-20221208032759-85de2813cf6b/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= @@ -602,6 +604,8 @@ github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+Ce github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs= +github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM= +github.com/blang/semver/v4 v4.0.0/go.mod h1:IbckMUScFkM3pff0VJDNKRiT6TG/YpiHIM2yvyW5YoQ= github.com/boltdb/bolt v1.3.1 h1:JQmyP4ZBrce+ZQu0dY660FMfatumYDLun9hBCUVIkF4= github.com/boltdb/bolt v1.3.1/go.mod h1:clJnj/oiGkjum5o1McbSZDSLxVThjynRyGBgiAx27Ps= github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8= @@ -725,6 +729,8 @@ github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:ma github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7DoTY= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/cyphar/filepath-securejoin v0.6.0 h1:BtGB77njd6SVO6VztOHfPxKitJvd/VPT+OFBFMOi1Is= +github.com/cyphar/filepath-securejoin v0.6.0/go.mod h1:A8hd4EnAeyujCJRrICiOWqjS1AX0a9kM5XL+NwKoYSc= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= @@ -1103,6 +1109,8 @@ github.com/hamba/avro/v2 v2.17.2/go.mod h1:Q9YK+qxAhtVrNqOhwlZTATLgLA8qxG2vtvkhK github.com/hashicorp/consul/api v1.3.0/go.mod h1:MmDNSzIMUjNpY/mQ398R4bk2FnqQLoPndWW5VkKPlCE= github.com/hashicorp/consul/sdk v0.3.0/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8= github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= +github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I= +github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= github.com/hashicorp/go-cleanhttp v0.5.0/go.mod h1:JpRdi6/HCYpAwUzNwuwqhbovhLtngrth3wmdIIUrZ80= github.com/hashicorp/go-cleanhttp v0.5.1/go.mod h1:JpRdi6/HCYpAwUzNwuwqhbovhLtngrth3wmdIIUrZ80= github.com/hashicorp/go-hclog v1.6.3 h1:Qr2kF+eVWjTiYmU7Y31tYlP1h0q/X3Nl3tPGdaB11/k= @@ -1118,6 +1126,8 @@ github.com/hashicorp/go-msgpack v0.5.5/go.mod h1:ahLV/dePpqEmjfWmKiqvPkv/twdG7iP github.com/hashicorp/go-msgpack/v2 v2.1.5 h1:Ue879bPnutj/hXfmUk6s/jtIK90XxgiUIcXRl656T44= github.com/hashicorp/go-msgpack/v2 v2.1.5/go.mod h1:bjCsRXpZ7NsJdk45PoCQnzRGDaK8TKm5ZnDI/9y3J4M= github.com/hashicorp/go-multierror v1.0.0/go.mod h1:dHtQlpGsu+cZNNAkkCN/P3hoUDHhCYQXV3UM06sGGrk= +github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo= +github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM= github.com/hashicorp/go-retryablehttp v0.5.3/go.mod h1:9B5zBasrRhHXnJnui7y6sL7es7NDiJgTc6Er0maI1Xs= github.com/hashicorp/go-rootcerts v1.0.0/go.mod h1:K6zTfqpRlCUIjkwsN4Z+hiSfzSTQa6eBIzfwKfwNnHU= github.com/hashicorp/go-sockaddr v1.0.0/go.mod h1:7Xibr9yA9JjQq1JpNB2Vw7kxv8xerXegt+ozgdvDeDU= @@ -1282,6 +1292,8 @@ github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= github.com/moby/locker v1.0.1 h1:fOXqR41zeveg4fFODix+1Ch4mj/gT0NE1XJbp/epuBg= github.com/moby/locker v1.0.1/go.mod h1:S7SDdo5zpBK84bzzVlKr2V0hz+7x9hWbYC/kq7oQppc= +github.com/moby/sys/capability v0.4.0 h1:4D4mI6KlNtWMCM1Z/K0i7RV1FkX+DBDHKVJpCndZoHk= +github.com/moby/sys/capability v0.4.0/go.mod h1:4g9IK291rVkms3LKCDOoYlnV8xKwoDTpIrNEE35Wq0I= github.com/moby/sys/mountinfo v0.7.2 h1:1shs6aH5s4o5H2zQLn796ADW1wMrIwHsyJ2v9KouLrg= github.com/moby/sys/mountinfo v0.7.2/go.mod h1:1YOa8w8Ih7uW0wALDUgT1dTTSBrZ+HiBLGws92L2RU4= github.com/moby/sys/sequential v0.6.0 h1:qrx7XFUd/5DxtqcoH1h438hF5TmOvzC/lspjy7zgvCU= @@ -1340,6 +1352,10 @@ github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJw github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= github.com/opencontainers/runtime-spec v1.3.0 h1:YZupQUdctfhpZy3TM39nN9Ika5CBWT5diQ8ibYCRkxg= github.com/opencontainers/runtime-spec v1.3.0/go.mod h1:jwyrGlmzljRJv/Fgzds9SsS/C5hL+LL3ko9hs6T5lQ0= +github.com/opencontainers/runtime-tools v0.9.1-0.20251114084447-edf4cb3d2116 h1:tAKu3NkKWZYpqBSOJKwTxT1wIGueiF7gcmcNgr5pNTY= +github.com/opencontainers/runtime-tools v0.9.1-0.20251114084447-edf4cb3d2116/go.mod h1:DKDEfzxvRkoQ6n9TGhxQgg2IM1lY4aM0eaQP4e3oElw= +github.com/opencontainers/selinux v1.13.1 h1:A8nNeceYngH9Ow++M+VVEwJVpdFmrlxsN22F+ISDCJE= +github.com/opencontainers/selinux v1.13.1/go.mod h1:S10WXZ/osk2kWOYKy1x2f/eXF5ZHJoUs8UU/2caNRbg= github.com/opentracing-contrib/go-observer v0.0.0-20170622124052-a52f23424492/go.mod h1:Ngi6UdF0k5OKD5t5wlmGhe/EDKPoUM3BXZSSfIuJbis= github.com/opentracing/basictracer-go v1.0.0/go.mod h1:QfBfYuafItcjQuMwinw9GhYKwFXS9KnPs5lxoYwgW74= github.com/opentracing/opentracing-go v1.0.2/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o= @@ -1557,6 +1573,12 @@ github.com/u-root/uio v0.0.0-20230220225925-ffce2a382923 h1:tHNk7XK9GkmKUR6Gh8gV github.com/u-root/uio v0.0.0-20230220225925-ffce2a382923/go.mod h1:eLL9Nub3yfAho7qB0MzZizFhTU2QkLeoVsWdHtDW264= github.com/urfave/cli v1.20.0/go.mod h1:70zkFmudgCuE/ngEzBv17Jvp/497gISqfk5gWijbERA= github.com/urfave/cli v1.22.1/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0= +github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb h1:zGWFAtiMcyryUHoUjUJX0/lt1H2+i2Ka2n+D3DImSNo= +github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU= +github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHovont7NscjpAxXsDA8S8BMYve8Y5+7cuRE7R0= +github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:GwrjFmJcFw6At/Gs6z4yjiIwzuJ1/+UwLxMQDVQXShQ= +github.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17UxZ74= +github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y= github.com/xhit/go-str2duration v1.2.0/go.mod h1:3cPSlfZlUHVlneIVfePFWcJZsuwf+P1v2SRTV4cUmp4= github.com/xhit/go-str2duration/v2 v2.1.0/go.mod h1:ohY8p+0f07DiV6Em5LKB0s2YpLtXVyJfNt1+BlmyAsU= github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU= @@ -3192,4 +3214,10 @@ rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0= rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA= sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o= +sigs.k8s.io/yaml v1.6.0 h1:G8fkbMSAFqgEFgh4b1wmtzDnioxFCUgTZhlbj5P9QYs= +sigs.k8s.io/yaml v1.6.0/go.mod h1:796bPqUfzR/0jLAl6XjHl3Ck7MiyVv8dbTdyT3/pMf4= sourcegraph.com/sourcegraph/appdash v0.0.0-20190731080439-ebfcffb1b5c0/go.mod h1:hI742Nqp5OhwiqlzhgfbWU4mW4yO10fP+LoT9WOswdU= +tags.cncf.io/container-device-interface v1.1.0 h1:RnxNhxF1JOu6CJUVpetTYvrXHdxw9j9jFYgZpI+anSY= +tags.cncf.io/container-device-interface v1.1.0/go.mod h1:76Oj0Yqp9FwTx/pySDc8Bxjpg+VqXfDb50cKAXVJ34Q= +tags.cncf.io/container-device-interface/specs-go v1.1.0 h1:QRZVeAceQM+zTZe12eyfuJuuzp524EKYwhmvLd+h+yQ= +tags.cncf.io/container-device-interface/specs-go v1.1.0/go.mod h1:u86hoFWqnh3hWz3esofRFKbI261bUlvUfLKGrDhJkgQ= diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go index 3e3fba6f9..881137780 100644 --- a/spinifex/daemon/daemon.go +++ b/spinifex/daemon/daemon.go @@ -1014,6 +1014,7 @@ func (d *Daemon) subscribeAll() error { natsSub{"ecs.ListServices", d.handleECSListServices, "spinifex-workers"}, natsSub{"ecs.SubmitTaskStateChange", d.handleECSSubmitTaskStateChange, "spinifex-workers"}, natsSub{"ecs.PollAssignments", d.handleECSPollAssignments, "spinifex-workers"}, + natsSub{"ecs.ReportTaskGPU", d.handleECSReportTaskGPU, "spinifex-workers"}, natsSub{"ecs.ProvisionCapacity", d.handleECSProvisionCapacity, "spinifex-workers"}, natsSub{"ecs.TagResource", d.handleECSTagResource, "spinifex-workers"}, natsSub{"ecs.UntagResource", d.handleECSUntagResource, "spinifex-workers"}, diff --git a/spinifex/daemon/daemon_handlers_ecs.go b/spinifex/daemon/daemon_handlers_ecs.go index 859452edf..7b5a25922 100644 --- a/spinifex/daemon/daemon_handlers_ecs.go +++ b/spinifex/daemon/daemon_handlers_ecs.go @@ -112,6 +112,10 @@ func (d *Daemon) handleECSPollAssignments(msg *nats.Msg) { handleNATSRequest(msg, d.ecsService.PollAssignments) } +func (d *Daemon) handleECSReportTaskGPU(msg *nats.Msg) { + handleNATSRequest(msg, d.ecsService.ReportTaskGPU) +} + // --- Capacity --- func (d *Daemon) handleECSProvisionCapacity(msg *nats.Msg) { diff --git a/spinifex/gateway/ecs/actions.go b/spinifex/gateway/ecs/actions.go index 5fd005f7b..69085098d 100644 --- a/spinifex/gateway/ecs/actions.go +++ b/spinifex/gateway/ecs/actions.go @@ -244,6 +244,17 @@ func PollAssignments(ctx context.Context, nc *nats.Conn, accountID string, body return handlers_ecs.NewNATSECSService(nc).PollAssignments(ctx, input, accountID) } +// ReportTaskGPU carries the agent's local ledger's per-task GPU device +// assignment (agent → gateway). Internal action, not an AWS SDK shape: real +// AWS's SubmitTaskStateChange has no gpuIds field. +func ReportTaskGPU(ctx context.Context, nc *nats.Conn, accountID string, body []byte) (any, error) { + input := new(handlers_ecs.ReportTaskGPUInput) + if err := unmarshalIfBody(body, input); err != nil { + return nil, err + } + return handlers_ecs.NewNATSECSService(nc).ReportTaskGPU(ctx, input, accountID) +} + // --- Tags --- func TagResource(ctx context.Context, nc *nats.Conn, accountID string, body []byte) (any, error) { diff --git a/spinifex/gateway/ecs/handler.go b/spinifex/gateway/ecs/handler.go index 0c1790efa..a63fe0682 100644 --- a/spinifex/gateway/ecs/handler.go +++ b/spinifex/gateway/ecs/handler.go @@ -76,6 +76,9 @@ var Actions = map[string]Handler{ // Agent assignment polling (agent drains its inbox over the gateway). "PollAssignments": PollAssignments, + // Agent per-task GPU device-assignment report (agent → gateway → task record). + "ReportTaskGPU": ReportTaskGPU, + // Service. ListServicesByNamespace is a no-op stub in v1. "CreateService": CreateService, "UpdateService": UpdateService, diff --git a/spinifex/gateway/ecs/handler_test.go b/spinifex/gateway/ecs/handler_test.go index b3123ed32..7ba9289b2 100644 --- a/spinifex/gateway/ecs/handler_test.go +++ b/spinifex/gateway/ecs/handler_test.go @@ -33,6 +33,7 @@ var wiredActions = map[string]bool{ "DescribeServices": true, "ListServices": true, "SubmitTaskStateChange": true, "PollAssignments": true, + "ReportTaskGPU": true, "TagResource": true, "UntagResource": true, "ListTagsForResource": true, "PutClusterCapacityProviders": true, "CreateCapacityProvider": true, "DescribeCapacityProviders": true, "DeleteCapacityProvider": true, diff --git a/spinifex/handlers/ecs/bus/messages.go b/spinifex/handlers/ecs/bus/messages.go index af346d38e..668ff0547 100644 --- a/spinifex/handlers/ecs/bus/messages.go +++ b/spinifex/handlers/ecs/bus/messages.go @@ -6,6 +6,11 @@ import "time" type InstanceCapacity struct { CPU int `json:"cpu"` // total vCPU units (1 vCPU = 1024) MemoryMiB int `json:"memoryMiB"` // total memory in MiB + // GPU is the total whole-GPU count from the agent's nvidia-smi discovery. + GPU int `json:"gpu,omitempty"` + // GPUIDs holds the discovered device UUIDs backing GPU; nil on a non-GPU + // host or a build with no nvidia-smi. + GPUIDs []string `json:"gpuIds,omitempty"` } // RegisterInstance is published on RegisterSubject when an ecs-agent boots and @@ -57,10 +62,13 @@ type PortMapping struct { // AssignContainer is one container the agent must pull and run for a task. type AssignContainer struct { - Name string `json:"name"` - Image string `json:"image"` - CPU int `json:"cpu,omitempty"` - MemoryMiB int `json:"memoryMiB,omitempty"` + Name string `json:"name"` + Image string `json:"image"` + CPU int `json:"cpu,omitempty"` + MemoryMiB int `json:"memoryMiB,omitempty"` + // GPU is the whole-GPU count requested via resourceRequirements type=GPU. + // Device pinning/CDI injection is not wired yet (later Epic C tasks). + GPU int `json:"gpu,omitempty"` Essential bool `json:"essential"` Command []string `json:"command,omitempty"` Environment map[string]string `json:"environment,omitempty"` @@ -116,6 +124,9 @@ type ContainerStatus struct { Status string `json:"status"` // PENDING | RUNNING | STOPPED ContainerID string `json:"containerId,omitempty"` ExitCode *int `json:"exitCode,omitempty"` + // GPUIDs are the device UUIDs the agent's local ledger pinned to this + // container (empty when the container requested no GPU). + GPUIDs []string `json:"gpuIds,omitempty"` } // TaskState is published on TaskStateSubject as the agent drives a task through diff --git a/spinifex/handlers/ecs/conv.go b/spinifex/handlers/ecs/conv.go index 3ff2b2a69..513e0ed85 100644 --- a/spinifex/handlers/ecs/conv.go +++ b/spinifex/handlers/ecs/conv.go @@ -1,6 +1,8 @@ package handlers_ecs import ( + "log/slog" + "strconv" "strings" "github.com/aws/aws-sdk-go/aws" @@ -58,6 +60,27 @@ func capacityProviderShortName(ref string) string { return ref } +// gpuCountFromResourceRequirements sums the whole-GPU count from a container's +// resourceRequirements (AWS ECS GPU semantics: type=GPU, value=stringified +// int). Entries with a non-numeric or negative value are skipped and warned — +// AWS itself rejects those at RegisterTaskDefinition, so this only guards +// against malformed input reaching the scheduler. +func gpuCountFromResourceRequirements(in []*ecs.ResourceRequirement) int { + total := 0 + for _, rr := range in { + if rr == nil || aws.StringValue(rr.Type) != ecs.ResourceTypeGpu { + continue + } + n, err := strconv.Atoi(aws.StringValue(rr.Value)) + if err != nil || n < 0 { + slog.Warn("ECS: ignoring invalid GPU resourceRequirement value", "value", aws.StringValue(rr.Value)) + continue + } + total += n + } + return total +} + // containerDefsFromAWS maps SDK container definitions to the persisted subset. func containerDefsFromAWS(in []*ecs.ContainerDefinition) []ContainerDef { out := make([]ContainerDef, 0, len(in)) @@ -70,6 +93,7 @@ func containerDefsFromAWS(in []*ecs.ContainerDefinition) []ContainerDef { Image: aws.StringValue(c.Image), CPU: int(aws.Int64Value(c.Cpu)), MemoryMiB: int(aws.Int64Value(c.Memory)), + GPU: gpuCountFromResourceRequirements(c.ResourceRequirements), Essential: aws.BoolValue(c.Essential), Command: awsStringSlice(c.Command), } @@ -118,6 +142,12 @@ func (c ContainerDef) toAWS() *ecs.ContainerDefinition { if c.MemoryMiB > 0 { cd.Memory = aws.Int64(int64(c.MemoryMiB)) } + if c.GPU > 0 { + cd.ResourceRequirements = []*ecs.ResourceRequirement{{ + Type: aws.String(ecs.ResourceTypeGpu), + Value: aws.String(strconv.Itoa(c.GPU)), + }} + } for _, cmd := range c.Command { cd.Command = append(cd.Command, aws.String(cmd)) } @@ -154,6 +184,7 @@ func (c ContainerDef) toAssignContainer() bus.AssignContainer { Image: c.Image, CPU: c.CPU, MemoryMiB: c.MemoryMiB, + GPU: c.GPU, Essential: c.Essential, Command: c.Command, Environment: c.Environment, diff --git a/spinifex/handlers/ecs/conv_test.go b/spinifex/handlers/ecs/conv_test.go index 7fed6b93d..03c4dc9b3 100644 --- a/spinifex/handlers/ecs/conv_test.go +++ b/spinifex/handlers/ecs/conv_test.go @@ -81,3 +81,68 @@ func TestContainerDefRoundTrip(t *testing.T) { func TestContainerDefsFromAWS_SkipsNil(t *testing.T) { assert.Empty(t, containerDefsFromAWS([]*ecs.ContainerDefinition{nil})) } + +// TestContainerDefsFromAWS_GPU covers mulga-11opz (Epic C task C1): a +// resourceRequirements entry of type=GPU is parsed as a whole-GPU count and +// carried onto ContainerDef, its AWS round trip, and the bus AssignContainer. +func TestContainerDefsFromAWS_GPU(t *testing.T) { + in := []*ecs.ContainerDefinition{{ + Name: aws.String("gpu-app"), Image: aws.String("registry/gpu-app:1"), Essential: aws.Bool(true), + ResourceRequirements: []*ecs.ResourceRequirement{ + {Type: aws.String(ecs.ResourceTypeGpu), Value: aws.String("2")}, + }, + }} + defs := containerDefsFromAWS(in) + require.Len(t, defs, 1) + assert.Equal(t, 2, defs[0].GPU) + + back := defs[0].toAWS() + require.Len(t, back.ResourceRequirements, 1) + assert.Equal(t, ecs.ResourceTypeGpu, aws.StringValue(back.ResourceRequirements[0].Type)) + assert.Equal(t, "2", aws.StringValue(back.ResourceRequirements[0].Value)) + + ac := defs[0].toAssignContainer() + assert.Equal(t, 2, ac.GPU) +} + +// TestContainerDefsFromAWS_NoGPU_Regression covers the non-GPU path: no +// resourceRequirements means GPU stays 0 and toAWS omits the field. +func TestContainerDefsFromAWS_NoGPU_Regression(t *testing.T) { + in := []*ecs.ContainerDefinition{{Name: aws.String("web"), Image: aws.String("registry/web:1"), Essential: aws.Bool(true)}} + defs := containerDefsFromAWS(in) + require.Len(t, defs, 1) + assert.Zero(t, defs[0].GPU) + assert.Empty(t, defs[0].toAWS().ResourceRequirements) + assert.Zero(t, defs[0].toAssignContainer().GPU) +} + +// TestGPUCountFromResourceRequirements covers the extraction edge cases: nil +// entries, non-GPU types, multiple GPU entries summed, and an invalid value +// (non-numeric or negative) skipped rather than erroring. +func TestGPUCountFromResourceRequirements(t *testing.T) { + assert.Zero(t, gpuCountFromResourceRequirements(nil)) + assert.Zero(t, gpuCountFromResourceRequirements([]*ecs.ResourceRequirement{nil})) + assert.Zero(t, gpuCountFromResourceRequirements([]*ecs.ResourceRequirement{ + {Type: aws.String(ecs.ResourceTypeInferenceAccelerator), Value: aws.String("device1")}, + })) + assert.Equal(t, 3, gpuCountFromResourceRequirements([]*ecs.ResourceRequirement{ + {Type: aws.String(ecs.ResourceTypeGpu), Value: aws.String("1")}, + {Type: aws.String(ecs.ResourceTypeGpu), Value: aws.String("2")}, + })) + assert.Zero(t, gpuCountFromResourceRequirements([]*ecs.ResourceRequirement{ + {Type: aws.String(ecs.ResourceTypeGpu), Value: aws.String("not-a-number")}, + })) + assert.Zero(t, gpuCountFromResourceRequirements([]*ecs.ResourceRequirement{ + {Type: aws.String(ecs.ResourceTypeGpu), Value: aws.String("-1")}, + })) +} + +// TestTaskDefReservedGPU covers the task-level GPU aggregate (sum across +// container defs), mirroring TestTaskDefReservedSums for CPU/memory. +func TestTaskDefReservedGPU(t *testing.T) { + td := &TaskDefRecord{Containers: []ContainerDef{{GPU: 1}, {GPU: 3}}} + assert.Equal(t, 4, td.reservedGPU()) + + empty := &TaskDefRecord{Containers: []ContainerDef{{CPU: 128}}} + assert.Zero(t, empty.reservedGPU()) +} diff --git a/spinifex/handlers/ecs/gpu_report.go b/spinifex/handlers/ecs/gpu_report.go new file mode 100644 index 000000000..dcbbb12fb --- /dev/null +++ b/spinifex/handlers/ecs/gpu_report.go @@ -0,0 +1,63 @@ +package handlers_ecs + +import "context" + +// ReportTaskGPUInput is the agent's per-task GPU device-assignment report +// (agent → gateway). Internal agent↔gateway shape, not an AWS SDK action: +// real AWS's SubmitTaskStateChange carries no gpuIds field, so the agent's +// local ledger assignment is reported out-of-band and merged onto the task +// record here, ahead of DescribeTasks. +type ReportTaskGPUInput struct { + Cluster string `json:"cluster"` + Task string `json:"task"` + Containers []ContainerGPUReport `json:"containers,omitempty"` +} + +// ContainerGPUReport is one container's pinned device UUIDs. +type ContainerGPUReport struct { + Name string `json:"name"` + GPUIDs []string `json:"gpuIds,omitempty"` +} + +// ReportTaskGPUOutput acknowledges the report. +type ReportTaskGPUOutput struct { + Acknowledgment string `json:"acknowledgment"` +} + +// ReportTaskGPU merges the agent-reported device UUIDs onto the task's +// container records by name, so DescribeTasks can surface gpuIds. A missing +// task or container is a silent no-op — the state-report path already owns +// the task's lifecycle; this only enriches it. +func (s *Service) ReportTaskGPU(_ context.Context, input *ReportTaskGPUInput, accountID string) (*ReportTaskGPUOutput, error) { + cluster := clusterShortName(input.Cluster) + taskID := taskShortID(input.Task) + kv, err := s.bucket(accountID) + if err != nil { + return nil, err + } + var task TaskRecord + found, err := getJSON(kv, TaskKey(cluster, taskID), &task) + if err != nil { + return nil, err + } + if !found { + return &ReportTaskGPUOutput{Acknowledgment: "OK"}, nil + } + byName := make(map[string][]string, len(input.Containers)) + for _, c := range input.Containers { + byName[c.Name] = c.GPUIDs + } + changed := false + for i := range task.Containers { + if ids, ok := byName[task.Containers[i].Name]; ok { + task.Containers[i].GPUIDs = ids + changed = true + } + } + if changed { + if err := putJSON(kv, TaskKey(cluster, taskID), &task); err != nil { + return nil, err + } + } + return &ReportTaskGPUOutput{Acknowledgment: "OK"}, nil +} diff --git a/spinifex/handlers/ecs/gpu_report_test.go b/spinifex/handlers/ecs/gpu_report_test.go new file mode 100644 index 000000000..4658280ab --- /dev/null +++ b/spinifex/handlers/ecs/gpu_report_test.go @@ -0,0 +1,191 @@ +package handlers_ecs + +import ( + "context" + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ecs" + "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ReportTaskGPU merges the agent-reported device UUIDs onto the matching +// container(s) of an existing task record, by container name. +func TestReportTaskGPU_MergesOntoTaskContainers(t *testing.T) { + svc, _ := newTestService(t) + kv, err := svc.bucket(testAccountID) + require.NoError(t, err) + require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{ + TaskID: "t-1", Cluster: "web", + Containers: []ContainerState{ + {Name: "web", Status: "RUNNING"}, + {Name: "trainer", Status: "RUNNING"}, + }, + })) + + out, err := svc.ReportTaskGPU(context.Background(), &ReportTaskGPUInput{ + Cluster: "web", Task: "t-1", + Containers: []ContainerGPUReport{{Name: "trainer", GPUIDs: []string{"GPU-aaa"}}}, + }, testAccountID) + require.NoError(t, err) + assert.Equal(t, "OK", out.Acknowledgment) + + var rec TaskRecord + found, err := getJSON(kv, TaskKey("web", "t-1"), &rec) + require.NoError(t, err) + require.True(t, found) + for _, c := range rec.Containers { + if c.Name == "trainer" { + assert.Equal(t, []string{"GPU-aaa"}, c.GPUIDs) + } else { + assert.Empty(t, c.GPUIDs) + } + } +} + +// An unknown task is a silent no-op (still acknowledged) — the state-report +// path already owns the task's lifecycle. +func TestReportTaskGPU_UnknownTaskNoop(t *testing.T) { + svc, _ := newTestService(t) + out, err := svc.ReportTaskGPU(context.Background(), &ReportTaskGPUInput{ + Cluster: "web", Task: "t-missing", + Containers: []ContainerGPUReport{{Name: "trainer", GPUIDs: []string{"GPU-aaa"}}}, + }, testAccountID) + require.NoError(t, err) + assert.Equal(t, "OK", out.Acknowledgment) +} + +// DescribeTasks surfaces the reported UUIDs as the real AWS Container.gpuIds +// field, once ReportTaskGPU has merged them onto the task record. +func TestDescribeTasks_GPUIDsPopulatedFromReport(t *testing.T) { + svc, _ := newTestService(t) + kv, err := svc.bucket(testAccountID) + require.NoError(t, err) + require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{ + TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web", + Containers: []ContainerState{{Name: "trainer", Status: "RUNNING", GPUIDs: []string{"GPU-aaa", "GPU-bbb"}}}, + })) + + dt, err := svc.DescribeTasks(context.Background(), &ecs.DescribeTasksInput{ + Cluster: aws.String("web"), Tasks: []*string{aws.String("t-1")}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, dt.Tasks, 1) + require.Len(t, dt.Tasks[0].Containers, 1) + assert.Equal(t, []string{"GPU-aaa", "GPU-bbb"}, awsStringSlice(dt.Tasks[0].Containers[0].GpuIds)) +} + +// recordTaskState must not wipe a container's pinned GPU UUIDs when a later +// state-change report omits them — the STOPPED report for a fast one-shot GPU +// task carries no GPUIDs, and DescribeTasks must still surface the UUIDs the +// RUNNING report (or ReportTaskGPU) previously set. +func TestRecordTaskState_PreservesGPUIDsAcrossStoppedReport(t *testing.T) { + svc, _ := newTestService(t) + kv, err := svc.bucket(testAccountID) + require.NoError(t, err) + require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{ + TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web", + ContainerInstanceID: "i-1", + Containers: []ContainerState{{Name: "trainer", Status: "PENDING"}}, + })) + + // RUNNING report pins the GPU UUIDs onto the container. + require.NoError(t, svc.recordTaskState(context.Background(), &bus.TaskState{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: "t-1", + LastStatus: bus.TaskStatusRunning, + Containers: []bus.ContainerStatus{{Name: "trainer", Status: bus.TaskStatusRunning, GPUIDs: []string{"GPU-aaa", "GPU-bbb"}}}, + })) + + var rec TaskRecord + found, err := getJSON(kv, TaskKey("web", "t-1"), &rec) + require.NoError(t, err) + require.True(t, found) + require.Len(t, rec.Containers, 1) + assert.Equal(t, []string{"GPU-aaa", "GPU-bbb"}, rec.Containers[0].GPUIDs) + + // STOPPED report arrives without GPUIDs (the agent doesn't re-send them) — + // the previously pinned UUIDs must survive the transition. + exit := 0 + require.NoError(t, svc.recordTaskState(context.Background(), &bus.TaskState{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: "t-1", + LastStatus: bus.TaskStatusStopped, Reason: "exited", + Containers: []bus.ContainerStatus{{Name: "trainer", Status: bus.TaskStatusStopped, ExitCode: &exit}}, + })) + + found, err = getJSON(kv, TaskKey("web", "t-1"), &rec) + require.NoError(t, err) + require.True(t, found) + require.Len(t, rec.Containers, 1) + assert.Equal(t, []string{"GPU-aaa", "GPU-bbb"}, rec.Containers[0].GPUIDs, + "gpuIds must survive the STOPPED transition so DescribeTasks still reports them") + + // DescribeTasks still surfaces the preserved UUIDs after STOPPED. + dt, err := svc.DescribeTasks(context.Background(), &ecs.DescribeTasksInput{ + Cluster: aws.String("web"), Tasks: []*string{aws.String("t-1")}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, dt.Tasks, 1) + require.Len(t, dt.Tasks[0].Containers, 1) + assert.Equal(t, []string{"GPU-aaa", "GPU-bbb"}, awsStringSlice(dt.Tasks[0].Containers[0].GpuIds)) +} + +// A state-change report that DOES carry GPUIDs still overwrites/replaces the +// pinned set normally (e.g. a corrected report, or GPUs reassigned on a retry). +func TestRecordTaskState_ReportedGPUIDsOverwritePinned(t *testing.T) { + svc, _ := newTestService(t) + kv, err := svc.bucket(testAccountID) + require.NoError(t, err) + require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{ + TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web", + ContainerInstanceID: "i-1", + Containers: []ContainerState{{Name: "trainer", Status: "RUNNING", GPUIDs: []string{"GPU-aaa"}}}, + })) + + require.NoError(t, svc.recordTaskState(context.Background(), &bus.TaskState{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: "t-1", + LastStatus: bus.TaskStatusRunning, + Containers: []bus.ContainerStatus{{Name: "trainer", Status: bus.TaskStatusRunning, GPUIDs: []string{"GPU-ccc"}}}, + })) + + var rec TaskRecord + found, err := getJSON(kv, TaskKey("web", "t-1"), &rec) + require.NoError(t, err) + require.True(t, found) + require.Len(t, rec.Containers, 1) + assert.Equal(t, []string{"GPU-ccc"}, rec.Containers[0].GPUIDs) +} + +// A non-GPU container's GPUIDs stay empty across state transitions — there is +// nothing to preserve, and the fallback must not fabricate a value. +func TestRecordTaskState_NonGPUContainerStaysEmpty(t *testing.T) { + svc, _ := newTestService(t) + kv, err := svc.bucket(testAccountID) + require.NoError(t, err) + require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{ + TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web", + ContainerInstanceID: "i-1", + Containers: []ContainerState{{Name: "web", Status: "PENDING"}}, + })) + + require.NoError(t, svc.recordTaskState(context.Background(), &bus.TaskState{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: "t-1", + LastStatus: bus.TaskStatusRunning, + Containers: []bus.ContainerStatus{{Name: "web", Status: bus.TaskStatusRunning}}, + })) + + exit := 0 + require.NoError(t, svc.recordTaskState(context.Background(), &bus.TaskState{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: "t-1", + LastStatus: bus.TaskStatusStopped, + Containers: []bus.ContainerStatus{{Name: "web", Status: bus.TaskStatusStopped, ExitCode: &exit}}, + })) + + var rec TaskRecord + found, err := getJSON(kv, TaskKey("web", "t-1"), &rec) + require.NoError(t, err) + require.True(t, found) + require.Len(t, rec.Containers, 1) + assert.Empty(t, rec.Containers[0].GPUIDs) +} diff --git a/spinifex/handlers/ecs/instances.go b/spinifex/handlers/ecs/instances.go index 93fa4c8aa..be1774ca0 100644 --- a/spinifex/handlers/ecs/instances.go +++ b/spinifex/handlers/ecs/instances.go @@ -53,6 +53,12 @@ func (s *Service) RegisterContainerInstance(_ context.Context, input *ecs.Regist r.TotalCPU = int(aws.Int64Value(res.IntegerValue)) case "MEMORY": r.TotalMemoryMiB = int(aws.Int64Value(res.IntegerValue)) + case "GPU": + // AWS reports GPU as a STRINGSET of device UUIDs; the count is the + // capacity, the UUIDs are the placeholder inventory (Epic C3 populates + // them for real from the agent's nvidia-smi discovery). + r.GPUIDs = aws.StringValueSlice(res.StringSetValue) + r.TotalGPU = len(r.GPUIDs) } } if len(input.Tags) > 0 { @@ -151,6 +157,17 @@ func (s *Service) instanceToAWS(r *InstanceRecord) *ecs.ContainerInstance { {Name: aws.String("CPU"), Type: aws.String("INTEGER"), IntegerValue: aws.Int64(int64(r.TotalCPU - r.ReservedCPU))}, {Name: aws.String("MEMORY"), Type: aws.String("INTEGER"), IntegerValue: aws.Int64(int64(r.TotalMemoryMiB - r.ReservedMemoryMiB))}, } + // GPU is only reported for GPU-capable instances (AWS parity); the STRINGSET + // values are the device UUIDs when known, empty pending the agent's Epic C3 + // report-back — TotalGPU/ReservedGPU remain the authoritative counts either way. + if r.TotalGPU > 0 { + registered = append(registered, &ecs.Resource{ + Name: aws.String("GPU"), Type: aws.String("STRINGSET"), StringSetValue: aws.StringSlice(r.GPUIDs), + }) + remaining = append(remaining, &ecs.Resource{ + Name: aws.String("GPU"), Type: aws.String("STRINGSET"), StringSetValue: aws.StringSlice(r.remainingGPUIDs()), + }) + } return &ecs.ContainerInstance{ ContainerInstanceArn: aws.String(r.ARN), Ec2InstanceId: aws.String(r.InstanceID), @@ -178,6 +195,8 @@ func (s *Service) recordRegister(msg *bus.RegisterInstance) error { r.AgentVersion = msg.AgentVersion r.TotalCPU = msg.Capacity.CPU r.TotalMemoryMiB = msg.Capacity.MemoryMiB + r.TotalGPU = msg.Capacity.GPU + r.GPUIDs = msg.Capacity.GPUIDs r.Status = InstanceStatusActive }) return err @@ -217,10 +236,23 @@ func (s *Service) recordTaskState(ctx context.Context, msg *bus.TaskState) error prev := task.LastStatus task.LastStatus = msg.LastStatus if len(msg.Containers) > 0 { + // Preserve pinned GPU UUIDs across a state change that omits them — the + // STOPPED report carries no GPUIDs and would otherwise wipe the device + // IDs the RUNNING report set, leaving DescribeTasks gpuIds empty. + prevGPUIDs := make(map[string][]string, len(task.Containers)) + for _, c := range task.Containers { + if len(c.GPUIDs) > 0 { + prevGPUIDs[c.Name] = c.GPUIDs + } + } task.Containers = task.Containers[:0] for _, c := range msg.Containers { + gpuIDs := c.GPUIDs + if len(gpuIDs) == 0 { + gpuIDs = prevGPUIDs[c.Name] + } task.Containers = append(task.Containers, ContainerState{ - Name: c.Name, Status: c.Status, ContainerID: c.ContainerID, ExitCode: c.ExitCode, + Name: c.Name, Status: c.Status, ContainerID: c.ContainerID, ExitCode: c.ExitCode, GPUIDs: gpuIDs, }) } } @@ -261,7 +293,7 @@ func (s *Service) recordTaskState(ctx context.Context, msg *bus.TaskState) error if perr := putJSON(kv, TaskKey(msg.ClusterName, msg.TaskID), &task); perr != nil { slog.ErrorContext(ctx, "ECS task STOPPED: persist after EIP release failed", "task", msg.TaskID, "err", perr) } - return s.releaseReservation(kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID, task.ReservedCPU, task.ReservedMemoryMiB) + return s.releaseReservation(kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID, task.ReservedCPU, task.ReservedMemoryMiB, task.GPU) } return nil } @@ -324,7 +356,7 @@ func (s *Service) recordDeploymentFailure(ctx context.Context, kv nats.KeyValue, } // releaseReservation returns a stopped task's capacity to its instance under CAS. -func (s *Service) releaseReservation(kv nats.KeyValue, cluster, instanceID, taskID string, cpu, mem int) error { +func (s *Service) releaseReservation(kv nats.KeyValue, cluster, instanceID, taskID string, cpu, mem, gpu int) error { for range reservePlacementRetries { entry, err := kv.Get(InstanceKey(cluster, instanceID)) if err != nil { @@ -336,6 +368,7 @@ func (s *Service) releaseReservation(kv nats.KeyValue, cluster, instanceID, task } rec.ReservedCPU = max(rec.ReservedCPU-cpu, 0) rec.ReservedMemoryMiB = max(rec.ReservedMemoryMiB-mem, 0) + rec.ReservedGPU = max(rec.ReservedGPU-gpu, 0) rec.PlacedTasks = slices.DeleteFunc(rec.PlacedTasks, func(v string) bool { return v == taskID }) data, merr := json.Marshal(&rec) if merr != nil { diff --git a/spinifex/handlers/ecs/placement.go b/spinifex/handlers/ecs/placement.go index 668becbd8..6e838e31c 100644 --- a/spinifex/handlers/ecs/placement.go +++ b/spinifex/handlers/ecs/placement.go @@ -16,29 +16,43 @@ const ( // ErrNoCapacity is returned when no ACTIVE instance can fit the task. var ErrNoCapacity = errors.New("no container instance has capacity for the task") -// remainingCPU/remainingMemory report an instance's unreserved capacity. +// remainingCPU/remainingMemory/remainingGPU report an instance's unreserved capacity. func (r *InstanceRecord) remainingCPU() int { return r.TotalCPU - r.ReservedCPU } func (r *InstanceRecord) remainingMemory() int { return r.TotalMemoryMiB - r.ReservedMemoryMiB } +func (r *InstanceRecord) remainingGPU() int { return r.TotalGPU - r.ReservedGPU } + +// remainingGPUIDs returns the instance's free GPU device UUIDs: GPUIDs with the +// first ReservedGPU entries dropped. Which UUID a given task actually holds is +// the agent's local ledger's call (reported back per-task); this is only a +// count-consistent view of the instance's total inventory. +func (r *InstanceRecord) remainingGPUIDs() []string { + if r.ReservedGPU >= len(r.GPUIDs) { + return nil + } + return r.GPUIDs[r.ReservedGPU:] +} // fits reports whether the instance is ACTIVE and has room for the reservation. -func (r *InstanceRecord) fits(cpu, mem int) bool { +func (r *InstanceRecord) fits(cpu, mem, gpu int) bool { if r.Status != InstanceStatusActive { return false } - return r.remainingCPU() >= cpu && r.remainingMemory() >= mem + return r.remainingCPU() >= cpu && r.remainingMemory() >= mem && r.remainingGPU() >= gpu } -// placeTask selects a container instance for a task reserving (cpu, mem) using -// the requested strategy. Candidates are filtered to ACTIVE instances that fit; -// ties broken by instance ID for determinism. strategy "" defaults to binpack. +// placeTask selects a container instance for a task reserving (cpu, mem, gpu) +// using the requested strategy. Candidates are filtered to ACTIVE instances that +// fit; ties broken by instance ID for determinism. strategy "" defaults to binpack. // // binpack: pick the instance with the LEAST remaining memory that still fits // (tightest pack). spread: pick the MOST remaining memory (widest spread). -// random: caller-stable first fit by instance ID. -func placeTask(instances []InstanceRecord, cpu, mem int, strategy string) (*InstanceRecord, error) { +// random: caller-stable first fit by instance ID. GPU is a fit gate only; it does +// not participate in the memory-based sort (non-GPU tasks request gpu=0 and are +// unaffected). +func placeTask(instances []InstanceRecord, cpu, mem, gpu int, strategy string) (*InstanceRecord, error) { candidates := make([]InstanceRecord, 0, len(instances)) for _, inst := range instances { - if inst.fits(cpu, mem) { + if inst.fits(cpu, mem, gpu) { candidates = append(candidates, inst) } } diff --git a/spinifex/handlers/ecs/placement_test.go b/spinifex/handlers/ecs/placement_test.go index 9be1e144f..ecbd39543 100644 --- a/spinifex/handlers/ecs/placement_test.go +++ b/spinifex/handlers/ecs/placement_test.go @@ -18,6 +18,15 @@ func inst(id string, total, reserved int) InstanceRecord { } } +// instGPU builds an instance with ample CPU/memory but a fixed GPU capacity, for +// tests exercising the GPU placement dimension in isolation. +func instGPU(id string, totalGPU, reservedGPU int) InstanceRecord { + r := inst(id, 100000, 0) + r.TotalGPU = totalGPU + r.ReservedGPU = reservedGPU + return r +} + func TestPlaceTask_Binpack_PicksTightest(t *testing.T) { // remaining mem: a=900, b=300, c=600 → binpack picks b (least remaining). instances := []InstanceRecord{ @@ -25,7 +34,7 @@ func TestPlaceTask_Binpack_PicksTightest(t *testing.T) { inst("b", 1000, 700), inst("c", 1000, 400), } - got, err := placeTask(instances, 100, 200, "binpack:memory") + got, err := placeTask(instances, 100, 200, 0, "binpack:memory") require.NoError(t, err) assert.Equal(t, "b", got.InstanceID) } @@ -36,28 +45,28 @@ func TestPlaceTask_Spread_PicksWidest(t *testing.T) { inst("b", 1000, 700), // 300 inst("c", 1000, 400), // 600 } - got, err := placeTask(instances, 100, 200, StrategySpread) + got, err := placeTask(instances, 100, 200, 0, StrategySpread) require.NoError(t, err) assert.Equal(t, "a", got.InstanceID) } func TestPlaceTask_Random_StableByID(t *testing.T) { instances := []InstanceRecord{inst("z", 1000, 0), inst("a", 1000, 0), inst("m", 1000, 0)} - got, err := placeTask(instances, 100, 100, StrategyRandom) + got, err := placeTask(instances, 100, 100, 0, StrategyRandom) require.NoError(t, err) assert.Equal(t, "a", got.InstanceID) } func TestPlaceTask_NoCapacity(t *testing.T) { instances := []InstanceRecord{inst("a", 100, 90)} - _, err := placeTask(instances, 0, 50, StrategyBinpack) + _, err := placeTask(instances, 0, 50, 0, StrategyBinpack) assert.ErrorIs(t, err, ErrNoCapacity) } func TestPlaceTask_SkipsDraining(t *testing.T) { d := inst("drain", 1000, 0) d.Status = InstanceStatusDraining - _, err := placeTask([]InstanceRecord{d}, 0, 100, StrategyBinpack) + _, err := placeTask([]InstanceRecord{d}, 0, 100, 0, StrategyBinpack) assert.ErrorIs(t, err, ErrNoCapacity) } @@ -72,8 +81,55 @@ func TestNormalizeStrategy(t *testing.T) { func TestInstanceRecord_Fits(t *testing.T) { r := inst("a", 1000, 600) // remaining mem 400, cpu 10000 - assert.True(t, r.fits(100, 400)) - assert.False(t, r.fits(100, 401)) + assert.True(t, r.fits(100, 400, 0)) + assert.False(t, r.fits(100, 401, 0)) r.Status = InstanceStatusDraining - assert.False(t, r.fits(0, 0)) + assert.False(t, r.fits(0, 0, 0)) +} + +// --- GPU placement dimension (Epic C2) --- + +func TestInstanceRecord_Fits_GPU(t *testing.T) { + r := instGPU("a", 2, 1) // 1 GPU free + assert.True(t, r.fits(0, 0, 1)) + assert.False(t, r.fits(0, 0, 2)) +} + +func TestInstanceRecord_Fits_GPU_IgnoredForNonGPUTasks(t *testing.T) { + // An instance with zero free GPU still fits a task that requests none. + r := instGPU("a", 1, 1) + assert.True(t, r.fits(0, 0, 0)) +} + +func TestPlaceTask_GPU_RejectsInsufficientInstances(t *testing.T) { + instances := []InstanceRecord{instGPU("a", 1, 1), instGPU("b", 2, 2)} + _, err := placeTask(instances, 0, 0, 1, StrategyBinpack) + assert.ErrorIs(t, err, ErrNoCapacity) +} + +func TestPlaceTask_GPU_PicksInstanceWithCapacity(t *testing.T) { + instances := []InstanceRecord{instGPU("a", 1, 1), instGPU("b", 2, 0)} + got, err := placeTask(instances, 0, 0, 2, StrategyBinpack) + require.NoError(t, err) + assert.Equal(t, "b", got.InstanceID) +} + +func TestPlaceTask_GPU_NonGPUTaskUnaffectedByExhaustedGPU(t *testing.T) { + // A task requesting no GPU still places on an instance with zero free GPU. + instances := []InstanceRecord{instGPU("a", 1, 1)} + got, err := placeTask(instances, 0, 0, 0, StrategyBinpack) + require.NoError(t, err) + assert.Equal(t, "a", got.InstanceID) +} + +func TestInstanceRecord_RemainingGPUIDs(t *testing.T) { + r := InstanceRecord{TotalGPU: 3, ReservedGPU: 1, GPUIDs: []string{"gpu-0", "gpu-1", "gpu-2"}} + assert.Equal(t, []string{"gpu-1", "gpu-2"}, r.remainingGPUIDs()) + + r.ReservedGPU = 3 + assert.Nil(t, r.remainingGPUIDs()) + + // Empty GPUIDs (UUIDs not yet reported by the agent) is a valid placeholder. + r2 := InstanceRecord{TotalGPU: 2, ReservedGPU: 1} + assert.Nil(t, r2.remainingGPUIDs()) } diff --git a/spinifex/handlers/ecs/records.go b/spinifex/handlers/ecs/records.go index e77d4374b..ec29cf95c 100644 --- a/spinifex/handlers/ecs/records.go +++ b/spinifex/handlers/ecs/records.go @@ -176,10 +176,14 @@ type CapacityProviderRecord struct { // ContainerDef is the persisted subset of an ecs.ContainerDefinition needed to // pull and run a container (bridge mode v1). type ContainerDef struct { - Name string `json:"name"` - Image string `json:"image"` - CPU int `json:"cpu,omitempty"` - MemoryMiB int `json:"memoryMiB,omitempty"` + Name string `json:"name"` + Image string `json:"image"` + CPU int `json:"cpu,omitempty"` + MemoryMiB int `json:"memoryMiB,omitempty"` + // GPU is the whole-GPU count from a resourceRequirements entry of type GPU + // (AWS ECS semantics; the value is a stringified integer). Device pinning and + // placement accounting land in later Epic C tasks. + GPU int `json:"gpu,omitempty"` Essential bool `json:"essential"` Command []string `json:"command,omitempty"` Environment map[string]string `json:"environment,omitempty"` @@ -230,6 +234,17 @@ func (t *TaskDefRecord) reservedMemory() int { return total } +// reservedGPU sums the task definition's per-container whole-GPU counts. +// Placement/reservation against instance capacity is a later Epic C task; this +// is the task-level total carried onto the task record and the bus assign. +func (t *TaskDefRecord) reservedGPU() int { + total := 0 + for _, c := range t.Containers { + total += c.GPU + } + return total +} + // InstanceRecord is the persisted container-instance state at InstanceKey. The // scheduler writes it from the Layer-2 bus (register/heartbeat) and reserves // capacity by appending placed task IDs. @@ -243,15 +258,25 @@ type InstanceRecord struct { Tags map[string]string `json:"tags,omitempty"` TotalCPU int `json:"totalCpu"` TotalMemoryMiB int `json:"totalMemoryMiB"` + // TotalGPU is the instance's whole-GPU count, mirroring TotalCPU/TotalMemoryMiB. + // It is the placement/capacity source of truth even before GPUIDs is populated. + TotalGPU int `json:"totalGpu"` // ReservedCPU/ReservedMemoryMiB track capacity committed to placed tasks; // placement increments them under a KV CAS and the task-state STOPPED path // releases them. Remaining = Total - Reserved. - ReservedCPU int `json:"reservedCpu"` - ReservedMemoryMiB int `json:"reservedMemoryMiB"` - AgentVersion string `json:"agentVersion,omitempty"` - PlacedTasks []string `json:"placedTasks,omitempty"` - RegisteredAt time.Time `json:"registeredAt"` - LastSeen time.Time `json:"lastSeen"` + ReservedCPU int `json:"reservedCpu"` + ReservedMemoryMiB int `json:"reservedMemoryMiB"` + // ReservedGPU mirrors ReservedCPU/ReservedMemoryMiB for whole-GPU counts. + ReservedGPU int `json:"reservedGpu"` + // GPUIDs holds the instance's total GPU device UUIDs, as reported at + // registration (AWS "GPU" STRINGSET resource) from the agent's nvidia-smi + // discovery. Empty on a non-GPU host; TotalGPU (not len(GPUIDs)) stays the + // authoritative capacity count. + GPUIDs []string `json:"gpuIds,omitempty"` + AgentVersion string `json:"agentVersion,omitempty"` + PlacedTasks []string `json:"placedTasks,omitempty"` + RegisteredAt time.Time `json:"registeredAt"` + LastSeen time.Time `json:"lastSeen"` // Reaped marks a DRAINING caused by the heartbeat reaper (involuntary), as // opposed to an operator UpdateContainerInstancesState drain. A reaped // instance is restored to ACTIVE when its agent re-registers; an operator @@ -265,6 +290,9 @@ type ContainerState struct { Status string `json:"status"` ContainerID string `json:"containerId,omitempty"` ExitCode *int `json:"exitCode,omitempty"` + // GPUIDs are the agent-reported device UUIDs pinned to this container, + // surfaced verbatim as the AWS Container.gpuIds field on DescribeTasks. + GPUIDs []string `json:"gpuIds,omitempty"` } // TaskRecord is the persisted task state at TaskKey; source of truth for @@ -276,20 +304,23 @@ type TaskRecord struct { // Group / StartedBy mirror the AWS task fields. A service's tasks carry // Group="service:{name}" so the reconciler counts them and the task-state // hook resolves a RUNNING/STOPPED task back to its owning service. - Group string `json:"group,omitempty"` - StartedBy string `json:"startedBy,omitempty"` - TaskDefFamily string `json:"taskDefFamily"` - TaskDefRevision int `json:"taskDefRevision"` - TaskDefARN string `json:"taskDefArn"` - ContainerInstanceID string `json:"containerInstanceId,omitempty"` - ContainerInstanceARN string `json:"containerInstanceArn,omitempty"` - DesiredStatus string `json:"desiredStatus"` - LastStatus string `json:"lastStatus"` - StoppedReason string `json:"stoppedReason,omitempty"` - ReservedCPU int `json:"reservedCpu"` - ReservedMemoryMiB int `json:"reservedMemoryMiB"` - Tags map[string]string `json:"tags,omitempty"` - Containers []ContainerState `json:"containers,omitempty"` + Group string `json:"group,omitempty"` + StartedBy string `json:"startedBy,omitempty"` + TaskDefFamily string `json:"taskDefFamily"` + TaskDefRevision int `json:"taskDefRevision"` + TaskDefARN string `json:"taskDefArn"` + ContainerInstanceID string `json:"containerInstanceId,omitempty"` + ContainerInstanceARN string `json:"containerInstanceArn,omitempty"` + DesiredStatus string `json:"desiredStatus"` + LastStatus string `json:"lastStatus"` + StoppedReason string `json:"stoppedReason,omitempty"` + ReservedCPU int `json:"reservedCpu"` + ReservedMemoryMiB int `json:"reservedMemoryMiB"` + // GPU is the task-level whole-GPU total (sum of its container defs' GPU + // counts). Not yet reserved against instance capacity (Epic C placement). + GPU int `json:"gpu,omitempty"` + Tags map[string]string `json:"tags,omitempty"` + Containers []ContainerState `json:"containers,omitempty"` // NetworkMode is the resolved task network mode (awsvpc|bridge|host). The // STOPPED path consults it to decide whether an ENI must be reclaimed. NetworkMode string `json:"networkMode,omitempty"` diff --git a/spinifex/handlers/ecs/service.go b/spinifex/handlers/ecs/service.go index 67a39540a..b44563edb 100644 --- a/spinifex/handlers/ecs/service.go +++ b/spinifex/handlers/ecs/service.go @@ -57,6 +57,10 @@ type ECSService interface { // PollAssignments drains an instance's task-assignment inbox (agent → gateway). PollAssignments(ctx context.Context, input *PollAssignmentsInput, accountID string) (*PollAssignmentsOutput, error) + // ReportTaskGPU records the agent's local ledger's pinned device UUIDs for + // a task's containers (agent → gateway); not an AWS SDK shape. + ReportTaskGPU(ctx context.Context, input *ReportTaskGPUInput, accountID string) (*ReportTaskGPUOutput, error) + // ProvisionCapacity launches container-instance EC2 capacity into a cluster. ProvisionCapacity(ctx context.Context, input *ProvisionCapacityInput, accountID string) (*ProvisionCapacityOutput, error) diff --git a/spinifex/handlers/ecs/service_nats.go b/spinifex/handlers/ecs/service_nats.go index 6eaf7ef6b..75635d1b3 100644 --- a/spinifex/handlers/ecs/service_nats.go +++ b/spinifex/handlers/ecs/service_nats.go @@ -134,6 +134,10 @@ func (s *NATSECSService) PollAssignments(ctx context.Context, input *PollAssignm return utils.NATSRequest[PollAssignmentsOutput](ctx, s.natsConn, "ecs.PollAssignments", input, defaultTimeout, accountID) } +func (s *NATSECSService) ReportTaskGPU(ctx context.Context, input *ReportTaskGPUInput, accountID string) (*ReportTaskGPUOutput, error) { + return utils.NATSRequest[ReportTaskGPUOutput](ctx, s.natsConn, "ecs.ReportTaskGPU", input, defaultTimeout, accountID) +} + // --- Capacity --- func (s *NATSECSService) ProvisionCapacity(ctx context.Context, input *ProvisionCapacityInput, accountID string) (*ProvisionCapacityOutput, error) { diff --git a/spinifex/handlers/ecs/service_test.go b/spinifex/handlers/ecs/service_test.go index b8ef65e4c..cf7447a53 100644 --- a/spinifex/handlers/ecs/service_test.go +++ b/spinifex/handlers/ecs/service_test.go @@ -2,6 +2,7 @@ package handlers_ecs import ( "context" + "strconv" "testing" "time" @@ -448,6 +449,191 @@ func TestScheduler_AcquireLease_SingleLeader(t *testing.T) { assert.True(t, a.acquireOrRefresh()) // refresh keeps leadership } +// --- GPU placement dimension (Epic C2) --- + +// registerTaskDefGPU registers a single-container task def requesting gpu +// whole-GPUs via resourceRequirements (AWS ECS GPU semantics). +func registerTaskDefGPU(t *testing.T, svc *Service, family string, cpu, mem, gpu int) *ecs.RegisterTaskDefinitionOutput { + t.Helper() + out, err := svc.RegisterTaskDefinition(context.Background(), &ecs.RegisterTaskDefinitionInput{ + Family: aws.String(family), + ContainerDefinitions: []*ecs.ContainerDefinition{{ + Name: aws.String("app"), + Image: aws.String("registry/app:1"), + Cpu: aws.Int64(int64(cpu)), + Memory: aws.Int64(int64(mem)), + Essential: aws.Bool(true), + ResourceRequirements: []*ecs.ResourceRequirement{ + {Type: aws.String(ecs.ResourceTypeGpu), Value: aws.String(strconv.Itoa(gpu))}, + }, + }}, + }, testAccountID) + require.NoError(t, err) + return out +} + +// registerInstanceGPU seeds an ACTIVE container instance with CPU/memory/GPU +// capacity via the bus register path. +func registerInstanceGPU(t *testing.T, svc *Service, cluster, id string, cpu, mem, gpu int) { + t.Helper() + require.NoError(t, svc.recordRegister(&bus.RegisterInstance{ + AccountID: testAccountID, + ClusterName: cluster, + InstanceID: id, + Capacity: bus.InstanceCapacity{CPU: cpu, MemoryMiB: mem, GPU: gpu}, + })) +} + +func TestService_RunTask_GPU_ReservesOnPlacement(t *testing.T) { + svc, _ := newTestService(t) + _, err := svc.CreateCluster(context.Background(), &ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) + require.NoError(t, err) + registerTaskDefGPU(t, svc, "gpu-app", 128, 256, 1) + registerInstanceGPU(t, svc, "web", "i-1", 1024, 2048, 2) + + out, err := svc.RunTask(context.Background(), &ecs.RunTaskInput{ + Cluster: aws.String("web"), TaskDefinition: aws.String("gpu-app"), Count: aws.Int64(1), + }, testAccountID) + require.NoError(t, err) + require.Len(t, out.Tasks, 1) + assert.Empty(t, out.Failures) + taskID := containerInstanceShortID(aws.StringValue(out.Tasks[0].TaskArn)) + + di, err := svc.DescribeContainerInstances(context.Background(), &ecs.DescribeContainerInstancesInput{ + Cluster: aws.String("web"), ContainerInstances: []*string{aws.String("i-1")}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, di.ContainerInstances, 1) + + // Registered GPU capacity is reported; remaining reflects the 1-GPU reservation. + var registeredGPU, remainingGPU *ecs.Resource + for _, r := range di.ContainerInstances[0].RegisteredResources { + if aws.StringValue(r.Name) == "GPU" { + registeredGPU = r + } + } + for _, r := range di.ContainerInstances[0].RemainingResources { + if aws.StringValue(r.Name) == "GPU" { + remainingGPU = r + } + } + require.NotNil(t, registeredGPU) + require.NotNil(t, remainingGPU) + assert.Equal(t, "STRINGSET", aws.StringValue(registeredGPU.Type)) + assert.Equal(t, "STRINGSET", aws.StringValue(remainingGPU.Type)) + // No agent-reported UUIDs yet (Epic C3): both stringSets are placeholder-empty, + // but the underlying instance record's counts prove the reservation happened. + assert.Empty(t, registeredGPU.StringSetValue) + assert.Empty(t, remainingGPU.StringSetValue) + + kv, err := svc.bucket(testAccountID) + require.NoError(t, err) + var rec InstanceRecord + found, err := getJSON(kv, InstanceKey("web", "i-1"), &rec) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, 2, rec.TotalGPU) + assert.Equal(t, 1, rec.ReservedGPU) + + // Stopping the task releases the GPU reservation. + require.NoError(t, svc.recordTaskState(context.Background(), &bus.TaskState{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: taskID, + LastStatus: bus.TaskStatusStopped, Reason: "exited", + })) + found, err = getJSON(kv, InstanceKey("web", "i-1"), &rec) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, 0, rec.ReservedGPU) +} + +func TestService_RunTask_GPU_NoCapacityFails(t *testing.T) { + svc, _ := newTestService(t) + _, err := svc.CreateCluster(context.Background(), &ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) + require.NoError(t, err) + registerTaskDefGPU(t, svc, "gpu-app", 128, 256, 2) + registerInstanceGPU(t, svc, "web", "i-1", 1024, 2048, 1) // only 1 GPU available, task needs 2 + + out, err := svc.RunTask(context.Background(), &ecs.RunTaskInput{ + Cluster: aws.String("web"), TaskDefinition: aws.String("gpu-app"), Count: aws.Int64(1), + }, testAccountID) + require.NoError(t, err) + assert.Empty(t, out.Tasks) + require.Len(t, out.Failures, 1) +} + +func TestService_DescribeContainerInstances_NoGPUOmitsResource(t *testing.T) { + svc, _ := newTestService(t) + _, err := svc.CreateCluster(context.Background(), &ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) + require.NoError(t, err) + registerInstance(t, svc, "web", "i-1", 1024, 2048) + + di, err := svc.DescribeContainerInstances(context.Background(), &ecs.DescribeContainerInstancesInput{ + Cluster: aws.String("web"), ContainerInstances: []*string{aws.String("i-1")}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, di.ContainerInstances, 1) + for _, r := range di.ContainerInstances[0].RegisteredResources { + assert.NotEqual(t, "GPU", aws.StringValue(r.Name)) + } + for _, r := range di.ContainerInstances[0].RemainingResources { + assert.NotEqual(t, "GPU", aws.StringValue(r.Name)) + } +} + +// RegisterContainerInstance (the AWS-API path) reports GPU as a STRINGSET of +// device UUIDs, matching real ECS's TotalResources shape; the count becomes the +// instance's placement capacity. +func TestRegisterContainerInstance_GPU_StringSetResource(t *testing.T) { + svc, _ := newTestService(t) + _, err := svc.CreateCluster(context.Background(), &ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) + require.NoError(t, err) + + _, err = svc.RegisterContainerInstance(context.Background(), &ecs.RegisterContainerInstanceInput{ + Cluster: aws.String("web"), + InstanceIdentityDocument: aws.String("i-gpu"), + TotalResources: []*ecs.Resource{ + {Name: aws.String("CPU"), Type: aws.String("INTEGER"), IntegerValue: aws.Int64(1024)}, + {Name: aws.String("MEMORY"), Type: aws.String("INTEGER"), IntegerValue: aws.Int64(2048)}, + {Name: aws.String("GPU"), Type: aws.String("STRINGSET"), StringSetValue: aws.StringSlice([]string{"GPU-aaa", "GPU-bbb"})}, + }, + }, testAccountID) + require.NoError(t, err) + + rec := instanceStatus(t, svc, "web", "i-gpu") + assert.Equal(t, 2, rec.TotalGPU) + assert.Equal(t, []string{"GPU-aaa", "GPU-bbb"}, rec.GPUIDs) + + di, err := svc.DescribeContainerInstances(context.Background(), &ecs.DescribeContainerInstancesInput{ + Cluster: aws.String("web"), ContainerInstances: []*string{aws.String("i-gpu")}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, di.ContainerInstances, 1) + var gpuRes *ecs.Resource + for _, r := range di.ContainerInstances[0].RegisteredResources { + if aws.StringValue(r.Name) == "GPU" { + gpuRes = r + } + } + require.NotNil(t, gpuRes) + assert.Equal(t, "STRINGSET", aws.StringValue(gpuRes.Type)) + assert.ElementsMatch(t, []string{"GPU-aaa", "GPU-bbb"}, aws.StringValueSlice(gpuRes.StringSetValue)) +} + +// The Layer-2 bus register path (recordRegister) also carries the agent's +// discovered device UUIDs, mirroring the AWS-API RegisterContainerInstance +// path (Epic C3 register report-back). +func TestRecordRegister_GPU_CarriesDeviceUUIDs(t *testing.T) { + svc, _ := newTestService(t) + require.NoError(t, svc.recordRegister(&bus.RegisterInstance{ + AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", + Capacity: bus.InstanceCapacity{CPU: 1024, MemoryMiB: 2048, GPU: 2, GPUIDs: []string{"GPU-aaa", "GPU-bbb"}}, + })) + + rec := instanceStatus(t, svc, "web", "i-1") + assert.Equal(t, 2, rec.TotalGPU) + assert.Equal(t, []string{"GPU-aaa", "GPU-bbb"}, rec.GPUIDs) +} + func TestAccountIDFromBucket(t *testing.T) { id, ok := accountIDFromBucket(AccountBucketName(testAccountID)) assert.True(t, ok) diff --git a/spinifex/handlers/ecs/task_stop.go b/spinifex/handlers/ecs/task_stop.go index 7dc212658..5660cbedd 100644 --- a/spinifex/handlers/ecs/task_stop.go +++ b/spinifex/handlers/ecs/task_stop.go @@ -74,7 +74,7 @@ func (s *Service) StartTask(ctx context.Context, input *ecs.StartTaskInput, acco if mode == NetworkModeAwsvpc && netCfg.firstSubnet() == "" { return nil, errors.New(awserrors.ErrorInvalidParameterValue) } - cpu, mem := taskDef.reservedCPU(), taskDef.reservedMemory() + cpu, mem, gpu := taskDef.reservedCPU(), taskDef.reservedMemory(), taskDef.reservedGPU() group := aws.StringValue(input.Group) startedBy := aws.StringValue(input.StartedBy) @@ -82,14 +82,14 @@ func (s *Service) StartTask(ctx context.Context, input *ecs.StartTaskInput, acco for _, ref := range awsStringSlice(input.ContainerInstances) { instanceID := containerInstanceShortID(ref) taskID := uuid.NewString() - inst, rerr := s.reserveOnInstance(kv, cluster, instanceID, taskID, cpu, mem) + inst, rerr := s.reserveOnInstance(kv, cluster, instanceID, taskID, cpu, mem, gpu) if rerr != nil { out.Failures = append(out.Failures, &ecs.Failure{ Arn: aws.String(ref), Reason: aws.String("RESOURCE:placement"), Detail: aws.String(rerr.Error()), }) continue } - rec := s.newTaskRecord(accountID, cluster, taskID, taskDef, inst, cpu, mem) + rec := s.newTaskRecord(accountID, cluster, taskID, taskDef, inst, cpu, mem, gpu) rec.NetworkMode = mode rec.Group = group rec.StartedBy = startedBy @@ -114,7 +114,7 @@ func (s *Service) StartTask(ctx context.Context, input *ecs.StartTaskInput, acco // reserveOnInstance commits a task's capacity reservation onto a specific // container instance under a KV CAS. Unlike reservePlacement it does not bin-pack; // the instance is fixed by the caller (StartTask). -func (s *Service) reserveOnInstance(kv nats.KeyValue, cluster, instanceID, taskID string, cpu, mem int) (*InstanceRecord, error) { +func (s *Service) reserveOnInstance(kv nats.KeyValue, cluster, instanceID, taskID string, cpu, mem, gpu int) (*InstanceRecord, error) { for range reservePlacementRetries { entry, err := kv.Get(InstanceKey(cluster, instanceID)) if errors.Is(err, nats.ErrKeyNotFound) { @@ -127,11 +127,12 @@ func (s *Service) reserveOnInstance(kv nats.KeyValue, cluster, instanceID, taskI if uerr := json.Unmarshal(entry.Value(), &live); uerr != nil { return nil, uerr } - if live.Status != InstanceStatusActive || !live.fits(cpu, mem) { + if live.Status != InstanceStatusActive || !live.fits(cpu, mem, gpu) { return nil, errors.New("instance has insufficient capacity") } live.ReservedCPU += cpu live.ReservedMemoryMiB += mem + live.ReservedGPU += gpu live.PlacedTasks = append(live.PlacedTasks, taskID) data, merr := json.Marshal(&live) if merr != nil { @@ -195,7 +196,7 @@ func (s *Service) forceStopTask(ctx context.Context, kv nats.KeyValue, accountID s.reclaimAssignInbox(kv, task.Cluster, task.ContainerInstanceID, task.TaskID) s.reclaimStopInbox(kv, task.Cluster, task.ContainerInstanceID, task.TaskID) s.reclaimTaskENI(ctx, accountID, task) - if rerr := s.releaseReservation(kv, task.Cluster, task.ContainerInstanceID, task.TaskID, task.ReservedCPU, task.ReservedMemoryMiB); rerr != nil { + if rerr := s.releaseReservation(kv, task.Cluster, task.ContainerInstanceID, task.TaskID, task.ReservedCPU, task.ReservedMemoryMiB, task.GPU); rerr != nil { slog.ErrorContext(ctx, "ECS forceStopTask: release reservation failed", "task", task.TaskID, "err", rerr) } } diff --git a/spinifex/handlers/ecs/taskdef_validation_test.go b/spinifex/handlers/ecs/taskdef_validation_test.go index 8ce698e88..08cbe3c94 100644 --- a/spinifex/handlers/ecs/taskdef_validation_test.go +++ b/spinifex/handlers/ecs/taskdef_validation_test.go @@ -90,3 +90,53 @@ func TestRunTask_AssignCarriesExecutionRoleAndLogDriver(t *testing.T) { require.Len(t, poll.Assignments[0].Containers, 1) assert.Equal(t, LogDriverJSONFile, poll.Assignments[0].Containers[0].LogDriver) } + +// TestRunTask_AssignCarriesGPU covers mulga-11opz (Epic C task C1): a +// resourceRequirements GPU count on the task def is threaded end-to-end — +// conv -> ContainerDef -> task-level aggregate (TaskRecord.GPU) -> the +// AssignContainer the agent polls for. +func TestRunTask_AssignCarriesGPU(t *testing.T) { + svc, _ := newTestService(t) + _, err := svc.CreateCluster(context.Background(), &ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) + require.NoError(t, err) + _, err = svc.RegisterTaskDefinition(context.Background(), &ecs.RegisterTaskDefinitionInput{ + Family: aws.String("gpu-app"), + ContainerDefinitions: []*ecs.ContainerDefinition{ + { + Name: aws.String("trainer"), Image: aws.String("registry/trainer:1"), + Cpu: aws.Int64(512), Memory: aws.Int64(1024), Essential: aws.Bool(true), + ResourceRequirements: []*ecs.ResourceRequirement{ + {Type: aws.String(ecs.ResourceTypeGpu), Value: aws.String("1")}, + }, + }, + { + Name: aws.String("sidecar"), Image: aws.String("registry/sidecar:1"), + Cpu: aws.Int64(64), Memory: aws.Int64(128), Essential: aws.Bool(false), + }, + }, + }, testAccountID) + require.NoError(t, err) + registerInstanceGPU(t, svc, "web", "i-1", 1024, 2048, 1) + + out, err := svc.RunTask(context.Background(), &ecs.RunTaskInput{ + Cluster: aws.String("web"), TaskDefinition: aws.String("gpu-app"), Count: aws.Int64(1), + }, testAccountID) + require.NoError(t, err) + require.Len(t, out.Tasks, 1) + + poll, err := svc.PollAssignments(context.Background(), &PollAssignmentsInput{Cluster: "web", ContainerInstance: "i-1"}, testAccountID) + require.NoError(t, err) + require.Len(t, poll.Assignments, 1) + require.Len(t, poll.Assignments[0].Containers, 2) + assert.Equal(t, 1, poll.Assignments[0].Containers[0].GPU) + assert.Zero(t, poll.Assignments[0].Containers[1].GPU) + + taskID := taskShortID(aws.StringValue(out.Tasks[0].TaskArn)) + kv, err := svc.bucket(testAccountID) + require.NoError(t, err) + var rec TaskRecord + found, err := getJSON(kv, TaskKey("web", taskID), &rec) + require.NoError(t, err) + require.True(t, found) + assert.Equal(t, 1, rec.GPU) +} diff --git a/spinifex/handlers/ecs/tasks.go b/spinifex/handlers/ecs/tasks.go index 524e8e460..b6990d146 100644 --- a/spinifex/handlers/ecs/tasks.go +++ b/spinifex/handlers/ecs/tasks.go @@ -50,7 +50,7 @@ func (s *Service) RunTask(ctx context.Context, input *ecs.RunTaskInput, accountI count = 1 } strategy := placementStrategyFromAWS(input.PlacementStrategy) - cpu, mem := taskDef.reservedCPU(), taskDef.reservedMemory() + cpu, mem, gpu := taskDef.reservedCPU(), taskDef.reservedMemory(), taskDef.reservedGPU() mode := resolveNetworkMode(taskDef) netCfg, err := parseAwsvpcConfig(input, mode) @@ -61,7 +61,7 @@ func (s *Service) RunTask(ctx context.Context, input *ecs.RunTaskInput, accountI out := &ecs.RunTaskOutput{} for i := 0; i < count; i++ { taskID := uuid.NewString() - inst, err := s.reservePlacement(kv, cluster, taskID, cpu, mem, strategy) + inst, err := s.reservePlacement(kv, cluster, taskID, cpu, mem, gpu, strategy) if err != nil { out.Failures = append(out.Failures, &ecs.Failure{ Reason: aws.String("RESOURCE:placement"), @@ -70,7 +70,7 @@ func (s *Service) RunTask(ctx context.Context, input *ecs.RunTaskInput, accountI continue } - rec := s.newTaskRecord(accountID, cluster, taskID, taskDef, inst, cpu, mem) + rec := s.newTaskRecord(accountID, cluster, taskID, taskDef, inst, cpu, mem, gpu) rec.NetworkMode = mode rec.Group = aws.StringValue(input.Group) rec.StartedBy = aws.StringValue(input.StartedBy) @@ -99,7 +99,7 @@ func (s *Service) RunTask(ctx context.Context, input *ecs.RunTaskInput, accountI // the caller skips it without leaking the ENI or the reserved capacity. func (s *Service) provisionTaskENI(ctx context.Context, kv nats.KeyValue, accountID, cluster string, rec *TaskRecord, netCfg awsvpcConfig) *ecs.Failure { rollback := func(reason string, err error) *ecs.Failure { - if rerr := s.releaseReservation(kv, cluster, rec.ContainerInstanceID, rec.TaskID, rec.ReservedCPU, rec.ReservedMemoryMiB); rerr != nil { + if rerr := s.releaseReservation(kv, cluster, rec.ContainerInstanceID, rec.TaskID, rec.ReservedCPU, rec.ReservedMemoryMiB, rec.GPU); rerr != nil { slog.ErrorContext(ctx, "ECS RunTask: reservation rollback failed", "task", rec.TaskID, "err", rerr) } return &ecs.Failure{Reason: aws.String(reason), Detail: aws.String(err.Error())} @@ -126,13 +126,13 @@ func (s *Service) provisionTaskENI(ctx context.Context, kv nats.KeyValue, accoun // reservePlacement bin-packs a task onto an ACTIVE instance and commits the // reservation under a KV CAS, retrying on contention. -func (s *Service) reservePlacement(kv nats.KeyValue, cluster, taskID string, cpu, mem int, strategy string) (*InstanceRecord, error) { +func (s *Service) reservePlacement(kv nats.KeyValue, cluster, taskID string, cpu, mem, gpu int, strategy string) (*InstanceRecord, error) { for range reservePlacementRetries { instances, err := s.listInstanceRecords(kv, cluster) if err != nil { return nil, err } - chosen, err := placeTask(instances, cpu, mem, strategy) + chosen, err := placeTask(instances, cpu, mem, gpu, strategy) if err != nil { return nil, err } @@ -146,11 +146,12 @@ func (s *Service) reservePlacement(kv nats.KeyValue, cluster, taskID string, cpu if uerr := json.Unmarshal(entry.Value(), &live); uerr != nil { return nil, uerr } - if !live.fits(cpu, mem) { + if !live.fits(cpu, mem, gpu) { continue } live.ReservedCPU += cpu live.ReservedMemoryMiB += mem + live.ReservedGPU += gpu live.PlacedTasks = append(live.PlacedTasks, taskID) data, merr := json.Marshal(&live) if merr != nil { @@ -165,7 +166,7 @@ func (s *Service) reservePlacement(kv nats.KeyValue, cluster, taskID string, cpu } // newTaskRecord builds a PENDING task record for a placed task. -func (s *Service) newTaskRecord(accountID, cluster, taskID string, td *TaskDefRecord, inst *InstanceRecord, cpu, mem int) *TaskRecord { +func (s *Service) newTaskRecord(accountID, cluster, taskID string, td *TaskDefRecord, inst *InstanceRecord, cpu, mem, gpu int) *TaskRecord { now := time.Now().UTC() rec := &TaskRecord{ TaskID: taskID, @@ -180,6 +181,7 @@ func (s *Service) newTaskRecord(accountID, cluster, taskID string, td *TaskDefRe LastStatus: TaskStatusPending, ReservedCPU: cpu, ReservedMemoryMiB: mem, + GPU: gpu, CreatedAt: now, } for _, c := range td.Containers { @@ -306,6 +308,9 @@ func (s *Service) taskToAWS(accountID string, r *TaskRecord) *ecs.Task { if c.ExitCode != nil { ctr.ExitCode = aws.Int64(int64(*c.ExitCode)) } + if len(c.GPUIDs) > 0 { + ctr.GpuIds = aws.StringSlice(c.GPUIDs) + } t.Containers = append(t.Containers, ctr) } return t diff --git a/tests/e2e/gpu/ecs_gpu_test.go b/tests/e2e/gpu/ecs_gpu_test.go new file mode 100644 index 000000000..d2b34f936 --- /dev/null +++ b/tests/e2e/gpu/ecs_gpu_test.go @@ -0,0 +1,418 @@ +//go:build e2e + +package gpu + +import ( + "encoding/base64" + "errors" + "fmt" + "os" + "sort" + "strings" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/awserr" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/aws/aws-sdk-go/service/ecs" + "github.com/aws/aws-sdk-go/service/iam" + "github.com/mulgadc/spinifex/spinifex/instancetypes" + "github.com/mulgadc/spinifex/spinifex/tags" + "github.com/mulgadc/spinifex/tests/e2e/harness" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + ecsGPUClusterPfx = "ecs-gpu-e2e" + + // ecsGPUInstanceProfile is the account-global instance profile ECS nodes + // assume for IMDS task-credential vending, matching tests/e2e/ecs/ecs_test.go's + // ensureInstanceProfile (account-global, left in place across suites). + ecsGPUInstanceProfile = "ecsInstanceRole" + + // nvidiaSmiContainerName names the single container in the GPU task def. + nvidiaSmiContainerName = "nvidia-smi" +) + +// ecsGPUFixture bundles the shared package fixture's GPU instance type with +// the ECS-tagged GPU node AMI resolved for it. +type ecsGPUFixture struct { + Env *harness.Env + AWS *harness.AWSClient + Harness *harness.Fixture + InstanceType string + AMIID string +} + +// requireECSGPUFixture skips the calling test unless every precondition for +// exercising Epic C (ECS: expose GPU to task containers) is met. It reuses +// requireGPUFixture's discovery (SPINIFEX_E2E, SPINIFEX_MODE=single, GPU +// instance type advertised) and additionally requires the ECS-tagged GPU node +// AMI (spinifex:managed-by=ecs, gpu-vendor=) to be imported. +func requireECSGPUFixture(t *testing.T) *ecsGPUFixture { + t.Helper() + fix := requireGPUFixture(t) + vendor := instancetypes.GPUVendorForType(fix.GPUInstanceType) + amiID, reason := discoverECSGPUNodeAMI(fix.AWS, vendor) + if reason != "" { + t.Skip(reason) + } + return &ecsGPUFixture{ + Env: fix.Env, AWS: fix.AWS, Harness: fix.Harness, + InstanceType: fix.GPUInstanceType, AMIID: amiID, + } +} + +// discoverECSGPUNodeAMI returns the newest AMI tagged spinifex:managed-by=ecs +// + gpu-vendor=vendor (the spinifex-ecs-node-gpu image), or a skip reason. +func discoverECSGPUNodeAMI(c *harness.AWSClient, vendor string) (amiID, reason string) { + out, err := c.EC2.DescribeImages(&ec2.DescribeImagesInput{ + Filters: []*ec2.Filter{ + {Name: aws.String("tag:" + tags.ManagedByKey), Values: aws.StringSlice([]string{tags.ManagedByECS})}, + {Name: aws.String("tag:" + tags.GPUVendorKey), Values: aws.StringSlice([]string{vendor})}, + {Name: aws.String("state"), Values: aws.StringSlice([]string{"available"})}, + }, + }) + if err != nil { + return "", "DescribeImages (ecs gpu ami): " + err.Error() + } + if len(out.Images) == 0 { + return "", fmt.Sprintf("ECS GPU node AMI not imported (managed-by=ecs, gpu-vendor=%s) — run: spx admin images import --name spinifex-ecs-node-gpu", vendor) + } + sort.Slice(out.Images, func(i, j int) bool { + return aws.StringValue(out.Images[i].CreationDate) > aws.StringValue(out.Images[j].CreationDate) + }) + return aws.StringValue(out.Images[0].ImageId), "" +} + +// TestECSGPUTaskExposure drives Epic C end-to-end: an ECS cluster plus a +// single GPU-capable container instance, then a task definition carrying a +// GPU resourceRequirement running nvidia-smi — asserting the task reaches +// RUNNING then STOPPED cleanly (nvidia-smi is one-shot), the container exits +// 0, and DescribeTasks reports the pinned GPU device UUID(s) (the C3 +// report-back). Container stdout/stderr are not retrievable through this +// harness (no log-retrieval API wired for ECS), so exit code + gpuIds stand +// in for asserting on the nvidia-smi table text directly. +func TestECSGPUTaskExposure(t *testing.T) { + fx := requireECSGPUFixture(t) + artifacts := harness.ArtifactDir(t, fx.Env) + + harness.Phase(t, "Ensuring ecsInstanceRole instance profile") + ensureECSGPUInstanceProfile(t, fx.AWS) + + vpc := harness.EnsureDefaultVPC(t, fx.Harness) + harness.AuthorizeSSHIngress(t, fx.AWS, vpc.SGID) + keyName, _ := harness.EnsureKeyPair(t, fx.Harness, artifacts) + + clusterName := fmt.Sprintf("%s-%d", ecsGPUClusterPfx, time.Now().Unix()) + harness.Phase(t, "Creating ECS cluster %q", clusterName) + // e2e:allow-create — the GPU cluster is the subject under test. + _, err := fx.AWS.ECS.CreateCluster(&ecs.CreateClusterInput{ClusterName: aws.String(clusterName)}) + require.NoError(t, err, "create-cluster") + t.Cleanup(func() { deleteECSGPUClusterBestEffort(t, fx.AWS, clusterName) }) + waitECSGPUClusterActive(t, fx.AWS, clusterName) + + harness.Phase(t, "Launching GPU container instance (%s, %s)", fx.InstanceType, fx.AMIID) + instanceID := launchECSGPUContainerInstance(t, fx, vpc, clusterName, keyName) + harness.Detail(t, "instance", instanceID) + + ci := waitECSGPUContainerInstanceActive(t, fx.AWS, clusterName, instanceID) + assert.True(t, aws.BoolValue(ci.AgentConnected), "GPU container instance agent must be connected") + assert.Greater(t, registeredGPUResource(ci), int64(0), "container instance must register GPU capacity") + + harness.Phase(t, "Registering nvidia-smi GPU task definition") + tdArn := registerECSGPUTaskDef(t, fx.AWS, clusterName) + + harness.Phase(t, "Running nvidia-smi GPU task") + taskArn := runECSGPUTask(t, fx.AWS, clusterName, tdArn) + + task := waitECSGPUTaskStopped(t, fx.AWS, clusterName, taskArn, 6*time.Minute) + require.Len(t, task.Containers, 1, "task must report exactly one container") + ctr := task.Containers[0] + assert.Equal(t, int64(0), aws.Int64Value(ctr.ExitCode), + "nvidia-smi container exit code must be 0 (reason=%s, stoppedReason=%s)", + aws.StringValue(ctr.Reason), aws.StringValue(task.StoppedReason)) + assert.NotEmpty(t, ctr.GpuIds, "container must report pinned GPU device UUID(s) (C3 report-back)") + harness.Detail(t, "gpuIds", aws.StringValueSlice(ctr.GpuIds)) +} + +// registeredGPUResource returns the container instance's registered GPU +// count (STRINGSET length), or 0 if the resource is absent. +func registeredGPUResource(ci *ecs.ContainerInstance) int64 { + for _, r := range ci.RegisteredResources { + if aws.StringValue(r.Name) == "GPU" { + return int64(len(aws.StringValueSlice(r.StringSetValue))) + } + } + return 0 +} + +// ensureECSGPUInstanceProfile makes the account-global ecsInstanceRole exist +// so the GPU node's agent can vend task credentials from IMDS. Idempotent and +// account-global (shared with tests/e2e/ecs), so it is not torn down. +func ensureECSGPUInstanceProfile(t *testing.T, c *harness.AWSClient) { + t.Helper() + if _, err := c.IAM.GetInstanceProfile(&iam.GetInstanceProfileInput{ + InstanceProfileName: aws.String(ecsGPUInstanceProfile), + }); err == nil { + return + } + + const trust = `{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"ec2.amazonaws.com"},"Action":"sts:AssumeRole"}]}` + _, err := c.IAM.CreateRole(&iam.CreateRoleInput{ + RoleName: aws.String(ecsGPUInstanceProfile), + AssumeRolePolicyDocument: aws.String(trust), + }) + tolerateEntityExists(t, err, "create-role ecsInstanceRole") + + const policy = `{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Action":"ecs:*","Resource":"*"}]}` + _, err = c.IAM.PutRolePolicy(&iam.PutRolePolicyInput{ + RoleName: aws.String(ecsGPUInstanceProfile), PolicyName: aws.String("ecsInstanceRolePolicy"), + PolicyDocument: aws.String(policy), + }) + require.NoError(t, err, "put-role-policy ecsInstanceRole") + + _, err = c.IAM.CreateInstanceProfile(&iam.CreateInstanceProfileInput{ + InstanceProfileName: aws.String(ecsGPUInstanceProfile), + }) + tolerateEntityExists(t, err, "create-instance-profile ecsInstanceRole") + + _, err = c.IAM.AddRoleToInstanceProfile(&iam.AddRoleToInstanceProfileInput{ + InstanceProfileName: aws.String(ecsGPUInstanceProfile), RoleName: aws.String(ecsGPUInstanceProfile), + }) + tolerateEntityExists(t, err, "add-role-to-instance-profile ecsInstanceRole") +} + +func tolerateEntityExists(t *testing.T, err error, what string) { + t.Helper() + if err == nil { + return + } + var aerr awserr.Error + if errors.As(err, &aerr) && strings.Contains(aerr.Code(), "EntityAlreadyExists") { + return + } + require.NoError(t, err, what) +} + +// launchECSGPUContainerInstance boots the GPU AMI as an ECS container +// instance, wired to register with clusterName over the local gateway. +func launchECSGPUContainerInstance(t *testing.T, fx *ecsGPUFixture, vpc harness.VPCInfo, clusterName, keyName string) string { + t.Helper() + userData := ecsGPUUserData(t, fx.Env, clusterName) + out, err := fx.AWS.EC2.RunInstances(&ec2.RunInstancesInput{ // e2e:allow-create — the GPU container instance is the subject under test. + ImageId: aws.String(fx.AMIID), + InstanceType: aws.String(fx.InstanceType), + KeyName: aws.String(keyName), + SubnetId: aws.String(vpc.SubnetID), + SecurityGroupIds: aws.StringSlice([]string{vpc.SGID}), + MinCount: aws.Int64(1), + MaxCount: aws.Int64(1), + IamInstanceProfile: &ec2.IamInstanceProfileSpecification{ + Name: aws.String(ecsGPUInstanceProfile), + }, + UserData: aws.String(base64.StdEncoding.EncodeToString([]byte(userData))), + }) + require.NoError(t, err, "run-instances (ecs gpu node)") + require.NotEmpty(t, out.Instances, "run-instances returned no instance") + id := aws.StringValue(out.Instances[0].InstanceId) + t.Cleanup(func() { + _, _ = fx.AWS.EC2.TerminateInstances(&ec2.TerminateInstancesInput{ + InstanceIds: aws.StringSlice([]string{id}), + }) + harness.WaitForInstanceState(t, fx.AWS, id, "terminated") + }) + // GPU driver init on first boot adds latency beyond a plain node boot. + harness.WaitForInstanceRunning(t, fx.AWS, id, 6*time.Minute) + return id +} + +// ecsGPUUserData builds the cloud-config the GPU node boots with: the +// agent's gateway URL and CA, matching tests/e2e/ecs/ecs_test.go's +// guestUserData (reimplemented locally — separate test package/binary). +func ecsGPUUserData(t *testing.T, env *harness.Env, cluster string) string { + t.Helper() + caPath, err := harness.ResolveCACert(env) + require.NoError(t, err, "resolve CA cert") + caBytes, err := os.ReadFile(caPath) //nolint:gosec // CA path from harness env + require.NoError(t, err, "read CA cert") + + var ca strings.Builder + for _, line := range strings.Split(strings.TrimRight(string(caBytes), "\n"), "\n") { + ca.WriteString(" ") + ca.WriteString(line) + ca.WriteString("\n") + } + + return fmt.Sprintf(`#cloud-config +write_files: + - path: /etc/spinifex-ecs/agent.env + permissions: '0600' + content: | + ECS_GATEWAY_URL=%s + ECS_GATEWAY_CA=/etc/spinifex-ecs/gateway-ca.pem + ECS_REGION=%s + ECS_CLUSTER=%s + - path: /etc/spinifex-ecs/gateway-ca.pem + permissions: '0644' + content: | +%s`, ecsGPUGatewayURL(env), ecsGPURegion(), cluster, ca.String()) +} + +// ecsGPUGatewayURL resolves a gateway URL the guest VM can reach, preferring +// an explicit override then the harness WAN host, matching +// tests/e2e/ecs/ecs_test.go's guestGatewayURL. +func ecsGPUGatewayURL(env *harness.Env) string { + if u := os.Getenv("SPINIFEX_ECS_GATEWAY_URL"); u != "" { + return u + } + candidates := append([]string{env.WANHost}, env.NodeIPs...) + candidates = append(candidates, env.ServiceIPs...) + for _, ip := range candidates { + if ip != "" && !strings.HasPrefix(ip, "127.") { + return fmt.Sprintf("https://%s:%d", ip, env.AWSGWPort) + } + } + return fmt.Sprintf("https://127.0.0.1:%d", env.AWSGWPort) +} + +func ecsGPURegion() string { + if v := os.Getenv("SPINIFEX_AWS_REGION"); v != "" { + return v + } + return "ap-southeast-2" +} + +func waitECSGPUClusterActive(t *testing.T, c *harness.AWSClient, name string) { + t.Helper() + harness.EventuallyErr(t, func() error { + out, err := c.ECS.DescribeClusters(&ecs.DescribeClustersInput{Clusters: aws.StringSlice([]string{name})}) + if err != nil { + return fmt.Errorf("describe-clusters: %w", err) + } + if len(out.Clusters) == 0 { + return fmt.Errorf("cluster %s not found", name) + } + if s := aws.StringValue(out.Clusters[0].Status); s != "ACTIVE" { + return fmt.Errorf("cluster status %s, want ACTIVE", s) + } + return nil + }, 1*time.Minute, 2*time.Second) +} + +func waitECSGPUContainerInstanceActive(t *testing.T, c *harness.AWSClient, cluster, ec2ID string) *ecs.ContainerInstance { + t.Helper() + harness.Step(t, "waiting for GPU container instance %s to register in %s", ec2ID, cluster) + var ci *ecs.ContainerInstance + harness.EventuallyErr(t, func() error { + list, err := c.ECS.ListContainerInstances(&ecs.ListContainerInstancesInput{Cluster: aws.String(cluster)}) + if err != nil { + return fmt.Errorf("list-container-instances: %w", err) + } + if len(list.ContainerInstanceArns) == 0 { + return fmt.Errorf("no container instances yet") + } + desc, err := c.ECS.DescribeContainerInstances(&ecs.DescribeContainerInstancesInput{ + Cluster: aws.String(cluster), ContainerInstances: list.ContainerInstanceArns, + }) + if err != nil { + return fmt.Errorf("describe-container-instances: %w", err) + } + for _, inst := range desc.ContainerInstances { + if aws.StringValue(inst.Ec2InstanceId) == ec2ID && aws.StringValue(inst.Status) == "ACTIVE" { + ci = inst + return nil + } + } + return fmt.Errorf("%s not ACTIVE yet", ec2ID) + }, 6*time.Minute, 5*time.Second) // GPU driver init + agent registration + harness.Step(t, "%s registered", ec2ID) + return ci +} + +// registerECSGPUTaskDef registers a single-container task def whose container +// carries a GPU resourceRequirement and runs nvidia-smi once. +func registerECSGPUTaskDef(t *testing.T, c *harness.AWSClient, clusterName string) string { + t.Helper() + family := clusterName + "-nvidia-smi" + out, err := c.ECS.RegisterTaskDefinition(&ecs.RegisterTaskDefinitionInput{ + Family: aws.String(family), + NetworkMode: aws.String("bridge"), + RequiresCompatibilities: aws.StringSlice([]string{"EC2"}), + ContainerDefinitions: []*ecs.ContainerDefinition{{ + Name: aws.String(nvidiaSmiContainerName), + Image: aws.String(nvidiaSmiPodImage), + Cpu: aws.Int64(128), + Memory: aws.Int64(512), + Essential: aws.Bool(true), + Command: aws.StringSlice([]string{"nvidia-smi"}), + ResourceRequirements: []*ecs.ResourceRequirement{{ + Type: aws.String("GPU"), Value: aws.String("1"), + }}, + }}, + }) + require.NoError(t, err, "register-task-definition (nvidia-smi GPU)") + arn := aws.StringValue(out.TaskDefinition.TaskDefinitionArn) + t.Cleanup(func() { + _, _ = c.ECS.DeregisterTaskDefinition(&ecs.DeregisterTaskDefinitionInput{TaskDefinition: aws.String(arn)}) + }) + return arn +} + +func runECSGPUTask(t *testing.T, c *harness.AWSClient, cluster, tdArn string) string { + t.Helper() + out, err := c.ECS.RunTask(&ecs.RunTaskInput{ + Cluster: aws.String(cluster), + TaskDefinition: aws.String(tdArn), + Count: aws.Int64(1), + }) + require.NoError(t, err, "run-task (nvidia-smi GPU)") + require.NotEmpty(t, out.Tasks, "run-task returned no task") + arn := aws.StringValue(out.Tasks[0].TaskArn) + t.Cleanup(func() { + _, _ = c.ECS.StopTask(&ecs.StopTaskInput{Cluster: aws.String(cluster), Task: aws.String(arn)}) + }) + return arn +} + +// waitECSGPUTaskStopped polls until the task reaches STOPPED, tolerating (and +// noting) a RUNNING observation along the way — nvidia-smi is a one-shot +// command, so the task may transition RUNNING -> STOPPED between polls. +func waitECSGPUTaskStopped(t *testing.T, c *harness.AWSClient, cluster, taskArn string, timeout time.Duration) *ecs.Task { + t.Helper() + harness.Step(t, "waiting for task %s: RUNNING then STOPPED (nvidia-smi is one-shot)", taskArn) + var task *ecs.Task + sawRunning := false + harness.EventuallyErr(t, func() error { + out, err := c.ECS.DescribeTasks(&ecs.DescribeTasksInput{ + Cluster: aws.String(cluster), Tasks: aws.StringSlice([]string{taskArn}), + }) + if err != nil { + return fmt.Errorf("describe-tasks: %w", err) + } + if len(out.Tasks) == 0 { + return fmt.Errorf("task not found") + } + task = out.Tasks[0] + switch s := aws.StringValue(task.LastStatus); s { + case "STOPPED": + return nil + case "RUNNING": + sawRunning = true + return fmt.Errorf("task RUNNING, waiting for STOPPED") + default: + return fmt.Errorf("task status %s", s) + } + }, timeout, 5*time.Second) + harness.Detail(t, "sawRunning", sawRunning, "stoppedReason", aws.StringValue(task.StoppedReason)) + return task +} + +func deleteECSGPUClusterBestEffort(t *testing.T, c *harness.AWSClient, cluster string) { + t.Helper() + if _, err := c.ECS.DeleteCluster(&ecs.DeleteClusterInput{Cluster: aws.String(cluster)}); err != nil { + t.Logf("cleanup delete-cluster %s: %v", cluster, err) + } +} diff --git a/tests/e2e/gpu/gpu_test.go b/tests/e2e/gpu/gpu_test.go index 866875d20..44913108f 100644 --- a/tests/e2e/gpu/gpu_test.go +++ b/tests/e2e/gpu/gpu_test.go @@ -18,6 +18,7 @@ import ( // type succeeds and reaches running state, confirming the VFIO bind path. func TestGPUPassthrough_Launch(t *testing.T) { fix := requireGPUFixture(t) + requireBaseGPUAMI(t, fix) harness.Phase(t, "GPU — Launch: RunInstances with %s", fix.GPUInstanceType) vpc := harness.EnsureDefaultVPC(t, fix.Harness) @@ -35,6 +36,7 @@ func TestGPUPassthrough_Launch(t *testing.T) { // wired the device into the guest. func TestGPUPassthrough_DeviceVisible(t *testing.T) { fix := requireGPUFixture(t) + requireBaseGPUAMI(t, fix) harness.Phase(t, "GPU — DeviceVisible: NVIDIA PCI device present in VM sysfs") vpc := harness.EnsureDefaultVPC(t, fix.Harness) @@ -60,6 +62,7 @@ func TestGPUPassthrough_DeviceVisible(t *testing.T) { // proving the guest OS can communicate with the passed-through GPU. func TestGPUPassthrough_DriverAccess(t *testing.T) { fix := requireGPUFixture(t) + requireBaseGPUAMI(t, fix) harness.Phase(t, "GPU — DriverAccess: nvidia-smi succeeds in VM") vpc := harness.EnsureDefaultVPC(t, fix.Harness) @@ -88,6 +91,7 @@ func TestGPUPassthrough_DriverAccess(t *testing.T) { // insufficient-capacity error. func TestGPUPassthrough_Release(t *testing.T) { fix := requireGPUFixture(t) + requireBaseGPUAMI(t, fix) harness.Phase(t, "GPU — Release: GPU returns to pool after termination") vpc := harness.EnsureDefaultVPC(t, fix.Harness) diff --git a/tests/e2e/gpu/main_test.go b/tests/e2e/gpu/main_test.go index d05038fee..02361efb2 100644 --- a/tests/e2e/gpu/main_test.go +++ b/tests/e2e/gpu/main_test.go @@ -51,7 +51,10 @@ func (f *Fixture) ArtifactDir(t *testing.T) string { // on first call. Skips the calling test when: // - SPINIFEX_E2E is unset // - no GPU instance types are advertised (node has no GPU or gpu_passthrough=false) -// - the NVIDIA GPU AMI has not been imported yet +// +// The ubuntu base GPU AMI is resolved best-effort into Fixture.AMIID; tests that +// launch from it must additionally call requireBaseGPUAMI. The ECS test resolves +// its own spinifex-ecs-node-gpu AMI and does not need the base image. func requireGPUFixture(t *testing.T) *Fixture { t.Helper() pkgFixOnce.Do(func() { @@ -73,11 +76,12 @@ func requireGPUFixture(t *testing.T) *Fixture { } awsCli := harness.NewAWSClient(t, env) - gpuType, amiID, reason := discoverGPUResources(awsCli) + gpuType, reason := discoverGPUInstanceType(awsCli) if reason != "" { pkgSkipReason = reason return } + amiID := discoverBaseGPUAMI(awsCli) // empty if not imported; gated per-test h, err := harness.NewProcessFixture(awsCli) if err != nil { @@ -111,24 +115,35 @@ func requireGPUFixture(t *testing.T) *Fixture { return pkgFix } -// discoverGPUResources returns the first GPU instance type and the NVIDIA GPU -// AMI ID available on this node. A non-empty reason means one or both are -// absent and the suite should be skipped. -func discoverGPUResources(c *harness.AWSClient) (gpuType, amiID, reason string) { +// discoverGPUInstanceType returns the SMALLEST GPU instance type advertised by +// the node (fewest vCPUs, then least memory) so it fits on a single dev host — +// picking an arbitrary/large type (e.g. g5.8xlarge) fails with +// InsufficientInstanceCapacity. A non-empty reason means none are advertised. +func discoverGPUInstanceType(c *harness.AWSClient) (gpuType, reason string) { typesOut, err := c.EC2.DescribeInstanceTypes(&ec2.DescribeInstanceTypesInput{}) if err != nil { - return "", "", "DescribeInstanceTypes: " + err.Error() + return "", "DescribeInstanceTypes: " + err.Error() } + var bestVCPU, bestMem int64 for _, it := range typesOut.InstanceTypes { - if it.GpuInfo != nil && len(it.GpuInfo.Gpus) > 0 { - gpuType = aws.StringValue(it.InstanceType) - break + if it.GpuInfo == nil || len(it.GpuInfo.Gpus) == 0 { + continue + } + vcpu := aws.Int64Value(it.VCpuInfo.DefaultVCpus) + mem := aws.Int64Value(it.MemoryInfo.SizeInMiB) + if gpuType == "" || vcpu < bestVCPU || (vcpu == bestVCPU && mem < bestMem) { + gpuType, bestVCPU, bestMem = aws.StringValue(it.InstanceType), vcpu, mem } } if gpuType == "" { - return "", "", "no GPU instance types advertised — node has no GPU or gpu_passthrough is disabled" + return "", "no GPU instance types advertised — node has no GPU or gpu_passthrough is disabled" } + return gpuType, "" +} +// discoverBaseGPUAMI returns the ubuntu-26.04-nvidia-gpu-x86_64 image ID, or "" +// if it has not been imported. Only the raw-EC2 GPU tests launch from it. +func discoverBaseGPUAMI(c *harness.AWSClient) string { imgsOut, err := c.EC2.DescribeImages(&ec2.DescribeImagesInput{ Filters: []*ec2.Filter{ {Name: aws.String("name"), Values: []*string{aws.String("ubuntu-26.04-nvidia-gpu-x86_64")}}, @@ -136,7 +151,17 @@ func discoverGPUResources(c *harness.AWSClient) (gpuType, amiID, reason string) }, }) if err != nil || len(imgsOut.Images) == 0 { - return "", "", "ubuntu-26.04-nvidia-gpu-x86_64 AMI not imported — run: spx admin images import --name ubuntu-26.04-nvidia-gpu-x86_64" + return "" + } + return aws.StringValue(imgsOut.Images[0].ImageId) +} + +// requireBaseGPUAMI skips the calling test unless the ubuntu base GPU image is +// imported. Tests that launch a VM directly from the base AMI call this after +// requireGPUFixture; the ECS test uses its own node AMI and does not. +func requireBaseGPUAMI(t *testing.T, fix *Fixture) { + t.Helper() + if fix.AMIID == "" { + t.Skip("ubuntu-26.04-nvidia-gpu-x86_64 AMI not imported — run: spx admin images import --name ubuntu-26.04-nvidia-gpu-x86_64") } - return gpuType, aws.StringValue(imgsOut.Images[0].ImageId), "" }