From 0b159d8bae3552ea09e57151ad9653a27b25ed44 Mon Sep 17 00:00:00 2001 From: Brandon Li Date: Wed, 28 Oct 2026 23:33:48 -0400 Subject: [PATCH 1/5] Add fake clock and fake executor (M1 step 2) clock.Fake fires timers and AfterFunc callbacks synchronously in time order during Advance, so executor and scheduler tests are deterministic. executor/fake runs scripted jobs on any clock and injects the failures real executors must survive: lost and late create responses, OOM, launch errors and timeouts, the executor backstop, SIGTERM vs SIGKILL, lost and recovered nodes, failed cleanup and observation outages. Its durable receipts are separate from an event stream that tests can hold, drop, duplicate, reorder or lose on restart. Co-Authored-By: Claude Opus 5.5 --- docs/agents/plans/jobs-and-scheduling.md | 8 + internal/clock/clock.go | 7 + internal/clock/fake.go | 119 +++++ internal/clock/fake_test.go | 63 +++ internal/executor/fake/delivery.go | 114 +++++ internal/executor/fake/fake.go | 583 +++++++++++++++++++++++ internal/executor/fake/fake_test.go | 421 ++++++++++++++++ 7 files changed, 1315 insertions(+) create mode 100644 internal/clock/fake.go create mode 100644 internal/clock/fake_test.go create mode 100644 internal/executor/fake/delivery.go create mode 100644 internal/executor/fake/fake.go create mode 100644 internal/executor/fake/fake_test.go diff --git a/docs/agents/plans/jobs-and-scheduling.md b/docs/agents/plans/jobs-and-scheduling.md index 0303345..c15f0ae 100644 --- a/docs/agents/plans/jobs-and-scheduling.md +++ b/docs/agents/plans/jobs-and-scheduling.md @@ -74,6 +74,14 @@ requirements. Keep secrets and actual research data out of test fixtures. duplicate/out-of-order events, observation outages and failed cleanup. Persist simulated receipts separately from notification delivery. - *Check:* scheduler tests run without a cluster. + - *Status (28 Oct 2026):* `internal/clock` has `Fake` (timers and + `AfterFunc` callbacks fire synchronously, in time order, inside `Advance`). + `internal/executor/fake` runs scripted jobs on that clock and injects lost or + late create responses, OOM, launch errors and timeouts, the backstop, SIGTERM + versus SIGKILL, lost and recovered nodes, failed cleanup, and observation + outages. Its durable receipts are kept separate from an event stream that + tests can hold, drop, duplicate, reorder or lose on `Restart`. The scheduler + tests that use it come with step 3. 3. **Scheduler core.** - Tables: `partitions`, `qos`, `limits`, `jobs`, `job_specs`, `job_events`, `allocations`, `nodes`, `partition_nodes`, operation intents, consumed diff --git a/internal/clock/clock.go b/internal/clock/clock.go index 40e0a66..663345f 100644 --- a/internal/clock/clock.go +++ b/internal/clock/clock.go @@ -8,6 +8,9 @@ import "time" type Clock interface { Now() time.Time NewTimer(d time.Duration) Timer + // AfterFunc calls f after d. On a Fake clock, f runs synchronously + // inside Advance, in time order. The returned Timer's C is nil. + AfterFunc(d time.Duration, f func()) Timer } type Timer interface { @@ -22,6 +25,10 @@ func (Real) Now() time.Time { return time.Now() } func (Real) NewTimer(d time.Duration) Timer { return realTimer{time.NewTimer(d)} } +func (Real) AfterFunc(d time.Duration, f func()) Timer { + return realTimer{time.AfterFunc(d, f)} +} + type realTimer struct{ t *time.Timer } func (r realTimer) C() <-chan time.Time { return r.t.C } diff --git a/internal/clock/fake.go b/internal/clock/fake.go new file mode 100644 index 0000000..0b59fb1 --- /dev/null +++ b/internal/clock/fake.go @@ -0,0 +1,119 @@ +package clock + +import ( + "slices" + "sync" + "time" +) + +// Fake is a virtual clock. Time moves only when Advance or Set is called. +// Timers due at the same instant fire in creation order. +type Fake struct { + mu sync.Mutex + now time.Time + seq int64 + timers []*fakeTimer +} + +// NewFake returns a fake clock starting at start. +func NewFake(start time.Time) *Fake { return &Fake{now: start} } + +type fakeTimer struct { + clock *Fake + when time.Time + seq int64 + ch chan time.Time // NewTimer + fn func() // AfterFunc +} + +func (f *Fake) Now() time.Time { + f.mu.Lock() + defer f.mu.Unlock() + return f.now +} + +func (f *Fake) NewTimer(d time.Duration) Timer { + return f.add(d, make(chan time.Time, 1), nil) +} + +func (f *Fake) AfterFunc(d time.Duration, fn func()) Timer { + return f.add(d, nil, fn) +} + +func (f *Fake) add(d time.Duration, ch chan time.Time, fn func()) *fakeTimer { + f.mu.Lock() + defer f.mu.Unlock() + f.seq++ + t := &fakeTimer{clock: f, when: f.now.Add(d), seq: f.seq, ch: ch, fn: fn} + f.timers = append(f.timers, t) + return t +} + +func (t *fakeTimer) C() <-chan time.Time { return t.ch } + +func (t *fakeTimer) Stop() bool { + f := t.clock + f.mu.Lock() + defer f.mu.Unlock() + i := slices.Index(f.timers, t) + if i < 0 { + return false + } + f.timers = slices.Delete(f.timers, i, i+1) + return true +} + +// Advance moves time forward by d, firing every timer that comes due on the +// way, including timers created by callbacks it runs. +func (f *Fake) Advance(d time.Duration) { f.Set(f.Now().Add(d)) } + +// Set moves time forward to t. Moving backwards panics: virtual time never +// rewinds, just as Mist assumes monotonic time within a boot. +func (f *Fake) Set(target time.Time) { + for { + f.mu.Lock() + if target.Before(f.now) { + f.mu.Unlock() + panic("clock: Fake moved backwards") + } + next := f.nextDue(target) + if next == nil { + f.now = target + f.mu.Unlock() + return + } + f.timers = slices.DeleteFunc(f.timers, func(t *fakeTimer) bool { return t == next }) + f.now = next.when + f.mu.Unlock() + + if next.fn != nil { + next.fn() + } else { + select { + case next.ch <- next.when: + default: + } + } + } +} + +func (f *Fake) nextDue(target time.Time) *fakeTimer { + var best *fakeTimer + for _, t := range f.timers { + if t.when.After(target) { + continue + } + if best == nil || t.when.Before(best.when) || (t.when.Equal(best.when) && t.seq < best.seq) { + best = t + } + } + return best +} + +// Pending returns how many timers are waiting. Tests use it to know that +// a goroutine has reached its wait before advancing time. +func (f *Fake) Pending() int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.timers) +} diff --git a/internal/clock/fake_test.go b/internal/clock/fake_test.go new file mode 100644 index 0000000..6107595 --- /dev/null +++ b/internal/clock/fake_test.go @@ -0,0 +1,63 @@ +package clock + +import ( + "testing" + "time" +) + +var t0 = time.Date(2026, 11, 2, 12, 0, 0, 0, time.UTC) + +func TestFakeFiresInTimeOrder(t *testing.T) { + c := NewFake(t0) + var got []string + c.AfterFunc(3*time.Second, func() { got = append(got, "c") }) + c.AfterFunc(time.Second, func() { + got = append(got, "a") + // Created during Advance and due before the target: fires too. + c.AfterFunc(time.Second, func() { got = append(got, "b@"+c.Now().Sub(t0).String()) }) + }) + stopped := c.AfterFunc(2*time.Second, func() { got = append(got, "never") }) + if !stopped.Stop() || stopped.Stop() { + t.Fatal("Stop should succeed once") + } + c.Advance(5 * time.Second) + want := []string{"a", "b@2s", "c"} + if len(got) != len(want) { + t.Fatalf("got %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("got %v, want %v", got, want) + } + } + if !c.Now().Equal(t0.Add(5 * time.Second)) { + t.Fatalf("now = %v", c.Now()) + } +} + +func TestFakeTimerChannel(t *testing.T) { + c := NewFake(t0) + tm := c.NewTimer(time.Minute) + c.Advance(59 * time.Second) + select { + case <-tm.C(): + t.Fatal("fired early") + default: + } + c.Advance(time.Second) + if at := <-tm.C(); !at.Equal(t0.Add(time.Minute)) { + t.Fatalf("fired at %v", at) + } + if c.Pending() != 0 { + t.Fatal("fired timer still pending") + } +} + +func TestFakeRefusesToRewind(t *testing.T) { + defer func() { + if recover() == nil { + t.Fatal("want panic") + } + }() + NewFake(t0).Set(t0.Add(-time.Second)) +} diff --git a/internal/executor/fake/delivery.go b/internal/executor/fake/delivery.go new file mode 100644 index 0000000..c81b4e5 --- /dev/null +++ b/internal/executor/fake/delivery.go @@ -0,0 +1,114 @@ +package fake + +import ( + "slices" + "sync" + + "github.com/UTMIST/Mist/internal/executor" +) + +// delivery is an unbounded, controllable notification queue in front of an +// unbuffered channel, so emitting never blocks the executor. +type delivery struct { + ch chan executor.Event + done chan struct{} + + mu sync.Mutex + cond *sync.Cond + queue []executor.Event + held []executor.Event + holding bool + duplicate bool + closed bool +} + +func newDelivery() *delivery { + d := &delivery{ch: make(chan executor.Event), done: make(chan struct{})} + d.cond = sync.NewCond(&d.mu) + go d.pump() + return d +} + +func (d *delivery) pump() { + for { + d.mu.Lock() + for len(d.queue) == 0 && !d.closed { + d.cond.Wait() + } + if d.closed { + d.mu.Unlock() + return + } + ev := d.queue[0] + d.queue = d.queue[1:] + d.mu.Unlock() + select { + case d.ch <- ev: + case <-d.done: + return + } + } +} + +func (d *delivery) push(ev executor.Event) { + d.mu.Lock() + defer d.mu.Unlock() + evs := []executor.Event{ev} + if d.duplicate { + evs = append(evs, ev) + } + if d.holding { + d.held = append(d.held, evs...) + return + } + d.queue = append(d.queue, evs...) + d.cond.Signal() +} + +func (d *delivery) hold() { + d.mu.Lock() + defer d.mu.Unlock() + d.holding = true +} + +func (d *delivery) release(reversed bool) { + d.mu.Lock() + defer d.mu.Unlock() + if reversed { + slices.Reverse(d.held) + } + d.queue = append(d.queue, d.held...) + d.held, d.holding = nil, false + d.cond.Signal() +} + +func (d *delivery) dropHeld() { + d.mu.Lock() + defer d.mu.Unlock() + d.held, d.holding = nil, false +} + +// dropQueued loses every notification not yet handed to the channel. One +// event already taken by the pump may still arrive: delivery is +// at-least-once, never exactly-once. +func (d *delivery) dropQueued() { + d.mu.Lock() + defer d.mu.Unlock() + d.queue, d.held = nil, nil +} + +func (d *delivery) setDuplicate(on bool) { + d.mu.Lock() + defer d.mu.Unlock() + d.duplicate = on +} + +func (d *delivery) close() { + d.mu.Lock() + defer d.mu.Unlock() + if !d.closed { + d.closed = true + close(d.done) + d.cond.Signal() + } +} diff --git a/internal/executor/fake/fake.go b/internal/executor/fake/fake.go new file mode 100644 index 0000000..2ffb98a --- /dev/null +++ b/internal/executor/fake/fake.go @@ -0,0 +1,583 @@ +// Package fake is an in-memory Executor driven by a clock, for scheduler +// tests and the simulator. Jobs "run" for scripted durations; tests inject +// the failures that the real executors must survive: lost or late create +// responses, OOM, launch failures, lost nodes, failed cleanup, observation +// outages, and lost, duplicated or reordered notifications. +// +// Like a real executor it keeps two separate things: +// +// - durable receipts and execution records, which survive Restart and +// are what List reports, and +// - the at-least-once Events stream, which tests may drop, hold, +// duplicate or reorder without losing any receipt. +// +// With clock.Fake every transition runs synchronously inside Advance, so +// tests are deterministic. +package fake + +import ( + "context" + "errors" + "fmt" + "io" + "slices" + "strings" + "sync" + "time" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/executor" +) + +// Script is what one job does once launched. +type Script struct { + StartDelay time.Duration // admission to work start (image pull, mounts) + RunFor time.Duration // work duration if not stopped + ExitCode int // when RunFor elapses + Reason executor.ExitReason // Normal, OutOfMemory, ScratchExceeded... + Detail string + OutputIncomplete bool + LaunchError string // nonempty: definitive failure at StartDelay, before work starts + StopDelay time.Duration // SIGTERM to exit; capped at the grace period + IgnoreSIGTERM bool // exit only on SIGKILL after grace + CleanupDelay time.Duration // finish to node-confirmed cleanup + CleanupFails bool // cleanup never confirms until FixCleanup +} + +// DefaultScript is used for jobs without their own script. +var DefaultScript = Script{ + StartDelay: time.Second, + RunFor: time.Minute, + Reason: executor.ExitNormal, + StopDelay: time.Second, + CleanupDelay: time.Second, +} + +// LaunchFault is an injected failure of one create call. +type LaunchFault int + +const ( + // FaultTimeoutBeforeCreate: the request is lost; nothing was created. + FaultTimeoutBeforeCreate LaunchFault = iota + 1 + // FaultTimeoutAfterCreate: the object was created but the response lost. + FaultTimeoutAfterCreate + // FaultLateCreate: the response is lost and the create lands later + // (after the duration passed to FailNextLaunch). A retry in between + // still yields exactly one object, as a deterministic name would. + FaultLateCreate +) + +// ErrTimeout is returned for injected lost responses. The outcome is +// ambiguous: callers must observe, not assume failure. +var ErrTimeout = fmt.Errorf("fake: create response lost: %w", context.DeadlineExceeded) + +// ErrUnavailable is returned while the fake API or node is unreachable. +var ErrUnavailable = errors.New("fake: executor API unavailable") + +// Config sets up a Fake. +type Config struct { + Clock clock.Clock + Nodes []executor.Node + Default *Script // nil means DefaultScript + BackstopSlack time.Duration // beyond time limit + grace; default 60 s +} + +// Fake implements executor.Executor. +type Fake struct { + clock clock.Clock + slack time.Duration + events *delivery + + mu sync.Mutex + def Script + nodes map[string]*executor.Node + order []string // node names in config order + scripts map[int64]Script + faults map[int64]fault + execs map[int64]*execution + receipts map[string][]executor.Receipt + creates map[int64]int + uids int + outage bool +} + +type fault struct { + kind LaunchFault + lateBy time.Duration +} + +type execution struct { + spec executor.LaunchSpec + script Script + created bool + lateCreate bool // a lost create will still land + uid string + phase executor.Phase + startedAt *time.Time + finishedAt *time.Time + result *executor.Result + cleanup executor.CleanupStatus + evidence string + stopRequested string // reason, when stopped before the object existed + terminating bool + timers []clock.Timer + frozen bool // node lost: nothing progresses +} + +// New returns a Fake. Close it to stop event delivery. +func New(cfg Config) *Fake { + f := &Fake{ + clock: cfg.Clock, + slack: cfg.BackstopSlack, + events: newDelivery(), + def: DefaultScript, + nodes: map[string]*executor.Node{}, + scripts: map[int64]Script{}, + faults: map[int64]fault{}, + execs: map[int64]*execution{}, + receipts: map[string][]executor.Receipt{}, + creates: map[int64]int{}, + } + if f.clock == nil { + f.clock = clock.Real{} + } + if f.slack == 0 { + f.slack = 60 * time.Second + } + if cfg.Default != nil { + f.def = *cfg.Default + } + for _, n := range cfg.Nodes { + n := n + f.nodes[n.Name] = &n + f.order = append(f.order, n.Name) + } + return f +} + +// Close stops event delivery. +func (f *Fake) Close() { f.events.close() } + +func (f *Fake) Name() string { return "fake" } + +// SetScript sets what jobID does when launched. +func (f *Fake) SetScript(jobID int64, s Script) { + f.mu.Lock() + defer f.mu.Unlock() + f.scripts[jobID] = s +} + +// FailNextLaunch makes the next create for jobID fail with fault. lateBy +// applies to FaultLateCreate only. +func (f *Fake) FailNextLaunch(jobID int64, kind LaunchFault, lateBy time.Duration) { + f.mu.Lock() + defer f.mu.Unlock() + f.faults[jobID] = fault{kind, lateBy} +} + +// SetOutage makes Nodes, List and Launch fail as if the API were down. +func (f *Fake) SetOutage(down bool) { + f.mu.Lock() + defer f.mu.Unlock() + f.outage = down +} + +// Creates counts native objects created for jobID. Tests assert it is +// never more than one. +func (f *Fake) Creates(jobID int64) int { + f.mu.Lock() + defer f.mu.Unlock() + return f.creates[jobID] +} + +// Receipts returns the durable receipts for an execution. +func (f *Fake) Receipts(executionID string) []executor.Receipt { + f.mu.Lock() + defer f.mu.Unlock() + return slices.Clone(f.receipts[executionID]) +} + +// Restart simulates the executor process restarting: notifications not yet +// delivered are lost. Receipts and running work are unaffected. +func (f *Fake) Restart() { f.events.dropQueued() } + +// HoldEvents queues notifications instead of delivering them. +func (f *Fake) HoldEvents() { f.events.hold() } + +// ReleaseEvents delivers held notifications, optionally newest first. +func (f *Fake) ReleaseEvents(reversed bool) { f.events.release(reversed) } + +// DropHeldEvents discards held notifications; receipts remain. +func (f *Fake) DropHeldEvents() { f.events.dropHeld() } + +// DuplicateEvents makes every later notification arrive twice. +func (f *Fake) DuplicateEvents(on bool) { f.events.setDuplicate(on) } + +func (f *Fake) Events() <-chan executor.Event { return f.events.ch } + +func (f *Fake) Nodes(context.Context) ([]executor.Node, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.outage { + return nil, ErrUnavailable + } + out := make([]executor.Node, 0, len(f.order)) + for _, name := range f.order { + n := *f.nodes[name] + n.Partitions = slices.Clone(n.Partitions) + n.Devices = slices.Clone(n.Devices) + out = append(out, n) + } + return out, nil +} + +func (f *Fake) Launch(_ context.Context, spec executor.LaunchSpec) error { + if err := spec.Validate(); err != nil { + return err + } + f.mu.Lock() + defer f.mu.Unlock() + if f.outage { + return ErrUnavailable + } + node := f.nodes[spec.Node] + if node == nil { + return fmt.Errorf("fake: unknown node %q", spec.Node) + } + + e := f.execs[spec.JobID] + if e == nil { + script, ok := f.scripts[spec.JobID] + if !ok { + script = f.def + } + e = &execution{spec: spec, script: script, phase: executor.PhaseCreating} + f.execs[spec.JobID] = e + f.receipt(e, executor.ReceiptCreateIntent, nil, "") + } else if e.spec.ExecutionID != spec.ExecutionID || e.spec.SpecHash != spec.SpecHash { + return executor.ErrSpecConflict + } + + switch e.phase { + case executor.PhasePending, executor.PhaseRunning: + return nil // already exists + case executor.PhaseCreating: + // No object yet: creating is safe, and a pending late create will + // find it already exists. + default: + return executor.ErrExecutionStarted + } + if !node.Healthy { + return ErrUnavailable + } + + ft, faulted := f.faults[spec.JobID] + delete(f.faults, spec.JobID) + switch { + case !faulted: + f.create(e) + return nil + case ft.kind == FaultTimeoutBeforeCreate: + return ErrTimeout + case ft.kind == FaultTimeoutAfterCreate: + f.create(e) + return ErrTimeout + case ft.kind == FaultLateCreate: + e.lateCreate = true + f.clock.AfterFunc(ft.lateBy, func() { + f.mu.Lock() + defer f.mu.Unlock() + e.lateCreate = false + if e.phase == executor.PhaseCreating && !e.frozen { + f.create(e) + } + }) + return ErrTimeout + } + return fmt.Errorf("fake: unknown fault %d", ft.kind) +} + +// create makes the native object and schedules its lifecycle. f.mu held. +func (f *Fake) create(e *execution) { + f.uids++ + f.creates[e.spec.JobID]++ + e.created = true + e.uid = fmt.Sprintf("fake-uid-%d", f.uids) + e.phase = executor.PhasePending + f.receipt(e, executor.ReceiptCreated, nil, "") + + if e.stopRequested != "" { + f.finish(e, executor.EventFinished, executor.Result{Reason: executor.ExitSignaled, + Detail: "stopped before start: " + e.stopRequested}) + return + } + now := f.clock.Now() + s := e.script + switch { + case s.LaunchError != "": + f.after(e, s.StartDelay, func() { + f.finish(e, executor.EventLaunchFailed, executor.Result{Reason: executor.ExitLaunchError, Detail: s.LaunchError}) + }) + case now.Add(s.StartDelay).After(e.spec.LaunchDeadline): + f.after(e, max(e.spec.LaunchDeadline.Sub(now), 0), func() { + f.finish(e, executor.EventLaunchFailed, executor.Result{Reason: executor.ExitLaunchTimeout, + Detail: "work did not start before the launch deadline"}) + }) + default: + f.after(e, s.StartDelay, func() { f.start(e) }) + } +} + +func (f *Fake) start(e *execution) { + if e.phase != executor.PhasePending { + return + } + now := f.clock.Now() + e.phase = executor.PhaseRunning + e.startedAt = &now + f.receipt(e, executor.ReceiptStarted, nil, "") + f.emit(e, executor.EventStarted, nil, "", "", now) + + s := e.script + if s.Reason == "" { + s.Reason = executor.ExitNormal + } + f.after(e, s.RunFor, func() { + f.finish(e, executor.EventFinished, executor.Result{ExitCode: ptr(s.ExitCode), Reason: s.Reason, + Detail: s.Detail, OutputIncomplete: s.OutputIncomplete}) + }) + f.after(e, e.spec.TimeLimit+e.spec.SignalGrace+f.slack, func() { + f.finish(e, executor.EventFinished, executor.Result{ExitCode: ptr(137), Reason: executor.ExitDeadline, + Detail: "killed by executor backstop", OutputIncomplete: true}) + }) +} + +// finish records the process outcome and begins cleanup. The first outcome +// wins; later timers for the same execution are no-ops. +func (f *Fake) finish(e *execution, kind executor.EventKind, r executor.Result) { + if e.phase != executor.PhasePending && e.phase != executor.PhaseRunning && e.phase != executor.PhaseCreating { + return + } + f.stopTimers(e) + now := f.clock.Now() + e.phase = executor.PhaseFinished + e.finishedAt = &now + e.result = &r + e.cleanup = executor.CleanupPending + f.receipt(e, executor.ReceiptFinished, &r, "") + f.emit(e, kind, r.ExitCode, r.Reason, r.Detail, now) + + f.after(e, e.script.CleanupDelay, func() { + if e.script.CleanupFails { + e.cleanup = executor.CleanupFailed + e.evidence = "fake: cleanup reporter did not confirm" + return + } + f.cleaned(e, "fake: processes gone, scratch removed, devices free") + }) +} + +func (f *Fake) cleaned(e *execution, evidence string) { + e.phase = executor.PhaseCleaned + e.cleanup = executor.CleanupConfirmed + e.evidence = evidence + f.receipt(e, executor.ReceiptCleaned, nil, evidence) + var code *int + reason := executor.ExitNodeLost + if e.result != nil { + code, reason = e.result.ExitCode, e.result.Reason + } + f.emit(e, executor.EventExited, code, reason, evidence, f.clock.Now()) +} + +// FixCleanup lets a failed cleanup confirm now. +func (f *Fake) FixCleanup(jobID int64) { + f.mu.Lock() + defer f.mu.Unlock() + e := f.execs[jobID] + if e == nil { + return + } + e.script.CleanupFails = false + if e.cleanup == executor.CleanupFailed { + f.cleaned(e, "fake: cleanup confirmed on retry") + } +} + +func (f *Fake) Terminate(_ context.Context, jobID int64, grace time.Duration, reason string) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.outage { + return ErrUnavailable + } + e := f.execs[jobID] + if e == nil || e.frozen { + return nil + } + switch e.phase { + case executor.PhaseCreating: + if e.lateCreate { + e.stopRequested = reason // the late object stops as soon as it lands + return nil + } + // Nothing exists and nothing will: record a stopped outcome. + f.finish(e, executor.EventFinished, executor.Result{Reason: executor.ExitSignaled, + Detail: "stopped before creation: " + reason}) + case executor.PhasePending: + f.finish(e, executor.EventFinished, executor.Result{Reason: executor.ExitSignaled, + Detail: "stopped before start: " + reason}) + case executor.PhaseRunning: + if e.terminating { + return nil + } + e.terminating = true + if e.script.IgnoreSIGTERM { + f.after(e, grace, func() { + f.finish(e, executor.EventFinished, executor.Result{ExitCode: ptr(137), Reason: executor.ExitSignaled, + Detail: fmt.Sprintf("killed after %s grace: %s", grace, reason), OutputIncomplete: true}) + }) + } else { + f.after(e, min(e.script.StopDelay, grace), func() { + f.finish(e, executor.EventFinished, executor.Result{ExitCode: ptr(143), Reason: executor.ExitSignaled, + Detail: reason}) + }) + } + } + return nil +} + +// LoseNode makes a node unreachable. Every unfinished or uncleaned +// execution on it is frozen and reported Lost; nothing on it progresses. +func (f *Fake) LoseNode(name, detail string) { + f.mu.Lock() + defer f.mu.Unlock() + n := f.nodes[name] + n.Healthy, n.Detail = false, detail + for _, e := range f.sortedExecs() { + if e.spec.Node != name || e.phase == executor.PhaseCleaned || e.phase == executor.PhaseLost { + continue + } + f.stopTimers(e) + e.frozen = true + e.phase = executor.PhaseLost + f.emit(e, executor.EventLost, nil, executor.ExitNodeLost, "node "+name+" unreachable: "+detail, f.clock.Now()) + } +} + +// RecoverNode brings a node back. Recovery proves its lost executions are +// gone, so each gets a cleanup receipt and an Exited event. +func (f *Fake) RecoverNode(name string) { + f.mu.Lock() + defer f.mu.Unlock() + n := f.nodes[name] + n.Healthy, n.Detail = true, "" + for _, e := range f.sortedExecs() { + if e.spec.Node == name && e.phase == executor.PhaseLost { + e.frozen = false + f.cleaned(e, "fake: node "+name+" recovered; execution not running") + } + } +} + +func (f *Fake) List(context.Context) ([]executor.Execution, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.outage { + return nil, fmt.Errorf("fake: %w", executor.ErrIncompleteObservation) + } + var out []executor.Execution + for _, e := range f.sortedExecs() { + x := executor.Execution{ + JobID: e.spec.JobID, + ExecutionID: e.spec.ExecutionID, + SpecHash: e.spec.SpecHash, + Node: e.spec.Node, + Phase: e.phase, + ObjectUID: e.uid, + StartedAt: e.startedAt, + FinishedAt: e.finishedAt, + Cleanup: e.cleanup, + CleanupEvidence: e.evidence, + } + if e.result != nil { + r := *e.result + x.Result = &r + } + out = append(out, x) + } + return out, nil +} + +func (f *Fake) Logs(_ context.Context, jobID int64, _ bool) (io.ReadCloser, error) { + f.mu.Lock() + defer f.mu.Unlock() + e := f.execs[jobID] + if e == nil || !e.created { + return nil, executor.ErrNotFound + } + return io.NopCloser(strings.NewReader(fmt.Sprintf("fake output for job %d\n", jobID))), nil +} + +func (f *Fake) Attach(context.Context, int64, executor.Terminal) error { + return errors.New("fake: attach is not implemented until M2") +} + +// after runs fn under f.mu after d unless the execution is frozen. +func (f *Fake) after(e *execution, d time.Duration, fn func()) { + e.timers = append(e.timers, f.clock.AfterFunc(d, func() { + f.mu.Lock() + defer f.mu.Unlock() + if !e.frozen { + fn() + } + })) +} + +func (f *Fake) stopTimers(e *execution) { + for _, t := range e.timers { + t.Stop() + } + e.timers = nil +} + +func (f *Fake) sortedExecs() []*execution { + ids := make([]int64, 0, len(f.execs)) + for id := range f.execs { + ids = append(ids, id) + } + slices.Sort(ids) + out := make([]*execution, len(ids)) + for i, id := range ids { + out[i] = f.execs[id] + } + return out +} + +func (f *Fake) receipt(e *execution, kind executor.ReceiptKind, r *executor.Result, evidence string) { + f.receipts[e.spec.ExecutionID] = append(f.receipts[e.spec.ExecutionID], executor.Receipt{ + ExecutionID: e.spec.ExecutionID, + Kind: kind, + ObjectUID: e.uid, + At: f.clock.Now(), + Result: r, + Evidence: evidence, + }) +} + +func (f *Fake) emit(e *execution, kind executor.EventKind, code *int, reason executor.ExitReason, detail string, at time.Time) { + f.events.push(executor.Event{ + EventID: e.spec.ExecutionID + "/" + string(kind), + JobID: e.spec.JobID, + ExecutionID: e.spec.ExecutionID, + Kind: kind, + Node: e.spec.Node, + ExitCode: code, + Reason: reason, + Detail: detail, + At: at, + }) +} + +func ptr[T any](v T) *T { return &v } + +var _ executor.Executor = (*Fake)(nil) diff --git a/internal/executor/fake/fake_test.go b/internal/executor/fake/fake_test.go new file mode 100644 index 0000000..f813078 --- /dev/null +++ b/internal/executor/fake/fake_test.go @@ -0,0 +1,421 @@ +package fake + +import ( + "context" + "errors" + "slices" + "testing" + "time" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/executor" +) + +var t0 = time.Date(2026, 11, 2, 12, 0, 0, 0, time.UTC) + +const gpuNode = "utmist-z1opa08" + +type env struct { + t *testing.T + clock *clock.Fake + fake *Fake + ctx context.Context +} + +func newEnv(t *testing.T) *env { + c := clock.NewFake(t0) + f := New(Config{Clock: c, Nodes: []executor.Node{{ + Name: gpuNode, + Partitions: []string{"gpu"}, + Capacity: executor.Resources{MilliCPU: 48000, MemoryBytes: 240 << 30, GPUs: 2, ScratchBytes: 400 << 30}, + Healthy: true, + }}}) + t.Cleanup(f.Close) + return &env{t: t, clock: c, fake: f, ctx: context.Background()} +} + +func (e *env) spec(jobID int64) executor.LaunchSpec { + s := executor.LaunchSpec{ + JobID: jobID, + ExecutionID: executor.NewExecutionID(), + AuthGeneration: 1, + LaunchDeadline: e.clock.Now().Add(10 * time.Minute), + Kind: executor.KindBatch, + Node: gpuNode, + User: executor.User{Name: "alice", UID: 20001, PersonalGID: 20001}, + TeamSlug: "team-a", + TeamGID: 40001, + ExtraGIDs: []int{19000}, + Resources: executor.Resources{MilliCPU: 8000, MemoryBytes: 32 << 30, GPUs: 1, ScratchBytes: 100 << 30}, + Workdir: "/mist/teams/team-a/alice", + OutputPath: "/mist/teams/team-a/alice/mist.out", + Script: "true\n", + Runtime: "gpu@sha256:0", + TimeLimit: time.Hour, + SignalGrace: 60 * time.Second, + } + s.Seal() + return s +} + +func (e *env) launch(s executor.LaunchSpec) { + e.t.Helper() + if err := e.fake.Launch(e.ctx, s); err != nil { + e.t.Fatalf("launch: %v", err) + } +} + +// next returns the next delivered event. +func (e *env) next() executor.Event { + e.t.Helper() + select { + case ev := <-e.fake.Events(): + return ev + case <-time.After(2 * time.Second): + e.t.Fatal("no event delivered") + return executor.Event{} + } +} + +func (e *env) expect(kind executor.EventKind, reason executor.ExitReason) executor.Event { + e.t.Helper() + ev := e.next() + if ev.Kind != kind || (reason != "" && ev.Reason != reason) { + e.t.Fatalf("got %s/%s (%s), want %s/%s", ev.Kind, ev.Reason, ev.Detail, kind, reason) + } + return ev +} + +func (e *env) quiet() { + e.t.Helper() + select { + case ev := <-e.fake.Events(): + e.t.Fatalf("unexpected event %s/%s", ev.Kind, ev.Reason) + case <-time.After(50 * time.Millisecond): + } +} + +func (e *env) execution(jobID int64) executor.Execution { + e.t.Helper() + xs, err := e.fake.List(e.ctx) + if err != nil { + e.t.Fatal(err) + } + for _, x := range xs { + if x.JobID == jobID { + return x + } + } + e.t.Fatalf("job %d not listed", jobID) + return executor.Execution{} +} + +func receiptKinds(rs []executor.Receipt) []executor.ReceiptKind { + var out []executor.ReceiptKind + for _, r := range rs { + out = append(out, r.Kind) + } + return out +} + +func TestLifecycle(t *testing.T) { + e := newEnv(t) + s := e.spec(1) + e.launch(s) + e.launch(s) // repeat-safe + if x := e.execution(1); x.Phase != executor.PhasePending || x.ObjectUID == "" { + t.Fatalf("after launch: %+v", x) + } + + e.clock.Advance(time.Second) + if ev := e.expect(executor.EventStarted, ""); !ev.At.Equal(t0.Add(time.Second)) { + t.Fatalf("started at %v", ev.At) + } + e.clock.Advance(time.Minute) + if ev := e.expect(executor.EventFinished, executor.ExitNormal); *ev.ExitCode != 0 { + t.Fatalf("exit %d", *ev.ExitCode) + } + if x := e.execution(1); x.Phase != executor.PhaseFinished || x.Cleanup != executor.CleanupPending { + t.Fatalf("finish must not imply cleanup: %+v", x) + } + e.clock.Advance(time.Second) + e.expect(executor.EventExited, executor.ExitNormal) + + want := []executor.ReceiptKind{executor.ReceiptCreateIntent, executor.ReceiptCreated, + executor.ReceiptStarted, executor.ReceiptFinished, executor.ReceiptCleaned} + if got := receiptKinds(e.fake.Receipts(s.ExecutionID)); !slices.Equal(got, want) { + t.Fatalf("receipts %v, want %v", got, want) + } + if err := e.fake.Launch(e.ctx, s); !errors.Is(err, executor.ErrExecutionStarted) { + t.Fatalf("relaunch after finish: %v", err) + } + if n := e.fake.Creates(1); n != 1 { + t.Fatalf("%d objects created", n) + } +} + +func TestSpecConflict(t *testing.T) { + e := newEnv(t) + s := e.spec(1) + e.launch(s) + other := s + other.Resources.GPUs = 2 + other.Seal() + if err := e.fake.Launch(e.ctx, other); !errors.Is(err, executor.ErrSpecConflict) { + t.Fatalf("different spec: %v", err) + } + newID := e.spec(1) + if err := e.fake.Launch(e.ctx, newID); !errors.Is(err, executor.ErrSpecConflict) { + t.Fatalf("second execution for job: %v", err) + } +} + +func TestOutcomes(t *testing.T) { + cases := []struct { + name string + script Script + kind executor.EventKind + reason executor.ExitReason + after time.Duration + }{ + {"oom", Script{StartDelay: time.Second, RunFor: time.Minute, ExitCode: 137, Reason: executor.ExitOutOfMemory}, + executor.EventFinished, executor.ExitOutOfMemory, time.Minute}, + {"nonzero", Script{StartDelay: time.Second, RunFor: time.Minute, ExitCode: 1, Reason: executor.ExitNormal}, + executor.EventFinished, executor.ExitNormal, time.Minute}, + {"backstop", Script{StartDelay: time.Second, RunFor: 10 * time.Hour}, + executor.EventFinished, executor.ExitDeadline, time.Hour + 2*time.Minute}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + e := newEnv(t) + e.fake.SetScript(1, c.script) + e.launch(e.spec(1)) + e.clock.Advance(time.Second) + e.expect(executor.EventStarted, "") + e.clock.Advance(c.after - time.Second) + e.quiet() + e.clock.Advance(time.Second) + e.expect(c.kind, c.reason) + }) + } +} + +func TestLaunchFailures(t *testing.T) { + e := newEnv(t) + e.fake.SetScript(1, Script{StartDelay: 5 * time.Second, LaunchError: "image pull failed"}) + e.fake.SetScript(2, Script{StartDelay: 20 * time.Minute, RunFor: time.Minute}) + e.launch(e.spec(1)) + e.launch(e.spec(2)) + + e.clock.Advance(5 * time.Second) + e.expect(executor.EventLaunchFailed, executor.ExitLaunchError) + e.expect(executor.EventExited, executor.ExitLaunchError) // zero cleanup delay + e.clock.Advance(10*time.Minute - 5*time.Second) + ev := e.expect(executor.EventLaunchFailed, executor.ExitLaunchTimeout) + if ev.JobID != 2 || !ev.At.Equal(t0.Add(10*time.Minute)) { + t.Fatalf("launch timeout %+v", ev) + } + if x := e.execution(2); x.StartedAt != nil { + t.Fatal("timed-out launch has a start time") + } +} + +func TestTerminate(t *testing.T) { + e := newEnv(t) + e.fake.SetScript(2, Script{StartDelay: time.Second, RunFor: time.Hour, IgnoreSIGTERM: true}) + e.launch(e.spec(1)) + e.launch(e.spec(2)) + e.launch(e.spec(3)) + e.fake.Terminate(e.ctx, 3, time.Minute, "cancelled by alice") // before start + e.expect(executor.EventFinished, executor.ExitSignaled) + + e.clock.Advance(time.Second) + e.expect(executor.EventStarted, "") + e.expect(executor.EventStarted, "") + e.expect(executor.EventExited, executor.ExitSignaled) // job 3 cleaned up + + for range 2 { // idempotent + e.fake.Terminate(e.ctx, 1, time.Minute, "cancelled") + e.fake.Terminate(e.ctx, 2, time.Minute, "cancelled") + } + e.clock.Advance(time.Second) + if ev := e.expect(executor.EventFinished, executor.ExitSignaled); ev.JobID != 1 || *ev.ExitCode != 143 { + t.Fatalf("SIGTERM exit %+v", ev) + } + e.clock.Advance(time.Minute) + e.expect(executor.EventExited, executor.ExitSignaled) + if ev := e.expect(executor.EventFinished, executor.ExitSignaled); ev.JobID != 2 || *ev.ExitCode != 137 { + t.Fatalf("SIGKILL exit %+v", ev) + } +} + +func TestAmbiguousCreates(t *testing.T) { + e := newEnv(t) + + // Lost request: nothing created, retry creates exactly once. + s1 := e.spec(1) + e.fake.FailNextLaunch(1, FaultTimeoutBeforeCreate, 0) + if err := e.fake.Launch(e.ctx, s1); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("want timeout, got %v", err) + } + if x := e.execution(1); x.Phase != executor.PhaseCreating || e.fake.Creates(1) != 0 { + t.Fatalf("lost request: %+v creates=%d", x, e.fake.Creates(1)) + } + e.launch(s1) + + // Lost response: object exists, retry is a no-op. + s2 := e.spec(2) + e.fake.FailNextLaunch(2, FaultTimeoutAfterCreate, 0) + if err := e.fake.Launch(e.ctx, s2); err == nil { + t.Fatal("want timeout") + } + e.launch(s2) + + // Late create: retry wins, the late one finds it exists. + s3 := e.spec(3) + e.fake.FailNextLaunch(3, FaultLateCreate, 5*time.Second) + if err := e.fake.Launch(e.ctx, s3); err == nil { + t.Fatal("want timeout") + } + e.launch(s3) + e.clock.Advance(time.Minute) + + for job := int64(1); job <= 3; job++ { + if n := e.fake.Creates(job); n != 1 { + t.Errorf("job %d: %d objects", job, n) + } + } +} + +func TestStopBeforeLateCreateLands(t *testing.T) { + e := newEnv(t) + e.fake.FailNextLaunch(1, FaultLateCreate, 5*time.Second) + e.fake.Launch(e.ctx, e.spec(1)) + e.fake.Terminate(e.ctx, 1, time.Minute, "cancelled") + e.quiet() + e.clock.Advance(5 * time.Second) + e.expect(executor.EventFinished, executor.ExitSignaled) + e.clock.Advance(time.Hour) + e.expect(executor.EventExited, executor.ExitSignaled) + if x := e.execution(1); x.StartedAt != nil || e.fake.Creates(1) != 1 { + t.Fatalf("stopped late create ran: %+v", x) + } +} + +func TestLostNode(t *testing.T) { + e := newEnv(t) + s := e.spec(1) + e.launch(s) + e.clock.Advance(time.Second) + e.expect(executor.EventStarted, "") + + e.fake.LoseNode(gpuNode, "kubelet NotReady") + e.expect(executor.EventLost, executor.ExitNodeLost) + nodes, _ := e.fake.Nodes(e.ctx) + if nodes[0].Healthy { + t.Fatal("lost node reported healthy") + } + e.clock.Advance(24 * time.Hour) // nothing progresses, including cleanup + e.quiet() + if err := e.fake.Launch(e.ctx, s); !errors.Is(err, executor.ErrExecutionStarted) { + t.Fatalf("relaunch of lost execution: %v", err) + } + if x := e.execution(1); x.Cleanup == executor.CleanupConfirmed { + t.Fatal("lost execution reported cleaned") + } + + e.fake.RecoverNode(gpuNode) + e.expect(executor.EventExited, executor.ExitNodeLost) + if n := e.fake.Creates(1); n != 1 { + t.Fatalf("%d objects", n) + } +} + +func TestFailedCleanupHolds(t *testing.T) { + e := newEnv(t) + e.fake.SetScript(1, Script{StartDelay: time.Second, RunFor: time.Minute, CleanupDelay: time.Second, CleanupFails: true}) + e.launch(e.spec(1)) + e.clock.Advance(time.Second) + e.expect(executor.EventStarted, "") + e.clock.Advance(time.Minute) + e.expect(executor.EventFinished, executor.ExitNormal) + e.clock.Advance(time.Hour) + e.quiet() + if x := e.execution(1); x.Cleanup != executor.CleanupFailed { + t.Fatalf("cleanup %q", x.Cleanup) + } + e.fake.FixCleanup(1) + e.expect(executor.EventExited, executor.ExitNormal) +} + +func TestOutage(t *testing.T) { + e := newEnv(t) + e.fake.SetOutage(true) + if _, err := e.fake.List(e.ctx); !errors.Is(err, executor.ErrIncompleteObservation) { + t.Fatalf("List during outage: %v", err) + } + if _, err := e.fake.Nodes(e.ctx); err == nil { + t.Fatal("Nodes during outage succeeded") + } + if err := e.fake.Launch(e.ctx, e.spec(1)); err == nil { + t.Fatal("Launch during outage succeeded") + } + e.fake.SetOutage(false) + if xs, err := e.fake.List(e.ctx); err != nil || len(xs) != 0 { + t.Fatalf("after outage: %v %v", xs, err) + } +} + +func TestDelivery(t *testing.T) { + run := func(e *env, job int64) { + e.launch(e.spec(job)) + e.clock.Advance(time.Second) + e.clock.Advance(time.Minute) + e.clock.Advance(time.Second) + } + + t.Run("reordered", func(t *testing.T) { + e := newEnv(t) + e.fake.HoldEvents() + run(e, 1) + e.quiet() + e.fake.ReleaseEvents(true) + e.expect(executor.EventExited, "") + e.expect(executor.EventFinished, "") + e.expect(executor.EventStarted, "") + }) + + t.Run("duplicated", func(t *testing.T) { + e := newEnv(t) + e.fake.DuplicateEvents(true) + e.launch(e.spec(1)) + e.clock.Advance(time.Second) + a, b := e.next(), e.next() + if a.EventID != b.EventID || a.Kind != executor.EventStarted { + t.Fatalf("duplicate has different identity: %s vs %s", a.EventID, b.EventID) + } + }) + + t.Run("dropped", func(t *testing.T) { + e := newEnv(t) + e.fake.HoldEvents() + run(e, 1) + e.fake.DropHeldEvents() + e.quiet() + // Receipts and List still recover the outcome. + if x := e.execution(1); x.Phase != executor.PhaseCleaned || x.Result == nil || *x.Result.ExitCode != 0 { + t.Fatalf("not recoverable from List: %+v", x) + } + }) + + t.Run("restart", func(t *testing.T) { + e := newEnv(t) + e.fake.HoldEvents() + run(e, 1) + e.fake.Restart() + e.quiet() + if len(e.fake.Receipts(e.execution(1).ExecutionID)) != 5 { + t.Fatal("restart lost receipts") + } + }) +} From 7c0840e55ddf1b40fed2cb37684b2d28d8f4fe17 Mon Sep 17 00:00:00 2001 From: Brandon Li Date: Fri, 30 Oct 2026 09:21:15 -0400 Subject: [PATCH 2/5] Add scheduler core: ledger, admission, events and reconciliation (M1 step 3) Scheduler tables (partitions, qos, limits, jobs, job_specs, job_events, allocations, nodes, op_intents, consumed_events) with the proposed M1 defaults. The cycle refreshes nodes, applies durable revocations, reconciles List against the database, enforces walltime from started_at and launch deadlines, admits FIFO-within-QoS with user/team/partition limits and partition caps, then dispatches committed intents. Allocations are released only on a cleanup receipt; lost executions keep them quarantined. A leader epoch fences writes from a deposed scheduler. A never-sent launch intent proves nothing exists, so cancelling it frees capacity at once; sent ones are resent with the same ExecutionID. Tests run against Postgres with the fake executor and fake clock. Co-Authored-By: Claude Opus 5.5 --- cmd/mistd/main.go | 4 + docs/agents/architecture/scheduler.md | 9 + docs/agents/plans/jobs-and-scheduling.md | 13 + .../scheduler/0002_scheduler_core.sql | 206 ++++++++ internal/executor/fake/fake.go | 14 + internal/identity/identity.go | 12 + internal/scheduler/admit.go | 232 +++++++++ internal/scheduler/dispatch.go | 181 +++++++ internal/scheduler/events.go | 166 +++++++ internal/scheduler/harness_test.go | 232 +++++++++ internal/scheduler/ledger.go | 284 +++++++++++ internal/scheduler/reconcile.go | 343 +++++++++++++ internal/scheduler/scheduler.go | 191 ++++++- internal/scheduler/scheduler_test.go | 464 ++++++++++++++++++ internal/scheduler/store.go | 284 +++++++++++ internal/scheduler/submit.go | 139 ++++++ 16 files changed, 2770 insertions(+), 4 deletions(-) create mode 100644 internal/db/migrations/scheduler/0002_scheduler_core.sql create mode 100644 internal/scheduler/admit.go create mode 100644 internal/scheduler/dispatch.go create mode 100644 internal/scheduler/events.go create mode 100644 internal/scheduler/harness_test.go create mode 100644 internal/scheduler/ledger.go create mode 100644 internal/scheduler/reconcile.go create mode 100644 internal/scheduler/scheduler_test.go create mode 100644 internal/scheduler/store.go create mode 100644 internal/scheduler/submit.go diff --git a/cmd/mistd/main.go b/cmd/mistd/main.go index 2b7d7d5..a22f37f 100644 --- a/cmd/mistd/main.go +++ b/cmd/mistd/main.go @@ -21,6 +21,7 @@ import ( "time" "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/identity" "github.com/UTMIST/Mist/internal/scheduler" ) @@ -67,6 +68,9 @@ func run(ctx context.Context, args []string, stderr io.Writer) error { return scheduler.Run(ctx, scheduler.Config{ DatabaseURL: *dbURL, Pool: pool, + // Until Identity's implementation lands, fail closed: nothing is + // admitted (AUTH_UNAVAILABLE). No executors are wired yet either. + Identity: identity.Unavailable{}, Log: log, Tick: *tick, LeaderCheck: *leaderCheck, diff --git a/docs/agents/architecture/scheduler.md b/docs/agents/architecture/scheduler.md index 5a65b74..424a65d 100644 --- a/docs/agents/architecture/scheduler.md +++ b/docs/agents/architecture/scheduler.md @@ -148,6 +148,15 @@ reconciles unresolved operations before admissions on affected nodes. An advisor lock alone cannot fence an already in-flight Kubernetes request. Quarantine any node whose old operations cannot yet be proven complete. +Database writes are fenced by a **leader epoch**. A new leader increments +`scheduler_cursors.leader_epoch`. Every scheduler transaction first share-locks +that row and aborts unless the epoch is still its own. A transaction an old +leader started on another pooled connection therefore cannot commit after a +takeover. API-side writes (submit, cancel) are not fenced; job row locks +serialize them with admission. An intent's status distinguishes `pending` (the +call was never sent, which proves nothing exists on the node) from `in_flight` +(ambiguous until observed). + A cycle runs on every event (submit, cancel, job end, node change, config change), and also every 10 s: diff --git a/docs/agents/plans/jobs-and-scheduling.md b/docs/agents/plans/jobs-and-scheduling.md index c15f0ae..31f8b0c 100644 --- a/docs/agents/plans/jobs-and-scheduling.md +++ b/docs/agents/plans/jobs-and-scheduling.md @@ -100,6 +100,19 @@ requirements. Keep secrets and actual research data out of test fixtures. Cover M1-04/M1-05/M1-06: crash boundaries, split leadership, fast exit before Started, lost node with retained allocation and repeated cleanup receipts. Include shared-node and partition caps, not just accelerator counts. + - *Status (30 Oct 2026):* implemented in `internal/scheduler`, with tables in + `migrations/scheduler/0002`. `scheduler_test.go` covers every scenario + listed in the check above, plus revocation, PROVISIONING/AUTH_UNAVAILABLE, + stale generation before launch, ambiguous creates, and dropped or duplicated + events. Writes are fenced by a leader epoch (scheduler.md). Still open: + - **Unknown executions:** they are stopped, and their node is kept out of + admission in memory, but not yet charged as quarantined occupancy. + - **Executor receipt tables:** moved to step 5, since the fake keeps + receipts in memory. + - **Generation granularity:** the identity fixture's generation is global + rather than per account, so any membership change cancels unsent + launches. Identity's real implementation should scope it per account. + - **Pending detail:** refreshed only when the pending reason changes. 4. **`mist-job-entry`.** - umask 0027, symlink-safe private dot-folders, output header/footer, copying output to the file and stdout, SIGTERM forwarding, OOM detection from diff --git a/internal/db/migrations/scheduler/0002_scheduler_core.sql b/internal/db/migrations/scheduler/0002_scheduler_core.sql new file mode 100644 index 0000000..5805972 --- /dev/null +++ b/internal/db/migrations/scheduler/0002_scheduler_core.sql @@ -0,0 +1,206 @@ +-- Scheduler core (M1 step 3). See docs/agents/architecture/scheduler.md. +-- Times are written by the scheduler from its injected clock, never now(), +-- so tests and the simulator run on virtual time. + +CREATE TABLE scheduler_cursors ( + name text PRIMARY KEY, + value bigint NOT NULL +); +-- leader_epoch fences writes from a scheduler that lost leadership; +-- revocation_seq is the last Identity revocation applied. +INSERT INTO scheduler_cursors (name, value) VALUES ('leader_epoch', 0), ('revocation_seq', 0); + +CREATE TABLE qos ( + name text PRIMARY KEY, + priority integer NOT NULL, -- higher runs first + max_time_limit_seconds bigint +); +INSERT INTO qos (name, priority, max_time_limit_seconds) VALUES + ('normal', 0, NULL), + ('interactive', 1000, 14400); + +CREATE TABLE partitions ( + name text PRIMARY KEY, + executor text NOT NULL, + state text NOT NULL DEFAULT 'up' CHECK (state IN ('up', 'down')), + accelerator text NOT NULL CHECK (accelerator IN ('gpu', 'tt', 'none')), + runtime text NOT NULL, -- pinned image digest or host profile + min_accelerators integer NOT NULL DEFAULT 0, + max_accelerators integer NOT NULL DEFAULT 0, + -- Per accelerator (GPU/TT partitions) or per job (accelerator 'none'). + default_milli_cpu bigint NOT NULL, + default_memory_bytes bigint NOT NULL, + max_milli_cpu bigint NOT NULL, + max_memory_bytes bigint NOT NULL, + default_scratch_bytes bigint NOT NULL, + max_scratch_bytes bigint NOT NULL, + default_time_limit_seconds bigint NOT NULL, + max_time_limit_seconds bigint NOT NULL, + -- Partition-wide caps across all its nodes; NULL means node capacity. + cap_milli_cpu bigint, + cap_memory_bytes bigint +); +-- Proposed initial values from scheduler.md; admins change them in place. +-- The runtime digests are placeholders until the runtime owner pins images. +INSERT INTO partitions VALUES + ('gpu', 'k3s', 'up', 'gpu', 'unpinned-gpu-runtime', 1, 2, + 8000, 34359738368, 26000, 75161927680, + 107374182400, 214748364800, 3600, 86400, NULL, NULL), + ('cpu', 'k3s', 'up', 'none', 'unpinned-cpu-runtime', 0, 0, + 2000, 8589934592, 16000, 137438953472, + 53687091200, 536870912000, 3600, 86400, 16000, 206158430208); + +CREATE TABLE nodes ( + name text PRIMARY KEY, + executor text NOT NULL, + admin_state text NOT NULL DEFAULT 'up' CHECK (admin_state IN ('up', 'drain', 'down')), + -- Last report from Executor.Nodes(), net of executor reservations. + healthy boolean NOT NULL DEFAULT false, + health_detail text NOT NULL DEFAULT '', + observed_at timestamptz, + milli_cpu bigint NOT NULL DEFAULT 0, + memory_bytes bigint NOT NULL DEFAULT 0, + gpus bigint NOT NULL DEFAULT 0, + tt_boards bigint NOT NULL DEFAULT 0, + scratch_bytes bigint NOT NULL DEFAULT 0, + hugepages_bytes bigint NOT NULL DEFAULT 0, + -- Additional Mist host reserve, subtracted once. + reserve_milli_cpu bigint NOT NULL DEFAULT 0, + reserve_memory_bytes bigint NOT NULL DEFAULT 0, + -- Set when the scheduler cannot prove old operations finished. + quarantined boolean NOT NULL DEFAULT false, + quarantine_detail text NOT NULL DEFAULT '' +); + +CREATE TABLE partition_nodes ( + partition text NOT NULL REFERENCES partitions, + node text NOT NULL REFERENCES nodes, + PRIMARY KEY (partition, node) +); + +-- (scope, subject, partition, qos, resource) -> max. See scheduler.md#limits. +CREATE TABLE limits ( + scope text NOT NULL CHECK (scope IN ('user', 'team', 'partition')), + subject text NOT NULL, -- username, team slug or '*' + partition text NOT NULL, -- name or '*' + qos text NOT NULL, -- name or '*' + resource text NOT NULL CHECK (resource IN + ('gpus', 'tt_boards', 'milli_cpu', 'memory_bytes', 'jobs', 'submitted_jobs')), + max bigint NOT NULL CHECK (max >= 0), + PRIMARY KEY (scope, subject, partition, qos, resource) +); +INSERT INTO limits VALUES + ('user', '*', 'gpu', '*', 'gpus', 1), + ('team', '*', 'gpu', '*', 'gpus', 2), + ('user', '*', 'cpu', '*', 'jobs', 4), + ('user', '*', '*', '*', 'submitted_jobs', 50); + +CREATE TABLE jobs ( + id bigserial PRIMARY KEY, + username text NOT NULL, + team text NOT NULL, + partition text NOT NULL REFERENCES partitions, + qos text NOT NULL REFERENCES qos, + name text NOT NULL, + kind text NOT NULL CHECK (kind IN ('batch', 'command', 'interactive')), + state text NOT NULL, + pending_reason text, + pending_detail text NOT NULL DEFAULT '', + held boolean NOT NULL DEFAULT false, + -- Requested resources (already normalized). + milli_cpu bigint NOT NULL, + memory_bytes bigint NOT NULL, + gpus bigint NOT NULL, + tt_boards bigint NOT NULL, + scratch_bytes bigint NOT NULL, + time_limit_seconds bigint NOT NULL, + signal_grace_seconds bigint NOT NULL, + -- Set at admission and never changed. + node text REFERENCES nodes, + execution_id uuid UNIQUE, + spec_hash text, + auth_generation bigint, + admitted_at timestamptz, + launch_deadline timestamptz, + -- First committed stop cause wins. + stop_cause text, + stop_actor text, + stop_requested_at timestamptz, + -- Process outcome, separate from state and cleanup. + exit_code integer, + exit_reason text, + exit_detail text NOT NULL DEFAULT '', + output_incomplete boolean NOT NULL DEFAULT false, + submitted_at timestamptz NOT NULL, + started_at timestamptz, + finished_at timestamptz, + ended_at timestamptz, + cleanup_pending boolean NOT NULL DEFAULT false, + cleanup_detail text NOT NULL DEFAULT '' +); +CREATE INDEX jobs_active ON jobs (state) + WHERE state IN ('PENDING', 'LAUNCHING', 'RUNNING', 'STOPPING', 'CLEANING'); +CREATE INDEX jobs_user ON jobs (username, state); + +CREATE TABLE job_specs ( + job_id bigint PRIMARY KEY REFERENCES jobs, + request jsonb NOT NULL, -- normalized submission: script, args, argv, env, paths + launch_spec jsonb -- frozen LaunchSpec, written at admission +); + +CREATE TABLE job_events ( + id bigserial PRIMARY KEY, + job_id bigint NOT NULL REFERENCES jobs, + at timestamptz NOT NULL, + from_state text, + to_state text NOT NULL, + reason text NOT NULL DEFAULT '', + actor text NOT NULL DEFAULT '' +); +CREATE INDEX job_events_job ON job_events (job_id, id); + +-- The ledger. An allocation counts against capacity until released_at is +-- set, which only a matching cleanup receipt does. +CREATE TABLE allocations ( + job_id bigint PRIMARY KEY REFERENCES jobs, + node text NOT NULL REFERENCES nodes, + partition text NOT NULL REFERENCES partitions, + username text NOT NULL, + team text NOT NULL, + qos text NOT NULL, + milli_cpu bigint NOT NULL, + memory_bytes bigint NOT NULL, + gpus bigint NOT NULL, + tt_boards bigint NOT NULL, + scratch_bytes bigint NOT NULL, + hugepages_bytes bigint NOT NULL DEFAULT 0, + device_ids text[] NOT NULL DEFAULT '{}', + quarantined boolean NOT NULL DEFAULT false, -- execution lost; held until proven stopped + created_at timestamptz NOT NULL, + released_at timestamptz, + release_evidence text NOT NULL DEFAULT '' +); +CREATE INDEX allocations_unreleased ON allocations (node) WHERE released_at IS NULL; + +-- Committed before every executor call. status 'pending' proves the call +-- was never sent; 'in_flight' is ambiguous until observed. +CREATE TABLE op_intents ( + id bigserial PRIMARY KEY, + job_id bigint NOT NULL REFERENCES jobs, + execution_id uuid NOT NULL, + op text NOT NULL CHECK (op IN ('launch', 'terminate')), + status text NOT NULL CHECK (status IN ('pending', 'in_flight', 'done', 'failed')), + attempts integer NOT NULL DEFAULT 0, + last_error text NOT NULL DEFAULT '', + reason text NOT NULL DEFAULT '', + created_at timestamptz NOT NULL, + updated_at timestamptz NOT NULL, + UNIQUE (job_id, op) +); +CREATE INDEX op_intents_open ON op_intents (job_id) WHERE status IN ('pending', 'in_flight'); + +CREATE TABLE consumed_events ( + event_id text PRIMARY KEY, + job_id bigint NOT NULL, + consumed_at timestamptz NOT NULL +); diff --git a/internal/executor/fake/fake.go b/internal/executor/fake/fake.go index 2ffb98a..c786de2 100644 --- a/internal/executor/fake/fake.go +++ b/internal/executor/fake/fake.go @@ -390,6 +390,20 @@ func (f *Fake) cleaned(e *execution, evidence string) { f.emit(e, executor.EventExited, code, reason, evidence, f.clock.Now()) } +// Forget erases every record of jobID's execution, as if the executor's +// state were wiped. Real executors must never do this; the scheduler has +// to treat the resulting absence as a loss, not as permission to rerun. +func (f *Fake) Forget(jobID int64) { + f.mu.Lock() + defer f.mu.Unlock() + if e := f.execs[jobID]; e != nil { + f.stopTimers(e) + e.frozen = true + delete(f.receipts, e.spec.ExecutionID) + delete(f.execs, jobID) + } +} + // FixCleanup lets a failed cleanup confirm now. func (f *Fake) FixCleanup(jobID int64) { f.mu.Lock() diff --git a/internal/identity/identity.go b/internal/identity/identity.go index 103b86d..ec1917c 100644 --- a/internal/identity/identity.go +++ b/internal/identity/identity.go @@ -121,3 +121,15 @@ type Revocation struct { Username string Team string } + +// Unavailable is an Admitter that can never establish authorization. It +// keeps a scheduler without a real Identity backend failing closed. +type Unavailable struct{} + +func (Unavailable) LockGrant(context.Context, pgx.Tx, string, string, int64) (Grant, error) { + return Grant{}, ErrUnavailable +} + +func (Unavailable) Revocations(context.Context, pgx.Tx, int64) ([]Revocation, int64, error) { + return nil, 0, ErrUnavailable +} diff --git a/internal/scheduler/admit.go b/internal/scheduler/admit.go new file mode 100644 index 0000000..720aa5d --- /dev/null +++ b/internal/scheduler/admit.go @@ -0,0 +1,232 @@ +package scheduler + +import ( + "context" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity" +) + +// admit walks pending jobs in priority order (QoS weight, then FIFO) and +// admits every one that fits. M1 lets a lower-priority job start when the +// one ahead of it does not fit; this can starve large jobs until M3. +func (s *Scheduler) admit(ctx context.Context) error { + var l *ledger + var pending []*job + err := s.fenced(ctx, func(tx pgx.Tx) error { + var err error + if l, err = loadLedger(ctx, tx); err != nil { + return err + } + rows, err := tx.Query(ctx, "SELECT "+jobColumns+jobFrom+ + "WHERE j.state = 'PENDING' ORDER BY q.priority DESC, j.submitted_at, j.id") + if err != nil { + return err + } + pending, err = pgx.CollectRows(rows, func(r pgx.CollectableRow) (*job, error) { return scanJob(r) }) + return err + }) + if err != nil { + return err + } + + for _, j := range pending { + if err := s.leaderCheck(); err != nil { + return err + } + if j.Held { + if err := s.setPending(ctx, j, PendingHeld, "held by an administrator"); err != nil { + return err + } + continue + } + if reason, detail := l.checkLimits(j); reason != "" { + if err := s.setPending(ctx, j, reason, detail); err != nil { + return err + } + continue + } + node, reason, detail := l.place(j, s.usable) + if node == nil { + if err := s.setPending(ctx, j, reason, detail); err != nil { + return err + } + continue + } + admitted, err := s.admitOne(ctx, j, node.Name) + if err != nil { + return err + } + if admitted { + l.add(allocation{Node: node.Name, Partition: j.Partition, Username: j.Username, + Team: j.Team, QOS: j.QOS, Res: j.Res}) + } + } + return nil +} + +func (s *Scheduler) setPending(ctx context.Context, j *job, reason PendingReason, detail string) error { + if j.PendingReason != nil && *j.PendingReason == string(reason) { + return nil // pending_detail may lag; the reason is what users act on + } + return s.fenced(ctx, func(tx pgx.Tx) error { + _, err := tx.Exec(ctx, + "UPDATE jobs SET pending_reason = $2, pending_detail = $3 WHERE id = $1 AND state = 'PENDING'", + j.ID, string(reason), detail) + return err + }) +} + +// admitOne commits allocation + LAUNCHING + node + launch intent in one +// transaction. The executor is called later, by dispatch, after commit. +func (s *Scheduler) admitOne(ctx context.Context, j *job, nodeName string) (bool, error) { + now := s.clock.Now() + var grantErr error + admitted, mismatch := false, false + err := s.fenced(ctx, func(tx pgx.Tx) error { + // Lock order: account, membership (Identity), then job, then node. + grant, err := s.identity.LockGrant(ctx, tx, j.Username, j.Team, 0) + if err != nil { + grantErr = err + return nil + } + cur, err := lockJob(ctx, tx, j.ID) + if err != nil { + return err + } + if cur == nil || cur.State != StatePending || cur.Held { + return nil // cancelled or held since the ledger was read + } + var p partitionInfo + var cap executor.Resources + var reserveCPU, reserveMem int64 + if err := tx.QueryRow(ctx, ` + SELECT n.milli_cpu, n.memory_bytes, n.gpus, n.tt_boards, n.scratch_bytes, n.hugepages_bytes, + n.reserve_milli_cpu, n.reserve_memory_bytes, p.runtime + FROM nodes n, partitions p WHERE n.name = $1 AND p.name = $2 FOR UPDATE OF n`, + nodeName, j.Partition).Scan(&cap.MilliCPU, &cap.MemoryBytes, &cap.GPUs, &cap.TTBoards, + &cap.ScratchBytes, &cap.HugepagesBytes, &reserveCPU, &reserveMem, &p.Runtime); err != nil { + return err + } + // Recheck the node against the database, not only the cycle's copy. + var used executor.Resources + if err := tx.QueryRow(ctx, ` + SELECT coalesce(sum(milli_cpu), 0), coalesce(sum(memory_bytes), 0), coalesce(sum(gpus), 0), + coalesce(sum(tt_boards), 0), coalesce(sum(scratch_bytes), 0), coalesce(sum(hugepages_bytes), 0) + FROM allocations WHERE node = $1 AND released_at IS NULL`, nodeName).Scan( + &used.MilliCPU, &used.MemoryBytes, &used.GPUs, &used.TTBoards, &used.ScratchBytes, &used.HugepagesBytes); err != nil { + return err + } + free := cap.Sub(executor.Resources{MilliCPU: reserveCPU, MemoryBytes: reserveMem}).Sub(used) + if !free.Fits(j.Res) { + mismatch = true // the cycle's copy is stale; wait and recompute + return nil + } + + var req Request + var body []byte + if err := tx.QueryRow(ctx, "SELECT request FROM job_specs WHERE job_id = $1", j.ID).Scan(&body); err != nil { + return err + } + if err := json.Unmarshal(body, &req); err != nil { + return err + } + spec := executor.LaunchSpec{ + JobID: j.ID, + ExecutionID: executor.NewExecutionID(), + AuthGeneration: grant.Generation, + LaunchDeadline: now.Add(s.launchTimeout), + Kind: req.Kind, + Node: nodeName, + User: executor.User{Name: grant.Account.Username, UID: grant.Account.UID, PersonalGID: grant.Account.PersonalGID}, + TeamSlug: grant.Team.Slug, + TeamGID: grant.Team.GID, + ExtraGIDs: grant.ExtraGIDs, + Resources: j.Res, + Mounts: []executor.Mount{ + {VolumeID: grant.Team.VolumeID, Source: "/teams/" + grant.Team.Slug, Target: "/mist/teams/" + grant.Team.Slug}, + {VolumeID: grant.Team.VolumeID, Source: "/software", Target: "/mist/software", ReadOnly: true}, + }, + Workdir: req.Workdir, + OutputPath: req.OutputPath, + Script: req.Script, + Args: req.Args, + Argv: req.Argv, + Env: req.Env, + Runtime: p.Runtime, + TimeLimit: j.TimeLimit, + SignalGrace: j.Grace, + } + spec.Seal() + if err := spec.Validate(); err != nil { + // A normalization bug, not a capacity question. Leave it pending + // with an explanation an admin can act on. + _, err2 := tx.Exec(ctx, + "UPDATE jobs SET pending_detail = $2 WHERE id = $1", j.ID, "cannot build launch spec: "+err.Error()) + return err2 + } + frozen, err := json.Marshal(spec) + if err != nil { + return err + } + + r := j.Res + if _, err := tx.Exec(ctx, ` + INSERT INTO allocations (job_id, node, partition, username, team, qos, + milli_cpu, memory_bytes, gpus, tt_boards, scratch_bytes, hugepages_bytes, device_ids, created_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, + j.ID, nodeName, j.Partition, j.Username, j.Team, j.QOS, + r.MilliCPU, r.MemoryBytes, r.GPUs, r.TTBoards, r.ScratchBytes, r.HugepagesBytes, append([]string{}, spec.DeviceIDs...), now); err != nil { + return err + } + if _, err := tx.Exec(ctx, "UPDATE job_specs SET launch_spec = $2 WHERE job_id = $1", j.ID, frozen); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` + UPDATE jobs SET node = $2, execution_id = $3, spec_hash = $4, auth_generation = $5, + admitted_at = $6, launch_deadline = $7, pending_reason = NULL, pending_detail = '' + WHERE id = $1`, + j.ID, nodeName, spec.ExecutionID, spec.SpecHash, grant.Generation, now, spec.LaunchDeadline); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` + INSERT INTO op_intents (job_id, execution_id, op, status, created_at, updated_at) + VALUES ($1, $2, 'launch', 'pending', $3, $3)`, j.ID, spec.ExecutionID, now); err != nil { + return err + } + if err := transition(ctx, tx, cur, StateLaunching, now, "admitted to "+nodeName, "mist"); err != nil { + return err + } + admitted = true + return nil + }) + if err != nil { + return false, err + } + + switch { + case mismatch: + s.log.Warn("ledger disagrees with database", "job", j.ID, "node", nodeName) + err = s.setPending(ctx, j, PendingResources, waitingFor(j.Res)) + case grantErr == nil: + case errors.Is(grantErr, identity.ErrRevoked): + err = s.fenced(ctx, func(tx pgx.Tx) error { + cur, err := lockJob(ctx, tx, j.ID) + if err != nil || cur == nil { + return err + } + return stop(ctx, tx, cur, StopRevoke, "identity", "access to team "+j.Team+" revoked", now) + }) + case errors.Is(grantErr, identity.ErrNotProvisioned): + err = s.setPending(ctx, j, PendingProvisioning, "team folders are not provisioned yet") + case errors.Is(grantErr, identity.ErrUnavailable): + err = s.setPending(ctx, j, PendingAuthUnavailable, "cannot confirm current team membership") + default: + err = grantErr + } + return admitted, err +} diff --git a/internal/scheduler/dispatch.go b/internal/scheduler/dispatch.go new file mode 100644 index 0000000..ac60a36 --- /dev/null +++ b/internal/scheduler/dispatch.go @@ -0,0 +1,181 @@ +package scheduler + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/jackc/pgx/v5" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity" +) + +type openIntent struct { + JobID int64 + Op OpKind + Node string + Reason string +} + +// dispatch sends every open intent to its executor. Launch intents that +// were sent but not answered are resent with the same ExecutionID and +// SpecHash; the executor's repeat-safety makes that harmless. +func (s *Scheduler) dispatch(ctx context.Context) error { + var open []openIntent + err := s.fenced(ctx, func(tx pgx.Tx) error { + rows, err := tx.Query(ctx, ` + SELECT i.job_id, i.op, j.node, i.reason FROM op_intents i JOIN jobs j ON j.id = i.job_id + WHERE i.status IN ('pending', 'in_flight') ORDER BY i.id`) + if err != nil { + return err + } + open, err = pgx.CollectRows(rows, func(r pgx.CollectableRow) (openIntent, error) { + var o openIntent + var op string + err := r.Scan(&o.JobID, &op, &o.Node, &o.Reason) + o.Op = OpKind(op) + return o, err + }) + return err + }) + if err != nil { + return err + } + for _, o := range open { + if err := s.leaderCheck(); err != nil { + return err + } + ex := s.executorFor(o.Node) + if ex == nil { + continue // executor not observed; the intent waits + } + switch o.Op { + case OpLaunch: + err = s.dispatchLaunch(ctx, ex, o) + case OpTerminate: + err = s.dispatchTerminate(ctx, ex, o) + } + if err != nil { + return err + } + } + return nil +} + +func (s *Scheduler) dispatchLaunch(ctx context.Context, ex executor.Executor, o openIntent) error { + now := s.clock.Now() + var spec *executor.LaunchSpec + err := s.fenced(ctx, func(tx pgx.Tx) error { + // Read the job without locking to learn who to authorize, then take + // locks in order: Identity's rows first, then the job. + var user, team string + var gen int64 + if err := tx.QueryRow(ctx, "SELECT username, team, auth_generation FROM jobs WHERE id = $1", + o.JobID).Scan(&user, &team, &gen); err != nil { + return err + } + _, grantErr := s.identity.LockGrant(ctx, tx, user, team, gen) + j, err := lockJob(ctx, tx, o.JobID) + if err != nil { + return err + } + if j.State != StateLaunching || j.StopCause != nil { + return nil // a stop owns this execution now + } + switch { + case errors.Is(grantErr, identity.ErrRevoked), errors.Is(grantErr, identity.ErrStaleGeneration): + // Never replay an unstarted launch under an older authorization. + return stop(ctx, tx, j, StopRevoke, "identity", "authorization changed before launch", now) + case grantErr != nil: + return nil // unavailable or not provisioned: retry; the launch deadline bounds it + case !now.Before(*j.LaunchDeadline): + return stop(ctx, tx, j, StopLaunchTimeout, "mist", "launch deadline passed", now) + } + var body []byte + if err := tx.QueryRow(ctx, "SELECT launch_spec FROM job_specs WHERE job_id = $1", j.ID).Scan(&body); err != nil { + return err + } + spec = new(executor.LaunchSpec) + if err := json.Unmarshal(body, spec); err != nil { + return err + } + spec.SpecHash = *j.SpecHash + _, err = tx.Exec(ctx, ` + UPDATE op_intents SET status = 'in_flight', attempts = attempts + 1, updated_at = $2 + WHERE job_id = $1 AND op = 'launch'`, j.ID, now) + return err + }) + if err != nil || spec == nil { + return err + } + + callCtx, cancel := context.WithTimeout(ctx, s.callTimeout) + callErr := ex.Launch(callCtx, *spec) + cancel() + + return s.fenced(ctx, func(tx pgx.Tx) error { + now := s.clock.Now() + switch { + case callErr == nil, errors.Is(callErr, executor.ErrExecutionStarted): + // It exists, or existed: events and List say what happened. + _, err := tx.Exec(ctx, + "UPDATE op_intents SET status = 'done', updated_at = $2, last_error = '' WHERE job_id = $1 AND op = 'launch' AND status = 'in_flight'", + o.JobID, now) + return err + case errors.Is(callErr, executor.ErrSpecConflict): + // A different execution claims this job: a bug or a stale writer. + // Stop using the node until an admin looks. + if err := setIntent(ctx, tx, o.JobID, OpLaunch, OpFailed, now, callErr.Error()); err != nil { + return err + } + _, err := tx.Exec(ctx, + "UPDATE nodes SET quarantined = true, quarantine_detail = $2 WHERE name = $1", + spec.Node, "spec conflict for job execution "+spec.ExecutionID) + return err + default: + // Ambiguous: it may or may not exist. Stay in_flight and retry. + _, err := tx.Exec(ctx, + "UPDATE op_intents SET last_error = $2, updated_at = $3 WHERE job_id = $1 AND op = 'launch'", + o.JobID, callErr.Error(), now) + return err + } + }) +} + +func (s *Scheduler) dispatchTerminate(ctx context.Context, ex executor.Executor, o openIntent) error { + var grace time.Duration + send := false + err := s.fenced(ctx, func(tx pgx.Tx) error { + j, err := lockJob(ctx, tx, o.JobID) + if err != nil { + return err + } + if j.State != StateStopping { + return setIntent(ctx, tx, j.ID, OpTerminate, OpDone, s.clock.Now(), "no longer stopping") + } + grace, send = j.Grace, true + _, err = tx.Exec(ctx, ` + UPDATE op_intents SET status = 'in_flight', attempts = attempts + 1, updated_at = $2 + WHERE job_id = $1 AND op = 'terminate'`, j.ID, s.clock.Now()) + return err + }) + if err != nil || !send { + return err + } + + callCtx, cancel := context.WithTimeout(ctx, s.callTimeout) + callErr := ex.Terminate(callCtx, o.JobID, grace, o.Reason) + cancel() + + return s.fenced(ctx, func(tx pgx.Tx) error { + if callErr != nil { + return setIntent(ctx, tx, o.JobID, OpTerminate, OpPending, s.clock.Now(), callErr.Error()) + } + _, err := tx.Exec(ctx, + "UPDATE op_intents SET status = 'done', updated_at = $2 WHERE job_id = $1 AND op = 'terminate' AND status = 'in_flight'", + o.JobID, s.clock.Now()) + return err + }) +} diff --git a/internal/scheduler/events.go b/internal/scheduler/events.go new file mode 100644 index 0000000..f9e152a --- /dev/null +++ b/internal/scheduler/events.go @@ -0,0 +1,166 @@ +package scheduler + +import ( + "context" + + "github.com/jackc/pgx/v5" + + "github.com/UTMIST/Mist/internal/executor" +) + +// HandleEvent consumes one executor notification. The event ID is recorded +// in the same transaction as the state change, so a redelivered event is +// skipped. Handlers are also idempotent by state, because reconciliation +// replays observations under its own IDs, and order is not guaranteed. +func (s *Scheduler) HandleEvent(ctx context.Context, ev executor.Event) error { + return s.fenced(ctx, func(tx pgx.Tx) error { + tag, err := tx.Exec(ctx, + "INSERT INTO consumed_events (event_id, job_id, consumed_at) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING", + ev.EventID, ev.JobID, s.clock.Now()) + if err != nil || tag.RowsAffected() == 0 { + return err + } + return s.apply(ctx, tx, ev) + }) +} + +// observe applies an observation from reconciliation, without dedupe. +func (s *Scheduler) observe(ctx context.Context, ev executor.Event) error { + return s.fenced(ctx, func(tx pgx.Tx) error { return s.apply(ctx, tx, ev) }) +} + +func (s *Scheduler) apply(ctx context.Context, tx pgx.Tx, ev executor.Event) error { + j, err := lockJob(ctx, tx, ev.JobID) + if err != nil || j == nil { + return err + } + if j.ExecutionID == nil || *j.ExecutionID != ev.ExecutionID { + s.log.Warn("ignoring event for another execution", "job", ev.JobID, "event", ev.EventID) + return nil + } + switch ev.Kind { + case executor.EventStarted: + return s.onStarted(ctx, tx, j, ev) + case executor.EventFinished, executor.EventLaunchFailed: + return s.onFinished(ctx, tx, j, ev) + case executor.EventExited: + return s.onExited(ctx, tx, j, ev) + case executor.EventLost: + return s.onLost(ctx, tx, j, ev) + } + return nil +} + +func (s *Scheduler) onStarted(ctx context.Context, tx pgx.Tx, j *job, ev executor.Event) error { + // A fast job may have finished before this arrived: adopt the real + // start time whatever the state, but only move LAUNCHING forward. + if j.StartedAt == nil { + if _, err := tx.Exec(ctx, "UPDATE jobs SET started_at = $2 WHERE id = $1", j.ID, ev.At); err != nil { + return err + } + j.StartedAt = &ev.At + } + if j.State != StateLaunching { + return nil + } + if err := setIntent(ctx, tx, j.ID, OpLaunch, OpDone, s.clock.Now(), ""); err != nil { + return err + } + return transition(ctx, tx, j, StateRunning, s.clock.Now(), "started on "+ev.Node, "executor") +} + +func (s *Scheduler) onFinished(ctx context.Context, tx pgx.Tx, j *job, ev executor.Event) error { + switch j.State { + case StateLaunching, StateRunning, StateStopping: + case StateNodeFail: + // A lost execution reported back: keep the result for display. + if j.ExitReason == nil { + return recordResult(ctx, tx, j, ev) + } + return nil + default: + return nil // CLEANING or terminal: the first result stands + } + if err := recordResult(ctx, tx, j, ev); err != nil { + return err + } + if err := closeOpenIntents(ctx, tx, j.ID, s.clock.Now(), ""); err != nil { + return err + } + if _, err := tx.Exec(ctx, "UPDATE jobs SET cleanup_pending = true WHERE id = $1", j.ID); err != nil { + return err + } + return transition(ctx, tx, j, StateCleaning, s.clock.Now(), describe(ev), "executor") +} + +func recordResult(ctx context.Context, tx pgx.Tx, j *job, ev executor.Event) error { + reason := string(ev.Reason) + if _, err := tx.Exec(ctx, + "UPDATE jobs SET exit_code = $2, exit_reason = $3, exit_detail = $4, finished_at = $5 WHERE id = $1", + j.ID, ev.ExitCode, reason, ev.Detail, ev.At); err != nil { + return err + } + j.ExitCode, j.ExitReason, j.ExitDetail, j.FinishedAt = ev.ExitCode, &reason, ev.Detail, &ev.At + return nil +} + +func (s *Scheduler) onExited(ctx context.Context, tx pgx.Tx, j *job, ev executor.Event) error { + now := s.clock.Now() + switch j.State { + case StateLaunching, StateRunning, StateStopping: + // The Finished notification was lost or is late; Exited carries + // the same outcome. + if err := s.onFinished(ctx, tx, j, ev); err != nil { + return err + } + case StateCleaning: + default: + // Terminal. A quarantined allocation (NODE_FAIL) is released by + // this matching cleanup receipt; a repeated receipt changes nothing. + if _, err := tx.Exec(ctx, + "UPDATE jobs SET cleanup_pending = false, cleanup_detail = $2 WHERE id = $1 AND cleanup_pending", + j.ID, ev.Detail); err != nil { + return err + } + return releaseAllocation(ctx, tx, j.ID, now, ev.Detail) + } + if err := releaseAllocation(ctx, tx, j.ID, now, ev.Detail); err != nil { + return err + } + if _, err := tx.Exec(ctx, "UPDATE jobs SET cleanup_detail = $2 WHERE id = $1", j.ID, ev.Detail); err != nil { + return err + } + return endJob(ctx, tx, j, decide(j), now, "executor") +} + +func (s *Scheduler) onLost(ctx context.Context, tx pgx.Tx, j *job, ev executor.Event) error { + switch j.State { + case StateLaunching, StateRunning, StateStopping, StateCleaning: + default: + return nil + } + now := s.clock.Now() + // Keep the allocation, marked quarantined: the devices may still be in + // use until recovery or fencing proves the old execution stopped. + if _, err := tx.Exec(ctx, + "UPDATE allocations SET quarantined = true WHERE job_id = $1 AND released_at IS NULL", j.ID); err != nil { + return err + } + if err := closeOpenIntents(ctx, tx, j.ID, now, "execution lost"); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` + UPDATE jobs SET ended_at = $2, cleanup_pending = true, cleanup_detail = $3, + exit_reason = coalesce(exit_reason, $4), exit_detail = CASE WHEN exit_reason IS NULL THEN $3 ELSE exit_detail END + WHERE id = $1`, j.ID, now, ev.Detail, string(executor.ExitNodeLost)); err != nil { + return err + } + return transition(ctx, tx, j, StateNodeFail, now, ev.Detail, "executor") +} + +func describe(ev executor.Event) string { + if ev.Detail != "" { + return string(ev.Reason) + ": " + ev.Detail + } + return string(ev.Reason) +} diff --git a/internal/scheduler/harness_test.go b/internal/scheduler/harness_test.go new file mode 100644 index 0000000..a30ce07 --- /dev/null +++ b/internal/scheduler/harness_test.go @@ -0,0 +1,232 @@ +package scheduler + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/db/dbtest" + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/executor/fake" + "github.com/UTMIST/Mist/internal/identity/identitytest" +) + +func TestMain(m *testing.M) { dbtest.Main(m) } + +var t0 = time.Date(2026, 11, 2, 12, 0, 0, 0, time.UTC) + +const ( + gpuNode = "utmist-z1opa08" + cpuNode = "utmist-tt" +) + +// harness wires a real Postgres, the fake executor, the fake clock and +// the identity fixture. Nothing touches a cluster. +type harness struct { + t *testing.T + ctx context.Context + pool *pgxpool.Pool + clock *clock.Fake + exec *fake.Fake + id *identitytest.Fixture + s *Scheduler + jobs Jobs +} + +func newHarness(t *testing.T) *harness { + ctx := context.Background() + pool, err := db.Open(ctx, dbtest.URL(t)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := db.Migrate(ctx, pool); err != nil { + t.Fatal(err) + } + for _, q := range []string{ + "INSERT INTO nodes (name, executor) VALUES ($1, 'k3s'), ($2, 'k3s')", + "INSERT INTO partition_nodes VALUES ('gpu', $1), ('cpu', $2)", + } { + if _, err := pool.Exec(ctx, q, gpuNode, cpuNode); err != nil { + t.Fatal(err) + } + } + if _, err := pool.Exec(ctx, "UPDATE partitions SET runtime = 'runtime@sha256:test'"); err != nil { + t.Fatal(err) + } + c := clock.NewFake(t0) + ex := fake.New(fake.Config{Clock: c, Nodes: []executor.Node{ + {Name: gpuNode, Partitions: []string{"gpu"}, Healthy: true, Capacity: executor.Resources{ + MilliCPU: 32000, MemoryBytes: 128 << 30, GPUs: 2, ScratchBytes: 400 << 30}}, + {Name: cpuNode, Partitions: []string{"cpu"}, Healthy: true, Capacity: executor.Resources{ + MilliCPU: 32000, MemoryBytes: 256 << 30, ScratchBytes: 1 << 40}}, + }}) + t.Cleanup(ex.Close) + h := &harness{t: t, ctx: ctx, pool: pool, clock: c, exec: ex, id: identitytest.New(), + jobs: Jobs{Pool: pool, Clock: c}} + h.s = h.newScheduler() + return h +} + +// newScheduler takes leadership, fencing any earlier scheduler. +func (h *harness) newScheduler() *Scheduler { + s, err := New(h.ctx, Config{ + Pool: h.pool, + Identity: h.id, + Executors: map[string]executor.Executor{"k3s": h.exec}, + Clock: h.clock, + }, nil) + if err != nil { + h.t.Fatal(err) + } + return s +} + +func (h *harness) submit(user, team, partition string, mutate ...func(*Request)) int64 { + h.t.Helper() + r := Request{ + Username: user, Team: team, Partition: partition, QOS: "normal", Name: "test", + Kind: executor.KindBatch, + TimeLimit: time.Hour, + SignalGrace: time.Minute, + Workdir: "/mist/teams/" + team + "/" + user, + OutputPath: "/mist/teams/" + team + "/" + user + "/mist.out", + Script: "python train.py\n", + } + if partition == "gpu" { + r.Resources = executor.Resources{MilliCPU: 8000, MemoryBytes: 32 << 30, GPUs: 1, ScratchBytes: 100 << 30} + } else { + r.Resources = executor.Resources{MilliCPU: 2000, MemoryBytes: 8 << 30, ScratchBytes: 50 << 30} + } + for _, m := range mutate { + m(&r) + } + id, err := h.jobs.Submit(h.ctx, r) + if err != nil { + h.t.Fatal(err) + } + return id +} + +func (h *harness) cycle() { + h.t.Helper() + if err := h.s.Cycle(h.ctx); err != nil { + h.t.Fatalf("cycle: %v", err) + } +} + +// drain handles every event the fake has delivered, and returns how many. +func (h *harness) drain() int { + h.t.Helper() + n := 0 + for { + select { + case ev := <-h.exec.Events(): + if err := h.s.HandleEvent(h.ctx, ev); err != nil { + h.t.Fatalf("event %s: %v", ev.EventID, err) + } + n++ + case <-time.After(20 * time.Millisecond): + return n + } + } +} + +// settle cycles and handles events until nothing more happens at this +// instant of virtual time. +func (h *harness) settle() { + h.t.Helper() + for range 10 { + h.cycle() + if h.drain() == 0 { + return + } + } + h.t.Fatal("did not settle") +} + +func (h *harness) advance(d time.Duration) { + h.t.Helper() + h.clock.Advance(d) + h.settle() +} + +func (h *harness) job(id int64) *job { + h.t.Helper() + j, err := scanJob(h.pool.QueryRow(h.ctx, "SELECT "+jobColumns+jobFrom+"WHERE j.id = $1", id)) + if err != nil { + h.t.Fatal(err) + } + return j +} + +func (h *harness) want(id int64, state State, reason ...PendingReason) { + h.t.Helper() + j := h.job(id) + if j.State != state { + h.t.Fatalf("job %d: state %s, want %s (history %s)", id, j.State, state, h.history(id)) + } + if len(reason) > 0 && (j.PendingReason == nil || *j.PendingReason != string(reason[0])) { + got := "" + if j.PendingReason != nil { + got = *j.PendingReason + } + h.t.Fatalf("job %d: pending reason %s, want %s", id, got, reason[0]) + } +} + +// history returns transitions like ">PENDING PENDING>LAUNCHING ...". +func (h *harness) history(id int64) string { + h.t.Helper() + rows, err := h.pool.Query(h.ctx, + "SELECT coalesce(from_state, ''), to_state FROM job_events WHERE job_id = $1 ORDER BY id", id) + if err != nil { + h.t.Fatal(err) + } + defer rows.Close() + var out []string + for rows.Next() { + var from, to string + rows.Scan(&from, &to) + out = append(out, from+">"+to) + } + return strings.Join(out, " ") +} + +// allocation reports whether job id has an allocation and whether it is +// released and quarantined. +func (h *harness) allocation(id int64) (exists, released, quarantined bool) { + h.t.Helper() + var rel *time.Time + err := h.pool.QueryRow(h.ctx, + "SELECT released_at, quarantined FROM allocations WHERE job_id = $1", id).Scan(&rel, &quarantined) + if err != nil { + return false, false, false + } + return true, rel != nil, quarantined +} + +func (h *harness) sql(sql string, args ...any) { + h.t.Helper() + if _, err := h.pool.Exec(h.ctx, sql, args...); err != nil { + h.t.Fatalf("%s: %v", sql, err) + } +} + +func (h *harness) intent(id int64, op OpKind) (OpStatus, string) { + h.t.Helper() + var st, last string + if err := h.pool.QueryRow(h.ctx, "SELECT status, last_error FROM op_intents WHERE job_id = $1 AND op = $2", + id, string(op)).Scan(&st, &last); err != nil { + return "", fmt.Sprint(err) + } + return OpStatus(st), last +} + +const lifecycle = ">PENDING PENDING>LAUNCHING LAUNCHING>RUNNING RUNNING>CLEANING CLEANING>COMPLETED" diff --git a/internal/scheduler/ledger.go b/internal/scheduler/ledger.go new file mode 100644 index 0000000..bcab51e --- /dev/null +++ b/internal/scheduler/ledger.go @@ -0,0 +1,284 @@ +package scheduler + +import ( + "context" + "fmt" + "slices" + + "github.com/jackc/pgx/v5" + + "github.com/UTMIST/Mist/internal/executor" +) + +// The ledger is the scheduler's view of capacity: what each node reports, +// minus the Mist host reserve, minus every unreleased allocation. It never +// asks Kubernetes whether something fits. + +type nodeInfo struct { + Name, Executor, AdminState string + Healthy, Quarantined bool + Capacity, Reserve, Used executor.Resources +} + +func (n *nodeInfo) free() executor.Resources { return n.Capacity.Sub(n.Reserve).Sub(n.Used) } + +type partitionInfo struct { + Name, Executor, State, Accelerator, Runtime string + CapMilliCPU, CapMemory *int64 + Nodes []string + Used executor.Resources +} + +type allocation struct { + Node, Partition, Username, Team, QOS string + Res executor.Resources +} + +type limitRow struct { + Scope, Subject, Partition, QOS, Resource string + Max int64 +} + +type ledger struct { + nodes map[string]*nodeInfo + partitions map[string]*partitionInfo + allocs []allocation + limits []limitRow +} + +func loadLedger(ctx context.Context, tx pgx.Tx) (*ledger, error) { + l := &ledger{nodes: map[string]*nodeInfo{}, partitions: map[string]*partitionInfo{}} + + rows, err := tx.Query(ctx, ` + SELECT name, executor, admin_state, healthy, quarantined, + milli_cpu, memory_bytes, gpus, tt_boards, scratch_bytes, hugepages_bytes, + reserve_milli_cpu, reserve_memory_bytes + FROM nodes`) + if err != nil { + return nil, err + } + for rows.Next() { + var n nodeInfo + c := &n.Capacity + if err := rows.Scan(&n.Name, &n.Executor, &n.AdminState, &n.Healthy, &n.Quarantined, + &c.MilliCPU, &c.MemoryBytes, &c.GPUs, &c.TTBoards, &c.ScratchBytes, &c.HugepagesBytes, + &n.Reserve.MilliCPU, &n.Reserve.MemoryBytes); err != nil { + return nil, err + } + l.nodes[n.Name] = &n + } + if err := rows.Err(); err != nil { + return nil, err + } + + rows, err = tx.Query(ctx, ` + SELECT p.name, p.executor, p.state, p.accelerator, p.runtime, p.cap_milli_cpu, p.cap_memory_bytes, + coalesce(array_agg(pn.node ORDER BY pn.node) FILTER (WHERE pn.node IS NOT NULL), '{}') + FROM partitions p LEFT JOIN partition_nodes pn ON pn.partition = p.name + GROUP BY p.name`) + if err != nil { + return nil, err + } + for rows.Next() { + var p partitionInfo + if err := rows.Scan(&p.Name, &p.Executor, &p.State, &p.Accelerator, &p.Runtime, + &p.CapMilliCPU, &p.CapMemory, &p.Nodes); err != nil { + return nil, err + } + l.partitions[p.Name] = &p + } + if err := rows.Err(); err != nil { + return nil, err + } + + // Every unreleased allocation counts: launching, running, stopping, + // cleaning and quarantined terminal jobs alike. + rows, err = tx.Query(ctx, ` + SELECT node, partition, username, team, qos, + milli_cpu, memory_bytes, gpus, tt_boards, scratch_bytes, hugepages_bytes + FROM allocations WHERE released_at IS NULL`) + if err != nil { + return nil, err + } + for rows.Next() { + var a allocation + r := &a.Res + if err := rows.Scan(&a.Node, &a.Partition, &a.Username, &a.Team, &a.QOS, + &r.MilliCPU, &r.MemoryBytes, &r.GPUs, &r.TTBoards, &r.ScratchBytes, &r.HugepagesBytes); err != nil { + return nil, err + } + l.add(a) + } + if err := rows.Err(); err != nil { + return nil, err + } + + rows, err = tx.Query(ctx, "SELECT scope, subject, partition, qos, resource, max FROM limits") + if err != nil { + return nil, err + } + for rows.Next() { + var r limitRow + if err := rows.Scan(&r.Scope, &r.Subject, &r.Partition, &r.QOS, &r.Resource, &r.Max); err != nil { + return nil, err + } + l.limits = append(l.limits, r) + } + return l, rows.Err() +} + +// add charges an allocation to its node and partition. +func (l *ledger) add(a allocation) { + l.allocs = append(l.allocs, a) + if n := l.nodes[a.Node]; n != nil { + n.Used = n.Used.Add(a.Res) + } + if p := l.partitions[a.Partition]; p != nil { + p.Used = p.Used.Add(a.Res) + } +} + +// effectiveLimits picks, for each scope and resource, the most specific +// matching row: exact subject over '*', then exact partition, then exact +// QoS. User, team and partition results all apply together. +func (l *ledger) effectiveLimits(j *job) []limitRow { + best := map[[2]string]limitRow{} + score := map[[2]string]int{} + for _, r := range l.limits { + if r.Resource == "submitted_jobs" { + continue // checked at submit + } + var subject string + switch r.Scope { + case "user": + subject = j.Username + case "team": + subject = j.Team + case "partition": + subject = j.Partition + } + if (r.Subject != subject && r.Subject != "*") || + (r.Partition != j.Partition && r.Partition != "*") || + (r.QOS != j.QOS && r.QOS != "*") { + continue + } + s := 0 + if r.Subject != "*" { + s += 4 + } + if r.Partition != "*" { + s += 2 + } + if r.QOS != "*" { + s++ + } + key := [2]string{r.Scope, r.Resource} + if old, ok := score[key]; !ok || s > old { + best[key], score[key] = r, s + } + } + out := make([]limitRow, 0, len(best)) + for _, r := range best { + out = append(out, r) + } + slices.SortFunc(out, func(a, b limitRow) int { // deterministic reason order + return scopeRank(a.Scope) - scopeRank(b.Scope) + }) + return out +} + +func scopeRank(scope string) int { + return map[string]int{"user": 0, "team": 1, "partition": 2}[scope] +} + +func amount(r executor.Resources, resource string) int64 { + switch resource { + case "gpus": + return r.GPUs + case "tt_boards": + return r.TTBoards + case "milli_cpu": + return r.MilliCPU + case "memory_bytes": + return r.MemoryBytes + case "jobs": + return 1 + } + return 0 +} + +// checkLimits returns the pending reason and detail if j would exceed a +// concurrent limit, or "" if it is within all of them. +func (l *ledger) checkLimits(j *job) (PendingReason, string) { + for _, r := range l.effectiveLimits(j) { + var used int64 + for _, a := range l.allocs { + if (r.Partition != "*" && a.Partition != r.Partition) || (r.QOS != "*" && a.QOS != r.QOS) { + continue + } + switch r.Scope { + case "user": + if a.Username != j.Username { + continue + } + case "team": + if a.Team != j.Team { + continue + } + } + used += amount(a.Res, r.Resource) + } + want := amount(j.Res, r.Resource) + if want == 0 || used+want <= r.Max { + continue + } + detail := fmt.Sprintf("%s limit: %d of %d %s in use", r.Scope, used, r.Max, r.Resource) + switch r.Scope { + case "user": + return PendingUserLimit, detail + case "team": + return PendingTeamLimit, detail + default: + return PendingQOSLimit, detail + } + } + return "", "" +} + +// place finds a node for j, or explains why there is none. usable reports +// whether the executor observed the node completely this cycle. +func (l *ledger) place(j *job, usable func(*nodeInfo) bool) (*nodeInfo, PendingReason, string) { + p := l.partitions[j.Partition] + if p == nil || p.State != "up" { + return nil, PendingPartitionDown, "partition " + j.Partition + " is down" + } + if p.CapMilliCPU != nil && p.Used.MilliCPU+j.Res.MilliCPU > *p.CapMilliCPU || + p.CapMemory != nil && p.Used.MemoryBytes+j.Res.MemoryBytes > *p.CapMemory { + return nil, PendingResources, "partition " + j.Partition + " is at its CPU or memory cap" + } + anyUp := false + for _, name := range p.Nodes { + n := l.nodes[name] + if n == nil || n.AdminState != "up" || !n.Healthy || n.Quarantined || !usable(n) { + continue + } + anyUp = true + if n.free().Fits(j.Res) { + return n, "", "" + } + } + if !anyUp { + return nil, PendingPartitionDown, "no schedulable node in partition " + j.Partition + } + return nil, PendingResources, waitingFor(j.Res) +} + +func waitingFor(r executor.Resources) string { + switch { + case r.GPUs > 0: + return fmt.Sprintf("waiting for %d GPU(s)", r.GPUs) + case r.TTBoards > 0: + return fmt.Sprintf("waiting for %d TT board(s)", r.TTBoards) + } + return fmt.Sprintf("waiting for %d CPU(s) and %d GiB memory", r.MilliCPU/1000, r.MemoryBytes>>30) +} diff --git a/internal/scheduler/reconcile.go b/internal/scheduler/reconcile.go new file mode 100644 index 0000000..cf6ee80 --- /dev/null +++ b/internal/scheduler/reconcile.go @@ -0,0 +1,343 @@ +package scheduler + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity" +) + +// refreshNodes copies each executor's node report into the nodes table. +// Nodes must already be registered by an admin; unknown ones are ignored. +func (s *Scheduler) refreshNodes(ctx context.Context) error { + for name, ex := range s.executors { + nodes, err := ex.Nodes(ctx) + if err != nil { + s.observed[name] = false + s.log.Warn("executor nodes unavailable", "executor", name, "err", err) + continue + } + s.observed[name] = true + err = s.fenced(ctx, func(tx pgx.Tx) error { + for _, n := range nodes { + c := n.Capacity + if _, err := tx.Exec(ctx, ` + UPDATE nodes SET healthy = $3, health_detail = $4, observed_at = $5, + milli_cpu = $6, memory_bytes = $7, gpus = $8, tt_boards = $9, scratch_bytes = $10, hugepages_bytes = $11 + WHERE name = $1 AND executor = $2`, + n.Name, name, n.Healthy, n.Detail, s.clock.Now(), + c.MilliCPU, c.MemoryBytes, c.GPUs, c.TTBoards, c.ScratchBytes, c.HugepagesBytes); err != nil { + return err + } + s.nodeExec[n.Name] = name + } + return nil + }) + if err != nil { + return err + } + } + return nil +} + +// usable reports whether admission may place work on n this cycle: its +// executor answered completely, and it hosts no unexplained execution. +func (s *Scheduler) usable(n *nodeInfo) bool { + return s.observed[n.Executor] && !s.unknownOn[n.Name] +} + +func (s *Scheduler) executorFor(node string) executor.Executor { + name, ok := s.nodeExec[node] + if !ok || !s.observed[name] { + return nil + } + return s.executors[name] +} + +// reconcile compares each executor's complete List with the database, +// following the matrix in scheduler.md#reconciliation. +func (s *Scheduler) reconcile(ctx context.Context) error { + for name, ex := range s.executors { + if !s.observed[name] { + continue + } + xs, err := ex.List(ctx) + if err != nil { + // No absence inference from a failed or partial list, and no new + // admissions on this executor's nodes this cycle. + s.observed[name] = false + s.log.Warn("executor list incomplete", "executor", name, "err", err) + continue + } + if err := s.reconcileExecutor(ctx, name, xs); err != nil { + return err + } + } + return nil +} + +func (s *Scheduler) reconcileExecutor(ctx context.Context, name string, xs []executor.Execution) error { + seen := map[int64]bool{} + unknown := map[string]bool{} + for _, x := range xs { + seen[x.JobID] = true + known, err := s.reconcileOne(ctx, name, x) + if err != nil { + return err + } + if !known && x.Cleanup != executor.CleanupConfirmed { + unknown[x.Node] = true + } + } + for node, ex := range s.nodeExec { + if ex == name { + s.unknownOn[node] = unknown[node] + } + } + + // Jobs this executor should have but did not report. + var active []*job + err := s.fenced(ctx, func(tx pgx.Tx) error { + rows, err := tx.Query(ctx, "SELECT "+jobColumns+jobFrom+` + JOIN nodes n ON n.name = j.node + WHERE n.executor = $1 AND j.state IN ('LAUNCHING', 'RUNNING', 'STOPPING', 'CLEANING')`, name) + if err != nil { + return err + } + active, err = pgx.CollectRows(rows, func(r pgx.CollectableRow) (*job, error) { return scanJob(r) }) + return err + }) + if err != nil { + return err + } + for _, j := range active { + if seen[j.ID] { + delete(s.missing, j.ID) + continue + } + if j.State == StateLaunching { + continue // dispatch sends or resends the same launch + } + // Require two consecutive complete lists before treating absence + // as real, so one racy read never ends a job. + s.missing[j.ID]++ + if s.missing[j.ID] < 2 { + continue + } + delete(s.missing, j.ID) + if err := s.resolveMissing(ctx, j); err != nil { + return err + } + } + return nil +} + +// reconcileOne turns one observed execution into the events it implies. +// It reports whether the execution belongs to a job the scheduler knows. +func (s *Scheduler) reconcileOne(ctx context.Context, exName string, x executor.Execution) (bool, error) { + var j *job + err := s.fenced(ctx, func(tx pgx.Tx) error { + var err error + j, err = scanJob(tx.QueryRow(ctx, "SELECT "+jobColumns+jobFrom+"WHERE j.id = $1", x.JobID)) + if errors.Is(err, pgx.ErrNoRows) { + j, err = nil, nil + } + return err + }) + if err != nil { + return false, err + } + if j == nil || j.ExecutionID == nil || *j.ExecutionID != x.ExecutionID || j.State == StatePending { + // Not ours to adopt: stop it and keep its node out of admission. + if x.Phase != executor.PhaseCleaned { + s.log.Warn("stopping unexpected execution", "job", x.JobID, "execution", x.ExecutionID, "node", x.Node) + if err := s.executors[exName].Terminate(ctx, x.JobID, 10*time.Second, "unexpected execution"); err != nil { + s.log.Warn("terminate unexpected execution", "err", err) + } + } + return false, nil + } + + base := executor.Event{JobID: x.JobID, ExecutionID: x.ExecutionID, Node: x.Node} + emit := func(kind executor.EventKind, at time.Time, r *executor.Result) error { + ev := base + ev.Kind, ev.At = kind, at + if r != nil { + ev.ExitCode, ev.Reason, ev.Detail = r.ExitCode, r.Reason, r.Detail + } + return s.observe(ctx, ev) + } + if x.StartedAt != nil { + if err := emit(executor.EventStarted, *x.StartedAt, nil); err != nil { + return true, err + } + } + if x.Result != nil && x.FinishedAt != nil { + kind := executor.EventFinished + if x.Result.Reason == executor.ExitLaunchError || x.Result.Reason == executor.ExitLaunchTimeout { + kind = executor.EventLaunchFailed + } + if err := emit(kind, *x.FinishedAt, x.Result); err != nil { + return true, err + } + } + switch { + case x.Cleanup == executor.CleanupConfirmed: + r := x.Result + if r == nil { + r = &executor.Result{Reason: executor.ExitNodeLost} + } + r2 := *r + r2.Detail = x.CleanupEvidence + return true, emit(executor.EventExited, s.clock.Now(), &r2) + case x.Phase == executor.PhaseLost: + return true, emit(executor.EventLost, s.clock.Now(), + &executor.Result{Reason: executor.ExitNodeLost, Detail: "node " + x.Node + " unreachable"}) + case j.State == StateStopping && (x.Phase == executor.PhaseRunning || x.Phase == executor.PhasePending): + return true, s.reissueStop(ctx, j) + case j.State.Terminal() && x.Phase != executor.PhaseFinished: + // Terminal in the database but still alive on the node. + return true, s.executors[exName].Terminate(ctx, j.ID, j.Grace, "job already ended") + } + return true, nil +} + +// reissueStop resends Terminate if the last one was answered but the +// execution is still alive well after its grace period. +func (s *Scheduler) reissueStop(ctx context.Context, j *job) error { + return s.fenced(ctx, func(tx pgx.Tx) error { + _, err := tx.Exec(ctx, ` + UPDATE op_intents SET status = 'pending', updated_at = $2 + WHERE job_id = $1 AND op = 'terminate' AND status = 'done' + AND updated_at + make_interval(secs => $3) <= $2`, + j.ID, s.clock.Now(), (j.Grace + s.stopRetry).Seconds()) + return err + }) +} + +// resolveMissing handles a started or stopping job whose execution is +// confirmed absent from complete observations. +func (s *Scheduler) resolveMissing(ctx context.Context, j *job) error { + return s.fenced(ctx, func(tx pgx.Tx) error { + cur, err := lockJob(ctx, tx, j.ID) + if err != nil || cur == nil { + return err + } + now := s.clock.Now() + if cur.State == StateStopping && cur.StartedAt == nil { + // The executor has no record of the create at all, including no + // create intent, so nothing was ever started. + st, err := intentStatus(ctx, tx, cur.ID, OpLaunch) + if err != nil { + return err + } + if st == OpPending || st == OpInFlight { + return finishUnsent(ctx, tx, cur, now, "executor has no record of the launch") + } + } + switch cur.State { + case StateRunning, StateStopping, StateCleaning: + return s.onLost(ctx, tx, cur, executor.Event{ + Detail: fmt.Sprintf("execution missing from node %s", *cur.Node), + }) + } + return nil + }) +} + +// enforceDeadlines stops RUNNING jobs at started_at + time limit, and +// LAUNCHING jobs that have not started by their launch deadline. +func (s *Scheduler) enforceDeadlines(ctx context.Context) error { + now := s.clock.Now() + var ids []int64 + var causes []StopCause + err := s.fenced(ctx, func(tx pgx.Tx) error { + rows, err := tx.Query(ctx, ` + SELECT id, CASE WHEN state = 'RUNNING' THEN 'deadline' ELSE 'launch_timeout' END FROM jobs + WHERE stop_cause IS NULL AND ( + (state = 'RUNNING' AND started_at + make_interval(secs => time_limit_seconds) <= $1) + OR (state = 'LAUNCHING' AND started_at IS NULL AND launch_deadline <= $1)) + ORDER BY id`, now) + if err != nil { + return err + } + for rows.Next() { + var id int64 + var c string + if err := rows.Scan(&id, &c); err != nil { + return err + } + ids, causes = append(ids, id), append(causes, StopCause(c)) + } + return rows.Err() + }) + if err != nil { + return err + } + for i, id := range ids { + err := s.fenced(ctx, func(tx pgx.Tx) error { + j, err := lockJob(ctx, tx, id) + if err != nil || j == nil { + return err + } + reason := fmt.Sprintf("time limit %s reached", j.TimeLimit) + if causes[i] == StopLaunchTimeout { + reason = "work did not start before the launch deadline" + } + return stop(ctx, tx, j, causes[i], "mist", reason, now) + }) + if err != nil { + return err + } + } + return nil +} + +// applyRevocations stops every nonterminal job affected by a durable +// Identity revocation, then advances the cursor in the same transaction. +func (s *Scheduler) applyRevocations(ctx context.Context) error { + err := s.fenced(ctx, func(tx pgx.Tx) error { + var cursor int64 + if err := tx.QueryRow(ctx, + "SELECT value FROM scheduler_cursors WHERE name = 'revocation_seq' FOR UPDATE").Scan(&cursor); err != nil { + return err + } + revs, last, err := s.identity.Revocations(ctx, tx, cursor) + if err != nil { + return err + } + now := s.clock.Now() + for _, r := range revs { + rows, err := tx.Query(ctx, ` + SELECT id FROM jobs WHERE username = $1 AND ($2 = '' OR team = $2) + AND state IN ('PENDING', 'LAUNCHING', 'RUNNING') ORDER BY id`, r.Username, r.Team) + if err != nil { + return err + } + ids, err := pgx.CollectRows(rows, pgx.RowTo[int64]) + if err != nil { + return err + } + for _, id := range ids { + j, err := lockJob(ctx, tx, id) + if err != nil { + return err + } + if err := stop(ctx, tx, j, StopRevoke, "identity", "access revoked", now); err != nil { + return err + } + } + } + _, err = tx.Exec(ctx, "UPDATE scheduler_cursors SET value = $1 WHERE name = 'revocation_seq'", last) + return err + }) + if errors.Is(err, identity.ErrUnavailable) { + return nil // try again next cycle; admission reports AUTH_UNAVAILABLE meanwhile + } + return err +} diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 65dcd63..7807702 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -4,25 +4,127 @@ package scheduler import ( "context" + "errors" "log/slog" "time" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/UTMIST/Mist/internal/clock" "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity" ) type Config struct { DatabaseURL string Pool *pgxpool.Pool - Clock clock.Clock - Log *slog.Logger + Identity identity.Admitter + // Executors by name, matching partitions.executor and nodes.executor. + Executors map[string]executor.Executor + Clock clock.Clock + Log *slog.Logger // Tick is the cycle period when nothing triggers one sooner. Tick time.Duration // LeaderCheck is how often leadership is confirmed. LeaderCheck time.Duration + // LaunchTimeout is admission to work start, including image pulls. + LaunchTimeout time.Duration + // CallTimeout bounds each executor call. + CallTimeout time.Duration +} + +// Scheduler is the active leader's state. Create it with New after +// acquiring leadership; tests drive Cycle and HandleEvent directly. +type Scheduler struct { + pool *pgxpool.Pool + identity identity.Admitter + executors map[string]executor.Executor + clock clock.Clock + log *slog.Logger + epoch int64 + launchTimeout time.Duration + callTimeout time.Duration + stopRetry time.Duration + lost <-chan struct{} + + // Per-cycle observations, rebuilt every cycle. + observed map[string]bool // executor -> Nodes and List complete + nodeExec map[string]string // node -> executor + unknownOn map[string]bool // node hosts an execution Mist did not admit + missing map[int64]int // job -> consecutive complete lists without it +} + +// New takes leadership: it bumps the epoch, which fences every older +// scheduler's writes. lost, if not nil, ends dispatch when closed. +func New(ctx context.Context, cfg Config, lost <-chan struct{}) (*Scheduler, error) { + if cfg.Clock == nil { + cfg.Clock = clock.Real{} + } + if cfg.Log == nil { + cfg.Log = slog.New(slog.DiscardHandler) + } + if cfg.LaunchTimeout == 0 { + cfg.LaunchTimeout = 10 * time.Minute + } + if cfg.CallTimeout == 0 { + cfg.CallTimeout = 30 * time.Second + } + epoch, err := takeLeadership(ctx, cfg.Pool) + if err != nil { + return nil, err + } + return &Scheduler{ + pool: cfg.Pool, + identity: cfg.Identity, + executors: cfg.Executors, + clock: cfg.Clock, + log: cfg.Log, + epoch: epoch, + launchTimeout: cfg.LaunchTimeout, + callTimeout: cfg.CallTimeout, + stopRetry: 30 * time.Second, + lost: lost, + observed: map[string]bool{}, + nodeExec: map[string]string{}, + unknownOn: map[string]bool{}, + missing: map[int64]int{}, + }, nil +} + +func (s *Scheduler) leaderCheck() error { + if s.lost == nil { + return nil + } + select { + case <-s.lost: + return ErrNotLeader + default: + return nil + } +} + +// Cycle runs one scheduling pass. Reconciliation comes first, so a new +// leader resolves outstanding operations before admitting anything. +func (s *Scheduler) Cycle(ctx context.Context) error { + for _, step := range []func(context.Context) error{ + s.refreshNodes, + s.applyRevocations, + s.reconcile, + s.enforceDeadlines, + s.admit, + s.dispatch, + } { + if err := s.leaderCheck(); err != nil { + return err + } + if err := step(ctx); err != nil { + return err + } + } + return nil } // Run waits for the scheduler lock, then runs cycles until ctx ends or @@ -50,9 +152,41 @@ func Run(ctx context.Context, cfg Config) error { defer leader.Release(context.WithoutCancel(ctx)) cfg.Log.Info("holding scheduler lock") + s, err := New(ctx, cfg, leader.Lost()) + if err != nil { + return err + } + + runCtx, cancel := context.WithCancel(ctx) + defer cancel() + wake := make(chan struct{}, 1) + trigger := func() { + select { + case wake <- struct{}{}: + default: + } + } + go listen(runCtx, cfg.DatabaseURL, trigger, cfg.Clock, cfg.Log) + events := make(chan executor.Event) + for _, ex := range cfg.Executors { + go func() { + for { + select { + case ev := <-ex.Events(): + select { + case events <- ev: + case <-runCtx.Done(): + return + } + case <-runCtx.Done(): + return + } + } + }() + } + + trigger() for { - // The scheduling cycle (M1 step 3) goes here: reconcile intents, - // then admit. It must check leader.Lost() between external calls. t := cfg.Clock.NewTimer(cfg.Tick) select { case <-ctx.Done(): @@ -61,7 +195,56 @@ func Run(ctx context.Context, cfg Config) error { case <-leader.Lost(): t.Stop() return leader.Err() + case ev := <-events: + t.Stop() + if err := s.HandleEvent(ctx, ev); err != nil { + if errors.Is(err, ErrNotLeader) { + return err + } + cfg.Log.Error("handle event", "event", ev.EventID, "err", err) + } + trigger() + continue + case <-wake: + case <-t.C(): + } + t.Stop() + if err := s.Cycle(ctx); err != nil { + if errors.Is(err, ErrNotLeader) { + return err + } + if ctx.Err() != nil { + return nil + } + cfg.Log.Error("scheduling cycle", "err", err) + } + } +} + +// listen wakes the leader when the API submits or cancels, through +// Postgres NOTIFY. It is only a latency optimization: the tick still runs. +func listen(ctx context.Context, url string, trigger func(), clk clock.Clock, log *slog.Logger) { + for ctx.Err() == nil { + conn, err := pgx.Connect(ctx, url) + if err == nil { + _, err = conn.Exec(ctx, "LISTEN "+notifyChannel) + for err == nil { + _, err = conn.WaitForNotification(ctx) + if err == nil { + trigger() + } + } + conn.Close(context.WithoutCancel(ctx)) + } + if ctx.Err() != nil { + return + } + log.Warn("scheduler notify listener reconnecting", "err", err) + t := clk.NewTimer(time.Second) + select { case <-t.C(): + case <-ctx.Done(): + t.Stop() } } } diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go new file mode 100644 index 0000000..3bfc334 --- /dev/null +++ b/internal/scheduler/scheduler_test.go @@ -0,0 +1,464 @@ +package scheduler + +import ( + "errors" + "testing" + "time" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/executor/fake" +) + +// Two 1-GPU jobs run side by side; a third waits with RESOURCES until a +// cleanup receipt, not process exit, frees a GPU. +func TestTwoGPUJobsThirdWaits(t *testing.T) { + h := newHarness(t) + a := h.submit("alice", "team-a", "gpu") + c := h.submit("carol", "team-b", "gpu") + b := h.submit("bob", "team-a", "gpu") + h.settle() + h.want(a, StateLaunching) + h.want(c, StateLaunching) + h.want(b, StatePending, PendingResources) + + h.advance(time.Second) + h.want(a, StateRunning) + h.advance(time.Minute) + h.want(a, StateCleaning) + if _, released, _ := h.allocation(a); released { + t.Fatal("allocation released on process exit") + } + h.want(b, StatePending, PendingResources) + + h.advance(time.Second) + h.want(a, StateCompleted) + h.want(c, StateCompleted) + h.want(b, StateLaunching) + if got := h.history(a); got != lifecycle { + t.Fatalf("history %s", got) + } +} + +func TestUserAndTeamLimits(t *testing.T) { + h := newHarness(t) + h.sql("INSERT INTO limits VALUES ('team', 'team-a', 'gpu', '*', 'gpus', 1)") + a1 := h.submit("alice", "team-a", "gpu") + a2 := h.submit("alice", "team-a", "gpu") + b1 := h.submit("bob", "team-a", "gpu") + h.settle() + h.want(a1, StateLaunching) + h.want(a2, StatePending, PendingUserLimit) + h.want(b1, StatePending, PendingTeamLimit) + + // A user override cannot bypass the team cap. + h.sql("INSERT INTO limits VALUES ('user', 'alice', 'gpu', '*', 'gpus', 2)") + h.settle() + h.want(a2, StatePending, PendingTeamLimit) +} + +func TestPartitionCapAndSharedNode(t *testing.T) { + h := newHarness(t) + eight := func(r *Request) { r.Resources.MilliCPU = 8000 } + j1 := h.submit("alice", "team-a", "cpu", eight) + j2 := h.submit("bob", "team-a", "cpu", eight) + j3 := h.submit("carol", "team-b", "cpu", eight) + h.settle() + h.want(j1, StateLaunching) + h.want(j2, StateLaunching) + h.want(j3, StatePending, PendingResources) // 16-CPU partition cap, node has 32 + + // A second partition on the same node draws from the same pool. + h.sql(`INSERT INTO partitions SELECT 'cpu2', executor, state, accelerator, runtime, min_accelerators, + max_accelerators, default_milli_cpu, default_memory_bytes, max_milli_cpu, max_memory_bytes, + default_scratch_bytes, max_scratch_bytes, default_time_limit_seconds, max_time_limit_seconds, NULL, NULL + FROM partitions WHERE name = 'cpu'`) + h.sql("INSERT INTO partition_nodes VALUES ('cpu2', $1)", cpuNode) + big := h.submit("carol", "team-b", "cpu2", func(r *Request) { r.Resources.MilliCPU = 16000 }) + small := h.submit("bob", "team-a", "cpu2", func(r *Request) { r.Resources.MilliCPU = 2000 }) + h.settle() + h.want(big, StateLaunching) + h.want(small, StatePending, PendingResources) +} + +func TestOutcomes(t *testing.T) { + cases := []struct { + name string + script fake.Script + after time.Duration + want State + }{ + {"completed", fake.DefaultScript, 2 * time.Minute, StateCompleted}, + {"failed", fake.Script{StartDelay: time.Second, RunFor: time.Minute, ExitCode: 3}, 2 * time.Minute, StateFailed}, + {"oom", fake.Script{StartDelay: time.Second, RunFor: time.Minute, ExitCode: 137, Reason: executor.ExitOutOfMemory}, + 2 * time.Minute, StateOutOfMemory}, + {"launch error", fake.Script{StartDelay: time.Second, LaunchError: "image pull failed"}, time.Minute, StateLaunchFailed}, + {"launch timeout", fake.Script{StartDelay: time.Hour}, 11 * time.Minute, StateLaunchFailed}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.exec.SetScript(id, c.script) + h.settle() + for elapsed := time.Duration(0); elapsed < c.after; elapsed += 30 * time.Second { + h.advance(30 * time.Second) + } + h.want(id, c.want) + if _, released, _ := h.allocation(id); !released { + t.Fatal("allocation not released after cleanup") + } + }) + } +} + +// The time limit runs from started_at, and the stop cause decides the +// terminal state even though the process exits by signal. +func TestTimeLimit(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.exec.SetScript(id, fake.Script{StartDelay: time.Minute, RunFor: 10 * time.Hour, StopDelay: time.Second, CleanupDelay: time.Second}) + h.settle() + h.advance(time.Minute) + h.want(id, StateRunning) + h.advance(time.Hour - time.Second) + h.want(id, StateRunning) // an hour from submit, but not from start + h.advance(time.Second) + h.want(id, StateStopping) + h.advance(time.Second) + h.want(id, StateCleaning) + h.advance(time.Second) + h.want(id, StateTimeout) +} + +func TestCancelInEachState(t *testing.T) { + t.Run("pending", func(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + if st, err := h.jobs.Cancel(h.ctx, id, "alice"); err != nil || st != StateCancelled { + t.Fatalf("cancel: %s %v", st, err) + } + h.settle() + h.want(id, StateCancelled) + if exists, _, _ := h.allocation(id); exists { + t.Fatal("cancelled pending job was allocated") + } + }) + + t.Run("launching, never sent", func(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + if err := h.s.refreshNodes(h.ctx); err != nil { + t.Fatal(err) + } + if err := h.s.admit(h.ctx); err != nil { // committed, not yet dispatched + t.Fatal(err) + } + h.want(id, StateLaunching) + h.jobs.Cancel(h.ctx, id, "alice") + h.want(id, StateCancelled) + if _, released, _ := h.allocation(id); !released { + t.Fatal("unsent launch kept its allocation") + } + h.settle() + if n := h.exec.Creates(id); n != 0 { + t.Fatalf("cancelled unsent launch created %d objects", n) + } + }) + + t.Run("launching, sent", func(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.jobs.Cancel(h.ctx, id, "alice") + h.want(id, StateStopping) + h.settle() + h.want(id, StateCleaning) + h.advance(time.Second) + h.want(id, StateCancelled) + }) + + t.Run("running", func(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.advance(time.Second) + h.jobs.Cancel(h.ctx, id, "alice") + h.jobs.Cancel(h.ctx, id, "alice") // repeat is harmless + h.settle() + h.advance(time.Second) + h.advance(time.Second) + h.want(id, StateCancelled) + }) + + t.Run("stopping keeps first cause", func(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.exec.SetScript(id, fake.Script{StartDelay: time.Second, RunFor: 10 * time.Hour, IgnoreSIGTERM: true, CleanupDelay: time.Second}) + h.settle() + h.advance(time.Second) + h.advance(time.Hour) + h.want(id, StateStopping) + if st, _ := h.jobs.Cancel(h.ctx, id, "alice"); st != StateStopping { + t.Fatalf("cancel during stop: %s", st) + } + h.advance(time.Minute) + h.advance(time.Second) + h.want(id, StateTimeout) + }) + + t.Run("cleaning and terminal", func(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.exec.SetScript(id, fake.Script{StartDelay: time.Second, RunFor: time.Minute, CleanupDelay: time.Minute}) + h.settle() + h.advance(time.Second) + h.advance(time.Minute) + h.want(id, StateCleaning) + h.jobs.Cancel(h.ctx, id, "alice") + h.want(id, StateCleaning) + h.advance(time.Minute) + h.want(id, StateCompleted) + before := h.history(id) + h.jobs.Cancel(h.ctx, id, "alice") + h.settle() + h.want(id, StateCompleted) + if h.history(id) != before { + t.Fatal("cancel changed a terminal job's history") + } + }) +} + +// A lost node keeps the allocation quarantined until recovery proves the +// execution stopped; repeated cleanup receipts change nothing. +func TestNodeFailRetainsAllocation(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.advance(time.Second) + h.exec.LoseNode(gpuNode, "kubelet NotReady") + h.settle() + h.want(id, StateNodeFail) + if _, released, quarantined := h.allocation(id); released || !quarantined { + t.Fatalf("lost allocation: released=%v quarantined=%v", released, quarantined) + } + if !h.job(id).CleanupPending { + t.Fatal("NODE_FAIL without cleanup_pending") + } + next := h.submit("carol", "team-b", "gpu") + h.advance(time.Hour) + h.want(next, StatePending, PendingPartitionDown) + if n := h.exec.Creates(id); n != 1 { + t.Fatalf("lost work recreated: %d objects", n) + } + + h.exec.RecoverNode(gpuNode) + h.settle() + if _, released, _ := h.allocation(id); !released { + t.Fatal("recovered node did not release the allocation") + } + h.want(id, StateNodeFail) + h.want(next, StateLaunching) + + before := h.history(id) + ev := executor.Event{EventID: "repeat", JobID: id, ExecutionID: *h.job(id).ExecutionID, + Kind: executor.EventExited, Detail: "again", At: h.clock.Now()} + if err := h.s.HandleEvent(h.ctx, ev); err != nil { + t.Fatal(err) + } + if h.history(id) != before { + t.Fatal("repeated cleanup receipt changed history") + } +} + +// An execution that vanishes from complete observations, with no receipt, +// is a loss after confirmation, never a reason to launch again. +func TestMissingExecutionIsNodeFail(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.advance(time.Second) + h.exec.Forget(id) + h.cycle() + h.want(id, StateRunning) // one absent list is not proof + h.cycle() + h.want(id, StateNodeFail) + h.settle() + if n := h.exec.Creates(id); n != 1 { + t.Fatalf("missing execution relaunched: %d", n) + } +} + +// A job can finish before Started is delivered; events arriving in reverse +// order still converge on the right outcome and real start time. +func TestFastExitOutOfOrder(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.exec.HoldEvents() + h.clock.Advance(time.Second) + h.clock.Advance(time.Minute) + h.clock.Advance(time.Second) + h.exec.ReleaseEvents(true) + h.drain() + h.want(id, StateCompleted) + if j := h.job(id); j.StartedAt == nil || !j.StartedAt.Equal(t0.Add(time.Second)) { + t.Fatalf("started_at %v", j.StartedAt) + } + if got := h.history(id); got != ">PENDING PENDING>LAUNCHING LAUNCHING>CLEANING CLEANING>COMPLETED" { + t.Fatalf("history %s", got) + } + h.settle() + h.want(id, StateCompleted) +} + +func TestDuplicateEvents(t *testing.T) { + h := newHarness(t) + h.exec.DuplicateEvents(true) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.advance(time.Second) + h.advance(time.Minute) + h.advance(time.Second) + if got := h.history(id); got != lifecycle { + t.Fatalf("history %s", got) + } +} + +// Lost notifications are recovered from List; a restarted scheduler picks +// up where the old one stopped. +func TestDroppedEventsAndRestart(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.exec.HoldEvents() + h.clock.Advance(time.Second) + h.exec.DropHeldEvents() + h.exec.HoldEvents() + h.s = h.newScheduler() + h.cycle() + h.want(id, StateRunning) + h.clock.Advance(time.Minute + time.Second) + h.exec.Restart() + h.cycle() + h.want(id, StateCompleted) +} + +func TestAmbiguousLaunchRetriesSameExecution(t *testing.T) { + for _, fault := range []fake.LaunchFault{fake.FaultTimeoutBeforeCreate, fake.FaultTimeoutAfterCreate, fake.FaultLateCreate} { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + h.exec.FailNextLaunch(id, fault, 500*time.Millisecond) + h.cycle() + if st, last := h.intent(id, OpLaunch); st != OpInFlight || last == "" { + t.Fatalf("fault %d: intent %s %q", fault, st, last) + } + h.settle() // retries the same ExecutionID + h.advance(time.Second) + h.advance(time.Minute) + h.advance(time.Second) + h.want(id, StateCompleted) + if n := h.exec.Creates(id); n != 1 { + t.Fatalf("fault %d: %d objects", fault, n) + } + } +} + +// A scheduler that lost leadership cannot commit anything; the new one +// dispatches the admission the old one committed, with the same identity. +func TestSplitLeadershipAndCrashBeforeDispatch(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + old := h.s + if err := old.refreshNodes(h.ctx); err != nil { + t.Fatal(err) + } + if err := old.admit(h.ctx); err != nil { + t.Fatal(err) + } + execID := *h.job(id).ExecutionID + + h.s = h.newScheduler() + if err := old.Cycle(h.ctx); !errors.Is(err, ErrNotLeader) { + t.Fatalf("old leader cycle: %v", err) + } + other := h.submit("carol", "team-b", "gpu") + if err := old.admit(h.ctx); !errors.Is(err, ErrNotLeader) { + t.Fatalf("old leader admit: %v", err) + } + h.want(other, StatePending) + + h.settle() + h.advance(time.Second) + if j := h.job(id); *j.ExecutionID != execID || j.State != StateRunning { + t.Fatalf("after takeover: %s %s", *j.ExecutionID, j.State) + } + if n := h.exec.Creates(id); n != 1 { + t.Fatalf("%d objects", n) + } +} + +func TestRevocation(t *testing.T) { + h := newHarness(t) + running := h.submit("alice", "team-a", "gpu") + otherTeam := h.submit("alice", "team-b", "cpu") + h.settle() + h.advance(time.Second) + queued := h.submit("alice", "team-a", "cpu") + h.id.RemoveMember("alice", "team-a") + h.settle() + h.want(queued, StateCancelled) + h.want(running, StateStopping) + h.want(otherTeam, StateRunning) + h.advance(time.Second) + h.advance(time.Second) + h.want(running, StateCancelled) + h.want(otherTeam, StateRunning) +} + +func TestAuthorizationNotReady(t *testing.T) { + h := newHarness(t) + h.id.SetProvisioned("team-a", false) + id := h.submit("alice", "team-a", "gpu") + h.settle() + h.want(id, StatePending, PendingProvisioning) + h.id.SetProvisioned("team-a", true) + h.id.Unavailable = true + h.settle() + h.want(id, StatePending, PendingAuthUnavailable) + h.id.Unavailable = false + h.settle() + h.want(id, StateLaunching) +} + +// An unstarted launch is never replayed under an older authorization. +func TestStaleGenerationCancelsUnsentLaunch(t *testing.T) { + h := newHarness(t) + id := h.submit("alice", "team-a", "gpu") + if err := h.s.refreshNodes(h.ctx); err != nil { + t.Fatal(err) + } + if err := h.s.admit(h.ctx); err != nil { + t.Fatal(err) + } + h.id.AddMember("bob", "team-b", "member") // bumps the generation + if err := h.s.dispatch(h.ctx); err != nil { + t.Fatal(err) + } + h.want(id, StateCancelled) + if n := h.exec.Creates(id); n != 0 { + t.Fatalf("stale launch sent: %d", n) + } +} + +func TestSubmitQueueLimit(t *testing.T) { + h := newHarness(t) + h.sql("UPDATE limits SET max = 2 WHERE resource = 'submitted_jobs'") + h.submit("alice", "team-a", "gpu") + h.submit("alice", "team-a", "gpu") + _, err := h.jobs.Submit(h.ctx, Request{Username: "alice", Team: "team-a", Partition: "gpu", QOS: "normal", + Kind: executor.KindBatch, Script: "true", TimeLimit: time.Hour, SignalGrace: time.Minute, + Resources: executor.Resources{MilliCPU: 1000, MemoryBytes: 1 << 30, GPUs: 1, ScratchBytes: 1 << 30}}) + if !errors.Is(err, ErrLimitExceeded) { + t.Fatalf("want ErrLimitExceeded, got %v", err) + } +} diff --git a/internal/scheduler/store.go b/internal/scheduler/store.go new file mode 100644 index 0000000..2442df5 --- /dev/null +++ b/internal/scheduler/store.go @@ -0,0 +1,284 @@ +package scheduler + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/UTMIST/Mist/internal/executor" +) + +// ErrNotLeader is returned when a write is attempted by a scheduler whose +// leadership epoch is no longer current. +var ErrNotLeader = errors.New("scheduler: not the current leader") + +// takeLeadership bumps the epoch. Any transaction still running for an +// older leader holds the epoch row FOR SHARE, so this waits for it, and +// every later fenced transaction from that leader fails. +func takeLeadership(ctx context.Context, pool *pgxpool.Pool) (int64, error) { + var epoch int64 + err := pool.QueryRow(ctx, + "UPDATE scheduler_cursors SET value = value + 1 WHERE name = 'leader_epoch' RETURNING value").Scan(&epoch) + return epoch, err +} + +// fenced runs fn in a transaction that commits only while s is the current +// leader. Every scheduler-side write goes through it. The epoch row is +// share-locked first; it is outside the account→…→node order because only +// takeLeadership ever locks it exclusively, and that takes nothing else. +func (s *Scheduler) fenced(ctx context.Context, fn func(pgx.Tx) error) error { + return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { + var epoch int64 + if err := tx.QueryRow(ctx, + "SELECT value FROM scheduler_cursors WHERE name = 'leader_epoch' FOR SHARE").Scan(&epoch); err != nil { + return err + } + if epoch != s.epoch { + return ErrNotLeader + } + return fn(tx) + }) +} + +// job is one jobs row. +type job struct { + ID int64 + Username, Team string + Partition, QOS string + Name, Kind string + State State + PendingReason *string + Held bool + Res executor.Resources + TimeLimit, Grace time.Duration + Node *string + ExecutionID *string + SpecHash *string + AuthGeneration *int64 + AdmittedAt *time.Time + LaunchDeadline *time.Time + StopCause *string + ExitCode *int + ExitReason *string + ExitDetail string + SubmittedAt time.Time + StartedAt *time.Time + FinishedAt *time.Time + EndedAt *time.Time + CleanupPending bool + QOSPriority int +} + +const jobColumns = `j.id, j.username, j.team, j.partition, j.qos, j.name, j.kind, j.state, j.pending_reason, j.held, + j.milli_cpu, j.memory_bytes, j.gpus, j.tt_boards, j.scratch_bytes, j.time_limit_seconds, j.signal_grace_seconds, + j.node, j.execution_id::text, j.spec_hash, j.auth_generation, j.admitted_at, j.launch_deadline, j.stop_cause, + j.exit_code, j.exit_reason, j.exit_detail, j.submitted_at, j.started_at, j.finished_at, j.ended_at, + j.cleanup_pending, q.priority` + +const jobFrom = ` FROM jobs j JOIN qos q ON q.name = j.qos ` + +func scanJob(row pgx.Row) (*job, error) { + var j job + var limit, grace int64 + var state string + err := row.Scan(&j.ID, &j.Username, &j.Team, &j.Partition, &j.QOS, &j.Name, &j.Kind, &state, &j.PendingReason, &j.Held, + &j.Res.MilliCPU, &j.Res.MemoryBytes, &j.Res.GPUs, &j.Res.TTBoards, &j.Res.ScratchBytes, &limit, &grace, + &j.Node, &j.ExecutionID, &j.SpecHash, &j.AuthGeneration, &j.AdmittedAt, &j.LaunchDeadline, &j.StopCause, + &j.ExitCode, &j.ExitReason, &j.ExitDetail, &j.SubmittedAt, &j.StartedAt, &j.FinishedAt, &j.EndedAt, + &j.CleanupPending, &j.QOSPriority) + if err != nil { + return nil, err + } + j.State = State(state) + j.TimeLimit = time.Duration(limit) * time.Second + j.Grace = time.Duration(grace) * time.Second + return &j, nil +} + +// lockJob loads and row-locks a job. It returns (nil, nil) if there is none. +func lockJob(ctx context.Context, tx pgx.Tx, id int64) (*job, error) { + j, err := scanJob(tx.QueryRow(ctx, "SELECT "+jobColumns+jobFrom+"WHERE j.id = $1 FOR UPDATE OF j", id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil + } + return j, err +} + +// transition moves a locked job to state `to` and records the event. Every +// state change goes through here, so job_events is the complete history. +func transition(ctx context.Context, tx pgx.Tx, j *job, to State, at time.Time, reason, actor string) error { + if _, err := tx.Exec(ctx, "UPDATE jobs SET state = $2 WHERE id = $1", j.ID, string(to)); err != nil { + return err + } + if _, err := tx.Exec(ctx, + "INSERT INTO job_events (job_id, at, from_state, to_state, reason, actor) VALUES ($1, $2, $3, $4, $5, $6)", + j.ID, at, string(j.State), string(to), reason, actor); err != nil { + return err + } + j.State = to + return nil +} + +// releaseAllocation frees a job's capacity. Callers must hold proof that +// nothing of the execution remains on the node. +func releaseAllocation(ctx context.Context, tx pgx.Tx, jobID int64, at time.Time, evidence string) error { + _, err := tx.Exec(ctx, + "UPDATE allocations SET released_at = $2, release_evidence = $3 WHERE job_id = $1 AND released_at IS NULL", + jobID, at, evidence) + return err +} + +func setIntent(ctx context.Context, tx pgx.Tx, jobID int64, op OpKind, status OpStatus, at time.Time, lastErr string) error { + _, err := tx.Exec(ctx, + "UPDATE op_intents SET status = $3, updated_at = $4, last_error = $5 WHERE job_id = $1 AND op = $2", + jobID, string(op), string(status), at, lastErr) + return err +} + +// closeOpenIntents marks every unfinished intent of a job done: the +// execution reached an outcome, so there is nothing left to send. +func closeOpenIntents(ctx context.Context, tx pgx.Tx, jobID int64, at time.Time, note string) error { + _, err := tx.Exec(ctx, + "UPDATE op_intents SET status = 'done', updated_at = $2, last_error = $3 WHERE job_id = $1 AND status IN ('pending', 'in_flight')", + jobID, at, note) + return err +} + +func intentStatus(ctx context.Context, tx pgx.Tx, jobID int64, op OpKind) (OpStatus, error) { + var st string + err := tx.QueryRow(ctx, "SELECT status FROM op_intents WHERE job_id = $1 AND op = $2", jobID, string(op)).Scan(&st) + if errors.Is(err, pgx.ErrNoRows) { + return "", nil + } + return OpStatus(st), err +} + +// stop records a stop request for a locked job. The first committed cause +// wins; once a finish result is committed (CLEANING or terminal) nothing +// changes. It is used by API cancellation and by the scheduler itself +// (deadline, revocation, launch timeout), so it must not depend on being +// the leader. +func stop(ctx context.Context, tx pgx.Tx, j *job, cause StopCause, actor, reason string, now time.Time) error { + switch j.State { + case StatePending: + if cause != StopCancel && cause != StopRevoke { + return nil + } + if _, err := tx.Exec(ctx, + "UPDATE jobs SET stop_cause = $2, stop_actor = $3, stop_requested_at = $4, ended_at = $4, pending_reason = NULL WHERE id = $1", + j.ID, string(cause), actor, now); err != nil { + return err + } + return transition(ctx, tx, j, StateCancelled, now, reason, actor) + + case StateLaunching, StateRunning: + if j.StopCause != nil { + return nil + } + if _, err := tx.Exec(ctx, + "UPDATE jobs SET stop_cause = $2, stop_actor = $3, stop_requested_at = $4 WHERE id = $1", + j.ID, string(cause), actor, now); err != nil { + return err + } + c := string(cause) + j.StopCause = &c + if err := transition(ctx, tx, j, StateStopping, now, reason, actor); err != nil { + return err + } + + // A launch intent still 'pending' was never sent, so nothing can + // exist on the node: finish now and release the allocation. + if st, err := intentStatus(ctx, tx, j.ID, OpLaunch); err != nil { + return err + } else if st == OpPending { + return finishUnsent(ctx, tx, j, now, "launch was never sent; nothing to clean up") + } + + _, err := tx.Exec(ctx, ` + INSERT INTO op_intents (job_id, execution_id, op, status, reason, created_at, updated_at) + VALUES ($1, $2, 'terminate', 'pending', $3, $4, $4) + ON CONFLICT (job_id, op) DO UPDATE SET status = 'pending', reason = $3, updated_at = $4`, + j.ID, *j.ExecutionID, reason, now) + return err + } + return nil // STOPPING keeps its first cause; CLEANING and terminal already have an outcome +} + +// finishUnsent ends a STOPPING job whose execution provably never reached +// the executor. +func finishUnsent(ctx context.Context, tx pgx.Tx, j *job, now time.Time, why string) error { + if err := closeOpenIntents(ctx, tx, j.ID, now, why); err != nil { + return err + } + if _, err := tx.Exec(ctx, + "UPDATE jobs SET finished_at = $2, exit_detail = $3 WHERE id = $1", j.ID, now, why); err != nil { + return err + } + if err := transition(ctx, tx, j, StateCleaning, now, why, "mist"); err != nil { + return err + } + if err := releaseAllocation(ctx, tx, j.ID, now, why); err != nil { + return err + } + return endJob(ctx, tx, j, decide(j), now, "mist") +} + +// endJob moves a CLEANING job to its terminal state. +func endJob(ctx context.Context, tx pgx.Tx, j *job, to State, now time.Time, actor string) error { + if _, err := tx.Exec(ctx, + "UPDATE jobs SET ended_at = $2, cleanup_pending = false WHERE id = $1", j.ID, now); err != nil { + return err + } + return transition(ctx, tx, j, to, now, exitSummary(j, to), actor) +} + +// decide picks the terminal state for a job whose cleanup is confirmed. +// A committed stop cause outranks the process result, because stop() +// refuses to set one after a finish result is committed. +func decide(j *job) State { + reason := "" + if j.ExitReason != nil { + reason = *j.ExitReason + } + cause := "" + if j.StopCause != nil { + cause = *j.StopCause + } + switch { + case cause == string(StopLaunchTimeout): + return StateLaunchFailed + case cause == string(StopCancel), cause == string(StopRevoke): + return StateCancelled + case cause == string(StopDeadline): + return StateTimeout + case reason == string(executor.ExitLaunchError), reason == string(executor.ExitLaunchTimeout): + return StateLaunchFailed + case reason == string(executor.ExitOutOfMemory): + return StateOutOfMemory + case reason == string(executor.ExitDeadline): + return StateTimeout + case reason == string(executor.ExitNormal) && j.ExitCode != nil && *j.ExitCode == 0: + return StateCompleted + } + return StateFailed +} + +func exitSummary(j *job, to State) string { + switch to { + case StateTimeout: + return fmt.Sprintf("time limit %s reached", j.TimeLimit) + case StateOutOfMemory: + return fmt.Sprintf("exceeded %d MiB memory limit", j.Res.MemoryBytes>>20) + case StateCompleted: + return "exit 0" + case StateFailed: + if j.ExitCode != nil { + return fmt.Sprintf("exit %d %s", *j.ExitCode, j.ExitDetail) + } + } + return j.ExitDetail +} diff --git a/internal/scheduler/submit.go b/internal/scheduler/submit.go new file mode 100644 index 0000000..690fcbe --- /dev/null +++ b/internal/scheduler/submit.go @@ -0,0 +1,139 @@ +package scheduler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/executor" +) + +// Request is a normalized submission. Filling defaults and validating +// paths, sizes and authorization happen in the API (M1 step 6); the +// scheduler stores what it is given and freezes it at admission. +type Request struct { + Username string `json:"username"` + Team string `json:"team"` + Partition string `json:"partition"` + QOS string `json:"qos"` + Name string `json:"name"` + Kind executor.Kind `json:"kind"` + Resources executor.Resources `json:"resources"` + TimeLimit time.Duration `json:"time_limit"` + SignalGrace time.Duration `json:"signal_grace"` + Workdir string `json:"workdir"` + OutputPath string `json:"output_path"` + Script string `json:"script"` + Args []string `json:"args"` + Argv []string `json:"argv"` + Env map[string]string `json:"env"` +} + +// ErrLimitExceeded is a submit-time limit (queue length) refusal. +var ErrLimitExceeded = errors.New("scheduler: limit exceeded") + +// ErrNotFound: no such job. +var ErrNotFound = errors.New("scheduler: job not found") + +// Jobs is the API-side handle: submit and cancel. It does not need +// leadership; row locks serialize it with the scheduler. +type Jobs struct { + Pool *pgxpool.Pool + Clock clock.Clock +} + +// notifyChannel wakes the leader for an early cycle. +const notifyChannel = "mist_scheduler" + +// Submit inserts a PENDING job and returns its ID. +func (js Jobs) Submit(ctx context.Context, r Request) (int64, error) { + if r.Resources.MilliCPU <= 0 || r.Resources.MemoryBytes <= 0 || r.Resources.ScratchBytes <= 0 || + r.TimeLimit < time.Second || r.SignalGrace < executor.MinSignalGrace || r.SignalGrace > executor.MaxSignalGrace { + return 0, fmt.Errorf("scheduler: request is not normalized") + } + body, err := json.Marshal(r) + if err != nil { + return 0, err + } + now := js.Clock.Now() + var id int64 + err = pgx.BeginFunc(ctx, js.Pool, func(tx pgx.Tx) error { + // Serialize one user's submissions so the queue-length check holds. + if _, err := tx.Exec(ctx, "SELECT pg_advisory_xact_lock(hashtext('submit:' || $1))", r.Username); err != nil { + return err + } + var max *int64 + err := tx.QueryRow(ctx, ` + SELECT max FROM limits + WHERE scope = 'user' AND resource = 'submitted_jobs' + AND subject IN ($1, '*') AND partition IN ($2, '*') AND qos IN ($3, '*') + ORDER BY (subject <> '*') DESC, (partition <> '*') DESC, (qos <> '*') DESC LIMIT 1`, + r.Username, r.Partition, r.QOS).Scan(&max) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return err + } + if max != nil { + var n int64 + if err := tx.QueryRow(ctx, ` + SELECT count(*) FROM jobs + WHERE username = $1 AND state IN ('PENDING', 'LAUNCHING', 'RUNNING', 'STOPPING', 'CLEANING')`, + r.Username).Scan(&n); err != nil { + return err + } + if n >= *max { + return fmt.Errorf("%w: %s already has %d queued or running jobs", ErrLimitExceeded, r.Username, n) + } + } + res := r.Resources + if err := tx.QueryRow(ctx, ` + INSERT INTO jobs (username, team, partition, qos, name, kind, state, pending_reason, + milli_cpu, memory_bytes, gpus, tt_boards, scratch_bytes, time_limit_seconds, signal_grace_seconds, submitted_at) + VALUES ($1, $2, $3, $4, $5, $6, 'PENDING', 'PRIORITY', $7, $8, $9, $10, $11, $12, $13, $14) + RETURNING id`, + r.Username, r.Team, r.Partition, r.QOS, r.Name, string(r.Kind), + res.MilliCPU, res.MemoryBytes, res.GPUs, res.TTBoards, res.ScratchBytes, + int64(r.TimeLimit/time.Second), int64(r.SignalGrace/time.Second), now).Scan(&id); err != nil { + return err + } + if _, err := tx.Exec(ctx, "INSERT INTO job_specs (job_id, request) VALUES ($1, $2)", id, body); err != nil { + return err + } + if _, err := tx.Exec(ctx, + "INSERT INTO job_events (job_id, at, to_state, reason, actor) VALUES ($1, $2, 'PENDING', 'submitted', $3)", + id, now, r.Username); err != nil { + return err + } + _, err = tx.Exec(ctx, "SELECT pg_notify($1, '')", notifyChannel) + return err + }) + return id, err +} + +// Cancel records a stop request and returns the job's state afterwards. +// It never waits for termination and never frees an allocation that might +// still be in use. Repeating it is harmless. +func (js Jobs) Cancel(ctx context.Context, jobID int64, actor string) (State, error) { + var state State + err := pgx.BeginFunc(ctx, js.Pool, func(tx pgx.Tx) error { + j, err := lockJob(ctx, tx, jobID) + if err != nil { + return err + } + if j == nil { + return ErrNotFound + } + if err := stop(ctx, tx, j, StopCancel, actor, "cancelled by "+actor, js.Clock.Now()); err != nil { + return err + } + state = j.State + _, err = tx.Exec(ctx, "SELECT pg_notify($1, '')", notifyChannel) + return err + }) + return state, err +} From a598d3c981a8eaa4853dbbe6644935071244ce11 Mon Sep 17 00:00:00 2001 From: Brandon Li Date: Fri, 30 Oct 2026 09:44:22 -0400 Subject: [PATCH 3/5] Add mist-job-entry (M1 step 4) The in-job wrapper: umask 0027, symlink-safe private dot-folders, output opened with openat2(RESOLVE_BENEATH) under the team root, header and footer, output copied to the file and stdout, SIGTERM forwarding to the job's process group, a local deadline timer with SIGKILL after grace, OOM detection from the cgroup's memory.events, and a structured termination message. Setup failures exit 125 without running the work. Process tests re-execute the test binary as the wrapper; the OOM test runs inside a 64 MiB systemd user scope. Co-Authored-By: Claude Opus 5.5 --- cmd/mist-job-entry/main.go | 31 +++ cmd/mist-job-entry/main_test.go | 305 ++++++++++++++++++++++ docs/agents/architecture/executors.md | 2 + docs/agents/plans/jobs-and-scheduling.md | 17 ++ go.mod | 1 + go.sum | 2 + internal/jobentry/env.go | 88 +++++++ internal/jobentry/fs.go | 81 ++++++ internal/jobentry/jobentry_test.go | 179 +++++++++++++ internal/jobentry/run.go | 306 +++++++++++++++++++++++ internal/jobentry/spec.go | 98 ++++++++ 11 files changed, 1110 insertions(+) create mode 100644 cmd/mist-job-entry/main.go create mode 100644 cmd/mist-job-entry/main_test.go create mode 100644 internal/jobentry/env.go create mode 100644 internal/jobentry/fs.go create mode 100644 internal/jobentry/jobentry_test.go create mode 100644 internal/jobentry/run.go create mode 100644 internal/jobentry/spec.go diff --git a/cmd/mist-job-entry/main.go b/cmd/mist-job-entry/main.go new file mode 100644 index 0000000..6712eb3 --- /dev/null +++ b/cmd/mist-job-entry/main.go @@ -0,0 +1,31 @@ +// Command mist-job-entry is the first process of every Mist job. See +// internal/jobentry and docs/agents/architecture/executors.md. +// +// mist-job-entry [-spec /etc/mist/job.json] +package main + +import ( + "flag" + "fmt" + "os" + "os/signal" + "syscall" + + "github.com/UTMIST/Mist/internal/jobentry" +) + +func main() { + specPath := flag.String("spec", "/etc/mist/job.json", "job description written by the executor") + flag.Parse() + // Handle signals before anything else: as PID 1 in a container, + // unhandled SIGTERM would be ignored. + sigs := make(chan os.Signal, 4) + signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP) + + spec, err := jobentry.LoadSpec(*specPath) + if err != nil { + fmt.Fprintln(os.Stdout, "[mist] invalid job description:", err) + os.Exit(jobentry.ExitSetupFailed) + } + os.Exit(jobentry.Run(spec, os.Stdout, sigs)) +} diff --git a/cmd/mist-job-entry/main_test.go b/cmd/mist-job-entry/main_test.go new file mode 100644 index 0000000..bbbc57e --- /dev/null +++ b/cmd/mist-job-entry/main_test.go @@ -0,0 +1,305 @@ +package main + +import ( + "bytes" + "encoding/json" + "errors" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync" + "syscall" + "testing" + "time" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/jobentry" +) + +// The test binary re-executes itself as mist-job-entry, so these tests run +// the real process: real signals, real exit statuses, real cgroups. +func TestMain(m *testing.M) { + if os.Getenv("MIST_JOB_ENTRY_TEST_CHILD") == "1" { + os.Args = append(os.Args[:1], os.Args[2:]...) // drop -test.run + main() + return + } + os.Exit(m.Run()) +} + +type job struct { + t *testing.T + root string + spec jobentry.Spec + stdout *lockedBuffer + cmd *exec.Cmd +} + +type lockedBuffer struct { + mu sync.Mutex + b bytes.Buffer +} + +func (l *lockedBuffer) Write(p []byte) (int, error) { + l.mu.Lock() + defer l.mu.Unlock() + return l.b.Write(p) +} + +func (l *lockedBuffer) String() string { + l.mu.Lock() + defer l.mu.Unlock() + return l.b.String() +} + +func newJob(t *testing.T, script string) *job { + root := t.TempDir() + team := filepath.Join(root, "mist", "teams", "team-a") + home := filepath.Join(team, "alice") + for _, d := range []string{home, filepath.Join(root, "scratch")} { + if err := os.MkdirAll(d, 0o750); err != nil { + t.Fatal(err) + } + } + os.WriteFile(filepath.Join(root, "script.sh"), []byte(script), 0o644) + return &job{t: t, root: root, stdout: &lockedBuffer{}, spec: jobentry.Spec{ + JobID: 412, Name: "test", User: "alice", Team: "team-a", Partition: "cpu", QOS: "normal", Node: "local", + Kind: executor.KindBatch, + Resources: executor.Resources{MilliCPU: 2000, MemoryBytes: 64 << 20, ScratchBytes: 1 << 30}, + TimeLimit: 60, + SignalGrace: 5, + TeamRoot: team, + Home: home, + Workdir: home, + OutputPath: filepath.Join(home, "mist-412.out"), + Scratch: filepath.Join(root, "scratch"), + ScriptPath: filepath.Join(root, "script.sh"), + TerminationLog: filepath.Join(root, "termination-log"), + }} +} + +// start runs mist-job-entry, optionally under a wrapper such as systemd-run. +func (j *job) start(wrapper ...string) { + j.t.Helper() + b, _ := json.Marshal(j.spec) + specFile := filepath.Join(j.root, "job.json") + os.WriteFile(specFile, b, 0o644) + self, _ := os.Executable() + argv := append(wrapper, self, "-test.run=none", "-spec", specFile) + j.cmd = exec.Command(argv[0], argv[1:]...) + j.cmd.Env = append(os.Environ(), "MIST_JOB_ENTRY_TEST_CHILD=1") + j.cmd.Stdout, j.cmd.Stderr = j.stdout, j.stdout + if err := j.cmd.Start(); err != nil { + j.t.Fatal(err) + } +} + +func (j *job) wait() int { + j.t.Helper() + done := make(chan error, 1) + go func() { done <- j.cmd.Wait() }() + select { + case err := <-done: + var ee *exec.ExitError + if errors.As(err, &ee) { + return ee.ExitCode() + } + return 0 + case <-time.After(30 * time.Second): + j.cmd.Process.Kill() + j.t.Fatalf("job did not finish; output:\n%s", j.stdout.String()) + return -1 + } +} + +func (j *job) output() string { + b, _ := os.ReadFile(j.spec.OutputPath) + return string(b) +} + +func (j *job) result() *executor.Result { + b, err := os.ReadFile(j.spec.TerminationLog) + if err != nil { + return nil + } + var r executor.Result + if err := json.Unmarshal(b, &r); err != nil { + j.t.Fatal(err) + } + return &r +} + +func (j *job) waitForOutput(s string) { + j.t.Helper() + deadline := time.Now().Add(10 * time.Second) + for !strings.Contains(j.output(), s) { + if time.Now().After(deadline) { + j.t.Fatalf("output never contained %q:\n%s", s, j.output()) + } + time.Sleep(20 * time.Millisecond) + } +} + +func TestBatchArgsAndOutput(t *testing.T) { + j := newJob(t, `echo "argc=$#"; for a in "$@"; do echo "arg=[$a]"; done; echo err >&2; pwd; umask; echo "home=$HOME"`) + j.spec.Args = []string{"two words", "$(id)", "; rm -rf /"} + j.start() + if code := j.wait(); code != 0 { + t.Fatalf("exit %d:\n%s", code, j.stdout.String()) + } + out := j.output() + for _, want := range []string{ + "[mist] job 412 (test) started", "argc=3", "arg=[two words]", "arg=[$(id)]", "arg=[; rm -rf /]", + "err\n", j.spec.Workdir + "\n", "0027\n", "home=" + j.spec.Home, + "[mist] job 412 ended", "[mist] exit code 0", + } { + if !strings.Contains(out, want) { + t.Errorf("output missing %q", want) + } + } + if !strings.Contains(j.stdout.String(), "argc=3") { + t.Error("output not copied to stdout") + } + if st, _ := os.Stat(j.spec.OutputPath); st.Mode().Perm() != 0o640 { + t.Errorf("output mode %v", st.Mode().Perm()) + } + if st, _ := os.Stat(filepath.Join(j.spec.Home, ".ssh")); st == nil || st.Mode().Perm() != 0o700 { + t.Error(".ssh not created private") + } + if r := j.result(); r == nil || *r.ExitCode != 0 || r.Reason != executor.ExitNormal || r.OutputIncomplete { + t.Errorf("result %+v", r) + } +} + +func TestCommandAndExitCode(t *testing.T) { + j := newJob(t, "") + j.spec.Kind, j.spec.ScriptPath = executor.KindCommand, "" + j.spec.Argv = []string{"sh", "-c", "echo running; exit 3"} + j.start() + if code := j.wait(); code != 3 { + t.Fatalf("exit %d", code) + } + if r := j.result(); *r.ExitCode != 3 || r.Reason != executor.ExitNormal { + t.Fatalf("result %+v", r) + } + if !strings.Contains(j.output(), "running") { + t.Fatal("command output missing") + } +} + +func TestSIGTERMIsForwarded(t *testing.T) { + j := newJob(t, `trap 'echo caught; exit 7' TERM; echo ready; sleep 30 & wait`) + j.start() + j.waitForOutput("ready") + j.cmd.Process.Signal(syscall.SIGTERM) + if code := j.wait(); code != 7 { + t.Fatalf("exit %d:\n%s", code, j.output()) + } + if !strings.Contains(j.output(), "caught") || !strings.Contains(j.output(), "stopped by request") { + t.Fatalf("output:\n%s", j.output()) + } + if r := j.result(); r.Reason != executor.ExitSignaled { + t.Fatalf("reason %s", r.Reason) + } +} + +func TestLocalTimeLimit(t *testing.T) { + t.Run("SIGTERM", func(t *testing.T) { + j := newJob(t, "sleep 30") + j.spec.TimeLimit, j.spec.SignalGrace = 1, 5 + j.start() + start := time.Now() + if code := j.wait(); code != 143 { + t.Fatalf("exit %d", code) + } + if el := time.Since(start); el > 4*time.Second { + t.Fatalf("took %s", el) + } + if r := j.result(); r.Reason != executor.ExitDeadline { + t.Fatalf("reason %s", r.Reason) + } + }) + t.Run("SIGKILL after grace", func(t *testing.T) { + j := newJob(t, `trap '' TERM; echo ready; while :; do sleep 0.1; done`) + j.spec.TimeLimit, j.spec.SignalGrace = 1, 1 + j.start() + if code := j.wait(); code != 137 { + t.Fatalf("exit %d:\n%s", code, j.output()) + } + if !strings.Contains(j.output(), "sent SIGKILL") || j.result().Reason != executor.ExitDeadline { + t.Fatalf("output:\n%s", j.output()) + } + }) +} + +// If the wrapper itself is killed there is no footer and no result: the +// executor must use its own observations. Nothing fabricates success. +func TestWrapperKilledLeavesNoFooter(t *testing.T) { + j := newJob(t, `echo "pid=$$"; echo working; exec sleep 30`) + j.start() + j.waitForOutput("working") + j.cmd.Process.Kill() + j.wait() + out := j.output() + if strings.Contains(out, "ended") || strings.Contains(out, "exit code") { + t.Fatalf("footer written after wrapper SIGKILL:\n%s", out) + } + if j.result() != nil { + t.Fatal("termination message written after wrapper SIGKILL") + } + // The work dies with the wrapper (parent-death signal). + pid, _ := strconv.Atoi(strings.TrimSpace(strings.SplitN(strings.SplitN(out, "pid=", 2)[1], "\n", 2)[0])) + deadline := time.Now().Add(5 * time.Second) + for syscall.Kill(pid, 0) == nil { + if time.Now().After(deadline) { + syscall.Kill(pid, syscall.SIGKILL) + t.Fatal("work survived the wrapper") + } + time.Sleep(20 * time.Millisecond) + } +} + +func TestUnsafeOutputNeverRunsWork(t *testing.T) { + j := newJob(t, "touch "+"ran-marker") + outside := t.TempDir() + os.Symlink(outside, filepath.Join(j.spec.Home, "escape")) + j.spec.OutputPath = filepath.Join(j.spec.Home, "escape", "mist.out") + j.start() + if code := j.wait(); code != jobentry.ExitSetupFailed { + t.Fatalf("exit %d", code) + } + if r := j.result(); r == nil || r.Reason != executor.ExitLaunchError || !r.OutputIncomplete { + t.Fatalf("result %+v", r) + } + if _, err := os.Stat(filepath.Join(j.spec.Home, "ran-marker")); err == nil { + t.Fatal("work ran despite unsafe output path") + } + if !strings.Contains(j.stdout.String(), "could not start") { + t.Fatalf("stdout:\n%s", j.stdout.String()) + } +} + +// Runs the job in a transient systemd scope with a 64 MiB memory limit and +// no swap, as the kernel would enforce it inside a Pod. OOMPolicy=continue +// isolates the kernel OOM kill: systemd's default (stop) would also SIGTERM +// the whole scope, which the M5 node executor must account for. +func TestOutOfMemoryInCgroup(t *testing.T) { + if err := exec.Command("systemd-run", "--user", "--scope", "-q", "true").Run(); err != nil { + t.Skipf("UNVERIFIED: no user systemd scopes here (%v)", err) + } + j := newJob(t, "echo allocating; head -c 512M /dev/zero | tail > /dev/null") + j.start("systemd-run", "--user", "--scope", "-q", "-p", "MemoryMax=64M", "-p", "MemorySwapMax=0", "-p", "OOMPolicy=continue") + code := j.wait() + r := j.result() + if r == nil || r.Reason != executor.ExitOutOfMemory { + t.Fatalf("exit %d, result %+v, output:\n%s", code, r, j.output()) + } + if code != 137 || !strings.Contains(j.output(), "out of memory: exceeded memory limit 64.0 MiB") { + t.Fatalf("exit %d, output:\n%s", code, j.output()) + } + if r.PeakMemoryBytes < 32<<20 { + t.Fatalf("peak %d", r.PeakMemoryBytes) + } +} diff --git a/docs/agents/architecture/executors.md b/docs/agents/architecture/executors.md index 1e13f03..81f7ed0 100644 --- a/docs/agents/architecture/executors.md +++ b/docs/agents/architecture/executors.md @@ -412,6 +412,8 @@ The agent starts a systemd **transient unit** `mist-job-.service` (in | Time limit | `RuntimeMaxSec=` as backstop; `KillSignal=SIGTERM`, `TimeoutStopSec=` grace | | Exit reason | Unit `Result=` (`oom-kill`, `timeout`, `exit-code`) and `ExecMainStatus` | +systemd's default `OOMPolicy=stop` also sends SIGTERM to the rest of the unit after a kernel OOM kill (observed in the M1 `mist-job-entry` cgroup test). Pin `OOMPolicy` explicitly, and take the reason from `Result=oom-kill` or `memory.events`, not from the exit status. + ### Interactive sessions Use the same persistent holder/PTY protocol as K3s, inside the main job unit. diff --git a/docs/agents/plans/jobs-and-scheduling.md b/docs/agents/plans/jobs-and-scheduling.md index 31f8b0c..37e5014 100644 --- a/docs/agents/plans/jobs-and-scheduling.md +++ b/docs/agents/plans/jobs-and-scheduling.md @@ -121,6 +121,23 @@ requirements. Keep secrets and actual research data out of test fixtures. `memory.max`. Verify the fallback when the wrapper itself is killed; report incomplete output without fabricating a successful footer. Test batch arguments and safe output opening inside the restricted view. + - *Status (30 Oct 2026):* `cmd/mist-job-entry` (static binary) with logic in + `internal/jobentry`. + - **How it works:** the output file opens via `openat2(RESOLVE_BENEATH)` + under the team root, regular files only. Dot-folders are opened + `O_NOFOLLOW` and checked by `fstat`; an unsafe replacement fails setup. + Setup failures exit 125 with a `launch_error` result and never run the + work. Work runs in its own process group with a parent-death signal; + the local deadline timer sends SIGTERM, then SIGKILL after grace. + - **Tests:** the process tests re-execute the test binary as the real + wrapper. They cover batch args with shell metacharacters, command exit + codes, SIGTERM forwarding, both deadline paths, killing the wrapper (no + footer, no result, work dies) and an unsafe output path. OOM runs in a + 64 MiB systemd scope with swap off and detects `oom_kill` from + `memory.events`. + - **Still open:** grandchildren that call `setsid` escape the group and + rely on the executor tearing down the container or cgroup. + `/dev/termination-log` wiring belongs to the step-5 Pod builder. 5. **K3s executor.** - A pod builder written as a pure function (`LaunchSpec` → Pod and ConfigMap), tested against golden files. diff --git a/go.mod b/go.mod index 17e6521..f390167 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.27 require ( github.com/fergusstrange/embedded-postgres v1.34.0 github.com/jackc/pgx/v5 v5.11.0 + golang.org/x/sys v0.49.0 ) require ( diff --git a/go.sum b/go.sum index caad47e..183ad37 100644 --- a/go.sum +++ b/go.sum @@ -26,6 +26,8 @@ go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/sys v0.49.0 h1:XbzkgYJHdqh/8m2Uu0W/dQv8nktxx4BFHp1M0gROTXA= +golang.org/x/sys v0.49.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/jobentry/env.go b/internal/jobentry/env.go new file mode 100644 index 0000000..25dd9ce --- /dev/null +++ b/internal/jobentry/env.go @@ -0,0 +1,88 @@ +package jobentry + +import ( + "fmt" + "slices" + "strings" + "time" +) + +// passthrough are variables taken from the wrapper's own environment: the +// runtime profile's baseline and what the executor/device runtime sets. +var passthrough = []string{"PATH", "LANG", "LC_ALL", "TERM", "CUDA_VISIBLE_DEVICES", "NVIDIA_VISIBLE_DEVICES"} + +// protected reports whether a user --env value may not set name. +func protected(name string) bool { + switch name { + case "HOME", "USER", "LOGNAME", "PWD", "SHELL", "TMPDIR", + "CUDA_VISIBLE_DEVICES", "NVIDIA_VISIBLE_DEVICES", "PATH": + return true + } + return strings.HasPrefix(name, "MIST_") || strings.HasPrefix(name, "NVIDIA_") +} + +// childEnv builds the job's complete environment. Nothing else from the +// wrapper's environment leaks in, so jobs behave the same however they +// were submitted. started is the actual work start, for MIST_DEADLINE. +func childEnv(s Spec, own []string, started time.Time) []string { + env := map[string]string{} + for _, kv := range own { + k, v, _ := strings.Cut(kv, "=") + if slices.Contains(passthrough, k) { + env[k] = v + } + } + if env["PATH"] == "" { + env["PATH"] = "/usr/local/bin:/usr/bin:/bin" + } + cpus := fmt.Sprint(max(s.Resources.MilliCPU/1000, 1)) + teamCache := s.TeamRoot + "/.cache" + tool := map[string]string{ // user --env may override these + "OMP_NUM_THREADS": cpus, + "MKL_NUM_THREADS": cpus, + "PYTHONUNBUFFERED": "1", + "HF_HOME": teamCache + "/huggingface", + "HF_TOKEN_PATH": s.Home + "/.cache/huggingface/token", + "TORCH_HOME": teamCache + "/torch", + } + for k, v := range tool { + env[k] = v + } + for k, v := range s.Env { + if !protected(k) && k != "" && !strings.ContainsAny(k, "=\x00") && !strings.Contains(v, "\x00") { + env[k] = v + } + } + fixed := map[string]string{ // set last: nothing overrides these + "HOME": s.Home, + "USER": s.User, + "LOGNAME": s.User, + "SHELL": "/bin/bash", + "PWD": s.Workdir, + "TMPDIR": s.Scratch, + "MIST_JOB_ID": fmt.Sprint(s.JobID), + "MIST_JOB_NAME": s.Name, + "MIST_USER": s.User, + "MIST_TEAM": s.Team, + "MIST_PARTITION": s.Partition, + "MIST_QOS": s.QOS, + "MIST_NODE": s.Node, + "MIST_CPUS": cpus, + "MIST_MEMORY_BYTES": fmt.Sprint(s.Resources.MemoryBytes), + "MIST_GPUS": fmt.Sprint(s.Resources.GPUs), + "MIST_TT_BOARDS": fmt.Sprint(s.Resources.TTBoards), + "MIST_TIME_LIMIT_SECONDS": fmt.Sprint(s.TimeLimit), + "MIST_DEADLINE": fmt.Sprint(started.Add(s.timeLimit()).Unix()), + "MIST_SUBMIT_DIR": s.Workdir, + "MIST_SCRATCH": s.Scratch, + } + for k, v := range fixed { + env[k] = v + } + out := make([]string, 0, len(env)) + for k, v := range env { + out = append(out, k+"="+v) + } + slices.Sort(out) + return out +} diff --git a/internal/jobentry/fs.go b/internal/jobentry/fs.go new file mode 100644 index 0000000..9d34862 --- /dev/null +++ b/internal/jobentry/fs.go @@ -0,0 +1,81 @@ +package jobentry + +import ( + "errors" + "fmt" + "os" + "strings" + + "golang.org/x/sys/unix" +) + +// privateDirs live in $HOME and are reset to 0700 on every job (D8). +var privateDirs = []string{".cache", ".config", ".local", ".ssh"} + +// ensurePrivateDirs creates or resets the private dot-folders without +// following symlinks. A replacement that is not a real directory owned by +// the job's user fails setup; its target is never chmod-ed. +func ensurePrivateDirs(home string) error { + hfd, err := unix.Open(home, unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_RDONLY|unix.O_CLOEXEC, 0) + if err != nil { + return fmt.Errorf("open home %s: %w", home, err) + } + defer unix.Close(hfd) + uid := os.Geteuid() + for _, name := range privateDirs { + if err := unix.Mkdirat(hfd, name, 0o700); err != nil && !errors.Is(err, unix.EEXIST) { + return fmt.Errorf("create %s/%s: %w", home, name, err) + } + fd, err := unix.Openat(hfd, name, unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_RDONLY|unix.O_CLOEXEC, 0) + if err != nil { + return fmt.Errorf("%s/%s is not a real directory (%w); fix or remove it", home, name, err) + } + var st unix.Stat_t + err = unix.Fstat(fd, &st) + if err == nil && int(st.Uid) != uid { + err = fmt.Errorf("owned by uid %d, not %d", st.Uid, uid) + } + if err == nil { + err = unix.Fchmod(fd, 0o700) + } + unix.Close(fd) + if err != nil { + return fmt.Errorf("%s/%s: %w", home, name, err) + } + } + return nil +} + +// openOutput opens the output file for appending, resolving the path +// strictly beneath the team root: symlinks may not escape it, and the +// result must be a regular file (never a FIFO or device that could block +// or redirect output). +func openOutput(teamRoot, output string) (*os.File, error) { + rel := strings.TrimPrefix(strings.TrimPrefix(output, teamRoot), "/") + if rel == "" || rel == output { + return nil, fmt.Errorf("output %s is not inside %s", output, teamRoot) + } + root, err := unix.Open(teamRoot, unix.O_DIRECTORY|unix.O_RDONLY|unix.O_CLOEXEC, 0) + if err != nil { + return nil, fmt.Errorf("open team folder: %w", err) + } + defer unix.Close(root) + fd, err := unix.Openat2(root, rel, &unix.OpenHow{ + Flags: unix.O_WRONLY | unix.O_CREAT | unix.O_APPEND | unix.O_NONBLOCK | unix.O_NOFOLLOW | unix.O_CLOEXEC, + Mode: 0o666, // umask 0027 makes it 0640 + Resolve: unix.RESOLVE_BENEATH | unix.RESOLVE_NO_MAGICLINKS, + }) + if err != nil { + return nil, fmt.Errorf("open output %s: %w", output, err) + } + var st unix.Stat_t + if err := unix.Fstat(fd, &st); err != nil || st.Mode&unix.S_IFMT != unix.S_IFREG { + unix.Close(fd) + return nil, fmt.Errorf("output %s is not a regular file", output) + } + if err := unix.SetNonblock(fd, false); err != nil { + unix.Close(fd) + return nil, err + } + return os.NewFile(uintptr(fd), output), nil +} diff --git a/internal/jobentry/jobentry_test.go b/internal/jobentry/jobentry_test.go new file mode 100644 index 0000000..7157032 --- /dev/null +++ b/internal/jobentry/jobentry_test.go @@ -0,0 +1,179 @@ +package jobentry + +import ( + "os" + "path/filepath" + "slices" + "strings" + "syscall" + "testing" + "time" + + "golang.org/x/sys/unix" + + "github.com/UTMIST/Mist/internal/executor" +) + +func TestPrivateDirs(t *testing.T) { + home := t.TempDir() + os.Mkdir(filepath.Join(home, ".cache"), 0o755) // too open: reset + if err := ensurePrivateDirs(home); err != nil { + t.Fatal(err) + } + for _, d := range privateDirs { + st, err := os.Lstat(filepath.Join(home, d)) + if err != nil || !st.IsDir() || st.Mode().Perm() != 0o700 { + t.Fatalf("%s: %v %v", d, st.Mode(), err) + } + } + + // A symlink in place of .ssh fails setup and its target is untouched. + elsewhere := t.TempDir() + os.Chmod(elsewhere, 0o755) + os.RemoveAll(filepath.Join(home, ".ssh")) + os.Symlink(elsewhere, filepath.Join(home, ".ssh")) + if err := ensurePrivateDirs(home); err == nil { + t.Fatal("symlinked .ssh accepted") + } + if st, _ := os.Stat(elsewhere); st.Mode().Perm() != 0o755 { + t.Fatalf("symlink target chmod-ed to %v", st.Mode().Perm()) + } + + // A file in place of a directory also fails. + os.Remove(filepath.Join(home, ".ssh")) + os.WriteFile(filepath.Join(home, ".ssh"), nil, 0o600) + if err := ensurePrivateDirs(home); err == nil { + t.Fatal("file in place of .ssh accepted") + } +} + +func TestOpenOutput(t *testing.T) { + old := unix.Umask(0o027) + defer unix.Umask(old) + team := t.TempDir() + os.MkdirAll(filepath.Join(team, "alice", "proj"), 0o750) + + f, err := openOutput(team, filepath.Join(team, "alice", "proj", "mist-1.out")) + if err != nil { + t.Fatal(err) + } + f.Close() + if st, _ := os.Stat(filepath.Join(team, "alice", "proj", "mist-1.out")); st.Mode().Perm() != 0o640 { + t.Fatalf("output mode %v, want 0640", st.Mode().Perm()) + } + + // A symlink inside the team tree is fine. + os.Symlink("proj", filepath.Join(team, "alice", "latest")) + if f, err := openOutput(team, filepath.Join(team, "alice", "latest", "mist-2.out")); err != nil { + t.Fatalf("in-tree symlink rejected: %v", err) + } else { + f.Close() + } + + outside := t.TempDir() + os.Symlink(outside, filepath.Join(team, "alice", "escape")) + os.Symlink(filepath.Join(outside, "victim"), filepath.Join(team, "alice", "link.out")) + syscall.Mkfifo(filepath.Join(team, "alice", "fifo.out"), 0o600) + for _, bad := range []string{ + filepath.Join(team, "alice", "escape", "mist.out"), // directory symlink out of the tree + filepath.Join(team, "alice", "link.out"), // final component is a symlink + filepath.Join(team, "alice", "fifo.out"), // not a regular file; must not block + filepath.Join(team, "alice", "..", "..", "x.out"), // not beneath the team root + "/etc/passwd", + } { + done := make(chan error, 1) + go func() { + f, err := openOutput(team, bad) + if f != nil { + f.Close() + } + done <- err + }() + select { + case err := <-done: + if err == nil { + t.Errorf("%s: accepted", bad) + } + case <-time.After(2 * time.Second): + t.Fatalf("%s: open blocked", bad) + } + } + if entries, _ := os.ReadDir(outside); len(entries) != 0 { + t.Fatal("wrote outside the team tree") + } +} + +func testSpec(root string) Spec { + team := filepath.Join(root, "mist", "teams", "team-a") + return Spec{ + JobID: 7, Name: "t", User: "alice", Team: "team-a", Partition: "cpu", QOS: "normal", Node: "n1", + Kind: executor.KindBatch, + Resources: executor.Resources{MilliCPU: 4000, MemoryBytes: 8 << 30, ScratchBytes: 1 << 30}, + TimeLimit: 3600, + SignalGrace: 60, + TeamRoot: team, + Home: filepath.Join(team, "alice"), + Workdir: filepath.Join(team, "alice"), + OutputPath: filepath.Join(team, "alice", "mist-7.out"), + Scratch: filepath.Join(root, "scratch"), + ScriptPath: filepath.Join(root, "script.sh"), + } +} + +func TestChildEnv(t *testing.T) { + s := testSpec("/r") + s.Env = map[string]string{ + "HOME": "/evil", "MIST_JOB_ID": "1", "CUDA_VISIBLE_DEVICES": "0,1", "PATH": "/evil", + "OMP_NUM_THREADS": "2", "WANDB_PROJECT": "x", + } + start := time.Unix(1_800_000_000, 0) + env := childEnv(s, []string{"PATH=/usr/bin", "AWS_SECRET_ACCESS_KEY=leak", "CUDA_VISIBLE_DEVICES=GPU-abc"}, start) + get := func(k string) string { + for _, kv := range env { + if v, ok := strings.CutPrefix(kv, k+"="); ok { + return v + } + } + return "" + } + want := map[string]string{ + "HOME": s.Home, + "MIST_JOB_ID": "7", + "CUDA_VISIBLE_DEVICES": "GPU-abc", // from the device runtime, not the user + "PATH": "/usr/bin", + "OMP_NUM_THREADS": "2", // tuning values may be overridden + "WANDB_PROJECT": "x", + "MIST_CPUS": "4", + "MIST_DEADLINE": "1800003600", + "TMPDIR": s.Scratch, + "AWS_SECRET_ACCESS_KEY": "", + "HF_HOME": s.TeamRoot + "/.cache/huggingface", + } + for k, v := range want { + if got := get(k); got != v { + t.Errorf("%s = %q, want %q", k, got, v) + } + } + if !slices.IsSorted(env) { + t.Error("environment is not deterministic") + } +} + +func TestSpecValidate(t *testing.T) { + if err := testSpec("/r").Validate(); err != nil { + t.Fatal(err) + } + for name, mutate := range map[string]func(*Spec){ + "output outside team": func(s *Spec) { s.OutputPath = "/r/other/x.out" }, + "unclean path": func(s *Spec) { s.Workdir = s.TeamRoot + "/alice/../bob" }, + "interactive": func(s *Spec) { s.Kind = executor.KindInteractive; s.ScriptPath = "" }, + "command with script": func(s *Spec) { s.Kind = executor.KindCommand; s.Argv = []string{"ls"} }, + "no time limit": func(s *Spec) { s.TimeLimit = 0 }, + } { + s := testSpec("/r") + mutate(&s) + if s.Validate() == nil { + t.Errorf("%s: accepted", name) + } + } +} diff --git a/internal/jobentry/run.go b/internal/jobentry/run.go new file mode 100644 index 0000000..11580d5 --- /dev/null +++ b/internal/jobentry/run.go @@ -0,0 +1,306 @@ +package jobentry + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync" + "syscall" + "time" + + "golang.org/x/sys/unix" + + "github.com/UTMIST/Mist/internal/executor" +) + +// ExitSetupFailed is returned when the work never started (bad output +// path, unsafe dot-folder, missing program). The executor reports it as a +// launch failure, not as the job's own exit. +const ExitSetupFailed = 125 + +// footerMargin is how long before the executor's SIGKILL the wrapper kills +// the work itself, so it still has time to write the footer. +const footerMargin = 2 * time.Second + +// Run executes the job described by s and returns the wrapper's exit code. +// stdout receives a copy of everything written to the output file, so +// executor log streaming works without reading team files. Signals sent to +// the wrapper arrive on sigs. +func Run(s Spec, stdout io.Writer, sigs <-chan os.Signal) int { + unix.Umask(0o027) + started := time.Now() + out := &output{stdout: stdout} + cg := ownCgroup() + oomBefore := oomKills(cg) + + setupFailed := func(err error) int { + out.mist("job %d could not start: %v", s.JobID, err) + writeTermination(s.TerminationLog, executor.Result{ + ExitCode: ptr(ExitSetupFailed), Reason: executor.ExitLaunchError, Detail: err.Error(), + OutputIncomplete: out.file == nil || out.failed, + }) + return ExitSetupFailed + } + + if err := ensurePrivateDirs(s.Home); err != nil { + return setupFailed(err) + } + f, err := openOutput(s.TeamRoot, s.OutputPath) + if err != nil { + return setupFailed(err) + } + out.file = f + out.header(s, started) + + env := childEnv(s, os.Environ(), started) + var cmd *exec.Cmd + if s.Kind == executor.KindBatch { + cmd = exec.Command("/bin/bash", append([]string{s.ScriptPath}, s.Args...)...) + } else { + prog, err := lookPath(s.Argv[0], env) + if err != nil { + return setupFailed(err) + } + cmd = exec.Command(prog, s.Argv[1:]...) + cmd.Args[0] = s.Argv[0] + } + cmd.Dir, cmd.Env = s.Workdir, env + // Own process group so signals reach the whole job; die with the wrapper. + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, Pdeathsig: syscall.SIGKILL} + r, w, err := os.Pipe() + if err != nil { + return setupFailed(err) + } + cmd.Stdout, cmd.Stderr = w, w // one pipe keeps stdout/stderr interleaving + if err := cmd.Start(); err != nil { + w.Close() + r.Close() + return setupFailed(err) + } + w.Close() + copied := make(chan struct{}) + go func() { io.Copy(out, r); close(copied) }() + waited := make(chan error, 1) + go func() { waited <- cmd.Wait() }() + + group := -cmd.Process.Pid + deadline := time.NewTimer(s.timeLimit()) + defer deadline.Stop() + var kill <-chan time.Time + stopReason := executor.ExitNormal + stopWith := func(why string, after time.Duration) { + syscall.Kill(group, syscall.SIGTERM) + if kill == nil { + out.mist("%s; sent SIGTERM, SIGKILL in %s", why, after) + kill = time.After(after) + } + } +wait: + for { + select { + case sig := <-sigs: + if stopReason == executor.ExitNormal { + stopReason = executor.ExitSignaled + } + stopWith(fmt.Sprintf("received %s", sig), max(s.grace()-footerMargin, time.Second)) + case <-deadline.C: + if stopReason == executor.ExitNormal { // the first stop cause wins + stopReason = executor.ExitDeadline + } + stopWith(fmt.Sprintf("time limit %s reached", s.timeLimit()), s.grace()) + case <-kill: + out.mist("grace period over; sent SIGKILL") + syscall.Kill(group, syscall.SIGKILL) + case err = <-waited: + break wait + } + } + // Background processes left in the group would keep the pipe open. + syscall.Kill(group, syscall.SIGKILL) + select { + case <-copied: + case <-time.After(5 * time.Second): + out.mist("output still open after job exit; some output may be missing") + out.failed = true + } + r.Close() + + res := executor.Result{Reason: stopReason, PeakMemoryBytes: peakMemory(cg)} + code := 0 + var exitErr *exec.ExitError + switch { + case err == nil: + case errors.As(err, &exitErr): + ws := exitErr.Sys().(syscall.WaitStatus) + if ws.Signaled() { + code = 128 + int(ws.Signal()) + } else { + code = ws.ExitStatus() + } + default: + code = 1 + } + res.ExitCode = &code + if after := oomKills(cg); after > oomBefore { + res.Reason = executor.ExitOutOfMemory + res.Detail = fmt.Sprintf("exceeded memory limit %s (peak %s)", + bytes(s.Resources.MemoryBytes), bytes(res.PeakMemoryBytes)) + } else if after < 0 && code == 128+int(syscall.SIGKILL) && stopReason == executor.ExitNormal { + res.Detail = "killed by SIGKILL (memory events unavailable; possibly out of memory)" + } + + out.footer(s, started, res) + if err := f.Close(); err != nil { + out.failed = true + } + res.OutputIncomplete = out.failed + writeTermination(s.TerminationLog, res) + return code +} + +// output copies job output to the file and the wrapper's stdout. A failed +// file write is remembered rather than retried, and reported as +// output_incomplete; stdout keeps receiving everything. +type output struct { + mu sync.Mutex + stdout io.Writer + file *os.File + failed bool +} + +func (o *output) Write(p []byte) (int, error) { + o.mu.Lock() + defer o.mu.Unlock() + o.stdout.Write(p) + if o.file != nil && !o.failed { + if _, err := o.file.Write(p); err != nil { + o.failed = true + } + } + return len(p), nil +} + +func (o *output) mist(format string, a ...any) { + o.Write([]byte("[mist] " + fmt.Sprintf(format, a...) + "\n")) +} + +func (o *output) header(s Spec, at time.Time) { + r := s.Resources + o.mist("job %d (%s) started %s on %s", s.JobID, s.Name, at.UTC().Format(time.RFC3339), s.Node) + o.mist("%s in team %s, partition %s, qos %s", s.User, s.Team, s.Partition, s.QOS) + acc := "" + if r.GPUs > 0 { + acc = fmt.Sprintf(", %d GPU", r.GPUs) + } + if r.TTBoards > 0 { + acc += fmt.Sprintf(", %d TT board", r.TTBoards) + } + o.mist("%g CPUs, %s memory%s, %s scratch, time limit %s", + float64(r.MilliCPU)/1000, bytes(r.MemoryBytes), acc, bytes(r.ScratchBytes), s.timeLimit()) +} + +func (o *output) footer(s Spec, started time.Time, r executor.Result) { + end := time.Now() + o.mist("job %d ended %s on %s", s.JobID, end.UTC().Format(time.RFC3339), s.Node) + switch r.Reason { + case executor.ExitOutOfMemory: + o.mist("out of memory: %s; request more with --mem", r.Detail) + case executor.ExitDeadline: + o.mist("stopped at the time limit (%s)", s.timeLimit()) + case executor.ExitSignaled: + o.mist("stopped by request") + default: + if r.Detail != "" { + o.mist("%s", r.Detail) + } + } + peak := "unknown" + if r.PeakMemoryBytes > 0 { + peak = bytes(r.PeakMemoryBytes) + } + o.mist("exit code %d, elapsed %s, peak memory %s", *r.ExitCode, end.Sub(started).Round(time.Second), peak) +} + +func writeTermination(file string, r executor.Result) { + if file == "" { + return + } + if len(r.Detail) > 1024 { // Kubernetes keeps at most 4 KiB + r.Detail = r.Detail[:1024] + } + b, _ := json.Marshal(r) + os.WriteFile(file, b, 0o644) +} + +// ownCgroup returns this process's cgroup v2 directory, or "". +func ownCgroup() string { + b, err := os.ReadFile("/proc/self/cgroup") + if err != nil { + return "" + } + for _, line := range strings.Split(string(b), "\n") { + if p, ok := strings.CutPrefix(line, "0::"); ok { + return filepath.Join("/sys/fs/cgroup", p) + } + } + return "" +} + +// oomKills returns the cgroup's oom_kill count, or -1 if unreadable. +func oomKills(cg string) int64 { + b, err := os.ReadFile(filepath.Join(cg, "memory.events")) + if cg == "" || err != nil { + return -1 + } + for _, line := range strings.Split(string(b), "\n") { + if v, ok := strings.CutPrefix(line, "oom_kill "); ok { + n, _ := strconv.ParseInt(v, 10, 64) + return n + } + } + return -1 +} + +func peakMemory(cg string) int64 { + b, err := os.ReadFile(filepath.Join(cg, "memory.peak")) + if cg == "" || err != nil { + return 0 + } + n, _ := strconv.ParseInt(strings.TrimSpace(string(b)), 10, 64) + return n +} + +func lookPath(name string, env []string) (string, error) { + if strings.Contains(name, "/") { + return name, nil + } + for _, kv := range env { + if p, ok := strings.CutPrefix(kv, "PATH="); ok { + for _, dir := range filepath.SplitList(p) { + cand := filepath.Join(dir, name) + if st, err := os.Stat(cand); err == nil && st.Mode().IsRegular() && st.Mode()&0o111 != 0 { + return cand, nil + } + } + } + } + return "", fmt.Errorf("%s: command not found", name) +} + +func bytes(n int64) string { + switch { + case n >= 1<<30: + return fmt.Sprintf("%.1f GiB", float64(n)/(1<<30)) + case n >= 1<<20: + return fmt.Sprintf("%.1f MiB", float64(n)/(1<<20)) + } + return fmt.Sprintf("%d B", n) +} + +func ptr[T any](v T) *T { return &v } diff --git a/internal/jobentry/spec.go b/internal/jobentry/spec.go new file mode 100644 index 0000000..732ea64 --- /dev/null +++ b/internal/jobentry/spec.go @@ -0,0 +1,98 @@ +// Package jobentry is mist-job-entry: the first process of every job under +// every executor. It is the one place that implements the job contract in +// docs/agents/architecture/runtime-environment.md: umask, private +// dot-folders, the output file with header and footer, signal forwarding, +// the local time limit, the exit reason and the termination message. +// +// It runs unprivileged as the job's user. It cannot guarantee a footer +// after SIGKILL, OOM or node loss, so executors take authoritative reasons +// from kubelet/cgroup observations and treat its report as advisory. +package jobentry + +import ( + "encoding/json" + "fmt" + "os" + "path" + "strings" + "time" + + "github.com/UTMIST/Mist/internal/executor" +) + +// Spec is /etc/mist/job.json, written by the executor from the LaunchSpec. +type Spec struct { + JobID int64 `json:"job_id"` + Name string `json:"name"` + User string `json:"user"` + Team string `json:"team"` + Partition string `json:"partition"` + QOS string `json:"qos"` + Node string `json:"node"` + Kind executor.Kind `json:"kind"` + Resources executor.Resources `json:"resources"` + TimeLimit int64 `json:"time_limit_seconds"` + SignalGrace int64 `json:"signal_grace_seconds"` + + TeamRoot string `json:"team_root"` // /mist/teams/; outputs must stay beneath it + Home string `json:"home"` // /mist/teams// + Workdir string `json:"workdir"` // start directory + OutputPath string `json:"output_path"` // absolute, inside TeamRoot + Scratch string `json:"scratch"` // /scratch + + ScriptPath string `json:"script_path"` // batch: the script file from the ConfigMap + Args []string `json:"args"` // batch arguments + Argv []string `json:"argv"` // command + Env map[string]string `json:"env"` // user --env values; protected names are dropped + + // TerminationLog receives the structured Result (K3s: the Pod's + // terminationMessagePath). Empty disables it. + TerminationLog string `json:"termination_log"` +} + +// LoadSpec reads and checks a job.json file. +func LoadSpec(file string) (Spec, error) { + var s Spec + b, err := os.ReadFile(file) + if err != nil { + return s, err + } + if err := json.Unmarshal(b, &s); err != nil { + return s, fmt.Errorf("jobentry: parse %s: %w", file, err) + } + return s, s.Validate() +} + +func (s Spec) timeLimit() time.Duration { return time.Duration(s.TimeLimit) * time.Second } +func (s Spec) grace() time.Duration { return time.Duration(s.SignalGrace) * time.Second } + +// Validate rejects a spec the wrapper cannot run safely. +func (s Spec) Validate() error { + clean := func(p string) bool { return path.IsAbs(p) && path.Clean(p) == p } + under := func(p, root string) bool { return p == root || strings.HasPrefix(p, root+"/") } + switch { + case s.JobID <= 0: + return fmt.Errorf("jobentry: job_id must be positive") + case !clean(s.TeamRoot) || !clean(s.Home) || !clean(s.Workdir) || !clean(s.OutputPath) || !clean(s.Scratch): + return fmt.Errorf("jobentry: paths must be clean and absolute") + case !under(s.Home, s.TeamRoot) || !under(s.OutputPath, s.TeamRoot): + return fmt.Errorf("jobentry: home and output must be inside %s", s.TeamRoot) + case s.TimeLimit <= 0 || s.SignalGrace <= 0: + return fmt.Errorf("jobentry: time limit and signal grace must be positive") + } + switch s.Kind { + case executor.KindBatch: + if s.ScriptPath == "" || len(s.Argv) != 0 { + return fmt.Errorf("jobentry: batch needs script_path and no argv") + } + case executor.KindCommand: + if len(s.Argv) == 0 || s.ScriptPath != "" || len(s.Args) != 0 { + return fmt.Errorf("jobentry: command needs argv and no script or args") + } + case executor.KindInteractive: + return fmt.Errorf("jobentry: interactive jobs arrive in M2") + default: + return fmt.Errorf("jobentry: unknown kind %q", s.Kind) + } + return nil +} From 1f48661778fbcd8cfc5ba97b3ed4a292ebc4cf34 Mon Sep 17 00:00:00 2001 From: Brandon Li Date: Fri, 30 Oct 2026 10:14:18 -0400 Subject: [PATCH 4/5] Add K3s executor with k3d integration tests (M1 step 5) Pod builder (pure, golden-tested) for plain Pods and immutable ConfigMaps: Mist's nodeName, strict supplemental groups, no service-account token, read-only root, requests = limits, per-job passwd/group, NFS or hostPath team mounts, dnsPolicy None. Executor on client-go informers and a work queue, with durable receipts in k3s_executions: repeat-safe Launch that adopts ambiguous creates and never recreates a recorded Pod, Terminate by shortening activeDeadlineSeconds so the final status survives, results that prefer kubelet observations (OOMKilled, eviction, deadline) over the wrapper, log tail saved before deletion, Lost on outside deletion or 2 min NotReady, and Exited only on a matching node cleanup receipt. k3d testing found that NetworkPolicy is programmed about a second after containers start; every job Pod now waits in a network-gate init container until the node's kubelet is unreachable. Also: deploy/k3s namespace, NetworkPolicy and RBAC, mistd cleanup-report, scripts/k3d-test.sh, and evidence in docs/agents/evidence. Co-Authored-By: Claude Opus 5.5 --- README.md | 7 + cmd/mist-job-entry/main.go | 12 + cmd/mistd/main.go | 35 +- cmd/mistd/main_test.go | 3 +- deploy/k3s/mist-jobs.yaml | 32 + deploy/k3s/rbac.yaml | 53 ++ docs/agents/architecture/executors.md | 32 + docs/agents/evidence/m1-step5-k3d.md | 50 ++ docs/agents/plans/jobs-and-scheduling.md | 21 + go.mod | 42 +- go.sum | 110 ++- .../executor/0001_k3s_executions.sql | 31 + internal/executor/fake/fake_test.go | 2 + internal/executor/k3s/cleanup.go | 67 ++ internal/executor/k3s/cleanup_test.go | 40 + internal/executor/k3s/executor.go | 767 ++++++++++++++++++ internal/executor/k3s/executor_test.go | 387 +++++++++ internal/executor/k3s/k3d_test.go | 420 ++++++++++ internal/executor/k3s/pod.go | 295 +++++++ internal/executor/k3s/pod_test.go | 159 ++++ internal/executor/k3s/store.go | 129 +++ .../testdata/cpu-command-hostpath.golden.yaml | 197 +++++ .../k3s/testdata/gpu-batch-nfs.golden.yaml | 205 +++++ .../executor/k3s/testdata/runtime/Dockerfile | 4 + internal/executor/queue.go | 64 ++ internal/executor/spec.go | 7 +- internal/executor/spec_test.go | 3 + internal/jobentry/isolation.go | 38 + internal/jobentry/isolation_test.go | 35 + internal/scheduler/admit.go | 3 + scripts/k3d-test.sh | 38 + 31 files changed, 3274 insertions(+), 14 deletions(-) create mode 100644 deploy/k3s/mist-jobs.yaml create mode 100644 deploy/k3s/rbac.yaml create mode 100644 docs/agents/evidence/m1-step5-k3d.md create mode 100644 internal/db/migrations/executor/0001_k3s_executions.sql create mode 100644 internal/executor/k3s/cleanup.go create mode 100644 internal/executor/k3s/cleanup_test.go create mode 100644 internal/executor/k3s/executor.go create mode 100644 internal/executor/k3s/executor_test.go create mode 100644 internal/executor/k3s/k3d_test.go create mode 100644 internal/executor/k3s/pod.go create mode 100644 internal/executor/k3s/pod_test.go create mode 100644 internal/executor/k3s/store.go create mode 100644 internal/executor/k3s/testdata/cpu-command-hostpath.golden.yaml create mode 100644 internal/executor/k3s/testdata/gpu-batch-nfs.golden.yaml create mode 100644 internal/executor/k3s/testdata/runtime/Dockerfile create mode 100644 internal/executor/queue.go create mode 100644 internal/jobentry/isolation.go create mode 100644 internal/jobentry/isolation_test.go create mode 100755 scripts/k3d-test.sh diff --git a/README.md b/README.md index 63a9f22..d6b71d3 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,13 @@ Database tests start an embedded Postgres as the current user (downloaded to `~/.embedded-postgres-go` on first run). To use an existing server instead, set `MIST_TEST_DATABASE_URL` to a superuser URL; each test creates its own database. +Executor integration tests run on a local k3d cluster (needs Docker without +sudo, plus k3d and kubectl): + +```bash +scripts/k3d-test.sh +``` + Check documentation links with: ```bash diff --git a/cmd/mist-job-entry/main.go b/cmd/mist-job-entry/main.go index 6712eb3..5b18545 100644 --- a/cmd/mist-job-entry/main.go +++ b/cmd/mist-job-entry/main.go @@ -2,6 +2,7 @@ // internal/jobentry and docs/agents/architecture/executors.md. // // mist-job-entry [-spec /etc/mist/job.json] +// mist-job-entry -await-isolation host:port[,host:port] (init container) package main import ( @@ -9,14 +10,25 @@ import ( "fmt" "os" "os/signal" + "strings" "syscall" + "time" "github.com/UTMIST/Mist/internal/jobentry" ) func main() { specPath := flag.String("spec", "/etc/mist/job.json", "job description written by the executor") + isolation := flag.String("await-isolation", "", "comma-separated host:port targets that must become unreachable, then exit") flag.Parse() + if *isolation != "" { + if err := jobentry.AwaitIsolation(strings.Split(*isolation, ","), 60*time.Second); err != nil { + fmt.Fprintln(os.Stdout, "[mist]", err) + os.WriteFile("/dev/termination-log", []byte(err.Error()), 0o644) + os.Exit(1) + } + return + } // Handle signals before anything else: as PID 1 in a container, // unhandled SIGTERM would be ignored. sigs := make(chan os.Signal, 4) diff --git a/cmd/mistd/main.go b/cmd/mistd/main.go index a22f37f..940268d 100644 --- a/cmd/mistd/main.go +++ b/cmd/mistd/main.go @@ -2,6 +2,8 @@ // // mistd api HTTP API // mistd scheduler queue, admission and executor calls (one active leader) +// mistd cleanup-report on a k3s node, as root: print a cleanup receipt +// for a deleted Pod, or fail if anything of it remains // // Both connect to MIST_DATABASE_URL (or -database-url) and apply pending // migrations at start. @@ -9,6 +11,7 @@ package main import ( "context" + "encoding/json" "errors" "flag" "fmt" @@ -21,6 +24,7 @@ import ( "time" "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/executor/k3s" "github.com/UTMIST/Mist/internal/identity" "github.com/UTMIST/Mist/internal/scheduler" ) @@ -28,17 +32,20 @@ import ( func main() { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - if err := run(ctx, os.Args[1:], os.Stderr); err != nil { + if err := run(ctx, os.Args[1:], os.Stdout, os.Stderr); err != nil { fmt.Fprintln(os.Stderr, "mistd:", err) os.Exit(1) } } -func run(ctx context.Context, args []string, stderr io.Writer) error { +func run(ctx context.Context, args []string, stdout, stderr io.Writer) error { if len(args) == 0 { - return errors.New("usage: mistd [flags]") + return errors.New("usage: mistd [flags]") } cmd, args := args[0], args[1:] + if cmd == "cleanup-report" { + return cleanupReport(args, stdout) + } fs := flag.NewFlagSet("mistd "+cmd, flag.ContinueOnError) fs.SetOutput(stderr) dbURL := fs.String("database-url", os.Getenv("MIST_DATABASE_URL"), "Postgres URL") @@ -100,3 +107,25 @@ func run(ctx context.Context, args []string, stderr io.Writer) error { return fmt.Errorf("unknown subcommand %q", cmd) } } + +// cleanupReport checks a node for leftovers of one Pod and prints the +// receipt. Until the internal API exists (M1 step 6) the receipt is +// delivered by whoever runs this; it is evidence, never a request. +func cleanupReport(args []string, stdout io.Writer) error { + fs := flag.NewFlagSet("mistd cleanup-report", flag.ContinueOnError) + kubelet := fs.String("kubelet-root", "/var/lib/kubelet", "kubelet root directory") + cgroups := fs.String("cgroup-root", "/sys/fs/cgroup", "cgroup v2 mount") + node := fs.String("node", "", "this node's name") + execID := fs.String("execution-id", "", "execution ID") + podUID := fs.String("pod-uid", "", "recorded Pod UID") + if err := fs.Parse(args); err != nil { + return err + } + evidence, err := k3s.VerifyCleanup(*kubelet, *cgroups, *podUID) + if err != nil { + return err + } + return json.NewEncoder(stdout).Encode(k3s.CleanupReceipt{ + ExecutionID: *execID, PodUID: *podUID, Node: *node, Evidence: evidence, + }) +} diff --git a/cmd/mistd/main_test.go b/cmd/mistd/main_test.go index 4d41be2..b441b3b 100644 --- a/cmd/mistd/main_test.go +++ b/cmd/mistd/main_test.go @@ -3,6 +3,7 @@ package main import ( "bytes" "context" + "io" "strings" "sync" "testing" @@ -53,7 +54,7 @@ func startScheduler(t *testing.T, url string) *instance { ctx, cancel := context.WithCancel(context.Background()) in := &instance{log: &syncBuffer{}, cancel: cancel, done: make(chan error, 1)} go func() { - in.done <- run(ctx, []string{"scheduler", "-database-url", url, "-leader-check", "100ms"}, in.log) + in.done <- run(ctx, []string{"scheduler", "-database-url", url, "-leader-check", "100ms"}, io.Discard, in.log) }() t.Cleanup(func() { cancel(); <-in.done }) return in diff --git a/deploy/k3s/mist-jobs.yaml b/deploy/k3s/mist-jobs.yaml new file mode 100644 index 0000000..c0d055a --- /dev/null +++ b/deploy/k3s/mist-jobs.yaml @@ -0,0 +1,32 @@ +# Namespace and network boundary for job Pods (executors.md#network). +# NetworkPolicy alone is not the host boundary: deployment-and-network.md +# adds host filtering for job traffic to each node's own services. +apiVersion: v1 +kind: Namespace +metadata: + name: mist-jobs + # No Pod Security admission label: both "restricted" and "baseline" forbid + # hostPath, which jobs on the volume's host node need. The pod builder's + # invariant tests (internal/executor/k3s) are the guard instead, and only + # mistd may create Pods here. +--- +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: jobs-internet-egress-only + namespace: mist-jobs +spec: + podSelector: {} + policyTypes: [Ingress, Egress] + ingress: [] # nothing connects in + egress: + - to: # IPv4 internet only; IPv6 is not listed, so it is denied + - ipBlock: + cidr: 0.0.0.0/0 + except: + - 10.0.0.0/8 # LAN, cluster Pod/Service ranges + - 172.16.0.0/12 + - 192.168.0.0/16 # includes the direct link + - 100.64.0.0/10 # Tailscale + - 169.254.0.0/16 # link-local and metadata + - 127.0.0.0/8 diff --git a/deploy/k3s/rbac.yaml b/deploy/k3s/rbac.yaml new file mode 100644 index 0000000..1b78f66 --- /dev/null +++ b/deploy/k3s/rbac.yaml @@ -0,0 +1,53 @@ +# mistd's Kubernetes permissions (executors.md#rbac): Pods and ConfigMaps in +# mist-jobs only, plus read-only nodes cluster-wide. ResourceClaims and +# ResourceSlices join in M4. +apiVersion: v1 +kind: Namespace +metadata: + name: mist-system +--- +apiVersion: v1 +kind: ServiceAccount +metadata: + name: mistd + namespace: mist-system +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: mistd-jobs + namespace: mist-jobs +rules: + - apiGroups: [""] + resources: [pods] + verbs: [get, list, watch, create, patch, delete] + - apiGroups: [""] + resources: [pods/log] + verbs: [get] + - apiGroups: [""] + resources: [configmaps] + verbs: [get, list, watch, create, patch, delete] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: mistd-jobs + namespace: mist-jobs +roleRef: {apiGroup: rbac.authorization.k8s.io, kind: Role, name: mistd-jobs} +subjects: [{kind: ServiceAccount, name: mistd, namespace: mist-system}] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: mistd-nodes +rules: + - apiGroups: [""] + resources: [nodes] + verbs: [get, list, watch] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: mistd-nodes +roleRef: {apiGroup: rbac.authorization.k8s.io, kind: ClusterRole, name: mistd-nodes} +subjects: [{kind: ServiceAccount, name: mistd, namespace: mist-system}] diff --git a/docs/agents/architecture/executors.md b/docs/agents/architecture/executors.md index 81f7ed0..993b7b2 100644 --- a/docs/agents/architecture/executors.md +++ b/docs/agents/architecture/executors.md @@ -292,6 +292,20 @@ DNS is an explicit port-53 exception to approved resolvers, not a blanket cluste service exception. Disable IPv6 for job networks until equivalent filtering is verified. Test both local and remote host addresses and service IPs. +**Isolation gate (found in M1 step 5).** k3s's policy controller (kube-router) +programs a Pod's rules about a second *after* its containers start. In k3d, a +job's first connections to the API service, the node's API server and kubelet +all succeeded, and the same probes were blocked a second later. Every job Pod +therefore has an init container, `mist-job-entry -await-isolation +$(MIST_HOST_IP):10250`. It shares the Pod's network namespace and passes only +after three consecutive failed connections to the node's kubelet, which is +known to be listening because it is starting the Pod. After 60 s it fails +closed, and the job becomes `LAUNCH_FAILED`. Jobs use `dnsPolicy: None` with +approved public resolvers, so they never need cluster DNS. Once rules are in +place, kube-router also filters job traffic to its own node (INPUT and OUTPUT +chains) in k3d. That still has to be shown on the real hosts, together with +the host firewall. + ### Status mapping | Observed | Event | @@ -338,6 +352,24 @@ actual process groups, rather than trusting an accepted field. Reference: The relay/API preserve stream framing and backpressure without buffering unbounded output. +### Implementation notes (M1) + +- **Terminate** shortens the Pod's `activeDeadlineSeconds`, which may only + decrease, instead of deleting the Pod. Kubelet sends SIGTERM, then SIGKILL + after `terminationGracePeriodSeconds` (the job's signal grace), and keeps the + final status, so the result and log tail are saved before the executor + deletes the Pod. +- **`/tmp`** is the scratch emptyDir mounted a second time, not a `subPath`: + kubelet creates subPath directories root-owned, so a non-root job could not + write to them. +- **Pod Security admission** is not enabled on `mist-jobs`: the `restricted` + and `baseline` levels both forbid the hostPath that jobs on the volume's host + need. The pod builder's invariant tests guard the security fields instead, + and only mistd may create Pods there. +- **Receipts** live in `k3s_executions` and `k3s_receipts`. A recorded Pod + UID forbids recreating the Pod. A create whose response was lost is read + back by name and adopted only if its execution ID and spec hash match. + ### Cleanup and reconciliation diff --git a/docs/agents/evidence/m1-step5-k3d.md b/docs/agents/evidence/m1-step5-k3d.md new file mode 100644 index 0000000..3861814 --- /dev/null +++ b/docs/agents/evidence/m1-step5-k3d.md @@ -0,0 +1,50 @@ +# M1 step 5: K3s executor on k3d (30 October 2026) + +This is local integration evidence for the K3s executor. It is **not** evidence +for any M1 release gate: k3d cannot show GPU visibility, NFS, the host firewall +or the production cleanup reporter. Those are steps 7–8 on the real nodes. See +[implementation contracts](../architecture/implementation-contracts.md#m1-acceptance-matrix). + +## Environment + +| Component | Version | +|---|---| +| Host | WSL2, kernel 6.18.33.2-microsoft-standard-WSL2, cgroup v2 | +| Docker | 29.1.3 | +| k3d | v5.9.0 | +| k3s image | `rancher/k3s:v1.36.5-k3s1` | +| containerd (node) | 2.3.4-k3s1.36 | +| Network policy controller | kube-router v2.6.3-k3s1 (built into k3s) | +| Node features | `supplementalGroupsPolicy: true` | +| Runtime image | test stand-in: `ubuntu:22.04` + `mist-job-entry` | + +## Command + +```bash +scripts/k3d-test.sh +``` + +## Results + +All eight `TestK3d*` tests passed. The executor ran impersonating +`system:serviceaccount:mist-system:mistd`, so it had only the permissions in +`deploy/k3s/rbac.yaml`. + +| Test | What it showed | +|---|---| +| `BatchJobIsolation` | `uid=20001(alice) gid=40001(team-team-a) groups=40001(team-team-a),19000(mist-users),20001(alice)`. Only `team-a` under `/mist/teams`. No service-account token, read-only root, writable `/tmp` and `/scratch`, umask `0027`. The API service, the node's API server and kubelet were blocked; `1.1.1.1:53` was reachable. | +| `TerminateForwardsSIGTERM` | The script's TERM trap ran; reason `signaled` | +| `OutOfMemory` | 64 MiB limit; kubelet reported `OOMKilled`, giving reason `out_of_memory` | +| `LocalTimeLimit` | 3 s limit: the wrapper sent SIGTERM; exit 143, reason `deadline` | +| `SetupFailureIsLaunchFailed` | A missing output directory gave `launch_failed`/`launch_error`; the work never ran | +| `ExecutorRestart` | A new executor instance listed the running job from receipts, refused to relaunch it, and stopped it | +| `RBACIsNarrow` | Listing Secrets in `mist-jobs` or Pods in `kube-system` was forbidden | +| `SchedulerEndToEnd` | PENDING → LAUNCHING → RUNNING → CLEANING → COMPLETED; the allocation was held until the node's cleanup receipt | + +## Finding + +Before the isolation gate existed, a job's first connections (within about +1 s of container start) reached the API service, the node's API server and +kubelet. The same probes were blocked about 1.5 s later. kube-router programs +Pod rules after start, so the gate init container was added. See +[executors.md](../architecture/executors.md#network). diff --git a/docs/agents/plans/jobs-and-scheduling.md b/docs/agents/plans/jobs-and-scheduling.md index 37e5014..96b586a 100644 --- a/docs/agents/plans/jobs-and-scheduling.md +++ b/docs/agents/plans/jobs-and-scheduling.md @@ -151,6 +151,27 @@ requirements. Keep secrets and actual research data out of test fixtures. executor backstop; launch timeout is separate from work walltime. - *Check:* golden tests, plus a run on a local k3d cluster using the `cpu` path. + - *Status (30 Oct 2026):* `internal/executor/k3s`, plus `deploy/k3s/` (the + namespace, NetworkPolicy and RBAC) and `mistd cleanup-report`. + - **Unit tests:** golden Pod and ConfigMap files with separate security + invariants, and executor behaviour on client-go's fake clientset. + - **Integration:** `scripts/k3d-test.sh` runs on k3s v1.36.5 (containerd + 2.3.4), as mistd's service account. It covers strict groups (only the + team, `mist-users` and the personal group), other teams absent, no + service-account token, a read-only root, writable scratch, umask 0027, + the network blocked except the internet, SIGTERM forwarding, kernel + OOM, the local time limit, setup failure as LAUNCH_FAILED, executor + restart, narrow RBAC, and scheduler → k3s → node cleanup receipt → + COMPLETED. + - **Found:** NetworkPolicy is applied about a second after the containers + start. Fixed with the isolation gate init container (executors.md). + - **Still open:** + - Cleanup receipts reach the executor through `RecordCleanup`; the + authenticated node → API transport is step 6. + - Orphan ConfigMap collection. + - Throttling `Lost` until the node has been NotReady for 2 min + *and* the scheduler has confirmed it. + - GPU, NFS, host firewall and real-node cleanup, which are steps 7–8. 6. **Job API endpoints.** - Submit (normalise with defaults, validate), list, get, cancel, log (live and tail), partitions, and queue (with anonymisation). diff --git a/go.mod b/go.mod index f390167..935f1a4 100644 --- a/go.mod +++ b/go.mod @@ -6,14 +6,52 @@ require ( github.com/fergusstrange/embedded-postgres v1.34.0 github.com/jackc/pgx/v5 v5.11.0 golang.org/x/sys v0.49.0 + k8s.io/api v0.36.5 + k8s.io/apimachinery v0.36.5 + k8s.io/client-go v0.36.5 + sigs.k8s.io/yaml v1.6.0 ) require ( + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/emicklei/go-restful/v3 v3.13.0 // indirect + github.com/fxamacker/cbor/v2 v2.9.0 // indirect + github.com/go-logr/logr v1.4.3 // indirect + github.com/go-openapi/jsonpointer v0.21.0 // indirect + github.com/go-openapi/jsonreference v0.20.2 // indirect + github.com/go-openapi/swag v0.23.0 // indirect + github.com/google/gnostic-models v0.7.0 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/josharian/intern v1.0.0 // indirect + github.com/json-iterator/go v1.1.12 // indirect github.com/lib/pq v1.10.9 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect + github.com/spf13/pflag v1.0.9 // indirect + github.com/x448/float16 v0.8.4 // indirect github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect - golang.org/x/sync v0.17.0 // indirect - golang.org/x/text v0.29.0 // indirect + go.yaml.in/yaml/v2 v2.4.3 // indirect + go.yaml.in/yaml/v3 v3.0.4 // indirect + golang.org/x/net v0.56.0 // indirect + golang.org/x/oauth2 v0.34.0 // indirect + golang.org/x/sync v0.21.0 // indirect + golang.org/x/term v0.44.0 // indirect + golang.org/x/text v0.39.0 // indirect + golang.org/x/time v0.14.0 // indirect + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect + gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect + gopkg.in/inf.v0 v0.9.1 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect + k8s.io/klog/v2 v2.140.0 // indirect + k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a // indirect + k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 // indirect + sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect + sigs.k8s.io/randfill v1.0.0 // indirect + sigs.k8s.io/structured-merge-diff/v6 v6.3.3 // indirect ) diff --git a/go.sum b/go.sum index 183ad37..2f6546e 100644 --- a/go.sum +++ b/go.sum @@ -1,8 +1,31 @@ +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 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= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/emicklei/go-restful/v3 v3.13.0 h1:C4Bl2xDndpU6nJ4bc1jXd+uTmYPVUwkD6bFY/oTyCes= +github.com/emicklei/go-restful/v3 v3.13.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= github.com/fergusstrange/embedded-postgres v1.34.0 h1:c6RKhPKFsLVU+Tdxsx8q0UxCHsvZZ/iShAnljRBXs6s= github.com/fergusstrange/embedded-postgres v1.34.0/go.mod h1:w0YvnCgf19o6tskInrOOACtnqfVlOvluz3hlNLY7tRk= +github.com/fxamacker/cbor/v2 v2.9.0 h1:NpKPmjDBgUfBms6tr6JZkTHtfFGcMKsw3eGcmD/sapM= +github.com/fxamacker/cbor/v2 v2.9.0/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= +github.com/go-openapi/jsonpointer v0.21.0 h1:YgdVicSA9vH5RiHs9TZW5oyafXZFc6+2Vc1rr/O9oNQ= +github.com/go-openapi/jsonpointer v0.21.0/go.mod h1:IUyH9l/+uyhIYQ/PXVA41Rexl+kOkAPDdXEYns6fzUY= +github.com/go-openapi/jsonreference v0.20.2 h1:3sVjiK66+uXK/6oQ8xgcRKcFgQ5KXa2KvnJRumpMGbE= +github.com/go-openapi/jsonreference v0.20.2/go.mod h1:Bl1zwGIM8/wsvqjsOQLJ/SH+En5Ap4rVB5KVcIDZG2k= +github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= +github.com/go-openapi/swag v0.23.0 h1:vsEVJDUo2hPJ2tu0/Xc+4noaxyEffXNIs3cOULZ+GrE= +github.com/go-openapi/swag v0.23.0/go.mod h1:esZ8ITTYEsH1V2trKHjAN8Ai7xHb8RV+YSZ577vPjgQ= +github.com/google/gnostic-models v0.7.0 h1:qwTtogB15McXDaNqTZdzPJRHvaVJlAl+HVQnLmJEJxo= +github.com/google/gnostic-models v0.7.0/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -11,26 +34,101 @@ github.com/jackc/pgx/v5 v5.11.0 h1:IzBBtyK9AHqf98cctWFifYSci2hgQR/cd56wB4p+ogg= github.com/jackc/pgx/v5 v5.11.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= -github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee h1:W5t00kpgFdJifH4BDsTlE89Zl93FEloxaWZfGcifgq8= +github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= +github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 h1:nIPpBwaJSVYIxUFsDv3M8ofmx9yWTog9BfvIu0q41lo= github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8/go.mod h1:HUYIGzjTL3rfEspMxjDjgmT5uz5wzYJKVo23qUhYTos= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= -golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= -golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +go.yaml.in/yaml/v2 v2.4.3 h1:6gvOSjQoTB3vt1l+CU+tSyi/HOjfOjRLJ4YwYZGwRO0= +go.yaml.in/yaml/v2 v2.4.3/go.mod h1:zSxWcmIDjOzPXpjlTTbAsKokqkDNAVtZO0WOMiT90s8= +go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= +golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= +golang.org/x/oauth2 v0.34.0 h1:hqK/t4AKgbqWkdkcAeI8XLmbK+4m4G5YeQRrmiotGlw= +golang.org/x/oauth2 v0.34.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.49.0 h1:XbzkgYJHdqh/8m2Uu0W/dQv8nktxx4BFHp1M0gROTXA= golang.org/x/sys v0.49.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= -golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= -golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= +golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= +golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= +golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus= +golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM= +golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI= +golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/evanphx/json-patch.v4 v4.13.0 h1:czT3CmqEaQ1aanPc5SdlgQrrEIb8w/wwCvWWnfEbYzo= +gopkg.in/evanphx/json-patch.v4 v4.13.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M= +gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc= +gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +k8s.io/api v0.36.5 h1:vtL/ByHmw7suLt+SGVMPZnuj8QdfZ0kN1vvMvHRCY9c= +k8s.io/api v0.36.5/go.mod h1:erfc3/3d6z30KO2Kf+BmOewE6CajqZTz1Hcd+FQqg60= +k8s.io/apimachinery v0.36.5 h1:X5Xbg6G4om0ma6LOcGbmUmo+B5Nq7mmIpAWhk+tgCh8= +k8s.io/apimachinery v0.36.5/go.mod h1:oPTSicSaDHDdoGwOs5DFPo6YS2wdTnzFHIYMEglfJ7I= +k8s.io/client-go v0.36.5 h1:D79FgevPon0NIBcjUHMy+QnALtihLGTBqOvIGCgoZmw= +k8s.io/client-go v0.36.5/go.mod h1:JAGQ3N4Z8E3H0W0Nm6NBg1EvB9+/wtEoBPDe8XGkHsI= +k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= +k8s.io/klog/v2 v2.140.0/go.mod h1:o+/RWfJ6PwpnFn7OyAG3QnO47BFsymfEfrz6XyYSSp0= +k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a h1:xCeOEAOoGYl2jnJoHkC3hkbPJgdATINPMAxaynU2Ovg= +k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a/go.mod h1:uGBT7iTA6c6MvqUvSXIaYZo9ukscABYi2btjhvgKGZ0= +k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 h1:AZYQSJemyQB5eRxqcPky+/7EdBj0xi3g0ZcxxJ7vbWU= +k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2/go.mod h1:xDxuJ0whA3d0I4mf/C4ppKHxXynQ+fxnkmQH0vTHnuk= +sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 h1:IpInykpT6ceI+QxKBbEflcR5EXP7sU1kvOlxwZh5txg= +sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730/go.mod h1:mdzfpAEoE6DHQEN0uh9ZbOCuHbLK5wOm7dK4ctXE9Tg= +sigs.k8s.io/randfill v1.0.0 h1:JfjMILfT8A6RbawdsK2JXGBR5AQVfd+9TbzrlneTyrU= +sigs.k8s.io/randfill v1.0.0/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY= +sigs.k8s.io/structured-merge-diff/v6 v6.3.3 h1:u08YRbVUi59ri4YD6cg0UqNM4Dimn0sIl+wldcx5PYw= +sigs.k8s.io/structured-merge-diff/v6 v6.3.3/go.mod h1:M3W8sfWvn2HhQDIbGWj3S099YozAsymCo/wrT5ohRUE= +sigs.k8s.io/yaml v1.6.0 h1:G8fkbMSAFqgEFgh4b1wmtzDnioxFCUgTZhlbj5P9QYs= +sigs.k8s.io/yaml v1.6.0/go.mod h1:796bPqUfzR/0jLAl6XjHl3Ck7MiyVv8dbTdyT3/pMf4= diff --git a/internal/db/migrations/executor/0001_k3s_executions.sql b/internal/db/migrations/executor/0001_k3s_executions.sql new file mode 100644 index 0000000..25540c3 --- /dev/null +++ b/internal/db/migrations/executor/0001_k3s_executions.sql @@ -0,0 +1,31 @@ +-- K3s executor receipts (M1 step 5). Only internal/executor/k3s reads or +-- writes these tables. A row outlives its Pod: deleting a Pod never erases +-- what happened to the execution. +CREATE TABLE k3s_executions ( + execution_id uuid PRIMARY KEY, + job_id bigint NOT NULL UNIQUE, -- one execution per job, ever + spec_hash text NOT NULL, + node text NOT NULL, + create_intent_at timestamptz NOT NULL, -- written before the first create call + object_uid text, -- Pod UID once observed; forbids recreation + started_at timestamptz, + finished_at timestamptz, + result jsonb, + log_tail bytea, + terminate_requested_at timestamptz, + lost_at timestamptz, + lost_detail text NOT NULL DEFAULT '', + pod_deleted_at timestamptz, + cleaned_at timestamptz, + cleanup_evidence text NOT NULL DEFAULT '' +); + +CREATE TABLE k3s_receipts ( + id bigserial PRIMARY KEY, + execution_id uuid NOT NULL REFERENCES k3s_executions, + kind text NOT NULL, + at timestamptz NOT NULL, + object_uid text NOT NULL DEFAULT '', + detail text NOT NULL DEFAULT '' +); +CREATE INDEX k3s_receipts_execution ON k3s_receipts (execution_id, id); diff --git a/internal/executor/fake/fake_test.go b/internal/executor/fake/fake_test.go index f813078..6781c43 100644 --- a/internal/executor/fake/fake_test.go +++ b/internal/executor/fake/fake_test.go @@ -41,6 +41,8 @@ func (e *env) spec(jobID int64) executor.LaunchSpec { AuthGeneration: 1, LaunchDeadline: e.clock.Now().Add(10 * time.Minute), Kind: executor.KindBatch, + Partition: "gpu", + QOS: "normal", Node: gpuNode, User: executor.User{Name: "alice", UID: 20001, PersonalGID: 20001}, TeamSlug: "team-a", diff --git a/internal/executor/k3s/cleanup.go b/internal/executor/k3s/cleanup.go new file mode 100644 index 0000000..af68d3a --- /dev/null +++ b/internal/executor/k3s/cleanup.go @@ -0,0 +1,67 @@ +package k3s + +import ( + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" + "regexp" + "strings" +) + +var podUIDPattern = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`) + +// VerifyCleanup runs on the node, as root, and checks that kubelet has +// finished with a Pod: its pod directory (which holds emptyDir scratch and +// volume mounts) is gone, and no cgroup for it remains, so none of its +// processes can still run or hold devices. It only reads; it never deletes +// anything, and it accepts only a well-formed UID, never a path. +// +// Device-plugin release is implied by the Pod's removal from kubelet; TT +// boards (M4) additionally need a reset receipt. +func VerifyCleanup(kubeletRoot, cgroupRoot, podUID string) (string, error) { + if !podUIDPattern.MatchString(podUID) { + return "", fmt.Errorf("cleanup: %q is not a Pod UID", podUID) + } + podDir := filepath.Join(kubeletRoot, "pods", podUID) + if _, err := os.Lstat(podDir); err == nil { + return "", fmt.Errorf("cleanup: %s still exists", podDir) + } else if !errors.Is(err, fs.ErrNotExist) { + return "", err + } + if _, err := os.Stat(filepath.Join(kubeletRoot, "pods")); err != nil { + return "", fmt.Errorf("cleanup: kubelet root %s unreadable: %w", kubeletRoot, err) + } + + // cgroupfs driver: .../kubepods/[burstable/]pod + // systemd driver: .../kubepods-[burstable-]pod.slice + names := []string{"pod" + podUID, "pod" + strings.ReplaceAll(podUID, "-", "_")} + var found string + err := filepath.WalkDir(cgroupRoot, func(p string, d fs.DirEntry, err error) error { + if err != nil { + return nil // a cgroup vanishing mid-walk is fine + } + if !d.IsDir() { + return nil + } + rel, _ := filepath.Rel(cgroupRoot, p) + if strings.Count(rel, string(filepath.Separator)) > 4 { + return filepath.SkipDir // pod cgroups sit a few levels down at most + } + for _, n := range names { + if strings.Contains(d.Name(), n) { + found = p + return filepath.SkipAll + } + } + return nil + }) + if err != nil { + return "", err + } + if found != "" { + return "", fmt.Errorf("cleanup: cgroup %s still exists", found) + } + return fmt.Sprintf("kubelet pod directory and cgroups for Pod %s are gone", podUID), nil +} diff --git a/internal/executor/k3s/cleanup_test.go b/internal/executor/k3s/cleanup_test.go new file mode 100644 index 0000000..54186ce --- /dev/null +++ b/internal/executor/k3s/cleanup_test.go @@ -0,0 +1,40 @@ +package k3s + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestVerifyCleanup(t *testing.T) { + uid := "0b0e6a5c-1f2e-4d3c-9b8a-7f6e5d4c3b2a" + kubelet, cgroups := t.TempDir(), t.TempDir() + os.MkdirAll(filepath.Join(kubelet, "pods", uid, "volumes"), 0o755) + os.MkdirAll(filepath.Join(cgroups, "kubepods", "pod"+uid), 0o755) + + if _, err := VerifyCleanup(kubelet, cgroups, uid); err == nil { + t.Fatal("pod directory present, cleanup confirmed") + } + os.RemoveAll(filepath.Join(kubelet, "pods", uid)) + if _, err := VerifyCleanup(kubelet, cgroups, uid); err == nil || !strings.Contains(err.Error(), "cgroup") { + t.Fatalf("cgroup present: %v", err) + } + os.RemoveAll(filepath.Join(cgroups, "kubepods")) + os.MkdirAll(filepath.Join(cgroups, "kubepods.slice", "kubepods-pod"+strings.ReplaceAll(uid, "-", "_")+".slice"), 0o755) + if _, err := VerifyCleanup(kubelet, cgroups, uid); err == nil { + t.Fatal("systemd-style cgroup missed") + } + os.RemoveAll(filepath.Join(cgroups, "kubepods.slice")) + if ev, err := VerifyCleanup(kubelet, cgroups, uid); err != nil || ev == "" { + t.Fatalf("clean node: %q %v", ev, err) + } + for _, bad := range []string{"../../etc", "", uid + "/x"} { + if _, err := VerifyCleanup(kubelet, cgroups, bad); err == nil { + t.Errorf("%q accepted", bad) + } + } + if _, err := VerifyCleanup(filepath.Join(kubelet, "missing"), cgroups, uid); err == nil { + t.Error("unreadable kubelet root treated as proof") + } +} diff --git a/internal/executor/k3s/executor.go b/internal/executor/k3s/executor.go new file mode 100644 index 0000000..d3ee588 --- /dev/null +++ b/internal/executor/k3s/executor.go @@ -0,0 +1,767 @@ +package k3s + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "strconv" + "strings" + "sync" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/informers" + "k8s.io/client-go/kubernetes" + listersv1 "k8s.io/client-go/listers/core/v1" + "k8s.io/client-go/tools/cache" + "k8s.io/client-go/util/workqueue" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/executor" +) + +// Options are the runtime settings beyond the pod builder's Config. +type Options struct { + Config + Client kubernetes.Interface + Pool *pgxpool.Pool + Clock clock.Clock + Log *slog.Logger + // Resync re-examines every execution this often (default 60 s). + Resync time.Duration + // LostAfter is how long a node must be NotReady before its unfinished + // executions are reported Lost (default 2 min). + LostAfter time.Duration + // TailBytes of output kept in the database (default 64 KiB). + TailBytes int +} + +// Executor implements executor.Executor on k3s. +type Executor struct { + opt Options + client kubernetes.Interface + store store + clock clock.Clock + log *slog.Logger + events *executor.Queue + + pods listersv1.PodLister + nodes listersv1.NodeLister + synced []cache.InformerSynced + queue workqueue.TypedRateLimitingInterface[int64] + started chan struct{} + + jobLocks sync.Map // int64 -> *sync.Mutex; serializes launch/stop/sync per job +} + +// New wires informers but starts nothing; call Start. +func New(opt Options) *Executor { + if opt.Clock == nil { + opt.Clock = clock.Real{} + } + if opt.Log == nil { + opt.Log = slog.New(slog.DiscardHandler) + } + if opt.Resync == 0 { + opt.Resync = 60 * time.Second + } + if opt.LostAfter == 0 { + opt.LostAfter = 2 * time.Minute + } + if opt.TailBytes == 0 { + opt.TailBytes = 64 << 10 + } + e := &Executor{ + opt: opt, + client: opt.Client, + store: store{opt.Pool}, + clock: opt.Clock, + log: opt.Log, + events: executor.NewQueue(), + queue: workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[int64]()), + started: make(chan struct{}), + } + return e +} + +func (e *Executor) Name() string { return "k3s" } + +func (e *Executor) Events() <-chan executor.Event { return e.events.C() } + +// Start runs informers and the sync worker until ctx ends. It returns once +// caches have synced; before that List reports an incomplete observation. +func (e *Executor) Start(ctx context.Context) error { + podFactory := informers.NewSharedInformerFactoryWithOptions(e.client, e.opt.Resync, + informers.WithNamespace(e.opt.Namespace), + informers.WithTweakListOptions(func(o *metav1.ListOptions) { o.LabelSelector = LabelJobID })) + nodeFactory := informers.NewSharedInformerFactory(e.client, e.opt.Resync) + podInf := podFactory.Core().V1().Pods() + nodeInf := nodeFactory.Core().V1().Nodes() + e.pods, e.nodes = podInf.Lister(), nodeInf.Lister() + e.synced = []cache.InformerSynced{podInf.Informer().HasSynced, nodeInf.Informer().HasSynced} + + enqueuePod := func(obj any) { + if tomb, ok := obj.(cache.DeletedFinalStateUnknown); ok { + obj = tomb.Obj + } + if p, ok := obj.(*corev1.Pod); ok { + if id, err := strconv.ParseInt(p.Labels[LabelJobID], 10, 64); err == nil { + e.queue.Add(id) + } + } + } + podInf.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{ + AddFunc: enqueuePod, + UpdateFunc: func(_, obj any) { enqueuePod(obj) }, + DeleteFunc: enqueuePod, + }) + // Node changes can make executions Lost: re-examine everything. + nodeInf.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{ + UpdateFunc: func(_, _ any) { e.enqueueAll(ctx) }, + }) + + podFactory.Start(ctx.Done()) + nodeFactory.Start(ctx.Done()) + if !cache.WaitForCacheSync(ctx.Done(), e.synced...) { + return fmt.Errorf("k3s: informer caches did not sync") + } + close(e.started) + + go func() { + <-ctx.Done() + e.queue.ShutDown() + e.events.Close() + }() + go e.worker(ctx) + go func() { + for { + e.enqueueAll(ctx) + t := e.clock.NewTimer(e.opt.Resync) + select { + case <-ctx.Done(): + t.Stop() + return + case <-t.C(): + } + } + }() + return nil +} + +func (e *Executor) enqueueAll(ctx context.Context) { + recs, err := e.store.all(ctx) + if err != nil { + e.log.Warn("k3s: list records", "err", err) + return + } + for _, r := range recs { + if r.CleanedAt == nil { + e.queue.Add(r.JobID) + } + } +} + +func (e *Executor) worker(ctx context.Context) { + for { + id, quit := e.queue.Get() + if quit { + return + } + if err := e.sync(ctx, id); err != nil { + e.log.Warn("k3s: sync", "job", id, "err", err) + e.queue.AddRateLimited(id) + } else { + e.queue.Forget(id) + } + e.queue.Done(id) + } +} + +func (e *Executor) lock(jobID int64) func() { + m, _ := e.jobLocks.LoadOrStore(jobID, &sync.Mutex{}) + m.(*sync.Mutex).Lock() + return m.(*sync.Mutex).Unlock +} + +func (e *Executor) cachesSynced() bool { + select { + case <-e.started: + default: + return false + } + for _, s := range e.synced { + if !s() { + return false + } + } + return true +} + +// Launch ensures the one Pod for spec exists. See executor.Executor. +func (e *Executor) Launch(ctx context.Context, spec executor.LaunchSpec) error { + cm, pod, err := Build(spec, e.opt.Config) + if err != nil { + return err + } + defer e.lock(spec.JobID)() + + rec, err := e.store.byJob(ctx, spec.JobID) + if err != nil { + return err + } + if rec != nil { + if rec.ExecutionID != spec.ExecutionID || rec.SpecHash != spec.SpecHash { + return executor.ErrSpecConflict + } + if rec.StartedAt != nil || rec.FinishedAt != nil || rec.LostAt != nil { + return executor.ErrExecutionStarted + } + if rec.ObjectUID != "" { + // Read back from the API, not the cache: the Pod must be the very + // object we recorded. If it is gone it may have run; never recreate. + live, err := e.client.CoreV1().Pods(e.opt.Namespace).Get(ctx, pod.Name, metav1.GetOptions{}) + switch { + case apierrors.IsNotFound(err): + return executor.ErrExecutionStarted + case err != nil: + return err + case string(live.UID) != rec.ObjectUID: + return executor.ErrSpecConflict + } + return nil + } + } else if err := e.store.createIntent(ctx, spec, e.clock.Now()); err != nil { + return err + } + + // ConfigMap first, then Pod, then the owner reference. A crash between + // steps is repaired by the next Launch or by orphan collection. + if _, err := e.client.CoreV1().ConfigMaps(e.opt.Namespace).Create(ctx, cm, metav1.CreateOptions{}); err != nil { + if !apierrors.IsAlreadyExists(err) { + return err + } + live, err := e.client.CoreV1().ConfigMaps(e.opt.Namespace).Get(ctx, cm.Name, metav1.GetOptions{}) + if err != nil { + return err + } + if !sameExecution(live.ObjectMeta, spec) { + return executor.ErrSpecConflict // never overwrite another spec + } + } + created, err := e.client.CoreV1().Pods(e.opt.Namespace).Create(ctx, pod, metav1.CreateOptions{}) + if apierrors.IsAlreadyExists(err) { + created, err = e.client.CoreV1().Pods(e.opt.Namespace).Get(ctx, pod.Name, metav1.GetOptions{}) + if err == nil && !sameExecution(created.ObjectMeta, spec) { + return executor.ErrSpecConflict + } + } + if err != nil { + return err // ambiguous: the next Launch reads back by name + } + if err := e.recordUID(ctx, spec.ExecutionID, string(created.UID)); err != nil { + return err + } + owner := fmt.Sprintf(`{"metadata":{"ownerReferences":[{"apiVersion":"v1","kind":"Pod","name":%q,"uid":%q}]}}`, + created.Name, created.UID) + if _, err := e.client.CoreV1().ConfigMaps(e.opt.Namespace).Patch(ctx, cm.Name, types.MergePatchType, + []byte(owner), metav1.PatchOptions{}); err != nil { + e.log.Warn("k3s: set ConfigMap owner", "job", spec.JobID, "err", err) + } + return nil +} + +func sameExecution(m metav1.ObjectMeta, spec executor.LaunchSpec) bool { + return m.Labels[LabelExecutionID] == spec.ExecutionID && m.Annotations[AnnotSpecHash] == spec.SpecHash +} + +func (e *Executor) recordUID(ctx context.Context, execID, uid string) error { + _, err := e.store.update(ctx, execID, executor.ReceiptCreated, e.clock.Now(), uid, "", + "object_uid = $2", "object_uid IS NULL", uid) + return err +} + +// Terminate shortens the Pod's activeDeadlineSeconds rather than deleting +// it: kubelet then sends SIGTERM and SIGKILL after the Pod's grace period, +// and keeps the final status so the result can be saved before deletion. +// The grace is fixed in the Pod at creation (the job's signal grace). +func (e *Executor) Terminate(ctx context.Context, jobID int64, _ time.Duration, reason string) error { + defer e.lock(jobID)() + rec, err := e.store.byJob(ctx, jobID) + if err != nil || rec == nil { + return err + } + if _, err := e.store.update(ctx, rec.ExecutionID, "", e.clock.Now(), "", reason, + "terminate_requested_at = $2", "terminate_requested_at IS NULL", e.clock.Now()); err != nil { + return err + } + return e.stopPod(ctx, jobID) +} + +func (e *Executor) stopPod(ctx context.Context, jobID int64) error { + pods := e.client.CoreV1().Pods(e.opt.Namespace) + p, err := pods.Get(ctx, Name(jobID), metav1.GetOptions{}) + if apierrors.IsNotFound(err) { + return nil // not created yet; sync stops it as soon as it appears + } + if err != nil { + return err + } + if p.Status.Phase == corev1.PodSucceeded || p.Status.Phase == corev1.PodFailed { + return nil + } + secs := int64(1) + if p.Status.StartTime != nil { + secs = max(int64(e.clock.Now().Sub(p.Status.StartTime.Time)/time.Second)+1, 1) + } + if p.Spec.ActiveDeadlineSeconds != nil && *p.Spec.ActiveDeadlineSeconds <= secs { + return nil // already shortened (it may only decrease) + } + patch := fmt.Sprintf(`{"spec":{"activeDeadlineSeconds":%d}}`, secs) + _, err = pods.Patch(ctx, p.Name, types.StrategicMergePatchType, []byte(patch), + metav1.PatchOptions{}) + return err +} + +// sync drives one job's execution from its durable record and the +// observed Pod. It is safe to run any number of times: every fact is +// recorded once, before it is acted on. +func (e *Executor) sync(ctx context.Context, jobID int64) error { + defer e.lock(jobID)() + rec, err := e.store.byJob(ctx, jobID) + if err != nil { + return err + } + pod, err := e.pods.Pods(e.opt.Namespace).Get(Name(jobID)) + if apierrors.IsNotFound(err) { + pod, err = nil, nil + } + if err != nil { + return err + } + if rec == nil || rec.CleanedAt != nil { + return nil // unknown Pods are reported by List for the scheduler to stop + } + if pod != nil && (pod.Labels[LabelExecutionID] != rec.ExecutionID || + (rec.ObjectUID != "" && string(pod.UID) != rec.ObjectUID)) { + e.log.Warn("k3s: Pod does not match the recorded execution", "job", jobID) + return nil + } + + if pod == nil { + switch { + case rec.FinishedAt != nil || rec.LostAt != nil: + if rec.PodDeletedAt == nil { + _, err = e.store.update(ctx, rec.ExecutionID, "", e.clock.Now(), "", "", + "pod_deleted_at = $2", "pod_deleted_at IS NULL", e.clock.Now()) + } + return err // cleanup now waits for the node's receipt + case rec.ObjectUID != "": + // The recorded Pod vanished before we saw it finish. + return e.markLost(ctx, rec, "Pod deleted outside Mist before its result was recorded") + } + return nil // create not yet observed + } + + if rec.ObjectUID == "" { + if err := e.recordUID(ctx, rec.ExecutionID, string(pod.UID)); err != nil { + return err + } + rec.ObjectUID = string(pod.UID) + } + if rec.TerminateRequestedAt != nil || rec.LostAt != nil { + if err := e.stopPod(ctx, jobID); err != nil { + return err + } + } + + cs := jobContainer(pod) + if at := startedAt(cs); at != nil && rec.StartedAt == nil { + applied, err := e.store.update(ctx, rec.ExecutionID, executor.ReceiptStarted, e.clock.Now(), rec.ObjectUID, "", + "started_at = $2", "started_at IS NULL", *at) + if err != nil { + return err + } + if applied { + rec.StartedAt = at + e.emit(rec, executor.EventStarted, nil, *at) + } + } + // Independent backstop in case the Pod deadline and the scheduler both fail. + if rec.StartedAt != nil && rec.FinishedAt == nil && rec.TerminateRequestedAt == nil { + var limit, grace int64 + if pod.Spec.TerminationGracePeriodSeconds != nil { + grace = *pod.Spec.TerminationGracePeriodSeconds + } + if pod.Spec.ActiveDeadlineSeconds != nil { + limit = *pod.Spec.ActiveDeadlineSeconds + } + if limit > 0 && e.clock.Now().After(rec.StartedAt.Add(time.Duration(limit+grace)*time.Second)) { + if err := e.stopPod(ctx, jobID); err != nil { + return err + } + } + } + + if lost, detail := e.nodeLost(pod.Spec.NodeName); lost && rec.FinishedAt == nil && rec.LostAt == nil { + return e.markLost(ctx, rec, detail) + } + if !podDone(pod) { + return nil + } + if rec.FinishedAt == nil { + res, kind := e.result(pod, cs, rec) + at := e.clock.Now() + if cs != nil && cs.State.Terminated != nil && !cs.State.Terminated.FinishedAt.IsZero() { + at = cs.State.Terminated.FinishedAt.Time + } + resJSON, _ := json.Marshal(res) + applied, err := e.store.update(ctx, rec.ExecutionID, executor.ReceiptFinished, e.clock.Now(), rec.ObjectUID, res.Detail, + "finished_at = $2, result = $3, log_tail = $4", "finished_at IS NULL", at, resJSON, e.logTail(ctx, pod)) + if err != nil { + return err + } + if applied { + rec.FinishedAt, rec.Result = &at, &res + e.emit(rec, kind, &res, at) + } + } + // Result and tail are saved: delete the Pod (its ConfigMap follows). + err = e.client.CoreV1().Pods(e.opt.Namespace).Delete(ctx, pod.Name, metav1.DeleteOptions{ + Preconditions: &metav1.Preconditions{UID: &pod.UID}, + }) + if apierrors.IsNotFound(err) || apierrors.IsConflict(err) { + err = nil + } + return err +} + +func (e *Executor) markLost(ctx context.Context, rec *record, detail string) error { + applied, err := e.store.update(ctx, rec.ExecutionID, "", e.clock.Now(), "", detail, + "lost_at = $2, lost_detail = $3", "lost_at IS NULL AND finished_at IS NULL", e.clock.Now(), detail) + if err == nil && applied { + e.emit(rec, executor.EventLost, &executor.Result{Reason: executor.ExitNodeLost, Detail: detail}, e.clock.Now()) + } + return err +} + +// nodeLost reports whether a node has been NotReady for LostAfter. +func (e *Executor) nodeLost(name string) (bool, string) { + n, err := e.nodes.Get(name) + if err != nil { + return false, "" // unknown node: not proof of loss + } + for _, c := range n.Status.Conditions { + if c.Type == corev1.NodeReady && c.Status != corev1.ConditionTrue && + e.clock.Now().Sub(c.LastTransitionTime.Time) >= e.opt.LostAfter { + return true, fmt.Sprintf("node %s unreachable since %s", name, c.LastTransitionTime.UTC().Format(time.RFC3339)) + } + } + return false, "" +} + +func jobContainer(p *corev1.Pod) *corev1.ContainerStatus { + for i := range p.Status.ContainerStatuses { + if p.Status.ContainerStatuses[i].Name == "job" { + return &p.Status.ContainerStatuses[i] + } + } + return nil +} + +func startedAt(cs *corev1.ContainerStatus) *time.Time { + switch { + case cs == nil: + case cs.State.Running != nil && !cs.State.Running.StartedAt.IsZero(): + return &cs.State.Running.StartedAt.Time + case cs.State.Terminated != nil && !cs.State.Terminated.StartedAt.IsZero(): + return &cs.State.Terminated.StartedAt.Time + } + return nil +} + +func podDone(p *corev1.Pod) bool { + if p.Status.Phase == corev1.PodSucceeded || p.Status.Phase == corev1.PodFailed { + return true + } + cs := jobContainer(p) + return cs != nil && cs.State.Terminated != nil +} + +// result works out the outcome. Kubernetes' own observations (OOMKilled, +// eviction, deadline) outrank the wrapper's message, which a dying or +// hostile process could have written or failed to write. +func (e *Executor) result(p *corev1.Pod, cs *corev1.ContainerStatus, rec *record) (executor.Result, executor.EventKind) { + var res executor.Result + stopped := rec.TerminateRequestedAt != nil + term := (*corev1.ContainerStateTerminated)(nil) + if cs != nil { + term = cs.State.Terminated + } + if term == nil { + // The work never started. + switch { + case stopped: + return executor.Result{Reason: executor.ExitSignaled, Detail: "stopped before start"}, executor.EventFinished + case p.Status.Reason == "DeadlineExceeded": + return executor.Result{Reason: executor.ExitLaunchTimeout, Detail: "did not start before the Pod deadline"}, executor.EventLaunchFailed + } + detail := strings.TrimSpace(p.Status.Reason + " " + p.Status.Message) + for _, ic := range p.Status.InitContainerStatuses { + if t := ic.State.Terminated; t != nil && t.ExitCode != 0 { + detail = fmt.Sprintf("%s failed: %s", ic.Name, strings.TrimSpace(t.Message)) + } + } + return executor.Result{Reason: executor.ExitLaunchError, Detail: detail}, executor.EventLaunchFailed + } + + json.Unmarshal([]byte(term.Message), &res) // wrapper report, advisory + code := int(term.ExitCode) + res.ExitCode = &code + if term.Message == "" { + res.OutputIncomplete = true + } + switch { + case term.Reason == "OOMKilled": + res.Reason = executor.ExitOutOfMemory + if res.Detail == "" { + res.Detail = "killed by the kernel for exceeding the memory limit" + } + case p.Status.Reason == "Evicted": + res.Reason, res.Detail = executor.ExitEvicted, p.Status.Message + if strings.Contains(p.Status.Message, "ephemeral") { + res.Reason = executor.ExitScratchExceeded + } + case p.Status.Reason == "DeadlineExceeded" && !stopped: + res.Reason, res.Detail = executor.ExitDeadline, "killed by the Pod deadline backstop" + case stopped && res.Reason != executor.ExitDeadline: + res.Reason = executor.ExitSignaled + case res.Reason == "": + res.Reason = executor.ExitNormal + if code == 137 { + res.Detail = "killed by SIGKILL; output may be incomplete" + } + } + if res.Reason == executor.ExitLaunchError { + return res, executor.EventLaunchFailed + } + return res, executor.EventFinished +} + +func (e *Executor) logTail(ctx context.Context, p *corev1.Pod) []byte { + b, err := e.client.CoreV1().Pods(p.Namespace).GetLogs(p.Name, &corev1.PodLogOptions{ + Container: "job", TailLines: ptr[int64](2000), + }).DoRaw(ctx) + if err != nil { + e.log.Warn("k3s: log tail", "job", p.Labels[LabelJobID], "err", err) + return nil + } + if len(b) > e.opt.TailBytes { + b = b[len(b)-e.opt.TailBytes:] + } + return b +} + +func (e *Executor) emit(rec *record, kind executor.EventKind, r *executor.Result, at time.Time) { + ev := executor.Event{ + EventID: rec.ExecutionID + "/" + string(kind), JobID: rec.JobID, ExecutionID: rec.ExecutionID, + Kind: kind, Node: rec.Node, At: at, + } + if r != nil { + ev.ExitCode, ev.Reason, ev.Detail = r.ExitCode, r.Reason, r.Detail + } + e.events.Push(ev) +} + +// CleanupReceipt is a node's proof that an execution left nothing behind. +type CleanupReceipt struct { + ExecutionID string `json:"execution_id"` + PodUID string `json:"pod_uid"` + Node string `json:"node"` + Evidence string `json:"evidence"` +} + +// RecordCleanup accepts a node cleanup receipt. Only a receipt matching +// the recorded execution, Pod UID and node, for a Pod the API no longer +// has, confirms cleanup and lets the scheduler release resources. +func (e *Executor) RecordCleanup(ctx context.Context, r CleanupReceipt) error { + recs, err := e.store.all(ctx) + if err != nil { + return err + } + var rec *record + for _, x := range recs { + if x.ExecutionID == r.ExecutionID { + rec = x + } + } + if rec == nil || rec.ObjectUID == "" || rec.ObjectUID != r.PodUID || rec.Node != r.Node { + return fmt.Errorf("k3s: cleanup receipt does not match a recorded execution") + } + defer e.lock(rec.JobID)() + if rec.FinishedAt == nil && rec.LostAt == nil { + return fmt.Errorf("k3s: execution %s has no recorded outcome yet", rec.ExecutionID) + } + live, err := e.client.CoreV1().Pods(e.opt.Namespace).Get(ctx, Name(rec.JobID), metav1.GetOptions{}) + if err == nil && string(live.UID) == rec.ObjectUID { + return fmt.Errorf("k3s: Pod %s still exists", live.Name) + } else if err != nil && !apierrors.IsNotFound(err) { + return err + } + applied, err := e.store.update(ctx, rec.ExecutionID, executor.ReceiptCleaned, e.clock.Now(), r.PodUID, r.Evidence, + "cleaned_at = $2, cleanup_evidence = $3", "cleaned_at IS NULL", e.clock.Now(), r.Evidence) + if err == nil && applied { + res := rec.Result + if res == nil { + res = &executor.Result{Reason: executor.ExitNodeLost} + } + x := *res + x.Detail = r.Evidence + e.emit(rec, executor.EventExited, &x, e.clock.Now()) + } + return err +} + +// Nodes reports node allocatable capacity. A node is healthy only if it +// is Ready, schedulable, under no pressure and supports strict +// supplemental groups; otherwise jobs could keep image-defined groups. +func (e *Executor) Nodes(context.Context) ([]executor.Node, error) { + if !e.cachesSynced() { + return nil, executor.ErrIncompleteObservation + } + list, err := e.nodes.List(labels.Everything()) + if err != nil { + return nil, err + } + out := make([]executor.Node, 0, len(list)) + for _, n := range list { + a, c := n.Status.Allocatable, n.Status.Capacity + node := executor.Node{ + Name: n.Name, + Capacity: executor.Resources{ + MilliCPU: a.Cpu().MilliValue(), + MemoryBytes: a.Memory().Value(), + GPUs: quantity(a, "nvidia.com/gpu"), + ScratchBytes: a.StorageEphemeral().Value(), + }, + Reserved: executor.Resources{ + MilliCPU: c.Cpu().MilliValue() - a.Cpu().MilliValue(), + MemoryBytes: c.Memory().Value() - a.Memory().Value(), + }, + Healthy: true, + } + var why []string + for _, cond := range n.Status.Conditions { + bad := cond.Status == corev1.ConditionTrue + if cond.Type == corev1.NodeReady { + bad = cond.Status != corev1.ConditionTrue + } + if bad { + why = append(why, string(cond.Type)) + } + } + if n.Spec.Unschedulable { + why = append(why, "cordoned") + } + if f := n.Status.Features; f == nil || f.SupplementalGroupsPolicy == nil || !*f.SupplementalGroupsPolicy { + why = append(why, "no strict supplemental groups support") + } + if len(why) > 0 { + node.Healthy, node.Detail = false, strings.Join(why, ", ") + } + out = append(out, node) + } + return out, nil +} + +func quantity(l corev1.ResourceList, name corev1.ResourceName) int64 { + if q, ok := l[name]; ok { + return q.Value() + } + return 0 +} + +// List merges durable records with live Pods. Pods without a record are +// reported as unknown executions for the scheduler to stop. +func (e *Executor) List(ctx context.Context) ([]executor.Execution, error) { + if !e.cachesSynced() { + return nil, fmt.Errorf("k3s: caches not synced: %w", executor.ErrIncompleteObservation) + } + recs, err := e.store.all(ctx) + if err != nil { + return nil, fmt.Errorf("k3s: %w: %v", executor.ErrIncompleteObservation, err) + } + seen := map[int64]bool{} + var out []executor.Execution + for _, r := range recs { + seen[r.JobID] = true + x := executor.Execution{ + JobID: r.JobID, ExecutionID: r.ExecutionID, SpecHash: r.SpecHash, Node: r.Node, + ObjectUID: r.ObjectUID, StartedAt: r.StartedAt, FinishedAt: r.FinishedAt, Result: r.Result, + CleanupEvidence: r.CleanupEvidence, + } + switch { + case r.CleanedAt != nil: + x.Phase, x.Cleanup = executor.PhaseCleaned, executor.CleanupConfirmed + case r.LostAt != nil: + x.Phase, x.Cleanup = executor.PhaseLost, executor.CleanupPending + case r.FinishedAt != nil: + x.Phase, x.Cleanup = executor.PhaseFinished, executor.CleanupPending + case r.StartedAt != nil: + x.Phase = executor.PhaseRunning + case r.ObjectUID != "": + x.Phase = executor.PhasePending + default: + x.Phase = executor.PhaseCreating + } + out = append(out, x) + } + pods, err := e.pods.Pods(e.opt.Namespace).List(labels.Everything()) + if err != nil { + return nil, err + } + for _, p := range pods { + id, err := strconv.ParseInt(p.Labels[LabelJobID], 10, 64) + if err != nil || seen[id] { + continue + } + out = append(out, executor.Execution{ + JobID: id, ExecutionID: p.Labels[LabelExecutionID], Node: p.Spec.NodeName, + ObjectUID: string(p.UID), Phase: executor.PhaseRunning, + }) + } + return out, nil +} + +// Logs streams live output while the Pod exists, then the stored tail. +func (e *Executor) Logs(ctx context.Context, jobID int64, follow bool) (io.ReadCloser, error) { + rc, err := e.client.CoreV1().Pods(e.opt.Namespace).GetLogs(Name(jobID), &corev1.PodLogOptions{ + Container: "job", Follow: follow, + }).Stream(ctx) + if err == nil { + return rc, nil + } + rec, rerr := e.store.byJob(ctx, jobID) + if rerr != nil || rec == nil { + return nil, executor.ErrNotFound + } + return io.NopCloser(strings.NewReader(string(rec.LogTail))), nil +} + +func (e *Executor) Attach(context.Context, int64, executor.Terminal) error { + return errors.New("k3s: attach arrives in M2") +} + +// Receipts returns an execution's durable receipts, for diagnostics. +func (e *Executor) Receipts(ctx context.Context, executionID string) ([]executor.Receipt, error) { + return e.store.receipts(ctx, executionID) +} + +var _ executor.Executor = (*Executor)(nil) diff --git a/internal/executor/k3s/executor_test.go b/internal/executor/k3s/executor_test.go new file mode 100644 index 0000000..892d0f3 --- /dev/null +++ b/internal/executor/k3s/executor_test.go @@ -0,0 +1,387 @@ +package k3s + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/db/dbtest" + "github.com/UTMIST/Mist/internal/executor" +) + +func TestMain(m *testing.M) { dbtest.Main(m) } + +// kube plays kubelet: tests move Pods through their lifecycle by writing +// status, as kubelet would, and the executor reacts through its informers. +type kube struct { + t *testing.T + ctx context.Context + client *fake.Clientset + clock *clock.Fake + ex *Executor +} + +var now0 = time.Date(2026, 11, 2, 15, 0, 0, 0, time.UTC) + +func newKube(t *testing.T, nodes ...*corev1.Node) *kube { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + pool, err := db.Open(ctx, dbtest.URL(t)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := db.Migrate(ctx, pool); err != nil { + t.Fatal(err) + } + if len(nodes) == 0 { + nodes = []*corev1.Node{readyNode("utmist-z1opa08", true)} + } + objs := make([]runtime.Object, 0, len(nodes)) + for _, n := range nodes { + objs = append(objs, n) + } + client := fake.NewSimpleClientset(objs...) + // The real API server assigns UIDs; the fake does not. + uids := 0 + client.PrependReactor("create", "*", func(a k8stesting.Action) (bool, runtime.Object, error) { + if o, ok := a.(k8stesting.CreateAction).GetObject().(metav1.Object); ok && o.GetUID() == "" { + uids++ + o.SetUID(types.UID(fmt.Sprintf("uid-%d", uids))) + } + return false, nil, nil + }) + k := &kube{t: t, ctx: ctx, client: client, clock: clock.NewFake(now0)} + return k.start(pool) +} + +func (k *kube) start(pool *pgxpool.Pool) *kube { + k.ex = New(Options{Config: testConfig(), Client: k.client, Pool: pool, Clock: k.clock}) + if err := k.ex.Start(k.ctx); err != nil { + k.t.Fatal(err) + } + return k +} + +func readyNode(name string, strict bool) *corev1.Node { + n := &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{Name: name}, + Status: corev1.NodeStatus{ + Capacity: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("32"), corev1.ResourceMemory: resource.MustParse("128Gi")}, + Allocatable: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("30"), corev1.ResourceMemory: resource.MustParse("120Gi"), + corev1.ResourceEphemeralStorage: resource.MustParse("400Gi"), "nvidia.com/gpu": resource.MustParse("2")}, + Conditions: []corev1.NodeCondition{{Type: corev1.NodeReady, Status: corev1.ConditionTrue, + LastTransitionTime: metav1.NewTime(now0.Add(-time.Hour))}}, + Features: &corev1.NodeFeatures{SupplementalGroupsPolicy: ptr(strict)}, + }, + } + return n +} + +func (k *kube) next(kind executor.EventKind) executor.Event { + k.t.Helper() + select { + case ev := <-k.ex.Events(): + if ev.Kind != kind { + k.t.Fatalf("got %s (%s %s), want %s", ev.Kind, ev.Reason, ev.Detail, kind) + } + return ev + case <-time.After(5 * time.Second): + k.t.Fatalf("no %s event", kind) + return executor.Event{} + } +} + +func (k *kube) quiet() { + k.t.Helper() + select { + case ev := <-k.ex.Events(): + k.t.Fatalf("unexpected %s (%s)", ev.Kind, ev.Detail) + case <-time.After(100 * time.Millisecond): + } +} + +func (k *kube) pod(id int64) *corev1.Pod { + k.t.Helper() + p, err := k.client.CoreV1().Pods("mist-jobs").Get(k.ctx, Name(id), metav1.GetOptions{}) + if err != nil { + k.t.Fatal(err) + } + return p +} + +func (k *kube) setStatus(id int64, mutate func(*corev1.PodStatus)) { + k.t.Helper() + p := k.pod(id) + mutate(&p.Status) + if _, err := k.client.CoreV1().Pods("mist-jobs").UpdateStatus(k.ctx, p, metav1.UpdateOptions{}); err != nil { + k.t.Fatal(err) + } +} + +func running(at time.Time) func(*corev1.PodStatus) { + return func(s *corev1.PodStatus) { + s.Phase, s.StartTime = corev1.PodRunning, &metav1.Time{Time: at} + s.ContainerStatuses = []corev1.ContainerStatus{{Name: "job", State: corev1.ContainerState{ + Running: &corev1.ContainerStateRunning{StartedAt: metav1.NewTime(at)}}}} + } +} + +func terminated(code int32, reason string, msg *executor.Result, podReason string) func(*corev1.PodStatus) { + return func(s *corev1.PodStatus) { + var m []byte + if msg != nil { + m, _ = json.Marshal(msg) + } + s.Phase, s.Reason = corev1.PodFailed, podReason + if code == 0 { + s.Phase = corev1.PodSucceeded + } + s.ContainerStatuses = []corev1.ContainerStatus{{Name: "job", State: corev1.ContainerState{ + Terminated: &corev1.ContainerStateTerminated{ExitCode: code, Reason: reason, Message: string(m), + StartedAt: metav1.NewTime(now0.Add(time.Second)), FinishedAt: metav1.NewTime(now0.Add(time.Minute))}}}} + } +} + +func (k *kube) waitGone(id int64) { + k.t.Helper() + deadline := time.Now().Add(5 * time.Second) + for { + _, err := k.client.CoreV1().Pods("mist-jobs").Get(k.ctx, Name(id), metav1.GetOptions{}) + if err != nil { + return + } + if time.Now().After(deadline) { + k.t.Fatal("Pod was not deleted after its result was saved") + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestLaunchIsRepeatSafe(t *testing.T) { + k := newKube(t) + spec := gpuBatch() + if err := k.ex.Launch(k.ctx, spec); err != nil { + t.Fatal(err) + } + if err := k.ex.Launch(k.ctx, spec); err != nil { + t.Fatalf("repeat: %v", err) + } + if _, err := k.client.CoreV1().ConfigMaps("mist-jobs").Get(k.ctx, Name(spec.JobID), metav1.GetOptions{}); err != nil { + t.Fatal(err) + } + cm, _ := k.client.CoreV1().ConfigMaps("mist-jobs").Get(k.ctx, Name(spec.JobID), metav1.GetOptions{}) + if len(cm.OwnerReferences) != 1 || cm.OwnerReferences[0].UID != k.pod(spec.JobID).UID { + t.Fatal("ConfigMap not owned by its Pod") + } + other := spec + other.Resources.GPUs = 2 + other.Seal() + if err := k.ex.Launch(k.ctx, other); !errors.Is(err, executor.ErrSpecConflict) { + t.Fatalf("conflicting spec: %v", err) + } +} + +// A create whose response was lost: the Pod exists but no UID was recorded. +// The retry adopts it instead of failing or creating a second one. +func TestLaunchAdoptsAmbiguousCreate(t *testing.T) { + k := newKube(t) + spec := gpuBatch() + if err := k.ex.store.createIntent(k.ctx, spec, now0); err != nil { + t.Fatal(err) + } + cm, pod, _ := Build(spec, testConfig()) + k.client.CoreV1().ConfigMaps("mist-jobs").Create(k.ctx, cm, metav1.CreateOptions{}) + k.client.CoreV1().Pods("mist-jobs").Create(k.ctx, pod, metav1.CreateOptions{}) + if err := k.ex.Launch(k.ctx, spec); err != nil { + t.Fatal(err) + } + rec, _ := k.ex.store.byJob(k.ctx, spec.JobID) + if rec.ObjectUID != string(k.pod(spec.JobID).UID) { + t.Fatal("did not adopt the existing Pod") + } +} + +func TestLifecycleAndCleanup(t *testing.T) { + k := newKube(t) + spec := gpuBatch() + k.ex.Launch(k.ctx, spec) + k.setStatus(spec.JobID, running(now0.Add(time.Second))) + if ev := k.next(executor.EventStarted); !ev.At.Equal(now0.Add(time.Second)) { + t.Fatalf("start time %v is not the container's", ev.At) + } + uid := string(k.pod(spec.JobID).UID) + k.setStatus(spec.JobID, terminated(0, "Completed", &executor.Result{Reason: executor.ExitNormal}, "")) + ev := k.next(executor.EventFinished) + if *ev.ExitCode != 0 || ev.Reason != executor.ExitNormal { + t.Fatalf("finished %+v", ev) + } + k.waitGone(spec.JobID) + rec, _ := k.ex.store.byJob(k.ctx, spec.JobID) + if len(rec.LogTail) == 0 { + t.Fatal("log tail not saved before deletion") + } + k.quiet() // deletion alone is not cleanup + + for _, bad := range []CleanupReceipt{ + {ExecutionID: spec.ExecutionID, PodUID: "someone-else", Node: spec.Node}, + {ExecutionID: spec.ExecutionID, PodUID: uid, Node: "utmist-tt"}, + } { + if err := k.ex.RecordCleanup(k.ctx, bad); err == nil { + t.Fatalf("mismatched receipt accepted: %+v", bad) + } + } + good := CleanupReceipt{ExecutionID: spec.ExecutionID, PodUID: uid, Node: spec.Node, Evidence: "pod dir and cgroup gone"} + if err := k.ex.RecordCleanup(k.ctx, good); err != nil { + t.Fatal(err) + } + k.next(executor.EventExited) + k.ex.RecordCleanup(k.ctx, good) // repeated receipt: no second event + k.quiet() + + if err := k.ex.Launch(k.ctx, spec); !errors.Is(err, executor.ErrExecutionStarted) { + t.Fatalf("relaunch after finish: %v", err) + } +} + +func TestResults(t *testing.T) { + cases := []struct { + name string + terminate bool + status func(*corev1.PodStatus) + kind executor.EventKind + reason executor.ExitReason + }{ + {"kernel OOM outranks wrapper", false, + terminated(137, "OOMKilled", &executor.Result{Reason: executor.ExitNormal}, ""), + executor.EventFinished, executor.ExitOutOfMemory}, + {"wrapper killed, no message", false, terminated(137, "Error", nil, ""), + executor.EventFinished, executor.ExitNormal}, + {"scratch eviction", false, terminated(137, "Error", nil, "Evicted"), + executor.EventFinished, executor.ExitEvicted}, + {"Pod deadline backstop", false, terminated(143, "Error", nil, "DeadlineExceeded"), + executor.EventFinished, executor.ExitDeadline}, + {"terminated by Mist", true, terminated(143, "Error", &executor.Result{Reason: executor.ExitSignaled}, "DeadlineExceeded"), + executor.EventFinished, executor.ExitSignaled}, + {"wrapper setup failure", false, terminated(125, "Error", &executor.Result{Reason: executor.ExitLaunchError, Detail: "bad output path"}, ""), + executor.EventLaunchFailed, executor.ExitLaunchError}, + {"never started", false, func(s *corev1.PodStatus) { s.Phase, s.Reason = corev1.PodFailed, "DeadlineExceeded" }, + executor.EventLaunchFailed, executor.ExitLaunchTimeout}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + k := newKube(t) + spec := gpuBatch() + k.ex.Launch(k.ctx, spec) + if c.terminate { + if err := k.ex.Terminate(k.ctx, spec.JobID, time.Minute, "cancelled"); err != nil { + t.Fatal(err) + } + if d := k.pod(spec.JobID).Spec.ActiveDeadlineSeconds; d == nil || *d > 2 { + t.Fatalf("Terminate did not shorten the deadline: %v", d) + } + } + k.setStatus(spec.JobID, c.status) + if c.name != "never started" { // the container itself started + k.next(executor.EventStarted) + } + ev := k.next(c.kind) + if ev.Reason != c.reason { + t.Fatalf("reason %s (%s), want %s", ev.Reason, ev.Detail, c.reason) + } + }) + } +} + +func TestPodDeletedOutsideMistIsLost(t *testing.T) { + k := newKube(t) + spec := gpuBatch() + k.ex.Launch(k.ctx, spec) + k.setStatus(spec.JobID, running(now0)) + k.next(executor.EventStarted) + k.client.CoreV1().Pods("mist-jobs").Delete(k.ctx, Name(spec.JobID), metav1.DeleteOptions{}) + k.next(executor.EventLost) + if err := k.ex.Launch(k.ctx, spec); !errors.Is(err, executor.ErrExecutionStarted) { + t.Fatalf("relaunch of lost execution: %v", err) + } +} + +func TestNodeNotReadyIsLostAfterGrace(t *testing.T) { + k := newKube(t) + spec := gpuBatch() + k.ex.Launch(k.ctx, spec) + k.setStatus(spec.JobID, running(now0)) + k.next(executor.EventStarted) + + n, _ := k.client.CoreV1().Nodes().Get(k.ctx, spec.Node, metav1.GetOptions{}) + n.Status.Conditions[0].Status = corev1.ConditionUnknown + n.Status.Conditions[0].LastTransitionTime = metav1.NewTime(now0) + k.client.CoreV1().Nodes().UpdateStatus(k.ctx, n, metav1.UpdateOptions{}) + k.quiet() // not yet 2 minutes + nodes, _ := k.ex.Nodes(k.ctx) + if nodes[0].Healthy { + t.Fatal("NotReady node reported healthy") + } + + k.clock.Set(now0.Add(2 * time.Minute)) + n.Status.Conditions[0].Message = "still unreachable" // any update triggers a resync + k.client.CoreV1().Nodes().UpdateStatus(k.ctx, n, metav1.UpdateOptions{}) + k.next(executor.EventLost) +} + +func TestNodesRequireStrictGroups(t *testing.T) { + k := newKube(t, readyNode("utmist-z1opa08", true), readyNode("old-node", false)) + nodes, err := k.ex.Nodes(k.ctx) + if err != nil { + t.Fatal(err) + } + for _, n := range nodes { + if n.Name == "old-node" && n.Healthy { + t.Fatal("node without strict supplemental groups reported healthy") + } + if n.Name == "utmist-z1opa08" && (!n.Healthy || n.Capacity.GPUs != 2 || n.Capacity.MilliCPU != 30000 || n.Reserved.MilliCPU != 2000) { + t.Fatalf("node %+v", n) + } + } +} + +func TestListReportsUnknownPods(t *testing.T) { + k := newKube(t) + _, pod, _ := Build(cpuCommand(), testConfig()) + k.client.CoreV1().Pods("mist-jobs").Create(k.ctx, pod, metav1.CreateOptions{}) // nobody recorded it + deadline := time.Now().Add(5 * time.Second) + for { + xs, err := k.ex.List(k.ctx) + if err != nil { + t.Fatal(err) + } + if len(xs) == 1 && xs[0].JobID == 413 && xs[0].Phase == executor.PhaseRunning { + return + } + if time.Now().After(deadline) { + t.Fatalf("unknown Pod not listed: %+v", xs) + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestListBeforeSyncIsIncomplete(t *testing.T) { + ex := New(Options{Config: testConfig(), Client: fake.NewSimpleClientset()}) + if _, err := ex.List(context.Background()); !errors.Is(err, executor.ErrIncompleteObservation) { + t.Fatalf("want incomplete observation, got %v", err) + } +} diff --git a/internal/executor/k3s/k3d_test.go b/internal/executor/k3s/k3d_test.go new file mode 100644 index 0000000..a2e9130 --- /dev/null +++ b/internal/executor/k3s/k3d_test.go @@ -0,0 +1,420 @@ +//go:build k3d + +// Integration tests on a real k3s (k3d) cluster: run scripts/k3d-test.sh. +// They prove Pod construction, informers, status mapping, strict groups, +// read-only root, NetworkPolicy, OOM and cleanup receipts on real kubelet +// and containerd. They do not prove GPU, NFS, host firewall or production +// cleanup-reporter behaviour; those are checked on the real nodes (step 8). +package k3s + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "strings" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/tools/clientcmd" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/db/dbtest" + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity/identitytest" + "github.com/UTMIST/Mist/internal/scheduler" +) + +type cluster struct { + t *testing.T + ctx context.Context + cancel context.CancelFunc + node string + client kubernetes.Interface + pool *pgxpool.Pool + ex *Executor + + mu sync.Mutex + events map[int64][]executor.Event +} + +var nextJob = time.Now().Unix() % 1_000_000 * 100 // unique Pod names across runs + +func newCluster(t *testing.T) *cluster { + kctx, node := os.Getenv("MIST_K3D_CONTEXT"), os.Getenv("MIST_K3D_NODE") + if kctx == "" || node == "" { + t.Skip("run through scripts/k3d-test.sh") + } + cfg, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( + clientcmd.NewDefaultClientConfigLoadingRules(), + &clientcmd.ConfigOverrides{CurrentContext: kctx}).ClientConfig() + if err != nil { + t.Fatal(err) + } + // Act as mistd's service account, so deploy/k3s/rbac.yaml is what is tested. + cfg.Impersonate.UserName = "system:serviceaccount:mist-system:mistd" + client, err := kubernetes.NewForConfig(cfg) + if err != nil { + t.Fatal(err) + } + pool, err := db.Open(context.Background(), dbtest.URL(t)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := db.Migrate(context.Background(), pool); err != nil { + t.Fatal(err) + } + c := &cluster{t: t, node: node, client: client, pool: pool, events: map[int64][]executor.Event{}} + c.startExecutor() + return c +} + +func (c *cluster) startExecutor() { + if c.cancel != nil { + c.cancel() + } + c.ctx, c.cancel = context.WithCancel(context.Background()) + c.t.Cleanup(c.cancel) + c.ex = New(Options{ + Config: Config{ + Namespace: "mist-jobs", + EntryPath: "/usr/local/bin/mist-job-entry", + LaunchTimeout: 10 * time.Minute, + BackstopSlack: time.Minute, + Nameservers: []string{"1.1.1.1"}, + Volumes: map[string]Volume{"vol1": {HostNode: c.node, HostPath: "/mist"}}, + }, + Client: c.client, + Pool: c.pool, + Resync: 5 * time.Second, + }) + if err := c.ex.Start(c.ctx); err != nil { + c.t.Fatal(err) + } + ex := c.ex + go func() { + for ev := range ex.Events() { + c.mu.Lock() + c.events[ev.JobID] = append(c.events[ev.JobID], ev) + c.mu.Unlock() + } + }() +} + +func (c *cluster) spec(script string, mutate ...func(*executor.LaunchSpec)) executor.LaunchSpec { + nextJob++ + id := nextJob + s := executor.LaunchSpec{ + JobID: id, ExecutionID: executor.NewExecutionID(), AuthGeneration: 1, + LaunchDeadline: time.Now().Add(10 * time.Minute), + Kind: executor.KindBatch, JobName: c.t.Name(), Partition: "cpu", QOS: "normal", + Node: c.node, + User: executor.User{Name: "alice", UID: 20001, PersonalGID: 20001}, + TeamSlug: "team-a", TeamGID: 40001, ExtraGIDs: []int{19000}, + Resources: executor.Resources{MilliCPU: 500, MemoryBytes: 256 << 20, ScratchBytes: 1 << 30}, + Mounts: []executor.Mount{ + {VolumeID: "vol1", Source: "/teams/team-a", Target: "/mist/teams/team-a"}, + {VolumeID: "vol1", Source: "/software", Target: "/mist/software", ReadOnly: true}, + }, + Workdir: "/mist/teams/team-a/alice", + OutputPath: fmt.Sprintf("/mist/teams/team-a/alice/mist-%d.out", id), + Script: script, + Runtime: "mist-test-runtime:dev", + TimeLimit: 5 * time.Minute, + SignalGrace: 5 * time.Second, + } + for _, m := range mutate { + m(&s) + } + s.Seal() + return s +} + +func (c *cluster) launch(s executor.LaunchSpec) { + c.t.Helper() + if err := c.ex.Launch(c.ctx, s); err != nil { + c.t.Fatal(err) + } +} + +func (c *cluster) wait(job int64, kind executor.EventKind) executor.Event { + c.t.Helper() + deadline := time.Now().Add(3 * time.Minute) + for time.Now().Before(deadline) { + c.mu.Lock() + for _, ev := range c.events[job] { + if ev.Kind == kind { + c.mu.Unlock() + return ev + } + } + c.mu.Unlock() + time.Sleep(200 * time.Millisecond) + } + p, _ := c.client.CoreV1().Pods("mist-jobs").Get(c.ctx, Name(job), metav1.GetOptions{}) + c.t.Fatalf("job %d: no %s event; events %+v; pod %+v", job, kind, c.events[job], p.Status) + return executor.Event{} +} + +// cleanup runs the node-side check inside the k3d node and delivers its +// receipt, retrying while kubelet is still tearing the Pod down. +func (c *cluster) cleanup(s executor.LaunchSpec) { + c.t.Helper() + rec, err := c.ex.store.byJob(c.ctx, s.JobID) + if err != nil || rec == nil { + c.t.Fatalf("no record: %v", err) + } + var last string + for range 60 { + out, err := exec.Command("docker", "exec", c.node, "/mistd", "cleanup-report", + "-pod-uid", rec.ObjectUID, "-execution-id", s.ExecutionID, "-node", c.node).CombinedOutput() + if err == nil { + var r CleanupReceipt + if err := json.Unmarshal(out, &r); err != nil { + c.t.Fatal(err) + } + if err := c.ex.RecordCleanup(c.ctx, r); err != nil { + c.t.Fatal(err) + } + c.wait(s.JobID, executor.EventExited) + return + } + last = string(out) + time.Sleep(time.Second) + } + c.t.Fatalf("node never confirmed cleanup: %s", last) +} + +func (c *cluster) output(s executor.LaunchSpec) string { + out, err := exec.Command("docker", "exec", c.node, "cat", s.OutputPath).CombinedOutput() + if err != nil { + c.t.Fatalf("read output: %v %s", err, out) + } + return string(out) +} + +func TestK3dBatchJobIsolation(t *testing.T) { + c := newCluster(t) + nodeIP, _ := exec.Command("docker", "inspect", "-f", "{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}", c.node).Output() + s := c.spec(` +echo "id: $(id)" +echo "teams: $(ls /mist/teams)" +[ -e /var/run/secrets/kubernetes.io/serviceaccount ] || echo "no-sa-token" +touch /etc/probe 2>/dev/null || echo "rootfs-readonly" +echo x > /tmp/t && echo x > /scratch/s && echo "scratch-writable" +echo "umask $(umask)" +probe() { timeout 3 bash -c "/dev/null && echo "open $1:$2" || echo "blocked $1:$2"; } +probe 10.43.0.1 443 +probe "$NODE_IP" 6443 +probe "$NODE_IP" 10250 +probe 1.1.1.1 53 +`, func(s *executor.LaunchSpec) { s.Env = map[string]string{"NODE_IP": strings.TrimSpace(string(nodeIP))} }) + c.launch(s) + c.wait(s.JobID, executor.EventStarted) + fin := c.wait(s.JobID, executor.EventFinished) + if *fin.ExitCode != 0 { + t.Fatalf("exit %d: %s\n%s", *fin.ExitCode, fin.Detail, c.output(s)) + } + c.cleanup(s) + + out := c.output(s) + t.Logf("job output:\n%s", out) + ip := strings.TrimSpace(string(nodeIP)) + for _, want := range []string{ + // Strict supplemental groups: exactly team, mist-users and personal. + "id: uid=20001(alice) gid=40001(team-team-a) groups=40001(team-team-a),19000(mist-users),20001(alice)\n", + "teams: team-a\n", // other teams are absent, not merely denied + "no-sa-token", "rootfs-readonly", "scratch-writable", "umask 0027", + "blocked 10.43.0.1:443", "blocked " + ip + ":6443", "blocked " + ip + ":10250", + "open 1.1.1.1:53", + "[mist] job", "[mist] exit code 0", + } { + if !strings.Contains(out, want) { + t.Errorf("output lacks %q", want) + } + } + if err := c.ex.Launch(c.ctx, s); err == nil { + t.Error("finished execution relaunched") + } +} + +func TestK3dTerminateForwardsSIGTERM(t *testing.T) { + c := newCluster(t) + s := c.spec(`trap 'echo caught-term; exit 0' TERM; echo ready; sleep 300 & wait`) + c.launch(s) + c.wait(s.JobID, executor.EventStarted) + time.Sleep(time.Second) // let the trap install + if err := c.ex.Terminate(c.ctx, s.JobID, s.SignalGrace, "cancelled by test"); err != nil { + t.Fatal(err) + } + fin := c.wait(s.JobID, executor.EventFinished) + if fin.Reason != executor.ExitSignaled { + t.Fatalf("reason %s (%s)", fin.Reason, fin.Detail) + } + c.cleanup(s) + if !strings.Contains(c.output(s), "caught-term") { + t.Fatalf("SIGTERM did not reach the script:\n%s", c.output(s)) + } +} + +func TestK3dOutOfMemory(t *testing.T) { + c := newCluster(t) + s := c.spec(`echo allocating; head -c 1G /dev/zero | tail > /dev/null`, + func(s *executor.LaunchSpec) { s.Resources.MemoryBytes = 64 << 20 }) + c.launch(s) + fin := c.wait(s.JobID, executor.EventFinished) + if fin.Reason != executor.ExitOutOfMemory { + t.Fatalf("reason %s (%s), exit %v", fin.Reason, fin.Detail, *fin.ExitCode) + } + c.cleanup(s) +} + +func TestK3dLocalTimeLimit(t *testing.T) { + c := newCluster(t) + s := c.spec(`echo ready; sleep 120`, func(s *executor.LaunchSpec) { s.TimeLimit = 3 * time.Second }) + c.launch(s) + fin := c.wait(s.JobID, executor.EventFinished) + if fin.Reason != executor.ExitDeadline || *fin.ExitCode != 143 { + t.Fatalf("reason %s exit %d", fin.Reason, *fin.ExitCode) + } + c.cleanup(s) +} + +func TestK3dSetupFailureIsLaunchFailed(t *testing.T) { + c := newCluster(t) + s := c.spec(`echo should-not-run`, func(s *executor.LaunchSpec) { + s.OutputPath = "/mist/teams/team-a/alice/no-such-dir/mist.out" + }) + c.launch(s) + fin := c.wait(s.JobID, executor.EventLaunchFailed) + if fin.Reason != executor.ExitLaunchError || !strings.Contains(fin.Detail, "no-such-dir") { + t.Fatalf("reason %s (%s)", fin.Reason, fin.Detail) + } + c.cleanup(s) +} + +// A restarted executor rebuilds its view from receipts and the API, and +// can still stop the job it did not launch. +func TestK3dExecutorRestart(t *testing.T) { + c := newCluster(t) + s := c.spec(`echo ready; sleep 300`) + c.launch(s) + c.wait(s.JobID, executor.EventStarted) + c.startExecutor() + xs, err := c.ex.List(c.ctx) + if err != nil { + t.Fatal(err) + } + found := false + for _, x := range xs { + found = found || (x.JobID == s.JobID && x.Phase == executor.PhaseRunning && x.StartedAt != nil) + } + if !found { + t.Fatalf("restarted executor lost the running job: %+v", xs) + } + if err := c.ex.Launch(c.ctx, s); err == nil || !strings.Contains(err.Error(), "started") { + t.Fatalf("relaunch after restart: %v", err) + } + c.ex.Terminate(c.ctx, s.JobID, s.SignalGrace, "stop") + c.wait(s.JobID, executor.EventFinished) + c.cleanup(s) +} + +func TestK3dRBACIsNarrow(t *testing.T) { + c := newCluster(t) + if _, err := c.client.CoreV1().Secrets("mist-jobs").List(c.ctx, metav1.ListOptions{}); !apierrors.IsForbidden(err) { + t.Errorf("mistd can list secrets: %v", err) + } + if _, err := c.client.CoreV1().Pods("kube-system").List(c.ctx, metav1.ListOptions{}); !apierrors.IsForbidden(err) { + t.Errorf("mistd can list pods outside mist-jobs: %v", err) + } +} + +// The whole M1 path on a real cluster: submit, admit, launch on k3s, run, +// node cleanup receipt, COMPLETED, allocation released. +func TestK3dSchedulerEndToEnd(t *testing.T) { + c := newCluster(t) + ctx := c.ctx + for _, q := range []string{ + "INSERT INTO nodes (name, executor) VALUES ($1, 'k3s')", + "INSERT INTO partition_nodes VALUES ('cpu', $1)", + } { + if _, err := c.pool.Exec(ctx, q, c.node); err != nil { + t.Fatal(err) + } + } + c.pool.Exec(ctx, "UPDATE partitions SET runtime = 'mist-test-runtime:dev'") + // Keep job IDs clear of other tests' Pods in the shared namespace. + c.pool.Exec(ctx, "SELECT setval('jobs_id_seq', $1)", nextJob+50) + + s, err := scheduler.New(ctx, scheduler.Config{ + Pool: c.pool, Identity: identitytest.New(), + Executors: map[string]executor.Executor{"k3s": c.ex}, + }, nil) + if err != nil { + t.Fatal(err) + } + jobs := scheduler.Jobs{Pool: c.pool, Clock: clock.Real{}} + id, err := jobs.Submit(ctx, scheduler.Request{ + Username: "alice", Team: "team-a", Partition: "cpu", QOS: "normal", Name: "e2e", + Kind: executor.KindBatch, + Resources: executor.Resources{MilliCPU: 500, MemoryBytes: 256 << 20, ScratchBytes: 1 << 30}, + TimeLimit: 5 * time.Minute, SignalGrace: 5 * time.Second, + Workdir: "/mist/teams/team-a/alice", + OutputPath: "/mist/teams/team-a/alice/mist-e2e.out", + Script: "echo hello from $MIST_JOB_ID on $MIST_NODE\n", + }) + if err != nil { + t.Fatal(err) + } + + state := func() (st string, released bool) { + c.pool.QueryRow(ctx, "SELECT state FROM jobs WHERE id = $1", id).Scan(&st) + var rel *time.Time + c.pool.QueryRow(ctx, "SELECT released_at FROM allocations WHERE job_id = $1", id).Scan(&rel) + return st, rel != nil + } + handled, cleaned := 0, false + deadline := time.Now().Add(3 * time.Minute) + for time.Now().Before(deadline) { + if err := s.Cycle(ctx); err != nil { + t.Fatal(err) + } + c.mu.Lock() + evs := append([]executor.Event(nil), c.events[id]...) + c.mu.Unlock() + for ; handled < len(evs); handled++ { + if err := s.HandleEvent(ctx, evs[handled]); err != nil { + t.Fatal(err) + } + } + st, released := state() + if st == "CLEANING" && !cleaned { + if released { + t.Fatal("allocation released before the node confirmed cleanup") + } + var execID string + c.pool.QueryRow(ctx, "SELECT execution_id::text FROM jobs WHERE id = $1", id).Scan(&execID) + c.cleanup(executor.LaunchSpec{JobID: id, ExecutionID: execID}) + cleaned = true + } + if st == "COMPLETED" && released { + out := c.output(executor.LaunchSpec{OutputPath: "/mist/teams/team-a/alice/mist-e2e.out"}) + if !strings.Contains(out, fmt.Sprintf("hello from %d on %s", id, c.node)) { + t.Fatalf("output:\n%s", out) + } + return + } + time.Sleep(500 * time.Millisecond) + } + st, released := state() + t.Fatalf("job %d ended in %s (released=%v)", id, st, released) +} diff --git a/internal/executor/k3s/pod.go b/internal/executor/k3s/pod.go new file mode 100644 index 0000000..6a3bc0f --- /dev/null +++ b/internal/executor/k3s/pod.go @@ -0,0 +1,295 @@ +// Package k3s is the K3s executor (M1): one plain Pod and ConfigMap per +// job in the mist-jobs namespace, created after Mist has admitted the job +// and chosen its node. Contract: docs/agents/architecture/executors.md. +package k3s + +import ( + "encoding/json" + "fmt" + "path" + "slices" + "strconv" + "strings" + "time" + + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/jobentry" +) + +// Labels and annotations on every object this executor creates. +const ( + LabelJobID = "mist/job-id" + LabelExecutionID = "mist/execution-id" + LabelPartition = "mist/partition" + LabelTeam = "mist/team" + AnnotSpecHash = "mist/spec-hash" + + MistUsersGID = 19000 + + jobDir = "/etc/mist" + terminationLog = "/dev/termination-log" +) + +// Volume says how to reach one storage volume from a node. +type Volume struct { + Server string // NFS server address on the direct link + Export string // NFS export root + HostNode string // node that hosts the volume locally, if any + HostPath string // the volume root on that node +} + +// Config holds the cluster-specific settings the pod builder needs. +type Config struct { + Namespace string + EntryPath string // mist-job-entry inside every runtime image + LaunchTimeout time.Duration // admission to work start + BackstopSlack time.Duration // beyond launch timeout + time limit + grace + Nameservers []string // approved resolvers; jobs never use cluster DNS + Volumes map[string]Volume +} + +// Name is the Pod and ConfigMap name for a job. +func Name(jobID int64) string { return "job-" + strconv.FormatInt(jobID, 10) } + +// Build turns a LaunchSpec into the ConfigMap and Pod for it. It is a pure +// function: the same spec and config always give the same objects, which +// golden tests pin. +func Build(spec executor.LaunchSpec, cfg Config) (*corev1.ConfigMap, *corev1.Pod, error) { + if err := spec.Validate(); err != nil { + return nil, nil, err + } + if spec.Kind == executor.KindInteractive { + return nil, nil, fmt.Errorf("k3s: interactive jobs arrive in M2") + } + if spec.Resources.TTBoards > 0 { + return nil, nil, fmt.Errorf("k3s: TT boards need DRA claims (M4)") + } + name := Name(spec.JobID) + teamRoot := "/mist/teams/" + spec.TeamSlug + meta := metav1.ObjectMeta{ + Name: name, + Namespace: cfg.Namespace, + Labels: map[string]string{ + LabelJobID: strconv.FormatInt(spec.JobID, 10), + LabelExecutionID: spec.ExecutionID, + LabelPartition: spec.Partition, + LabelTeam: spec.TeamSlug, + }, + Annotations: map[string]string{AnnotSpecHash: spec.SpecHash}, + } + + entry := jobentry.Spec{ + JobID: spec.JobID, Name: spec.JobName, User: spec.User.Name, Team: spec.TeamSlug, + Partition: spec.Partition, QOS: spec.QOS, Node: spec.Node, Kind: spec.Kind, + Resources: spec.Resources, + TimeLimit: int64(spec.TimeLimit / time.Second), + SignalGrace: int64(spec.SignalGrace / time.Second), + TeamRoot: teamRoot, + Home: teamRoot + "/" + spec.User.Name, + Workdir: spec.Workdir, + OutputPath: spec.OutputPath, + Scratch: "/scratch", + Args: spec.Args, + Argv: spec.Argv, + Env: spec.Env, + TerminationLog: terminationLog, + } + data := map[string]string{"passwd": passwd(spec), "group": group(spec)} + if spec.Kind == executor.KindBatch { + entry.ScriptPath = jobDir + "/script" + data["script"] = spec.Script + } + if err := entry.Validate(); err != nil { + return nil, nil, err + } + jobJSON, err := json.MarshalIndent(entry, "", " ") + if err != nil { + return nil, nil, err + } + data["job.json"] = string(jobJSON) + "\n" + cm := &corev1.ConfigMap{ + TypeMeta: metav1.TypeMeta{APIVersion: "v1", Kind: "ConfigMap"}, + ObjectMeta: meta, + Data: data, + Immutable: ptr(true), + } + + mounts, volumes, err := durableMounts(spec, cfg) + if err != nil { + return nil, nil, err + } + r := spec.Resources + quantities := corev1.ResourceList{ + corev1.ResourceCPU: *resource.NewMilliQuantity(r.MilliCPU, resource.DecimalSI), + corev1.ResourceMemory: *resource.NewQuantity(r.MemoryBytes, resource.BinarySI), + corev1.ResourceEphemeralStorage: *resource.NewQuantity(r.ScratchBytes, resource.BinarySI), + } + if r.GPUs > 0 { + quantities["nvidia.com/gpu"] = *resource.NewQuantity(r.GPUs, resource.DecimalSI) + } + // The gate needs almost nothing; equal requests and limits keep the Pod + // in the Guaranteed QoS class. + gateQuantities := corev1.ResourceList{ + corev1.ResourceCPU: *resource.NewMilliQuantity(100, resource.DecimalSI), + corev1.ResourceMemory: *resource.NewQuantity(32<<20, resource.BinarySI), + } + scratchSize := resource.NewQuantity(r.ScratchBytes, resource.BinarySI) + shmSize := resource.NewQuantity(r.MemoryBytes/2, resource.BinarySI) // counts against memory + + volumes = append(volumes, + corev1.Volume{Name: "scratch", VolumeSource: corev1.VolumeSource{ + EmptyDir: &corev1.EmptyDirVolumeSource{SizeLimit: scratchSize}}}, + corev1.Volume{Name: "shm", VolumeSource: corev1.VolumeSource{ + EmptyDir: &corev1.EmptyDirVolumeSource{Medium: corev1.StorageMediumMemory, SizeLimit: shmSize}}}, + corev1.Volume{Name: "job", VolumeSource: corev1.VolumeSource{ConfigMap: &corev1.ConfigMapVolumeSource{ + LocalObjectReference: corev1.LocalObjectReference{Name: name}, + DefaultMode: ptr[int32](0o444), + }}}, + ) + mounts = append(mounts, + corev1.VolumeMount{Name: "scratch", MountPath: "/scratch"}, + // The same directory, not a subPath: kubelet creates subPath + // directories root-owned, which a non-root job could not write. + corev1.VolumeMount{Name: "scratch", MountPath: "/tmp"}, + corev1.VolumeMount{Name: "shm", MountPath: "/dev/shm"}, + corev1.VolumeMount{Name: "job", MountPath: jobDir + "/job.json", SubPath: "job.json", ReadOnly: true}, + corev1.VolumeMount{Name: "job", MountPath: "/etc/passwd", SubPath: "passwd", ReadOnly: true}, + corev1.VolumeMount{Name: "job", MountPath: "/etc/group", SubPath: "group", ReadOnly: true}, + ) + if spec.Kind == executor.KindBatch { + mounts = append(mounts, corev1.VolumeMount{Name: "job", MountPath: jobDir + "/script", SubPath: "script", ReadOnly: true}) + } + + groups := append(slices.Clone(spec.ExtraGIDs), spec.User.PersonalGID) + slices.Sort(groups) + groups = slices.Compact(groups) + sup := make([]int64, len(groups)) + for i, g := range groups { + sup[i] = int64(g) + } + strict := corev1.SupplementalGroupsPolicyStrict + // A coarse backstop that holds even with the scheduler down. Pod time + // includes startup, so it allows the launch timeout on top of walltime. + backstop := int64((cfg.LaunchTimeout + spec.TimeLimit + spec.SignalGrace + cfg.BackstopSlack) / time.Second) + + pod := &corev1.Pod{ + TypeMeta: metav1.TypeMeta{APIVersion: "v1", Kind: "Pod"}, + ObjectMeta: meta, + Spec: corev1.PodSpec{ + NodeName: spec.Node, // Mist placed it; the Kubernetes scheduler never does + RestartPolicy: corev1.RestartPolicyNever, + ActiveDeadlineSeconds: &backstop, + TerminationGracePeriodSeconds: ptr(int64(spec.SignalGrace / time.Second)), + AutomountServiceAccountToken: ptr(false), + EnableServiceLinks: ptr(false), + DNSPolicy: corev1.DNSNone, + DNSConfig: &corev1.PodDNSConfig{Nameservers: cfg.Nameservers}, + SecurityContext: &corev1.PodSecurityContext{ + RunAsUser: ptr(int64(spec.User.UID)), + RunAsGroup: ptr(int64(spec.TeamGID)), + SupplementalGroups: sup, + SupplementalGroupsPolicy: &strict, + RunAsNonRoot: ptr(true), + SeccompProfile: &corev1.SeccompProfile{Type: corev1.SeccompProfileTypeRuntimeDefault}, + // No fsGroup: it would rewrite ownership of team files. + }, + // Network policy is programmed only after the Pod starts. This + // gate holds user code back until the node's kubelet, which is + // certainly listening, has become unreachable from the Pod. + InitContainers: []corev1.Container{{ + Name: "network-gate", + Image: spec.Runtime, + ImagePullPolicy: corev1.PullIfNotPresent, + Command: []string{cfg.EntryPath, "-await-isolation", "$(MIST_HOST_IP):10250"}, + Env: []corev1.EnvVar{{Name: "MIST_HOST_IP", ValueFrom: &corev1.EnvVarSource{ + FieldRef: &corev1.ObjectFieldSelector{FieldPath: "status.hostIP"}}}}, + Resources: corev1.ResourceRequirements{Requests: gateQuantities, Limits: gateQuantities}, + TerminationMessagePath: terminationLog, + TerminationMessagePolicy: corev1.TerminationMessageReadFile, + SecurityContext: containerSecurity(), + }}, + Containers: []corev1.Container{{ + Name: "job", + Image: spec.Runtime, + ImagePullPolicy: corev1.PullIfNotPresent, + Command: []string{cfg.EntryPath, "-spec", jobDir + "/job.json"}, + Resources: corev1.ResourceRequirements{Requests: quantities, Limits: quantities}, + VolumeMounts: mounts, + TerminationMessagePath: terminationLog, + TerminationMessagePolicy: corev1.TerminationMessageReadFile, + SecurityContext: containerSecurity(), + }}, + Volumes: volumes, + }, + } + return cm, pod, nil +} + +// durableMounts resolves each LaunchSpec mount to NFS, or to hostPath when +// the job runs on the node that hosts the volume. +func durableMounts(spec executor.LaunchSpec, cfg Config) ([]corev1.VolumeMount, []corev1.Volume, error) { + var mounts []corev1.VolumeMount + var volumes []corev1.Volume + for i, m := range spec.Mounts { + v, ok := cfg.Volumes[m.VolumeID] + if !ok { + return nil, nil, fmt.Errorf("k3s: unknown volume %q", m.VolumeID) + } + name := fmt.Sprintf("durable-%d", i) + var src corev1.VolumeSource + if v.HostNode == spec.Node { + src.HostPath = &corev1.HostPathVolumeSource{ + Path: path.Join(v.HostPath, m.Source), + Type: ptr(corev1.HostPathDirectory), // never create missing folders + } + } else { + src.NFS = &corev1.NFSVolumeSource{Server: v.Server, Path: path.Join(v.Export, m.Source), ReadOnly: m.ReadOnly} + } + volumes = append(volumes, corev1.Volume{Name: name, VolumeSource: src}) + mounts = append(mounts, corev1.VolumeMount{Name: name, MountPath: m.Target, ReadOnly: m.ReadOnly}) + } + return mounts, volumes, nil +} + +// passwd and group project only what the job needs: system entries, the +// submitter and the charged team. No other team's membership is visible. +// Names were validated against ^[a-z][a-z0-9-]{1,30}$, so they cannot +// inject fields or lines. +func passwd(s executor.LaunchSpec) string { + home := "/mist/teams/" + s.TeamSlug + "/" + s.User.Name + return strings.Join([]string{ + "root:x:0:0:root:/root:/usr/sbin/nologin", + "nobody:x:65534:65534:nobody:/nonexistent:/usr/sbin/nologin", + fmt.Sprintf("%s:x:%d:%d::%s:/bin/bash", s.User.Name, s.User.UID, s.TeamGID, home), + }, "\n") + "\n" +} + +func group(s executor.LaunchSpec) string { + lines := []string{"root:x:0:", "nogroup:x:65534:", + fmt.Sprintf("team-%s:x:%d:%s", s.TeamSlug, s.TeamGID, s.User.Name), + fmt.Sprintf("%s:x:%d:%s", s.User.Name, s.User.PersonalGID, s.User.Name), + } + for _, g := range s.ExtraGIDs { + name := fmt.Sprintf("mist-gid-%d", g) + if g == MistUsersGID { + name = "mist-users" + } + lines = append(lines, fmt.Sprintf("%s:x:%d:%s", name, g, s.User.Name)) + } + return strings.Join(lines, "\n") + "\n" +} + +func containerSecurity() *corev1.SecurityContext { + return &corev1.SecurityContext{ + AllowPrivilegeEscalation: ptr(false), + Privileged: ptr(false), + ReadOnlyRootFilesystem: ptr(true), + Capabilities: &corev1.Capabilities{Drop: []corev1.Capability{"ALL"}}, + } +} + +func ptr[T any](v T) *T { return &v } diff --git a/internal/executor/k3s/pod_test.go b/internal/executor/k3s/pod_test.go new file mode 100644 index 0000000..ef869bd --- /dev/null +++ b/internal/executor/k3s/pod_test.go @@ -0,0 +1,159 @@ +package k3s + +import ( + "flag" + "os" + "path/filepath" + "strings" + "testing" + "time" + + corev1 "k8s.io/api/core/v1" + "sigs.k8s.io/yaml" + + "github.com/UTMIST/Mist/internal/executor" +) + +var update = flag.Bool("update", false, "rewrite golden files") + +func testConfig() Config { + return Config{ + Namespace: "mist-jobs", + EntryPath: "/usr/local/bin/mist-job-entry", + LaunchTimeout: 10 * time.Minute, + BackstopSlack: time.Minute, + Nameservers: []string{"1.1.1.1", "9.9.9.9"}, + Volumes: map[string]Volume{ + "vol1": {Server: "192.168.240.2", Export: "/mist", HostNode: "utmist-tt", HostPath: "/mist"}, + }, + } +} + +func gpuBatch() executor.LaunchSpec { + s := executor.LaunchSpec{ + JobID: 412, + ExecutionID: "3f1c2a9e-8b7d-4c6e-9a5b-1d2e3f4a5b6c", + AuthGeneration: 7, + LaunchDeadline: time.Date(2026, 11, 2, 15, 14, 5, 0, time.UTC), + Kind: executor.KindBatch, + JobName: "resnet-sweep-3", + Partition: "gpu", + QOS: "normal", + Node: "utmist-z1opa08", + User: executor.User{Name: "alice", UID: 20001, PersonalGID: 20001}, + TeamSlug: "vision", + TeamGID: 40001, + ExtraGIDs: []int{19000}, + Resources: executor.Resources{MilliCPU: 8000, MemoryBytes: 32 << 30, GPUs: 1, ScratchBytes: 100 << 30}, + Mounts: []executor.Mount{ + {VolumeID: "vol1", Source: "/teams/vision", Target: "/mist/teams/vision"}, + {VolumeID: "vol1", Source: "/software", Target: "/mist/software", ReadOnly: true}, + }, + Workdir: "/mist/teams/vision/alice/resnet", + OutputPath: "/mist/teams/vision/alice/resnet/mist-412.out", + Script: "#!/bin/bash\npython train.py --lr \"$1\"\n", + Args: []string{"3e-4"}, + Env: map[string]string{"WANDB_PROJECT": "resnet"}, + Runtime: "registry.mist.local/runtime/gpu@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + TimeLimit: 4 * time.Hour, + SignalGrace: 60 * time.Second, + } + s.Seal() + return s +} + +func cpuCommand() executor.LaunchSpec { + s := gpuBatch() + s.JobID, s.Kind, s.Partition, s.Node = 413, executor.KindCommand, "cpu", "utmist-tt" + s.Script, s.Args = "", nil + s.Argv = []string{"huggingface-cli", "download", "meta-llama/Llama-3-8B"} + s.Resources = executor.Resources{MilliCPU: 2000, MemoryBytes: 8 << 30, ScratchBytes: 50 << 30} + s.Runtime = "registry.mist.local/runtime/cpu@sha256:fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210" + s.Seal() + return s +} + +func TestBuildGolden(t *testing.T) { + for name, spec := range map[string]executor.LaunchSpec{"gpu-batch-nfs": gpuBatch(), "cpu-command-hostpath": cpuCommand()} { + t.Run(name, func(t *testing.T) { + cm, pod, err := Build(spec, testConfig()) + if err != nil { + t.Fatal(err) + } + var got []byte + for _, obj := range []any{cm, pod} { + b, err := yaml.Marshal(obj) + if err != nil { + t.Fatal(err) + } + got = append(append(got, "---\n"...), b...) + } + file := filepath.Join("testdata", name+".golden.yaml") + if *update { + os.MkdirAll("testdata", 0o755) + if err := os.WriteFile(file, got, 0o644); err != nil { + t.Fatal(err) + } + } + want, err := os.ReadFile(file) + if err != nil { + t.Fatalf("%v (run with -update to create)", err) + } + if string(got) != string(want) { + t.Fatalf("%s differs from golden; run go test ./internal/executor/k3s -update and review the diff", file) + } + }) + } +} + +// Security properties checked directly, so a golden regeneration cannot +// silently drop one. +func TestBuildInvariants(t *testing.T) { + for _, spec := range []executor.LaunchSpec{gpuBatch(), cpuCommand()} { + cm, pod, err := Build(spec, testConfig()) + if err != nil { + t.Fatal(err) + } + ps, c := pod.Spec, pod.Spec.Containers[0] + sc := ps.SecurityContext + check := func(ok bool, what string) { + if !ok { + t.Errorf("job %d: %s", spec.JobID, what) + } + } + check(ps.NodeName == spec.Node, "nodeName is Mist's chosen node") + check(ps.SchedulerName == "", "default scheduler fields untouched") + check(ps.RestartPolicy == corev1.RestartPolicyNever, "restartPolicy Never") + check(!*ps.AutomountServiceAccountToken && !*ps.EnableServiceLinks, "no service account token or service links") + check(ps.DNSPolicy == corev1.DNSNone, "no cluster DNS") + check(sc.FSGroup == nil, "no fsGroup") + check(*sc.RunAsUser == int64(spec.User.UID) && *sc.RunAsGroup == int64(spec.TeamGID), "uid and team gid") + check(*sc.SupplementalGroupsPolicy == corev1.SupplementalGroupsPolicyStrict, "strict supplemental groups") + check(len(sc.SupplementalGroups) == 2, "only mist-users and the personal group") + check(*sc.RunAsNonRoot, "runAsNonRoot") + check(!*c.SecurityContext.AllowPrivilegeEscalation && *c.SecurityContext.ReadOnlyRootFilesystem, "no escalation, read-only root") + check(len(c.SecurityContext.Capabilities.Drop) == 1 && c.SecurityContext.Capabilities.Drop[0] == "ALL", "drop ALL") + check(c.Resources.Requests.Cpu().Equal(*c.Resources.Limits.Cpu()) && + c.Resources.Requests.Memory().Equal(*c.Resources.Limits.Memory()), "requests = limits") + check(len(ps.InitContainers) == 1 && ps.InitContainers[0].Name == "network-gate", "network isolation gate before user code") + check(*ps.InitContainers[0].SecurityContext.ReadOnlyRootFilesystem, "gate is unprivileged too") + check(*ps.TerminationGracePeriodSeconds == int64(spec.SignalGrace/time.Second), "grace = signal grace") + check(pod.Annotations[AnnotSpecHash] == spec.SpecHash, "spec hash annotation") + for _, v := range ps.Volumes { + check(v.HostPath == nil || spec.Node == "utmist-tt", "hostPath only on the volume's host") + } + for _, f := range []string{"passwd", "group"} { + check(!strings.Contains(cm.Data[f], "carol") && strings.Count(cm.Data[f], "team-") <= 1, f+" has only the charged team") + } + } + + if _, _, err := Build(gpuBatch(), Config{Namespace: "mist-jobs"}); err == nil { + t.Error("unknown volume accepted") + } + tt := gpuBatch() + tt.Resources.GPUs, tt.Resources.TTBoards = 0, 1 + tt.Seal() + if _, _, err := Build(tt, testConfig()); err == nil { + t.Error("TT job built without DRA") + } +} diff --git a/internal/executor/k3s/store.go b/internal/executor/k3s/store.go new file mode 100644 index 0000000..c9e3000 --- /dev/null +++ b/internal/executor/k3s/store.go @@ -0,0 +1,129 @@ +package k3s + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/UTMIST/Mist/internal/executor" +) + +// record is one k3s_executions row: everything durable about an execution. +type record struct { + ExecutionID, SpecHash, Node string + JobID int64 + CreateIntentAt time.Time + ObjectUID string + StartedAt, FinishedAt *time.Time + Result *executor.Result + LogTail []byte + TerminateRequestedAt *time.Time + LostAt *time.Time + LostDetail string + PodDeletedAt, CleanedAt *time.Time + CleanupEvidence string +} + +type store struct{ pool *pgxpool.Pool } + +const recordColumns = `execution_id::text, spec_hash, node, job_id, create_intent_at, coalesce(object_uid, ''), + started_at, finished_at, result, log_tail, terminate_requested_at, lost_at, lost_detail, + pod_deleted_at, cleaned_at, cleanup_evidence` + +func scanRecord(row pgx.Row) (*record, error) { + var r record + var result []byte + err := row.Scan(&r.ExecutionID, &r.SpecHash, &r.Node, &r.JobID, &r.CreateIntentAt, &r.ObjectUID, + &r.StartedAt, &r.FinishedAt, &result, &r.LogTail, &r.TerminateRequestedAt, &r.LostAt, &r.LostDetail, + &r.PodDeletedAt, &r.CleanedAt, &r.CleanupEvidence) + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, err + } + if result != nil { + r.Result = new(executor.Result) + if err := json.Unmarshal(result, r.Result); err != nil { + return nil, err + } + } + return &r, nil +} + +func (s store) byJob(ctx context.Context, jobID int64) (*record, error) { + return scanRecord(s.pool.QueryRow(ctx, "SELECT "+recordColumns+" FROM k3s_executions WHERE job_id = $1", jobID)) +} + +func (s store) all(ctx context.Context) ([]*record, error) { + rows, err := s.pool.Query(ctx, "SELECT "+recordColumns+" FROM k3s_executions ORDER BY job_id") + if err != nil { + return nil, err + } + return pgx.CollectRows(rows, func(r pgx.CollectableRow) (*record, error) { return scanRecord(r) }) +} + +// createIntent records that a create is about to be attempted. It is a +// no-op if the same execution already has a row. +func (s store) createIntent(ctx context.Context, spec executor.LaunchSpec, at time.Time) error { + return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { + tag, err := tx.Exec(ctx, ` + INSERT INTO k3s_executions (execution_id, job_id, spec_hash, node, create_intent_at) + VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING`, + spec.ExecutionID, spec.JobID, spec.SpecHash, spec.Node, at) + if err != nil || tag.RowsAffected() == 0 { + return err + } + return receipt(ctx, tx, spec.ExecutionID, executor.ReceiptCreateIntent, at, "", "") + }) +} + +// update applies one change to a row and appends its receipt, atomically. +// cond guards the change (for example "started_at IS NULL"), so only the +// first observation of a fact is recorded. It reports whether it applied. +func (s store) update(ctx context.Context, execID string, kind executor.ReceiptKind, at time.Time, + objectUID, detail, set, cond string, args ...any) (bool, error) { + applied := false + err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { + q := "UPDATE k3s_executions SET " + set + " WHERE execution_id = $1" + if cond != "" { + q += " AND " + cond + } + tag, err := tx.Exec(ctx, q, append([]any{execID}, args...)...) + if err != nil || tag.RowsAffected() == 0 { + return err + } + applied = true + if kind == "" { + return nil + } + return receipt(ctx, tx, execID, kind, at, objectUID, detail) + }) + return applied, err +} + +func receipt(ctx context.Context, tx pgx.Tx, execID string, kind executor.ReceiptKind, at time.Time, uid, detail string) error { + _, err := tx.Exec(ctx, + "INSERT INTO k3s_receipts (execution_id, kind, at, object_uid, detail) VALUES ($1, $2, $3, $4, $5)", + execID, string(kind), at, uid, detail) + return err +} + +func (s store) receipts(ctx context.Context, execID string) ([]executor.Receipt, error) { + rows, err := s.pool.Query(ctx, + "SELECT kind, at, object_uid, detail FROM k3s_receipts WHERE execution_id = $1 ORDER BY id", execID) + if err != nil { + return nil, err + } + return pgx.CollectRows(rows, func(r pgx.CollectableRow) (executor.Receipt, error) { + var x executor.Receipt + var kind string + err := r.Scan(&kind, &x.At, &x.ObjectUID, &x.Evidence) + x.Kind, x.ExecutionID = executor.ReceiptKind(kind), execID + return x, err + }) +} diff --git a/internal/executor/k3s/testdata/cpu-command-hostpath.golden.yaml b/internal/executor/k3s/testdata/cpu-command-hostpath.golden.yaml new file mode 100644 index 0000000..5de8777 --- /dev/null +++ b/internal/executor/k3s/testdata/cpu-command-hostpath.golden.yaml @@ -0,0 +1,197 @@ +--- +apiVersion: v1 +data: + group: | + root:x:0: + nogroup:x:65534: + team-vision:x:40001:alice + alice:x:20001:alice + mist-users:x:19000:alice + job.json: | + { + "job_id": 413, + "name": "resnet-sweep-3", + "user": "alice", + "team": "vision", + "partition": "cpu", + "qos": "normal", + "node": "utmist-tt", + "kind": "command", + "resources": { + "milli_cpu": 2000, + "memory_bytes": 8589934592, + "gpus": 0, + "tt_boards": 0, + "scratch_bytes": 53687091200, + "hugepages_bytes": 0 + }, + "time_limit_seconds": 14400, + "signal_grace_seconds": 60, + "team_root": "/mist/teams/vision", + "home": "/mist/teams/vision/alice", + "workdir": "/mist/teams/vision/alice/resnet", + "output_path": "/mist/teams/vision/alice/resnet/mist-412.out", + "scratch": "/scratch", + "script_path": "", + "args": null, + "argv": [ + "huggingface-cli", + "download", + "meta-llama/Llama-3-8B" + ], + "env": { + "WANDB_PROJECT": "resnet" + }, + "termination_log": "/dev/termination-log" + } + passwd: | + root:x:0:0:root:/root:/usr/sbin/nologin + nobody:x:65534:65534:nobody:/nonexistent:/usr/sbin/nologin + alice:x:20001:40001::/mist/teams/vision/alice:/bin/bash +immutable: true +kind: ConfigMap +metadata: + annotations: + mist/spec-hash: sha256:2afd963d82ff2edd045185b45d2fb055673b2f8e4796c996c1415c1df0cfc028 + labels: + mist/execution-id: 3f1c2a9e-8b7d-4c6e-9a5b-1d2e3f4a5b6c + mist/job-id: "413" + mist/partition: cpu + mist/team: vision + name: job-413 + namespace: mist-jobs +--- +apiVersion: v1 +kind: Pod +metadata: + annotations: + mist/spec-hash: sha256:2afd963d82ff2edd045185b45d2fb055673b2f8e4796c996c1415c1df0cfc028 + labels: + mist/execution-id: 3f1c2a9e-8b7d-4c6e-9a5b-1d2e3f4a5b6c + mist/job-id: "413" + mist/partition: cpu + mist/team: vision + name: job-413 + namespace: mist-jobs +spec: + activeDeadlineSeconds: 15120 + automountServiceAccountToken: false + containers: + - command: + - /usr/local/bin/mist-job-entry + - -spec + - /etc/mist/job.json + image: registry.mist.local/runtime/cpu@sha256:fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210 + imagePullPolicy: IfNotPresent + name: job + resources: + limits: + cpu: "2" + ephemeral-storage: 50Gi + memory: 8Gi + requests: + cpu: "2" + ephemeral-storage: 50Gi + memory: 8Gi + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + privileged: false + readOnlyRootFilesystem: true + terminationMessagePath: /dev/termination-log + terminationMessagePolicy: File + volumeMounts: + - mountPath: /mist/teams/vision + name: durable-0 + - mountPath: /mist/software + name: durable-1 + readOnly: true + - mountPath: /scratch + name: scratch + - mountPath: /tmp + name: scratch + - mountPath: /dev/shm + name: shm + - mountPath: /etc/mist/job.json + name: job + readOnly: true + subPath: job.json + - mountPath: /etc/passwd + name: job + readOnly: true + subPath: passwd + - mountPath: /etc/group + name: job + readOnly: true + subPath: group + dnsConfig: + nameservers: + - 1.1.1.1 + - 9.9.9.9 + dnsPolicy: None + enableServiceLinks: false + initContainers: + - command: + - /usr/local/bin/mist-job-entry + - -await-isolation + - $(MIST_HOST_IP):10250 + env: + - name: MIST_HOST_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP + image: registry.mist.local/runtime/cpu@sha256:fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210 + imagePullPolicy: IfNotPresent + name: network-gate + resources: + limits: + cpu: 100m + memory: 32Mi + requests: + cpu: 100m + memory: 32Mi + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + privileged: false + readOnlyRootFilesystem: true + terminationMessagePath: /dev/termination-log + terminationMessagePolicy: File + nodeName: utmist-tt + restartPolicy: Never + securityContext: + runAsGroup: 40001 + runAsNonRoot: true + runAsUser: 20001 + seccompProfile: + type: RuntimeDefault + supplementalGroups: + - 19000 + - 20001 + supplementalGroupsPolicy: Strict + terminationGracePeriodSeconds: 60 + volumes: + - hostPath: + path: /mist/teams/vision + type: Directory + name: durable-0 + - hostPath: + path: /mist/software + type: Directory + name: durable-1 + - emptyDir: + sizeLimit: 50Gi + name: scratch + - emptyDir: + medium: Memory + sizeLimit: 4Gi + name: shm + - configMap: + defaultMode: 292 + name: job-413 + name: job +status: {} diff --git a/internal/executor/k3s/testdata/gpu-batch-nfs.golden.yaml b/internal/executor/k3s/testdata/gpu-batch-nfs.golden.yaml new file mode 100644 index 0000000..baba5b1 --- /dev/null +++ b/internal/executor/k3s/testdata/gpu-batch-nfs.golden.yaml @@ -0,0 +1,205 @@ +--- +apiVersion: v1 +data: + group: | + root:x:0: + nogroup:x:65534: + team-vision:x:40001:alice + alice:x:20001:alice + mist-users:x:19000:alice + job.json: | + { + "job_id": 412, + "name": "resnet-sweep-3", + "user": "alice", + "team": "vision", + "partition": "gpu", + "qos": "normal", + "node": "utmist-z1opa08", + "kind": "batch", + "resources": { + "milli_cpu": 8000, + "memory_bytes": 34359738368, + "gpus": 1, + "tt_boards": 0, + "scratch_bytes": 107374182400, + "hugepages_bytes": 0 + }, + "time_limit_seconds": 14400, + "signal_grace_seconds": 60, + "team_root": "/mist/teams/vision", + "home": "/mist/teams/vision/alice", + "workdir": "/mist/teams/vision/alice/resnet", + "output_path": "/mist/teams/vision/alice/resnet/mist-412.out", + "scratch": "/scratch", + "script_path": "/etc/mist/script", + "args": [ + "3e-4" + ], + "argv": null, + "env": { + "WANDB_PROJECT": "resnet" + }, + "termination_log": "/dev/termination-log" + } + passwd: | + root:x:0:0:root:/root:/usr/sbin/nologin + nobody:x:65534:65534:nobody:/nonexistent:/usr/sbin/nologin + alice:x:20001:40001::/mist/teams/vision/alice:/bin/bash + script: | + #!/bin/bash + python train.py --lr "$1" +immutable: true +kind: ConfigMap +metadata: + annotations: + mist/spec-hash: sha256:add67d194aa767e462365538aad75bf3714b005a47f14a9a5902f3737db49cb6 + labels: + mist/execution-id: 3f1c2a9e-8b7d-4c6e-9a5b-1d2e3f4a5b6c + mist/job-id: "412" + mist/partition: gpu + mist/team: vision + name: job-412 + namespace: mist-jobs +--- +apiVersion: v1 +kind: Pod +metadata: + annotations: + mist/spec-hash: sha256:add67d194aa767e462365538aad75bf3714b005a47f14a9a5902f3737db49cb6 + labels: + mist/execution-id: 3f1c2a9e-8b7d-4c6e-9a5b-1d2e3f4a5b6c + mist/job-id: "412" + mist/partition: gpu + mist/team: vision + name: job-412 + namespace: mist-jobs +spec: + activeDeadlineSeconds: 15120 + automountServiceAccountToken: false + containers: + - command: + - /usr/local/bin/mist-job-entry + - -spec + - /etc/mist/job.json + image: registry.mist.local/runtime/gpu@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + imagePullPolicy: IfNotPresent + name: job + resources: + limits: + cpu: "8" + ephemeral-storage: 100Gi + memory: 32Gi + nvidia.com/gpu: "1" + requests: + cpu: "8" + ephemeral-storage: 100Gi + memory: 32Gi + nvidia.com/gpu: "1" + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + privileged: false + readOnlyRootFilesystem: true + terminationMessagePath: /dev/termination-log + terminationMessagePolicy: File + volumeMounts: + - mountPath: /mist/teams/vision + name: durable-0 + - mountPath: /mist/software + name: durable-1 + readOnly: true + - mountPath: /scratch + name: scratch + - mountPath: /tmp + name: scratch + - mountPath: /dev/shm + name: shm + - mountPath: /etc/mist/job.json + name: job + readOnly: true + subPath: job.json + - mountPath: /etc/passwd + name: job + readOnly: true + subPath: passwd + - mountPath: /etc/group + name: job + readOnly: true + subPath: group + - mountPath: /etc/mist/script + name: job + readOnly: true + subPath: script + dnsConfig: + nameservers: + - 1.1.1.1 + - 9.9.9.9 + dnsPolicy: None + enableServiceLinks: false + initContainers: + - command: + - /usr/local/bin/mist-job-entry + - -await-isolation + - $(MIST_HOST_IP):10250 + env: + - name: MIST_HOST_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP + image: registry.mist.local/runtime/gpu@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + imagePullPolicy: IfNotPresent + name: network-gate + resources: + limits: + cpu: 100m + memory: 32Mi + requests: + cpu: 100m + memory: 32Mi + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + privileged: false + readOnlyRootFilesystem: true + terminationMessagePath: /dev/termination-log + terminationMessagePolicy: File + nodeName: utmist-z1opa08 + restartPolicy: Never + securityContext: + runAsGroup: 40001 + runAsNonRoot: true + runAsUser: 20001 + seccompProfile: + type: RuntimeDefault + supplementalGroups: + - 19000 + - 20001 + supplementalGroupsPolicy: Strict + terminationGracePeriodSeconds: 60 + volumes: + - name: durable-0 + nfs: + path: /mist/teams/vision + server: 192.168.240.2 + - name: durable-1 + nfs: + path: /mist/software + readOnly: true + server: 192.168.240.2 + - emptyDir: + sizeLimit: 100Gi + name: scratch + - emptyDir: + medium: Memory + sizeLimit: 16Gi + name: shm + - configMap: + defaultMode: 292 + name: job-412 + name: job +status: {} diff --git a/internal/executor/k3s/testdata/runtime/Dockerfile b/internal/executor/k3s/testdata/runtime/Dockerfile new file mode 100644 index 0000000..e2f58c7 --- /dev/null +++ b/internal/executor/k3s/testdata/runtime/Dockerfile @@ -0,0 +1,4 @@ +# Test-only stand-in for the runtime profiles (owned by the runtime owner). +# Ubuntu 22.04 matches the hosts; the real gpu/cpu images add the toolchain. +FROM ubuntu:22.04 +COPY mist-job-entry /usr/local/bin/mist-job-entry diff --git a/internal/executor/queue.go b/internal/executor/queue.go new file mode 100644 index 0000000..9a7ed19 --- /dev/null +++ b/internal/executor/queue.go @@ -0,0 +1,64 @@ +package executor + +import "sync" + +// Queue is an unbounded event queue in front of an unbuffered channel, so +// an executor never blocks on a slow consumer. Delivery is at-least-once +// from the consumer's point of view only because executors may re-emit; +// the queue itself neither drops nor duplicates. +type Queue struct { + ch chan Event + done chan struct{} + + mu sync.Mutex + cond *sync.Cond + items []Event + closed bool +} + +func NewQueue() *Queue { + q := &Queue{ch: make(chan Event), done: make(chan struct{})} + q.cond = sync.NewCond(&q.mu) + go q.pump() + return q +} + +func (q *Queue) C() <-chan Event { return q.ch } + +func (q *Queue) Push(ev Event) { + q.mu.Lock() + defer q.mu.Unlock() + q.items = append(q.items, ev) + q.cond.Signal() +} + +func (q *Queue) Close() { + q.mu.Lock() + defer q.mu.Unlock() + if !q.closed { + q.closed = true + close(q.done) + q.cond.Signal() + } +} + +func (q *Queue) pump() { + for { + q.mu.Lock() + for len(q.items) == 0 && !q.closed { + q.cond.Wait() + } + if q.closed { + q.mu.Unlock() + return + } + ev := q.items[0] + q.items = q.items[1:] + q.mu.Unlock() + select { + case q.ch <- ev: + case <-q.done: + return + } + } +} diff --git a/internal/executor/spec.go b/internal/executor/spec.go index 03fc475..5693867 100644 --- a/internal/executor/spec.go +++ b/internal/executor/spec.go @@ -22,6 +22,9 @@ type LaunchSpec struct { AuthGeneration int64 `json:"auth_generation"` LaunchDeadline time.Time `json:"launch_deadline"` // admitted_at + launch timeout Kind Kind `json:"kind"` + JobName string `json:"job_name"` + Partition string `json:"partition"` + QOS string `json:"qos"` Node string `json:"node"` User User `json:"user"` TeamSlug string `json:"team_slug"` @@ -111,8 +114,8 @@ func (s LaunchSpec) Validate() error { if s.LaunchDeadline.IsZero() { bad("launch_deadline is required") } - if s.Node == "" { - bad("node is required") + if s.Node == "" || s.Partition == "" || s.QOS == "" { + bad("node, partition and qos are required") } if !namePattern.MatchString(s.User.Name) || s.User.UID <= 0 || s.User.PersonalGID <= 0 { bad("user needs a valid name, uid and personal gid") diff --git a/internal/executor/spec_test.go b/internal/executor/spec_test.go index 3b94f7d..94b0ce8 100644 --- a/internal/executor/spec_test.go +++ b/internal/executor/spec_test.go @@ -13,6 +13,9 @@ func validSpec() LaunchSpec { AuthGeneration: 7, LaunchDeadline: time.Date(2026, 11, 2, 15, 14, 5, 0, time.UTC), Kind: KindBatch, + JobName: "resnet-sweep-3", + Partition: "gpu", + QOS: "normal", Node: "utmist-z1opa08", User: User{Name: "alice", UID: 20001, PersonalGID: 20001}, TeamSlug: "vision", diff --git a/internal/jobentry/isolation.go b/internal/jobentry/isolation.go new file mode 100644 index 0000000..980864c --- /dev/null +++ b/internal/jobentry/isolation.go @@ -0,0 +1,38 @@ +package jobentry + +import ( + "fmt" + "net" + "time" +) + +// AwaitIsolation blocks until every target refuses or drops connections +// several times in a row, proving the job's network filtering is in place. +// Kubernetes policy controllers program a Pod's rules only after it +// starts, so user code must not run until this passes. Targets must be +// services known to be listening (the node's kubelet is: it is starting +// this Pod), so a failed connection means filtering, not an absent service. +func AwaitIsolation(targets []string, timeout time.Duration) error { + const needBlocked = 3 + deadline := time.Now().Add(timeout) + blocked := 0 + for blocked < needBlocked { + if time.Now().After(deadline) { + return fmt.Errorf("network isolation not in place after %s: %v still reachable", timeout, targets) + } + open := false + for _, t := range targets { + if c, err := net.DialTimeout("tcp", t, time.Second); err == nil { + c.Close() + open = true + } + } + if open { + blocked = 0 + time.Sleep(200 * time.Millisecond) + } else { + blocked++ + } + } + return nil +} diff --git a/internal/jobentry/isolation_test.go b/internal/jobentry/isolation_test.go new file mode 100644 index 0000000..ba65368 --- /dev/null +++ b/internal/jobentry/isolation_test.go @@ -0,0 +1,35 @@ +package jobentry + +import ( + "net" + "testing" + "time" +) + +func TestAwaitIsolation(t *testing.T) { + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + open := l.Addr().String() + go func() { + for { + c, err := l.Accept() + if err != nil { + return + } + c.Close() + } + }() + if err := AwaitIsolation([]string{open}, time.Second); err == nil { + t.Fatal("reachable target passed the gate") + } + l.Close() // now refused, as REJECT rules make it + start := time.Now() + if err := AwaitIsolation([]string{open}, 5*time.Second); err != nil { + t.Fatalf("blocked target failed the gate: %v", err) + } + if time.Since(start) > 2*time.Second { + t.Fatal("gate slow on a blocked target") + } +} diff --git a/internal/scheduler/admit.go b/internal/scheduler/admit.go index 720aa5d..7b25509 100644 --- a/internal/scheduler/admit.go +++ b/internal/scheduler/admit.go @@ -141,6 +141,9 @@ func (s *Scheduler) admitOne(ctx context.Context, j *job, nodeName string) (bool AuthGeneration: grant.Generation, LaunchDeadline: now.Add(s.launchTimeout), Kind: req.Kind, + JobName: j.Name, + Partition: j.Partition, + QOS: j.QOS, Node: nodeName, User: executor.User{Name: grant.Account.Username, UID: grant.Account.UID, PersonalGID: grant.Account.PersonalGID}, TeamSlug: grant.Team.Slug, diff --git a/scripts/k3d-test.sh b/scripts/k3d-test.sh new file mode 100755 index 0000000..d04238e --- /dev/null +++ b/scripts/k3d-test.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# Runs the K3s executor integration tests on a local k3d cluster. +# Needs docker (without sudo), k3d and kubectl. Safe to re-run: it reuses +# the cluster if it exists. Delete it with: k3d cluster delete mist-test +set -euo pipefail + +cluster=mist-test +k3s_image=rancher/k3s:v1.36.5-k3s1 +node=k3d-$cluster-server-0 +root=$(cd "$(dirname "$0")/.." && pwd) +work=$(mktemp -d) +trap 'rm -rf "$work"' EXIT +cd "$root" + +if ! k3d cluster get "$cluster" >/dev/null 2>&1; then + k3d cluster create "$cluster" --image "$k3s_image" --servers 1 --agents 0 --no-lb \ + --k3s-arg "--disable=traefik,servicelb,metrics-server,local-storage@server:0" --wait --timeout 240s +fi +ctx=k3d-$cluster + +CGO_ENABLED=0 go build -o "$work/mist-job-entry" ./cmd/mist-job-entry +CGO_ENABLED=0 go build -o "$work/mistd" ./cmd/mistd +cp internal/executor/k3s/testdata/runtime/Dockerfile "$work/" +docker build -q -t mist-test-runtime:dev "$work" >/dev/null +k3d image import -c "$cluster" mist-test-runtime:dev >/dev/null +docker cp "$work/mistd" "$node:/mistd" + +kubectl --context "$ctx" apply -f deploy/k3s/mist-jobs.yaml -f deploy/k3s/rbac.yaml >/dev/null +# A miniature /mist on the node, with the storage contract's modes. +docker exec "$node" sh -c ' + mkdir -p /mist/software /mist/teams/team-a/alice /mist/teams/team-b/carol + chmod 755 /mist /mist/teams /mist/software + chown 0:40001 /mist/teams/team-a && chmod 2750 /mist/teams/team-a + chown 20001:40001 /mist/teams/team-a/alice && chmod 2750 /mist/teams/team-a/alice + chown 0:40002 /mist/teams/team-b && chmod 2750 /mist/teams/team-b' + +MIST_K3D_CONTEXT=$ctx MIST_K3D_NODE=$node \ + go test -tags k3d -count=1 -v -run K3d ./internal/executor/k3s "$@" From ef0cf5beb4fa2de47ecdcc2d17719ed8658ba9c8 Mon Sep 17 00:00:00 2001 From: Brandon Li Date: Fri, 30 Oct 2026 13:37:04 -0400 Subject: [PATCH 5/5] Add job API endpoints and cleanup receipt transport (M1 step 6) internal/api serves submit (database defaults, validation, required Idempotency-Key), list, get, cancel, log (tail and WebSocket follow), the anonymised queue and partitions, plus the host-token authenticated cleanup receipt endpoint. mistd cleanup-report -post delivers receipts. Production builds fail closed until Identity's authenticators land, and reject the development user header; -dev-fixture exists only with -tags mistdev. M1-09 is covered by unit tests on the fixture and by a k3d test that runs submit-to-COMPLETED through the API. Co-Authored-By: Claude Opus 5.5 --- README.md | 12 +- cmd/mistd/dev.go | 15 + cmd/mistd/main.go | 152 +++- cmd/mistd/main_test.go | 45 ++ cmd/mistd/nodev.go | 10 + cmd/mistd/nodev_test.go | 56 ++ docs/agents/architecture/api.md | 78 +++ docs/agents/architecture/executors.md | 5 + docs/agents/plans/jobs-and-scheduling.md | 40 ++ internal/api/api.go | 173 +++++ internal/api/api_test.go | 655 ++++++++++++++++++ internal/api/auth.go | 192 +++++ internal/api/cleanup.go | 43 ++ internal/api/dev.go | 24 + internal/api/dev_test.go | 16 + internal/api/jobs.go | 348 ++++++++++ internal/api/logs.go | 162 +++++ internal/api/submit.go | 439 ++++++++++++ internal/api/wire.go | 207 ++++++ .../db/migrations/scheduler/0003_job_api.sql | 21 + internal/e2e/k3d_test.go | 252 +++++++ internal/executor/k3s/executor.go | 10 +- internal/jobentry/env.go | 7 +- internal/scheduler/scheduler_test.go | 34 + internal/scheduler/submit.go | 83 ++- internal/scheduler/views.go | 306 ++++++++ scripts/k3d-test.sh | 2 +- 27 files changed, 3348 insertions(+), 39 deletions(-) create mode 100644 cmd/mistd/dev.go create mode 100644 cmd/mistd/nodev.go create mode 100644 cmd/mistd/nodev_test.go create mode 100644 internal/api/api.go create mode 100644 internal/api/api_test.go create mode 100644 internal/api/auth.go create mode 100644 internal/api/cleanup.go create mode 100644 internal/api/dev.go create mode 100644 internal/api/dev_test.go create mode 100644 internal/api/jobs.go create mode 100644 internal/api/logs.go create mode 100644 internal/api/submit.go create mode 100644 internal/api/wire.go create mode 100644 internal/db/migrations/scheduler/0003_job_api.sql create mode 100644 internal/e2e/k3d_test.go create mode 100644 internal/scheduler/views.go diff --git a/README.md b/README.md index d6b71d3..d2033df 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,17 @@ Database tests start an embedded Postgres as the current user (downloaded to `~/.embedded-postgres-go` on first run). To use an existing server instead, set `MIST_TEST_DATABASE_URL` to a superuser URL; each test creates its own database. -Executor integration tests run on a local k3d cluster (needs Docker without +For local development against the fixture accounts (alice, bob, carol, +root-admin), build with `-tags mistdev` and pass `-dev-fixture` to both `mistd +api` and `mistd scheduler`; requests then name the user in `X-Mist-Dev-User`. +Production builds contain neither and reject the header. Run the tests with the +tag too: + +```bash +go test -tags mistdev ./internal/api ./cmd/... +``` + +Executor and API integration tests run on a local k3d cluster (needs Docker without sudo, plus k3d and kubectl): ```bash diff --git a/cmd/mistd/dev.go b/cmd/mistd/dev.go new file mode 100644 index 0000000..c4b5339 --- /dev/null +++ b/cmd/mistd/dev.go @@ -0,0 +1,15 @@ +//go:build mistdev + +package main + +import ( + "github.com/UTMIST/Mist/internal/api" + "github.com/UTMIST/Mist/internal/identity/identitytest" +) + +// Each process gets its own in-memory fixture, identical at start; +// membership changes are not shared between api and scheduler. +func newDevIdentity() (*devIdentity, error) { + f := identitytest.New() + return &devIdentity{authz: f, admit: f, header: api.DevHeader{}}, nil +} diff --git a/cmd/mistd/main.go b/cmd/mistd/main.go index 940268d..4cc49d8 100644 --- a/cmd/mistd/main.go +++ b/cmd/mistd/main.go @@ -2,14 +2,19 @@ // // mistd api HTTP API // mistd scheduler queue, admission and executor calls (one active leader) -// mistd cleanup-report on a k3s node, as root: print a cleanup receipt -// for a deleted Pod, or fail if anything of it remains +// mistd cleanup-report on a k3s node, as root: check that nothing of a +// deleted Pod remains, then print or post the receipt // -// Both connect to MIST_DATABASE_URL (or -database-url) and apply pending -// migrations at start. +// api and scheduler connect to MIST_DATABASE_URL (or -database-url) and +// apply pending migrations at start. +// +// Binaries built with -tags mistdev also accept -dev-fixture: the +// identitytest accounts and the X-Mist-Dev-User header, for local +// development only. Production builds reject both. package main import ( + "bytes" "context" "encoding/json" "errors" @@ -18,11 +23,18 @@ import ( "io" "log/slog" "net/http" + "net/url" "os" "os/signal" + "strings" "syscall" "time" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/tools/clientcmd" + + "github.com/UTMIST/Mist/internal/api" + "github.com/UTMIST/Mist/internal/clock" "github.com/UTMIST/Mist/internal/db" "github.com/UTMIST/Mist/internal/executor/k3s" "github.com/UTMIST/Mist/internal/identity" @@ -44,7 +56,7 @@ func run(ctx context.Context, args []string, stdout, stderr io.Writer) error { } cmd, args := args[0], args[1:] if cmd == "cleanup-report" { - return cleanupReport(args, stdout) + return cleanupReport(ctx, args, stdout) } fs := flag.NewFlagSet("mistd "+cmd, flag.ContinueOnError) fs.SetOutput(stderr) @@ -52,9 +64,22 @@ func run(ctx context.Context, args []string, stdout, stderr io.Writer) error { listen := fs.String("listen", "127.0.0.1:8080", "API listen address (api)") tick := fs.Duration("tick", 10*time.Second, "scheduling cycle period (scheduler)") leaderCheck := fs.Duration("leader-check", 2*time.Second, "leadership check period (scheduler)") + hostTokens := fs.String("host-tokens", "", "file of hashed host agent tokens (api)") + kubeconfig := fs.String("kubeconfig", "", "kubeconfig for job logs and cleanup receipts; empty disables them (api)") + kubeContext := fs.String("kube-context", "", "kubeconfig context (api)") + namespace := fs.String("namespace", "mist-jobs", "job namespace (api)") + origins := fs.String("allowed-origins", "", "comma-separated browser origins allowed to open WebSockets (api)") + devFixture := fs.Bool("dev-fixture", false, "use the identitytest accounts (development builds only)") if err := fs.Parse(args); err != nil { return err } + var dev *devIdentity + if *devFixture { + var err error + if dev, err = newDevIdentity(); err != nil { + return err + } + } if *dbURL == "" { return errors.New("set MIST_DATABASE_URL or -database-url") } @@ -72,26 +97,66 @@ func run(ctx context.Context, args []string, stdout, stderr io.Writer) error { switch cmd { case "scheduler": + // Until Identity's implementation lands, fail closed: nothing is + // admitted (AUTH_UNAVAILABLE). No executors are wired yet either. + var admitter identity.Admitter = identity.Unavailable{} + if dev != nil { + admitter = dev.admit + } return scheduler.Run(ctx, scheduler.Config{ DatabaseURL: *dbURL, Pool: pool, - // Until Identity's implementation lands, fail closed: nothing is - // admitted (AUTH_UNAVAILABLE). No executors are wired yet either. - Identity: identity.Unavailable{}, + Identity: admitter, Log: log, Tick: *tick, LeaderCheck: *leaderCheck, }) case "api": - mux := http.NewServeMux() - mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) { - if err := pool.Ping(r.Context()); err != nil { - http.Error(w, "database unavailable", http.StatusServiceUnavailable) - return + s := &api.Server{ + Jobs: scheduler.Jobs{Pool: pool, Clock: clock.Real{}}, + // Fail closed until Identity provides the session and relay + // authenticators and the real Authorizer. + Authz: closedAuthorizer{}, + Log: log, + Ping: pool.Ping, + } + if dev != nil { + s.Authz, s.Accounts = dev.authz, []api.Authenticator{dev.header} + log.Warn("development identity fixture in use; never run this in production") + } + if *origins != "" { + s.AllowedOrigins = strings.Split(*origins, ",") + } + if *hostTokens != "" { + ht, err := api.LoadHostTokens(*hostTokens, "cleanup") + if err != nil { + return err } - io.WriteString(w, "ok\n") - }) - srv := &http.Server{Addr: *listen, Handler: mux, ReadHeaderTimeout: 10 * time.Second} + s.Hosts = ht + } + if *kubeconfig != "" { + cfg, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( + &clientcmd.ClientConfigLoadingRules{ExplicitPath: *kubeconfig}, + &clientcmd.ConfigOverrides{CurrentContext: *kubeContext}).ClientConfig() + if err != nil { + return err + } + client, err := kubernetes.NewForConfig(cfg) + if err != nil { + return err + } + // Not started: the API only reads logs and records cleanup + // receipts. The scheduler's executor sees recorded cleanups + // through List; the events this one emits are not needed. + ex := k3s.New(k3s.Options{Config: k3s.Config{Namespace: *namespace}, Client: client, + Pool: pool, Clock: clock.Real{}, Log: log}) + go func() { + for range ex.Events() { + } + }() + s.Logs, s.Cleanup = ex, ex + } + srv := &http.Server{Addr: *listen, Handler: s.Handler(), ReadHeaderTimeout: 10 * time.Second} go func() { <-ctx.Done() shut, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) @@ -108,16 +173,35 @@ func run(ctx context.Context, args []string, stdout, stderr io.Writer) error { } } +// devIdentity is the development identity (dev.go); production builds +// cannot construct one. +type devIdentity struct { + authz identity.Authorizer + admit identity.Admitter + header api.Authenticator +} + +// closedAuthorizer refuses everything as unavailable, so a server without +// a real Identity backend never grants access. +type closedAuthorizer struct{} + +func (closedAuthorizer) Authorize(context.Context, identity.Principal, identity.Action, identity.Target) error { + return identity.ErrUnavailable +} + // cleanupReport checks a node for leftovers of one Pod and prints the -// receipt. Until the internal API exists (M1 step 6) the receipt is -// delivered by whoever runs this; it is evidence, never a request. -func cleanupReport(args []string, stdout io.Writer) error { +// receipt, or with -post delivers it to the API using this host's cleanup +// token. The receipt is evidence for the executor to check, never a +// request it must obey. +func cleanupReport(ctx context.Context, args []string, stdout io.Writer) error { fs := flag.NewFlagSet("mistd cleanup-report", flag.ContinueOnError) kubelet := fs.String("kubelet-root", "/var/lib/kubelet", "kubelet root directory") cgroups := fs.String("cgroup-root", "/sys/fs/cgroup", "cgroup v2 mount") node := fs.String("node", "", "this node's name") execID := fs.String("execution-id", "", "execution ID") podUID := fs.String("pod-uid", "", "recorded Pod UID") + post := fs.String("post", "", "API base URL to deliver the receipt to, e.g. http://192.168.240.1:8080") + tokenFile := fs.String("token-file", "/etc/mist/cleanup-token", "this host's cleanup token (with -post)") if err := fs.Parse(args); err != nil { return err } @@ -125,7 +209,31 @@ func cleanupReport(args []string, stdout io.Writer) error { if err != nil { return err } - return json.NewEncoder(stdout).Encode(k3s.CleanupReceipt{ - ExecutionID: *execID, PodUID: *podUID, Node: *node, Evidence: evidence, - }) + receipt := k3s.CleanupReceipt{ExecutionID: *execID, PodUID: *podUID, Node: *node, Evidence: evidence} + if *post == "" { + return json.NewEncoder(stdout).Encode(receipt) + } + token, err := os.ReadFile(*tokenFile) + if err != nil { + return err + } + body, _ := json.Marshal(receipt) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + strings.TrimSuffix(*post, "/")+"/internal/v1/executions/"+url.PathEscape(*execID)+"/cleanup", bytes.NewReader(body)) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(string(token))) + req.Header.Set("Content-Type", "application/json") + resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + reply, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("cleanup receipt refused: %s: %s", resp.Status, bytes.TrimSpace(reply)) + } + fmt.Fprintf(stdout, "%s\n", bytes.TrimSpace(reply)) + return nil } diff --git a/cmd/mistd/main_test.go b/cmd/mistd/main_test.go index b441b3b..5581b06 100644 --- a/cmd/mistd/main_test.go +++ b/cmd/mistd/main_test.go @@ -3,7 +3,12 @@ package main import ( "bytes" "context" + "encoding/json" "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" "strings" "sync" "testing" @@ -12,6 +17,7 @@ import ( "github.com/jackc/pgx/v5" "github.com/UTMIST/Mist/internal/db/dbtest" + "github.com/UTMIST/Mist/internal/executor/k3s" ) func TestMain(m *testing.M) { dbtest.Main(m) } @@ -131,3 +137,42 @@ func TestSchedulerExitsWhenConnectionLost(t *testing.T) { t.Fatal("scheduler kept running after losing its lock connection") } } + +// cleanup-report -post delivers the receipt with the host's token. +func TestCleanupReportPosts(t *testing.T) { + kubelet, cgroups := t.TempDir(), t.TempDir() + os.Mkdir(filepath.Join(kubelet, "pods"), 0o755) + token := filepath.Join(t.TempDir(), "token") + os.WriteFile(token, []byte("tt-secret\n"), 0o600) + const execID = "3f1c2a9e-8b7d-4c6e-9a5b-1d2e3f4a5b6c" + const uid = "0b6e7f2c-1a2b-4c3d-8e9f-0123456789ab" + + var got k3s.CleanupReceipt + var auth, path string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + auth, path = r.Header.Get("Authorization"), r.URL.Path + json.NewDecoder(r.Body).Decode(&got) + io.WriteString(w, `{"status":"recorded"}`) + })) + defer srv.Close() + + var out bytes.Buffer + err := run(context.Background(), []string{"cleanup-report", "-kubelet-root", kubelet, "-cgroup-root", cgroups, + "-node", "utmist-tt", "-execution-id", execID, "-pod-uid", uid, "-post", srv.URL, "-token-file", token}, &out, io.Discard) + if err != nil { + t.Fatal(err) + } + if auth != "Bearer tt-secret" || path != "/internal/v1/executions/"+execID+"/cleanup" || + got.PodUID != uid || got.Node != "utmist-tt" || got.Evidence == "" { + t.Fatalf("posted %q %q %+v", auth, path, got) + } + + // Leftovers mean no receipt at all. + os.Mkdir(filepath.Join(kubelet, "pods", uid), 0o755) + auth = "" + err = run(context.Background(), []string{"cleanup-report", "-kubelet-root", kubelet, "-cgroup-root", cgroups, + "-node", "utmist-tt", "-execution-id", execID, "-pod-uid", uid, "-post", srv.URL, "-token-file", token}, &out, io.Discard) + if err == nil || auth != "" { + t.Fatalf("receipt sent with the Pod directory present: %v", err) + } +} diff --git a/cmd/mistd/nodev.go b/cmd/mistd/nodev.go new file mode 100644 index 0000000..eae9da5 --- /dev/null +++ b/cmd/mistd/nodev.go @@ -0,0 +1,10 @@ +//go:build !mistdev + +package main + +import "errors" + +// Production builds contain no development identity. +func newDevIdentity() (*devIdentity, error) { + return nil, errors.New("-dev-fixture needs a development build (go build -tags mistdev)") +} diff --git a/cmd/mistd/nodev_test.go b/cmd/mistd/nodev_test.go new file mode 100644 index 0000000..0f23481 --- /dev/null +++ b/cmd/mistd/nodev_test.go @@ -0,0 +1,56 @@ +//go:build !mistdev + +package main + +import ( + "context" + "io" + "net" + "net/http" + "strings" + "testing" + "time" + + "github.com/UTMIST/Mist/internal/db/dbtest" +) + +// M1 step 6 check: a production build refuses the development fixture, +// and its API rejects the development principal header. +func TestProductionRejectsDevIdentity(t *testing.T) { + url := dbtest.URL(t) + err := run(context.Background(), []string{"api", "-database-url", url, "-dev-fixture"}, io.Discard, io.Discard) + if err == nil || !strings.Contains(err.Error(), "development build") { + t.Fatalf("-dev-fixture in a production build: %v", err) + } + + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + addr := l.Addr().String() + l.Close() + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + done <- run(ctx, []string{"api", "-database-url", url, "-listen", addr}, io.Discard, io.Discard) + }() + t.Cleanup(func() { cancel(); <-done }) + + waitFor(t, "api to listen", func() bool { + resp, err := http.Get("http://" + addr + "/healthz") + if err == nil { + resp.Body.Close() + } + return err == nil && resp.StatusCode == 200 + }) + req, _ := http.NewRequest("GET", "http://"+addr+"/v1/jobs", nil) + req.Header.Set("X-Mist-Dev-User", "root-admin") + resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("dev header on a production API: %s", resp.Status) + } +} diff --git a/docs/agents/architecture/api.md b/docs/agents/architecture/api.md index 74698ad..52209a3 100644 --- a/docs/agents/architecture/api.md +++ b/docs/agents/architecture/api.md @@ -232,6 +232,84 @@ process group; callers cannot specify a PID, executable or alternate command. The holder owns grace timing, so API restart cannot reset it. The CLI's `mist shell` and later browser terminal use this same endpoint. +## Implementation notes (M1) + +The job endpoints are in `internal/api`. Where this section is more specific than +the contract above, the code follows this section. + +- **Status codes.** `unauthenticated` is 401, `forbidden` 403, `not_found` 404, + `invalid` 400, `limit_exceeded` 422, `conflict` 409 and `unavailable` 503. + Clients should switch on `code`, not on the status. +- **Authenticators.** `api.Authenticator` resolves a request to an + `identity.Principal`. It returns `api.ErrNoCredentials` when the request carries + none of its credentials, so the next one can try. Identity plugs the portal + session and relay authenticators into `Server.Accounts`. Until it does, + `mistd api` fails closed: every `/v1` request is `unauthenticated`, and + authorization is `unavailable`. Session-cookie CSRF and Origin checks belong + to the session authenticator. +- **Development identity.** `X-Mist-Dev-User` and `mistd -dev-fixture` exist only + in binaries built with `-tags mistdev`, which use the `identitytest` accounts. + - Production binaries contain neither, refuse `-dev-fixture`, and answer + `unauthenticated` to any request carrying the header, even alongside valid + credentials. + - `cmd/mistd/nodev_test.go` checks this on the built server. +- **Idempotency.** `Idempotency-Key` is required from every caller, not only the + CLI. The portal should send a fresh `crypto.randomUUID()` per submit click. + - Keys are per account. + - The hash covers the request as decoded: field order and spacing don't + matter. + - A replay answers `200` with header `Idempotent-Replayed: true` (a new job is + `201`), after rechecking that the caller may still view the job. + - The key is checked once before normalization and again inside the submit + transaction, under the per-user lock, so concurrent retries create one job. +- **Defaults and checks.** Defaults come from the `partitions` and `qos` rows, + including `partitions.is_default` and `default_signal_grace_seconds`. + - `%j` in `output` is expanded when the job ID is assigned. Other `%` + sequences are kept literally. + - Paths must be clean and absolute (no `.`, `..`, `//` or trailing `/`). + Output must lie strictly inside the charged team. + - Protected env names (`jobentry.ProtectedEnv`) are rejected, not dropped. + - A partition whose nodes have never been reported answers `unavailable` + rather than accepting a job that might never fit. + - `interactive` kind and QoS are rejected until M2. +- **`GET /v1/jobs`.** + - Without parameters it lists the caller's own jobs, keeping only teams they + still belong to. + - `?team=` lists a team's jobs for its members, and `?all=1` lists every job + for admins. + - `?state=` takes a comma list of states, or `active`. + - The cursor is the last job ID of the page; pages run newest first. +- **`GET /v1/jobs/{id}`** adds `events` (the state history) and `exit_detail`. + `?script=1` adds `script`, `args`, `argv` and `env`. +- **`GET /v1/jobs/{id}/log`** answers + `{state, output_path, output_incomplete, truncated, text}`, with at most the + last 64 KiB as text. Invalid UTF-8 becomes U+FFFD. + - `?follow=1` upgrades to a WebSocket of binary frames. It closes when output + ends, when the client sends anything, or when a recheck of view access + (every 10 s) fails. + - A browser `Origin` must be in `-allowed-origins`. Requests without an + Origin are not from a browser. +- **`GET /v1/queue`** entries carry `id`, `partition`, `qos`, `state`, + `resources`, `age_seconds`, `pending_reason`, `user`, `team`, `name`, + `workdir` and `output_path`. + - `user`, `team` and `name` are filled only for teams the caller belongs to. + - `workdir` and `output_path` are always null. + - `pending_reason` passes only known codes. +- **`GET /v1/partitions`** gives defaults, maximums (per accelerator when + `per_accelerator`), nodes with total and free capacity, and free capacity on + schedulable nodes. + - Only limits the caller can learn are listed: wildcard rows, rows naming the + caller, and rows naming teams the caller belongs to. + - Admins see every limit. +- **Host tokens.** `/internal` accepts only `Authorization: Bearer` host tokens, + and `/v1` never accepts them. `mistd api -host-tokens FILE` reads lines of + ` `; generate a hash with `api.HashToken`. + - Each server path accepts one scope (`cleanup`). + - A host may post cleanup receipts only for itself. The executor then checks + the receipt against its own records; a refusal is `conflict`. + - Issuing and rotating tokens is Identity's provisioning work. Until then, + an admin writes the file. + ## Versioning Additive changes (new fields, endpoints or enum values) don't change the diff --git a/docs/agents/architecture/executors.md b/docs/agents/architecture/executors.md index 993b7b2..fdd48f3 100644 --- a/docs/agents/architecture/executors.md +++ b/docs/agents/architecture/executors.md @@ -384,6 +384,11 @@ actual process groups, rather than trusting an accepted field. Reference: paths; it never deletes arbitrary paths supplied by jobs. Send authenticated, durable receipts via the internal API. Missing reporter/receipt keeps capacity reserved and alerts an admin. TT also requires a successful board reset. + M1: `mistd cleanup-report -post -token-file ` checks the node and + posts to `POST /internal/v1/executions/{id}/cleanup` with the node's + `cleanup` host token ([api.md](api.md#implementation-notes-m1)). The API's + executor records the receipt, and the scheduler sees it through `List`. What + runs the reporter after each Pod deletion is step 7. - Only then emit `Exited` and release resources. Keep diagnostics and tombstones in the database, including when API deletion or cleanup is retried. - Compare authoritative observations every 60 s and on restart. Follow the diff --git a/docs/agents/plans/jobs-and-scheduling.md b/docs/agents/plans/jobs-and-scheduling.md index 96b586a..27574a2 100644 --- a/docs/agents/plans/jobs-and-scheduling.md +++ b/docs/agents/plans/jobs-and-scheduling.md @@ -179,6 +179,46 @@ requirements. Keep secrets and actual research data out of test fixtures. allowlisted queue summaries and cleanup/incomplete-output fields. - *Check:* M1-09 through dev fixtures and real session/relay auth. Production build must reject the development principal header. + - *Status (30 Oct 2026):* `internal/api` and the `mistd api` wiring. The + details are in [api.md](../architecture/api.md#implementation-notes-m1). + - **Endpoints:** + - `POST /v1/jobs`, which normalizes with database defaults and requires + an `Idempotency-Key`; + - `GET /v1/jobs` and `GET /v1/jobs/{id}`; + - `POST /v1/jobs/{id}/cancel`; + - `GET /v1/jobs/{id}/log`, including the `?follow=1` WebSocket; + - `GET /v1/queue`, anonymised; + - `GET /v1/partitions`; + - `POST /internal/v1/executions/{id}/cleanup`, with host tokens. + + `mistd cleanup-report -post` delivers receipts. + - **M1-09 tests on the fixture:** + - A duplicate key gives one job, both through HTTP and inside the submit + transaction. A conflicting body is rejected, and a replay survives a + defaults change. + - Non-team detail, log and cancel requests get `not_found`. + - The lead, owner and admin rules follow the matrix, and former members + lose access. + - The queue and partition replies are free of other teams' identities, + names, env and paths. + - Hostile paths, env and argv are rejected and never stored. + - Log streams close on revocation. + - A production build refuses `-dev-fixture` and the dev header. + + Mutation checks confirm the tests catch removed authorization, queue + redaction, the idempotency check, dev-header rejection, path + confinement, protected env and node-scoped receipts. + - **k3d:** `TestK3dAPIEndToEnd` submits over HTTP, admits, runs on k3s, + reads the log through the API, and has the node post its own cleanup + receipt to `/internal`. The job then reaches `COMPLETED` with the + allocation released. + - **Still open:** + - Real session and relay auth (Identity) are not integrated, so the + second half of the check waits for Lucas's authenticators. + - Host-token issuing (Identity provisioning). + - TLS on the direct link (step 7). + - The scheduler role has no executors wired in `mistd` yet (step 7 + config). 7. **Cluster.** - Reinstall k3s on the direct link and install the GPU Operator with only the device plugin and toolkit. diff --git a/internal/api/api.go b/internal/api/api.go new file mode 100644 index 0000000..4132575 --- /dev/null +++ b/internal/api/api.go @@ -0,0 +1,173 @@ +// Package api is the HTTP API that `mistd api` serves: conventions, +// authentication middleware and the job endpoints. The contract is +// docs/agents/architecture/api.md. Identity, storage and admin endpoints +// are added by their owners on the same Server. +package api + +import ( + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "time" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/executor/k3s" + "github.com/UTMIST/Mist/internal/identity" + "github.com/UTMIST/Mist/internal/scheduler" +) + +// LogSource gives live output of a job, or the stored tail once the job +// has ended. executor.Executor satisfies it. +type LogSource interface { + Logs(ctx context.Context, jobID int64, follow bool) (io.ReadCloser, error) +} + +// CleanupRecorder accepts node cleanup receipts. *k3s.Executor satisfies it. +type CleanupRecorder interface { + RecordCleanup(ctx context.Context, r k3s.CleanupReceipt) error +} + +// Server holds the API's dependencies. Nil optional fields make their +// endpoints answer `unavailable` rather than guess. +type Server struct { + Jobs scheduler.Jobs + Authz identity.Authorizer + // Accounts resolve /v1 callers (portal session, relay, API token), in + // order. Hosts resolves /internal callers. Neither accepts the other's + // credentials. + Accounts []Authenticator + Hosts Authenticator + Logs LogSource // optional + Cleanup CleanupRecorder // optional + Clock clock.Clock + Log *slog.Logger + // AllowedOrigins are the browser origins allowed to open WebSockets. + AllowedOrigins []string + // StreamRecheck is how often an open stream re-checks authorization + // (default 10 s). + StreamRecheck time.Duration + // Ping checks the database for /healthz. + Ping func(context.Context) error +} + +// Handler returns the API's routes. +func (s *Server) Handler() http.Handler { + if s.Clock == nil { + s.Clock = clock.Real{} + } + if s.Log == nil { + s.Log = slog.New(slog.DiscardHandler) + } + if s.StreamRecheck == 0 { + s.StreamRecheck = 10 * time.Second + } + mux := http.NewServeMux() + mux.HandleFunc("GET /healthz", s.health) + mux.Handle("POST /v1/jobs", s.account(s.submit)) + mux.Handle("GET /v1/jobs", s.account(s.listJobs)) + mux.Handle("GET /v1/jobs/{id}", s.account(s.getJob)) + mux.Handle("POST /v1/jobs/{id}/cancel", s.account(s.cancelJob)) + mux.Handle("GET /v1/jobs/{id}/log", s.account(s.jobLog)) + mux.Handle("GET /v1/partitions", s.account(s.partitions)) + mux.Handle("GET /v1/queue", s.account(s.queue)) + mux.Handle("POST /internal/v1/executions/{execution_id}/cleanup", s.host(s.cleanup)) + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + writeError(w, &Error{http.StatusNotFound, CodeNotFound, "no such endpoint", nil}) + }) + return mux +} + +func (s *Server) health(w http.ResponseWriter, r *http.Request) { + if s.Ping != nil { + if err := s.Ping(r.Context()); err != nil { + http.Error(w, "database unavailable", http.StatusServiceUnavailable) + return + } + } + io.WriteString(w, "ok\n") +} + +// Error codes (api.md#conventions). +const ( + CodeUnauthenticated = "unauthenticated" + CodeForbidden = "forbidden" + CodeNotFound = "not_found" + CodeInvalid = "invalid" + CodeLimitExceeded = "limit_exceeded" + CodeConflict = "conflict" + CodeUnavailable = "unavailable" + CodeInternal = "internal" +) + +// Error is an API error reply. +type Error struct { + Status int `json:"-"` + Code string `json:"code"` + Message string `json:"message"` + Details map[string]any `json:"details"` +} + +func (e *Error) Error() string { return e.Code + ": " + e.Message } + +func invalid(field, msg string) *Error { + return &Error{http.StatusBadRequest, CodeInvalid, msg, map[string]any{"field": field}} +} + +func limitExceeded(field, msg string) *Error { + return &Error{http.StatusUnprocessableEntity, CodeLimitExceeded, msg, map[string]any{"field": field}} +} + +func notFound(msg string) *Error { return &Error{http.StatusNotFound, CodeNotFound, msg, nil} } + +func unavailable(msg string) *Error { + return &Error{http.StatusServiceUnavailable, CodeUnavailable, msg, nil} +} + +func writeError(w http.ResponseWriter, e *Error) { + if e.Details == nil { + e.Details = map[string]any{} + } + writeJSON(w, e.Status, map[string]*Error{"error": e}) +} + +func writeJSON(w http.ResponseWriter, status int, v any) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Cache-Control", "no-store") + w.WriteHeader(status) + json.NewEncoder(w).Encode(v) +} + +// fail writes err, which is an *Error or an unexpected failure. Internal +// details are logged, never returned. +func (s *Server) fail(w http.ResponseWriter, r *http.Request, err error) { + var e *Error + if errors.As(err, &e) { + writeError(w, e) + return + } + if errors.Is(err, context.Canceled) && r.Context().Err() != nil { + return + } + s.Log.Error("request failed", "method", r.Method, "path", r.URL.Path, "err", err) + writeError(w, &Error{http.StatusInternalServerError, CodeInternal, "internal error", nil}) +} + +// authzError turns an Authorizer error into a reply. With hide, refusals +// become not_found, so non-members learn nothing about a job. +func authzError(err error, hide bool, what string) error { + switch { + case err == nil: + return nil + case errors.Is(err, identity.ErrUnavailable): + return unavailable("cannot confirm current account and team membership; try again shortly") + case errors.Is(err, identity.ErrForbidden), errors.Is(err, identity.ErrRevoked): + if hide { + return notFound(what + " not found") + } + return &Error{http.StatusForbidden, CodeForbidden, "not allowed", nil} + } + return err +} diff --git a/internal/api/api_test.go b/internal/api/api_test.go new file mode 100644 index 0000000..7de8a6d --- /dev/null +++ b/internal/api/api_test.go @@ -0,0 +1,655 @@ +package api + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "golang.org/x/net/websocket" + + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/db/dbtest" + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/executor/k3s" + "github.com/UTMIST/Mist/internal/identity" + "github.com/UTMIST/Mist/internal/identity/identitytest" + "github.com/UTMIST/Mist/internal/scheduler" +) + +func TestMain(m *testing.M) { dbtest.Main(m) } + +var t0 = time.Date(2026, 11, 2, 12, 0, 0, 0, time.UTC) + +// testUser stands in for Identity's session and relay authenticators. +type testUser struct{} + +func (testUser) Authenticate(r *http.Request) (identity.Principal, error) { + u := r.Header.Get("X-Test-User") + if u == "" { + return identity.Principal{}, ErrNoCredentials + } + return identity.Principal{Kind: identity.PrincipalAccount, Username: u}, nil +} + +type fakeLogs struct { + mu sync.Mutex + text map[int64]string +} + +func (f *fakeLogs) Logs(_ context.Context, id int64, _ bool) (io.ReadCloser, error) { + f.mu.Lock() + defer f.mu.Unlock() + t, ok := f.text[id] + if !ok { + return nil, executor.ErrNotFound + } + return io.NopCloser(strings.NewReader(t)), nil +} + +type fakeCleanup struct { + mu sync.Mutex + got []k3s.CleanupReceipt + fail error +} + +func (f *fakeCleanup) RecordCleanup(_ context.Context, r k3s.CleanupReceipt) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.fail != nil { + return f.fail + } + f.got = append(f.got, r) + return nil +} + +type env struct { + t *testing.T + ctx context.Context + srv *httptest.Server + s *Server + id *identitytest.Fixture + logs *fakeLogs + cleanup *fakeCleanup + clock *clock.Fake +} + +func newEnv(t *testing.T) *env { + ctx := context.Background() + pool, err := db.Open(ctx, dbtest.URL(t)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := db.Migrate(ctx, pool); err != nil { + t.Fatal(err) + } + for _, q := range []string{ + `INSERT INTO nodes (name, executor, healthy, observed_at, milli_cpu, memory_bytes, gpus, scratch_bytes) + VALUES ('utmist-z1opa08', 'k3s', true, now(), 32000, 137438953472, 2, 429496729600), + ('utmist-tt', 'k3s', true, now(), 32000, 274877906944, 0, 1099511627776)`, + "INSERT INTO partition_nodes VALUES ('gpu', 'utmist-z1opa08'), ('cpu', 'utmist-tt')", + } { + if _, err := pool.Exec(ctx, q); err != nil { + t.Fatal(err) + } + } + tokens := filepath.Join(t.TempDir(), "tokens") + os.WriteFile(tokens, []byte("# host scope hash\n"+ + "utmist-tt cleanup "+HashToken("tt-secret")+"\n"+ + "utmist-tt provision "+HashToken("tt-provision")+"\n"), 0o600) + hosts, err := LoadHostTokens(tokens, "cleanup") + if err != nil { + t.Fatal(err) + } + c := clock.NewFake(t0) + e := &env{t: t, ctx: ctx, id: identitytest.New(), clock: c, + logs: &fakeLogs{text: map[int64]string{}}, cleanup: &fakeCleanup{}} + e.s = &Server{ + Jobs: scheduler.Jobs{Pool: pool, Clock: c}, Authz: e.id, + Accounts: []Authenticator{testUser{}}, Hosts: hosts, + Logs: e.logs, Cleanup: e.cleanup, Clock: c, + AllowedOrigins: []string{"https://mist.example"}, + StreamRecheck: 20 * time.Millisecond, + } + e.srv = httptest.NewServer(e.s.Handler()) + t.Cleanup(e.srv.Close) + return e +} + +type reply struct { + status int + header http.Header + body []byte +} + +func (r reply) decode(t *testing.T, v any) { + t.Helper() + if err := json.Unmarshal(r.body, v); err != nil { + t.Fatalf("decode %s: %v", r.body, err) + } +} + +func (r reply) code(t *testing.T) string { + t.Helper() + var e struct{ Error Error } + r.decode(t, &e) + return e.Error.Code +} + +func (e *env) do(method, path, user string, body any, hdr ...string) reply { + e.t.Helper() + var rd io.Reader + switch b := body.(type) { + case nil: + case string: + rd = strings.NewReader(b) + default: + buf, _ := json.Marshal(b) + rd = bytes.NewReader(buf) + } + req, _ := http.NewRequest(method, e.srv.URL+path, rd) + if user != "" { + req.Header.Set("X-Test-User", user) + } + for i := 0; i+1 < len(hdr); i += 2 { + req.Header.Set(hdr[i], hdr[i+1]) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + e.t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + return reply{resp.StatusCode, resp.Header, b} +} + +var keys int + +func (e *env) submit(user string, body any) reply { + e.t.Helper() + keys++ + return e.do("POST", "/v1/jobs", user, body, "Idempotency-Key", fmt.Sprintf("key-%d", keys)) +} + +func batch(team string) map[string]any { + return map[string]any{"team": team, "kind": "batch", "script": "#!/bin/bash\necho hi\n"} +} + +func (e *env) mustSubmit(user string, body any) Job { + e.t.Helper() + r := e.submit(user, body) + if r.status != http.StatusCreated { + e.t.Fatalf("submit: %d %s", r.status, r.body) + } + var j Job + r.decode(e.t, &j) + return j +} + +func TestSubmitFillsDefaults(t *testing.T) { + e := newEnv(t) + j := e.mustSubmit("alice", map[string]any{"team": "team-a", "kind": "batch", "script": "x", "resources": map[string]any{"gpus": 2}}) + want := Resources{CPUs: 16, MemoryBytes: 64 << 30, GPUs: 2, ScratchBytes: 100 << 30} + if j.State != "PENDING" || j.Partition != "gpu" || j.QOS != "normal" || j.Resources != want || + j.TimeLimitSeconds != 3600 || j.SignalGraceSeconds != 60 || j.Name != "batch" { + t.Fatalf("defaults: %+v", j) + } + if j.Workdir != "/mist/teams/team-a/alice" || j.OutputPath != fmt.Sprintf("/mist/teams/team-a/alice/mist-%d.out", j.ID) { + t.Fatalf("paths: %q %q", j.Workdir, j.OutputPath) + } + if !j.CanCancel || j.CanAttach { + t.Fatalf("can_cancel/can_attach: %+v", j) + } + + // Team derived from the workdir; CPU partition defaults are per job. + j = e.mustSubmit("bob", map[string]any{"partition": "cpu", "kind": "command", + "argv": []string{"/usr/bin/huggingface-cli", "download", "x"}, "workdir": "/mist/teams/team-a/bob/dl", "output": "logs/%j-%x.txt"}) + if j.Team != "team-a" || j.Resources.CPUs != 2 || j.Resources.MemoryBytes != 8<<30 || j.Name != "huggingface-cli" || + j.OutputPath != fmt.Sprintf("/mist/teams/team-a/bob/dl/logs/%d-%%x.txt", j.ID) { + t.Fatalf("command job: %+v", j) + } +} + +// M1-09: a duplicate submission key gives one job; a conflicting payload +// under the same key is rejected; a retry survives a defaults change. +func TestSubmitIdempotency(t *testing.T) { + e := newEnv(t) + body := batch("team-a") + first := e.do("POST", "/v1/jobs", "alice", body, "Idempotency-Key", "k1") + if first.status != http.StatusCreated { + t.Fatalf("%d %s", first.status, first.body) + } + var j1 Job + first.decode(t, &j1) + + if _, err := e.s.Jobs.Pool.Exec(e.ctx, "UPDATE partitions SET default_time_limit_seconds = 7200"); err != nil { + t.Fatal(err) + } + again := e.do("POST", "/v1/jobs", "alice", body, "Idempotency-Key", "k1") + var j2 Job + again.decode(t, &j2) + if again.status != http.StatusOK || again.header.Get("Idempotent-Replayed") != "true" || j2.ID != j1.ID || j2.TimeLimitSeconds != 3600 { + t.Fatalf("replay: %d %s", again.status, again.body) + } + + changed := batch("team-a") + changed["script"] = "#!/bin/bash\necho different\n" + if r := e.do("POST", "/v1/jobs", "alice", changed, "Idempotency-Key", "k1"); r.status != http.StatusConflict || r.code(t) != CodeConflict { + t.Fatalf("conflict: %d %s", r.status, r.body) + } + // Keys are per account: Bob's k1 is a different submission. + if r := e.do("POST", "/v1/jobs", "bob", body, "Idempotency-Key", "k1"); r.status != http.StatusCreated { + t.Fatalf("bob: %d %s", r.status, r.body) + } + if r := e.do("POST", "/v1/jobs", "alice", body); r.status != http.StatusBadRequest { + t.Fatalf("missing key: %d", r.status) + } + + // Concurrent retries of one key still give one job. + var wg sync.WaitGroup + ids := make(chan int64, 8) + for range 8 { + wg.Go(func() { + r := e.do("POST", "/v1/jobs", "carol", batch("team-b"), "Idempotency-Key", "race") + var j Job + json.Unmarshal(r.body, &j) + ids <- j.ID + }) + } + wg.Wait() + close(ids) + seen := map[int64]bool{} + for id := range ids { + seen[id] = true + } + if len(seen) != 1 || seen[0] { + t.Fatalf("concurrent retries made jobs %v", seen) + } + + // The replay re-checks access: after leaving the team, Alice's retry + // does not return the job. + e.id.RemoveMember("alice", "team-a") + if r := e.do("POST", "/v1/jobs", "alice", body, "Idempotency-Key", "k1"); r.status != http.StatusNotFound { + t.Fatalf("replay after removal: %d %s", r.status, r.body) + } +} + +// M1-09: hostile paths, env and argv are rejected before they are stored. +func TestSubmitValidation(t *testing.T) { + e := newEnv(t) + cases := []struct { + name string + edit func(map[string]any) + status int + }{ + {"dotdot workdir", func(b map[string]any) { b["workdir"] = "/mist/teams/team-a/../team-b/carol" }, 400}, + {"other team workdir", func(b map[string]any) { b["workdir"] = "/mist/teams/team-b/carol" }, 400}, + {"landing folder", func(b map[string]any) { b["workdir"] = "/mist/login/alice" }, 400}, + {"relative workdir", func(b map[string]any) { b["workdir"] = "team-a/alice" }, 400}, + {"trailing slash", func(b map[string]any) { b["workdir"] = "/mist/teams/team-a/alice/" }, 400}, + {"output escapes", func(b map[string]any) { b["output"] = "../../team-b/x.out" }, 400}, + {"output elsewhere", func(b map[string]any) { b["output"] = "/etc/cron.d/x" }, 400}, + {"output is team root", func(b map[string]any) { b["workdir"] = "/mist/teams/team-a"; b["output"] = "." }, 400}, + {"control char path", func(b map[string]any) { b["output"] = "a\nb.out" }, 400}, + {"protected env", func(b map[string]any) { b["env"] = map[string]string{"HOME": "/tmp"} }, 400}, + {"mist env", func(b map[string]any) { b["env"] = map[string]string{"MIST_JOB_ID": "1"} }, 400}, + {"device env", func(b map[string]any) { b["env"] = map[string]string{"CUDA_VISIBLE_DEVICES": "0,1"} }, 400}, + {"bad env name", func(b map[string]any) { b["env"] = map[string]string{"A=B": "c"} }, 400}, + {"nul env", func(b map[string]any) { b["env"] = map[string]string{"A": "b\x00c"} }, 400}, + {"nul arg", func(b map[string]any) { b["args"] = []string{"a\x00b"} }, 400}, + {"batch with argv", func(b map[string]any) { b["argv"] = []string{"ls"} }, 400}, + {"command with script", func(b map[string]any) { b["kind"] = "command"; b["argv"] = []string{"ls"} }, 400}, + {"interactive", func(b map[string]any) { b["kind"] = "interactive"; delete(b, "script") }, 400}, + {"interactive qos", func(b map[string]any) { b["qos"] = "interactive" }, 400}, + {"unknown field", func(b map[string]any) { b["uid"] = 0 }, 400}, + {"huge script", func(b map[string]any) { b["script"] = strings.Repeat("x", maxScript+1) }, 400}, + {"long name", func(b map[string]any) { b["name"] = strings.Repeat("n", 129) }, 400}, + {"fractional cpus", func(b map[string]any) { b["resources"] = map[string]any{"cpus": 1.5} }, 400}, + {"zero gpus on gpu", func(b map[string]any) { b["resources"] = map[string]any{"gpus": 0} }, 400}, + {"gpu on cpu", func(b map[string]any) { b["partition"] = "cpu"; b["resources"] = map[string]any{"gpus": 1} }, 400}, + {"grace too long", func(b map[string]any) { b["signal_grace_seconds"] = 121 }, 400}, + {"other runtime", func(b map[string]any) { b["runtime"] = "evil@sha256:00" }, 400}, + {"unknown partition", func(b map[string]any) { b["partition"] = "tpu" }, 400}, + {"three gpus", func(b map[string]any) { b["resources"] = map[string]any{"gpus": 3} }, 422}, + {"too many cpus", func(b map[string]any) { b["resources"] = map[string]any{"cpus": 27} }, 422}, + {"too long", func(b map[string]any) { b["time_limit_seconds"] = 86401 }, 422}, + {"bigger than node", func(b map[string]any) { b["resources"] = map[string]any{"gpus": 2, "cpus": 40} }, 422}, + {"not a member", func(b map[string]any) { b["team"] = "team-b"; b["workdir"] = "/mist/teams/team-b/bob" }, 403}, + } + for _, c := range cases { + b := batch("team-a") + c.edit(b) + if r := e.submit("bob", b); r.status != c.status { + t.Errorf("%s: got %d %s, want %d", c.name, r.status, r.body, c.status) + } + } + var n int + e.s.Jobs.Pool.QueryRow(e.ctx, "SELECT count(*) FROM jobs").Scan(&n) + if n != 0 { + t.Fatalf("%d rejected submissions were stored", n) + } + + e.id.Unavailable = true + if r := e.submit("bob", batch("team-a")); r.status != http.StatusServiceUnavailable { + t.Fatalf("identity outage: %d", r.status) + } +} + +// M1-09: non-team callers get not_found on detail, log and cancel; the +// authorization matrix follows api.md. +func TestAuthorizationMatrix(t *testing.T) { + e := newEnv(t) + aliceJob := e.mustSubmit("alice", batch("team-a")) + bobJob := e.mustSubmit("bob", batch("team-a")) + path := func(id int64, rest string) string { return fmt.Sprintf("/v1/jobs/%d%s", id, rest) } + + for _, u := range []string{"alice", "bob", "root-admin"} { + if r := e.do("GET", path(aliceJob.ID, "?script=1"), u, nil); r.status != 200 { + t.Errorf("%s view: %d", u, r.status) + } + } + for _, rest := range []string{"", "?script=1", "/log"} { + if r := e.do("GET", path(aliceJob.ID, rest), "carol", nil); r.status != 404 || r.code(t) != CodeNotFound { + t.Errorf("carol GET %s: %d %s", rest, r.status, r.body) + } + } + if r := e.do("GET", path(99999, ""), "carol", nil); r.status != 404 { + t.Errorf("missing job: %d", r.status) + } + if r := e.do("POST", path(aliceJob.ID, "/cancel"), "carol", nil); r.status != 404 { + t.Errorf("carol cancel: %d", r.status) + } + + // can_cancel is computed per caller. + var seen Job + e.do("GET", path(aliceJob.ID, ""), "bob", nil).decode(t, &seen) + if seen.CanCancel { + t.Error("bob may not cancel alice's job") + } + if r := e.do("POST", path(aliceJob.ID, "/cancel"), "bob", nil); r.status != 403 { + t.Errorf("bob cancel alice: %d", r.status) + } + // Alice leads team-a, so she can cancel Bob's job. + r := e.do("POST", path(bobJob.ID, "/cancel"), "alice", nil) + var cancelled Job + r.decode(t, &cancelled) + if r.status != 200 || cancelled.State != "CANCELLED" || cancelled.CanCancel { + t.Fatalf("lead cancel: %d %s", r.status, r.body) + } + if r := e.do("POST", path(bobJob.ID, "/cancel"), "bob", nil); r.status != 200 { + t.Errorf("repeat cancel: %d", r.status) + } + + // Former membership or ownership alone never grants access. + e.id.RemoveMember("bob", "team-a") + if r := e.do("GET", path(bobJob.ID, ""), "bob", nil); r.status != 404 { + t.Errorf("bob after removal: %d", r.status) + } + var mine List[Job] + e.do("GET", "/v1/jobs", "bob", nil).decode(t, &mine) + if len(mine.Items) != 0 { + t.Errorf("bob's list after removal: %+v", mine.Items) + } + e.id.Disable("alice") + if r := e.do("GET", path(aliceJob.ID, ""), "alice", nil); r.status != 404 { + t.Errorf("disabled alice: %d", r.status) + } + + if r := e.do("GET", "/v1/jobs", "", nil); r.status != 401 || r.code(t) != CodeUnauthenticated { + t.Errorf("anonymous: %d", r.status) + } + if r := e.do("GET", "/v1/jobs?all=1", "carol", nil); r.status != 403 { + t.Errorf("non-admin all: %d", r.status) + } + var all List[Job] + e.do("GET", "/v1/jobs?all=1", "root-admin", nil).decode(t, &all) + if len(all.Items) != 2 { + t.Errorf("admin all: %d jobs", len(all.Items)) + } +} + +func TestListPagination(t *testing.T) { + e := newEnv(t) + for range 5 { + e.mustSubmit("alice", batch("team-a")) + } + e.mustSubmit("carol", batch("team-b")) + var ids []int64 + cursor := "" + for { + var page List[Job] + e.do("GET", "/v1/jobs?team=team-a&limit=2"+cursor, "bob", nil).decode(t, &page) + for _, j := range page.Items { + ids = append(ids, j.ID) + } + if page.NextCursor == nil { + break + } + cursor = "&cursor=" + *page.NextCursor + } + if len(ids) != 5 || ids[0] < ids[4] { + t.Fatalf("pages: %v", ids) + } + if r := e.do("GET", "/v1/jobs?team=team-b", "bob", nil); r.status != 403 { + t.Fatalf("other team list: %d", r.status) + } + var running List[Job] + e.do("GET", "/v1/jobs?state=RUNNING", "alice", nil).decode(t, &running) + if len(running.Items) != 0 { + t.Fatal("state filter") + } +} + +// M1-09: the queue shows other teams' jobs without embedded identities. +func TestQueueAnonymised(t *testing.T) { + e := newEnv(t) + secret := batch("team-a") + secret["name"] = "alice-secret-project" + secret["env"] = map[string]string{"WANDB_API_KEY": "hunter2"} + a := e.mustSubmit("alice", secret) + c := e.mustSubmit("carol", batch("team-b")) + + r := e.do("GET", "/v1/queue", "carol", nil) + for _, leak := range []string{"alice", "team-a", "secret", "hunter2", "/mist/"} { + if bytes.Contains(r.body, []byte(leak)) { + t.Errorf("queue leaks %q: %s", leak, r.body) + } + } + var q List[QueueEntry] + r.decode(t, &q) + if len(q.Items) != 2 { + t.Fatalf("queue: %s", r.body) + } + for _, it := range q.Items { + switch it.ID { + case a.ID: + if it.User != nil || it.Team != nil || it.Name != nil || it.Resources.GPUs != 1 { + t.Errorf("other team entry: %+v", it) + } + case c.ID: + if it.User == nil || *it.User != "carol" || it.Workdir != nil { + t.Errorf("own entry: %+v", it) + } + } + } + // Bob is in team-a and sees names there. + e.do("GET", "/v1/queue?partition=gpu", "bob", nil).decode(t, &q) + for _, it := range q.Items { + if it.ID == a.ID && (it.User == nil || *it.User != "alice") { + t.Errorf("team member entry: %+v", it) + } + } +} + +func TestPartitions(t *testing.T) { + e := newEnv(t) + if _, err := e.s.Jobs.Pool.Exec(e.ctx, + "INSERT INTO limits VALUES ('user', 'carol', 'gpu', '*', 'gpus', 2), ('team', 'team-b', 'gpu', '*', 'gpus', 3)"); err != nil { + t.Fatal(err) + } + r := e.do("GET", "/v1/partitions", "bob", nil) + if bytes.Contains(r.body, []byte("carol")) || bytes.Contains(r.body, []byte("team-b")) { + t.Fatalf("other subjects' limits shown: %s", r.body) + } + var ps List[Partition] + r.decode(t, &ps) + var gpu *Partition + for i := range ps.Items { + if ps.Items[i].Name == "gpu" { + gpu = &ps.Items[i] + } + } + if gpu == nil || !gpu.Default || gpu.Free.GPUs != 2 || gpu.DefaultCPUs != 8 || !gpu.PerAccelerator || len(gpu.Limits) != 3 { + t.Fatalf("gpu partition: %+v", gpu) + } + e.do("GET", "/v1/partitions", "carol", nil).decode(t, &ps) + for _, p := range ps.Items { + if p.Name == "gpu" && len(p.Limits) != 5 { + t.Fatalf("carol's limits: %+v", p.Limits) + } + } +} + +func TestLogs(t *testing.T) { + e := newEnv(t) + j := e.mustSubmit("alice", batch("team-a")) + var l Log + e.do("GET", fmt.Sprintf("/v1/jobs/%d/log", j.ID), "bob", nil).decode(t, &l) + if l.Text != "" || l.State != "PENDING" || l.OutputPath != j.OutputPath { + t.Fatalf("pending log: %+v", l) + } + + // Pretend it was admitted and wrote a lot of output. + if _, err := e.s.Jobs.Pool.Exec(e.ctx, "UPDATE jobs SET node = 'utmist-z1opa08', state = 'RUNNING' WHERE id = $1", j.ID); err != nil { + t.Fatal(err) + } + e.logs.text[j.ID] = strings.Repeat("a", 100<<10) + "the end\n" + e.do("GET", fmt.Sprintf("/v1/jobs/%d/log", j.ID), "alice", nil).decode(t, &l) + if len(l.Text) != tailBytes || !l.Truncated || !strings.HasSuffix(l.Text, "the end\n") { + t.Fatalf("tail: %d bytes, truncated %v", len(l.Text), l.Truncated) + } + + // Follow over a WebSocket. + wsURL := "ws" + strings.TrimPrefix(e.srv.URL, "http") + fmt.Sprintf("/v1/jobs/%d/log?follow=1", j.ID) + cfg, _ := websocket.NewConfig(wsURL, "https://mist.example") + cfg.Header.Set("X-Test-User", "alice") + conn, err := websocket.DialConfig(cfg) + if err != nil { + t.Fatal(err) + } + var got []byte + for { + var frame []byte + if err := websocket.Message.Receive(conn, &frame); err != nil { + break + } + got = append(got, frame...) + } + conn.Close() + if string(got) != e.logs.text[j.ID] { + t.Fatalf("followed %d bytes", len(got)) + } + + cfg, _ = websocket.NewConfig(wsURL, "https://evil.example") + cfg.Header.Set("X-Test-User", "alice") + if _, err := websocket.DialConfig(cfg); err == nil { + t.Fatal("disallowed origin accepted") + } +} + +// An open stream ends when the caller loses access. +func TestFollowRevoked(t *testing.T) { + e := newEnv(t) + j := e.mustSubmit("alice", batch("team-a")) + e.s.Jobs.Pool.Exec(e.ctx, "UPDATE jobs SET node = 'utmist-z1opa08', state = 'RUNNING' WHERE id = $1", j.ID) + e.s.Logs = blockingLogs{} + wsURL := "ws" + strings.TrimPrefix(e.srv.URL, "http") + fmt.Sprintf("/v1/jobs/%d/log?follow=1", j.ID) + cfg, _ := websocket.NewConfig(wsURL, "https://mist.example") + cfg.Header.Set("X-Test-User", "bob") + conn, err := websocket.DialConfig(cfg) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + e.id.RemoveMember("bob", "team-a") + conn.SetReadDeadline(time.Now().Add(5 * time.Second)) + var frame []byte + if err := websocket.Message.Receive(conn, &frame); err == nil || strings.Contains(err.Error(), "timeout") { + t.Fatalf("stream still open after revocation: %v", err) + } +} + +type blockingLogs struct{} + +func (blockingLogs) Logs(ctx context.Context, _ int64, _ bool) (io.ReadCloser, error) { + pr, pw := io.Pipe() + go func() { <-ctx.Done(); pw.Close() }() + return pr, nil +} + +// The production build has no development authenticator and rejects its +// header outright, even alongside valid credentials. +func TestProductionRejectsDevHeader(t *testing.T) { + e := newEnv(t) + r := e.do("GET", "/v1/jobs", "alice", nil, DevUserHeader, "root-admin") + if r.status != 401 || !bytes.Contains(r.body, []byte("development")) { + t.Fatalf("dev header: %d %s", r.status, r.body) + } + r = e.do("GET", "/v1/jobs", "", nil, DevUserHeader, "root-admin") + if r.status != 401 { + t.Fatalf("dev header alone: %d", r.status) + } +} + +func TestCleanupEndpoint(t *testing.T) { + e := newEnv(t) + receipt := k3s.CleanupReceipt{ExecutionID: "3f1c2a9e-8b7d-4c6e-9a5b-1d2e3f4a5b6c", PodUID: "u", Node: "utmist-tt", Evidence: "gone"} + url := "/internal/v1/executions/" + receipt.ExecutionID + "/cleanup" + post := func(token string, rc k3s.CleanupReceipt, path string) reply { + hdr := []string{} + if token != "" { + hdr = []string{"Authorization", "Bearer " + token} + } + return e.do("POST", path, "", rc, hdr...) + } + if r := post("", receipt, url); r.status != 401 { + t.Errorf("no token: %d", r.status) + } + if r := post("tt-provision", receipt, url); r.status != 401 { + t.Errorf("token for another scope: %d", r.status) + } + // Account credentials do not reach /internal. + if r := e.do("POST", url, "root-admin", receipt); r.status != 401 { + t.Errorf("account on internal: %d", r.status) + } + other := receipt + other.Node = "utmist-z1opa08" + if r := post("tt-secret", other, url); r.status != 403 { + t.Errorf("report for another node: %d", r.status) + } + if r := post("tt-secret", receipt, "/internal/v1/executions/other/cleanup"); r.status != 400 { + t.Errorf("path mismatch: %d", r.status) + } + if r := post("tt-secret", receipt, url); r.status != 200 || len(e.cleanup.got) != 1 { + t.Errorf("valid receipt: %d %s", r.status, r.body) + } + e.cleanup.fail = fmt.Errorf("%w: Pod job-1 still exists", k3s.ErrCleanupRejected) + if r := post("tt-secret", receipt, url); r.status != 409 { + t.Errorf("rejected receipt: %d", r.status) + } + // Host tokens do not reach /v1. + if r := e.do("GET", "/v1/jobs", "", nil, "Authorization", "Bearer tt-secret"); r.status != 401 { + t.Errorf("host on v1: %d", r.status) + } +} diff --git a/internal/api/auth.go b/internal/api/auth.go new file mode 100644 index 0000000..258269c --- /dev/null +++ b/internal/api/auth.go @@ -0,0 +1,192 @@ +package api + +import ( + "bufio" + "context" + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "errors" + "fmt" + "net/http" + "os" + "strings" + + "github.com/UTMIST/Mist/internal/identity" +) + +// Authenticator resolves the caller of a request. It returns +// ErrNoCredentials when the request carries no credential of its kind, so +// the next authenticator can try, and any other error when a credential is +// present but invalid, which ends the request as unauthenticated. +// +// Identity provides the portal session and relay authenticators; this +// package provides host tokens for /internal and, in development builds +// only, a header naming the user. +type Authenticator interface { + Authenticate(r *http.Request) (identity.Principal, error) +} + +// ErrNoCredentials: the request carries no credential this authenticator +// understands. +var ErrNoCredentials = errors.New("api: no credentials") + +// DevUserHeader names the account in development builds (tag mistdev). +// Production binaries do not contain the code that reads it, and reject +// any request that sends it, so a misconfigured proxy cannot let it +// through silently. +const DevUserHeader = "X-Mist-Dev-User" + +// devAuthenticator is implemented only by the development authenticator. +type devAuthenticator interface{ acceptsDevHeader() } + +type principalKey struct{} + +func principalFrom(ctx context.Context) identity.Principal { + p, _ := ctx.Value(principalKey{}).(identity.Principal) + return p +} + +type handler func(w http.ResponseWriter, r *http.Request, p identity.Principal) error + +func unauthenticated(msg string) *Error { + return &Error{http.StatusUnauthorized, CodeUnauthenticated, msg, nil} +} + +// account authenticates a /v1 request as an account. +func (s *Server) account(h handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + p, err := s.authenticate(r, s.Accounts) + if err == nil && (p.Kind != identity.PrincipalAccount || p.Username == "") { + err = unauthenticated("account credentials required") + } + if err == nil { + err = h(w, r.WithContext(context.WithValue(r.Context(), principalKey{}, p)), p) + } + if err != nil { + s.fail(w, r, err) + } + }) +} + +// host authenticates an /internal request as a host agent. +func (s *Server) host(h handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var list []Authenticator + if s.Hosts != nil { + list = []Authenticator{s.Hosts} + } + p, err := s.authenticate(r, list) + if err == nil && (p.Kind != identity.PrincipalHost || p.Host == "") { + err = unauthenticated("host credentials required") + } + if err == nil { + err = h(w, r, p) + } + if err != nil { + s.fail(w, r, err) + } + }) +} + +func (s *Server) authenticate(r *http.Request, list []Authenticator) (identity.Principal, error) { + if r.Header.Get(DevUserHeader) != "" { + allowed := false + for _, a := range list { + if _, ok := a.(devAuthenticator); ok { + allowed = true + } + } + if !allowed { + return identity.Principal{}, unauthenticated("the development user header is not accepted by this server") + } + } + for _, a := range list { + p, err := a.Authenticate(r) + if errors.Is(err, ErrNoCredentials) { + continue + } + if err != nil { + s.Log.Info("authentication failed", "path", r.URL.Path, "err", err) + return identity.Principal{}, unauthenticated("invalid credentials") + } + return p, nil + } + return identity.Principal{}, unauthenticated("sign in or use the mist command on the login node") +} + +// HostTokens authenticates host agents by bearer token. Only SHA-256 +// hashes of tokens are stored, one per line: " ". +// Each instance accepts one scope, so a token issued for one operation +// cannot be used for another. +type HostTokens struct { + scope string + entries []hostToken +} + +type hostToken struct { + host string + hash [sha256.Size]byte +} + +// LoadHostTokens reads a token file, keeping the entries for scope. +// Blank lines and lines starting with # are ignored. +func LoadHostTokens(path, scope string) (*HostTokens, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + ht := &HostTokens{scope: scope} + sc := bufio.NewScanner(f) + for n := 1; sc.Scan(); n++ { + line := strings.TrimSpace(sc.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + fields := strings.Fields(line) + if len(fields) != 3 { + return nil, fmt.Errorf("%s:%d: want \" \"", path, n) + } + raw, err := hex.DecodeString(fields[2]) + if err != nil || len(raw) != sha256.Size { + return nil, fmt.Errorf("%s:%d: token hash is not 64 hex digits", path, n) + } + if fields[1] != scope { + continue + } + var t hostToken + t.host = fields[0] + copy(t.hash[:], raw) + ht.entries = append(ht.entries, t) + } + return ht, sc.Err() +} + +// HashToken is the form stored in a host token file. +func HashToken(token string) string { + h := sha256.Sum256([]byte(token)) + return hex.EncodeToString(h[:]) +} + +func (ht *HostTokens) Authenticate(r *http.Request) (identity.Principal, error) { + auth := r.Header.Get("Authorization") + if auth == "" { + return identity.Principal{}, ErrNoCredentials + } + token, ok := strings.CutPrefix(auth, "Bearer ") + if !ok || token == "" { + return identity.Principal{}, errors.New("malformed Authorization header") + } + sum := sha256.Sum256([]byte(token)) + host := "" + for _, t := range ht.entries { // compare every entry, in constant time each + if subtle.ConstantTimeCompare(sum[:], t.hash[:]) == 1 { + host = t.host + } + } + if host == "" { + return identity.Principal{}, fmt.Errorf("unknown %s token", ht.scope) + } + return identity.Principal{Kind: identity.PrincipalHost, Host: host}, nil +} diff --git a/internal/api/cleanup.go b/internal/api/cleanup.go new file mode 100644 index 0000000..4a0d501 --- /dev/null +++ b/internal/api/cleanup.go @@ -0,0 +1,43 @@ +package api + +import ( + "encoding/json" + "errors" + "net/http" + + "github.com/UTMIST/Mist/internal/executor/k3s" + "github.com/UTMIST/Mist/internal/identity" +) + +// cleanup handles POST /internal/v1/executions/{execution_id}/cleanup: a +// node's cleanup reporter delivering the receipt from `mistd +// cleanup-report`. A host may only report for itself; the executor then +// checks the receipt against its own records before anything is released. +func (s *Server) cleanup(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + if s.Cleanup == nil { + return unavailable("cleanup receipts are not accepted by this API server") + } + var rc k3s.CleanupReceipt + dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, 64<<10)) + dec.DisallowUnknownFields() + if err := dec.Decode(&rc); err != nil { + return invalid("", "receipt: "+err.Error()) + } + if rc.ExecutionID != r.PathValue("execution_id") { + return invalid("execution_id", "receipt is for a different execution") + } + if rc.Node != p.Host { + return &Error{http.StatusForbidden, CodeForbidden, "a host can only report cleanup on itself", nil} + } + if err := s.Cleanup.RecordCleanup(r.Context(), rc); errors.Is(err, k3s.ErrCleanupRejected) { + return &Error{http.StatusConflict, CodeConflict, err.Error(), nil} + } else if err != nil { + return err + } + s.Log.Info("cleanup receipt recorded", "execution", rc.ExecutionID, "node", rc.Node) + if err := s.Jobs.Wake(r.Context()); err != nil { + s.Log.Warn("wake scheduler", "err", err) + } + writeJSON(w, http.StatusOK, map[string]string{"status": "recorded"}) + return nil +} diff --git a/internal/api/dev.go b/internal/api/dev.go new file mode 100644 index 0000000..f3360fe --- /dev/null +++ b/internal/api/dev.go @@ -0,0 +1,24 @@ +//go:build mistdev + +package api + +import ( + "net/http" + + "github.com/UTMIST/Mist/internal/identity" +) + +// DevHeader trusts DevUserHeader as the account name. It exists only in +// binaries built with -tags mistdev, for local development against the +// identitytest fixture. +type DevHeader struct{} + +func (DevHeader) acceptsDevHeader() {} + +func (DevHeader) Authenticate(r *http.Request) (identity.Principal, error) { + u := r.Header.Get(DevUserHeader) + if u == "" { + return identity.Principal{}, ErrNoCredentials + } + return identity.Principal{Kind: identity.PrincipalAccount, Username: u}, nil +} diff --git a/internal/api/dev_test.go b/internal/api/dev_test.go new file mode 100644 index 0000000..7ac75f7 --- /dev/null +++ b/internal/api/dev_test.go @@ -0,0 +1,16 @@ +//go:build mistdev + +package api + +import ( + "net/http" + "testing" +) + +func TestDevHeader(t *testing.T) { + e := newEnv(t) + e.s.Accounts = []Authenticator{DevHeader{}} + if r := e.do("GET", "/v1/jobs", "", nil, DevUserHeader, "alice"); r.status != http.StatusOK { + t.Fatalf("dev header in a development build: %d %s", r.status, r.body) + } +} diff --git a/internal/api/jobs.go b/internal/api/jobs.go new file mode 100644 index 0000000..8d3904e --- /dev/null +++ b/internal/api/jobs.go @@ -0,0 +1,348 @@ +package api + +import ( + "context" + "errors" + "net/http" + "slices" + "strconv" + "strings" + "time" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity" + "github.com/UTMIST/Mist/internal/scheduler" +) + +func jobID(r *http.Request) (int64, error) { + id, err := strconv.ParseInt(r.PathValue("id"), 10, 64) + if err != nil || id < 1 { + return 0, notFound("job not found") + } + return id, nil +} + +func target(j *scheduler.JobInfo) identity.Target { + return identity.Target{Team: j.Team, Owner: j.Username} +} + +// viewableJob loads a job the caller may view under current membership. +// Anything else, including a job that does not exist, is not_found. +func (s *Server) viewableJob(ctx context.Context, p identity.Principal, id int64) (*scheduler.JobInfo, error) { + j, err := s.Jobs.Get(ctx, id) + if errors.Is(err, scheduler.ErrNotFound) { + return nil, notFound("job not found") + } + if err != nil { + return nil, err + } + if err := s.Authz.Authorize(ctx, p, identity.ActionViewJob, target(j)); err != nil { + return nil, authzError(err, true, "job") + } + return j, nil +} + +// cancellable: the job can still be stopped by a cancel request. +func cancellable(st scheduler.State) bool { + return st == scheduler.StatePending || st == scheduler.StateLaunching || st == scheduler.StateRunning +} + +// jobObject builds the job object, computing can_cancel from state and +// current authorization so clients never decide it. +func (s *Server) jobObject(ctx context.Context, p identity.Principal, j *scheduler.JobInfo) (Job, error) { + can := false + if cancellable(j.State) { + err := s.Authz.Authorize(ctx, p, identity.ActionCancelJob, target(j)) + if errors.Is(err, identity.ErrUnavailable) { + return Job{}, authzError(err, false, "") + } + can = err == nil + } + return wireJob(j, can), nil +} + +// getJob handles GET /v1/jobs/{id}[?script=1]. +func (s *Server) getJob(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + id, err := jobID(r) + if err != nil { + return err + } + j, err := s.viewableJob(r.Context(), p, id) + if err != nil { + return err + } + out, err := s.jobObject(r.Context(), p, j) + if err != nil { + return err + } + events, err := s.Jobs.Events(r.Context(), id) + if err != nil { + return err + } + out.Events = make([]Event, len(events)) + for i, e := range events { + out.Events[i] = Event{At: e.At.UTC(), From: e.From, To: e.To, Reason: e.Reason, Actor: e.Actor} + } + if r.URL.Query().Get("script") == "1" { + spec, err := s.Jobs.Spec(r.Context(), id) + if err != nil { + return err + } + out.Script, out.Args, out.Argv, out.Env = &spec.Script, spec.Args, spec.Argv, spec.Env + } + writeJSON(w, http.StatusOK, out) + return nil +} + +// listJobs handles GET /v1/jobs. Without parameters it lists the caller's +// own jobs in teams they still belong to. ?team= lists a team's jobs for +// its members; ?all=1 lists every job for admins. ?state= takes a comma +// list of states, or "active". +func (s *Server) listJobs(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + q := r.URL.Query() + f := scheduler.ListFilter{Team: q.Get("team"), Limit: 50} + if v := q.Get("limit"); v != "" { + n, err := strconv.Atoi(v) + if err != nil || n < 1 || n > 200 { + return invalid("limit", "limit must be 1–200") + } + f.Limit = n + } + if v := q.Get("cursor"); v != "" { + n, err := strconv.ParseInt(v, 10, 64) + if err != nil || n < 1 { + return invalid("cursor", "invalid cursor") + } + f.Before = n + } + for _, st := range strings.Split(q.Get("state"), ",") { + switch { + case st == "": + case st == "active": + f.States = append(f.States, scheduler.StatePending, scheduler.StateLaunching, + scheduler.StateRunning, scheduler.StateStopping, scheduler.StateCleaning) + case knownState(scheduler.State(st)): + f.States = append(f.States, scheduler.State(st)) + default: + return invalid("state", "unknown state "+st) + } + } + + switch { + case q.Get("all") == "1": + if err := s.Authz.Authorize(r.Context(), p, identity.ActionAdmin, identity.Target{}); err != nil { + return authzError(err, false, "") + } + case f.Team != "": + if err := s.Authz.Authorize(r.Context(), p, identity.ActionViewJob, identity.Target{Team: f.Team}); err != nil { + return authzError(err, false, "") + } + default: + f.Username = p.Username + } + + rows, err := s.Jobs.List(r.Context(), f) + if err != nil { + return err + } + out := List[Job]{Items: []Job{}} + viewable := map[string]bool{} // own jobs: only teams the caller still belongs to + for _, j := range rows { + if f.Username != "" { + ok, seen := viewable[j.Team] + if !seen { + err := s.Authz.Authorize(r.Context(), p, identity.ActionViewJob, target(j)) + if errors.Is(err, identity.ErrUnavailable) { + return authzError(err, false, "") + } + ok = err == nil + viewable[j.Team] = ok + } + if !ok { + continue + } + } + job, err := s.jobObject(r.Context(), p, j) + if err != nil { + return err + } + out.Items = append(out.Items, job) + } + if len(rows) == f.Limit { + c := strconv.FormatInt(rows[len(rows)-1].ID, 10) + out.NextCursor = &c + } + writeJSON(w, http.StatusOK, out) + return nil +} + +func knownState(st scheduler.State) bool { + return slices.Contains([]scheduler.State{ + scheduler.StatePending, scheduler.StateLaunching, scheduler.StateRunning, scheduler.StateStopping, + scheduler.StateCleaning, scheduler.StateCompleted, scheduler.StateFailed, scheduler.StateOutOfMemory, + scheduler.StateLaunchFailed, scheduler.StateCancelled, scheduler.StateTimeout, scheduler.StateNodeFail, + }, st) +} + +// cancelJob handles POST /v1/jobs/{id}/cancel. It records the stop +// request and returns the job; it never waits for the job to end. +// Cancelling an ended job is not an error. +func (s *Server) cancelJob(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + id, err := jobID(r) + if err != nil { + return err + } + j, err := s.viewableJob(r.Context(), p, id) + if err != nil { + return err + } + if err := s.Authz.Authorize(r.Context(), p, identity.ActionCancelJob, target(j)); err != nil { + if !errors.Is(err, identity.ErrUnavailable) { + return &Error{http.StatusForbidden, CodeForbidden, + "only the job's owner, a lead of team " + j.Team + " or an admin can cancel it", nil} + } + return authzError(err, false, "") + } + if !j.State.Terminal() { + if _, err := s.Jobs.Cancel(r.Context(), id, p.Username); err != nil { + return err + } + if j, err = s.Jobs.Get(r.Context(), id); err != nil { + return err + } + } + out, err := s.jobObject(r.Context(), p, j) + if err != nil { + return err + } + writeJSON(w, http.StatusOK, out) + return nil +} + +// queue handles GET /v1/queue[?partition=]. Every caller sees every +// queued and running job, but only as an allowlisted summary; identities +// are filled in only for jobs the caller could view anyway (D15). +func (s *Server) queue(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + rows, err := s.Jobs.Queue(r.Context(), r.URL.Query().Get("partition")) + if err != nil { + return err + } + now := s.Clock.Now() + viewable := map[string]bool{} + out := List[QueueEntry]{Items: make([]QueueEntry, 0, len(rows))} + for _, j := range rows { + ok, seen := viewable[j.Team] + if !seen { + err := s.Authz.Authorize(r.Context(), p, identity.ActionViewJob, identity.Target{Team: j.Team}) + if errors.Is(err, identity.ErrUnavailable) { + return authzError(err, false, "") + } + ok = err == nil + viewable[j.Team] = ok + } + e := QueueEntry{ + ID: j.ID, Partition: j.Partition, QOS: j.QOS, State: string(j.State), + Resources: wireResources(j.Resources), + AgeSeconds: max(int64(now.Sub(j.SubmittedAt)/time.Second), 0), + PendingReason: publicReason(j.PendingReason), + } + if ok { + e.User, e.Team, e.Name = &j.Username, &j.Team, &j.Name + } + out.Items = append(out.Items, e) + } + writeJSON(w, http.StatusOK, out) + return nil +} + +// publicReason passes only known pending-reason codes, never free text. +func publicReason(r *string) *string { + if r == nil { + return nil + } + switch scheduler.PendingReason(*r) { + case scheduler.PendingPriority, scheduler.PendingResources, scheduler.PendingUserLimit, + scheduler.PendingTeamLimit, scheduler.PendingQOSLimit, scheduler.PendingPartitionDown, + scheduler.PendingProvisioning, scheduler.PendingAuthUnavailable, scheduler.PendingHeld: + return r + } + return nil +} + +// partitions handles GET /v1/partitions: configuration, defaults, live +// free capacity and the limits that apply to the caller. +func (s *Server) partitions(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + cat, err := s.Jobs.Catalog(r.Context()) + if err != nil { + return err + } + admin := s.Authz.Authorize(r.Context(), p, identity.ActionAdmin, identity.Target{}) == nil + member := map[string]bool{} + applies := func(l scheduler.Limit) (bool, error) { + switch { + case admin || l.Subject == "*": + return true, nil + case l.Scope == "user": + return l.Subject == p.Username, nil + case l.Scope == "team": + ok, seen := member[l.Subject] + if !seen { + err := s.Authz.Authorize(r.Context(), p, identity.ActionViewJob, identity.Target{Team: l.Subject}) + if errors.Is(err, identity.ErrUnavailable) { + return false, authzError(err, false, "") + } + ok = err == nil + member[l.Subject] = ok + } + return ok, nil + } + return false, nil + } + + out := List[Partition]{Items: []Partition{}} + for _, c := range cat.Partitions { + pw := Partition{ + Name: c.Name, State: c.State, Default: c.IsDefault, Accelerator: c.Accelerator, + MinAccelerators: c.MinAccelerators, MaxAccelerators: c.MaxAccelerators, + PerAccelerator: c.Accelerator != "none", + DefaultCPUs: c.DefaultMilliCPU / 1000, MaxCPUs: c.MaxMilliCPU / 1000, + DefaultMemoryBytes: c.DefaultMemory, MaxMemoryBytes: c.MaxMemory, + DefaultScratchBytes: c.DefaultScratch, MaxScratchBytes: c.MaxScratch, + DefaultTimeLimitSeconds: int64(c.DefaultTimeLimit / time.Second), + MaxTimeLimitSeconds: int64(c.MaxTimeLimit / time.Second), + DefaultSignalGraceSeconds: int64(c.DefaultSignalGrace / time.Second), + CapMemoryBytes: c.CapMemory, + Nodes: []Node{}, + Limits: []Limit{}, + } + if c.CapMilliCPU != nil { + v := *c.CapMilliCPU / 1000 + pw.CapCPUs = &v + } + var free executor.Resources + for _, n := range c.Nodes { + pw.Nodes = append(pw.Nodes, Node{Name: n.Name, State: n.AdminState, Healthy: n.Healthy, + Quarantined: n.Quarantined, Total: wireResources(n.Total), Free: wireResources(n.Free)}) + if n.AdminState == "up" && n.Healthy && !n.Quarantined && c.State == "up" { + free = free.Add(n.Free) + } + } + pw.Free = wireResources(free) + for _, l := range cat.Limits { + if l.Partition != "*" && l.Partition != c.Name { + continue + } + ok, err := applies(l) + if err != nil { + return err + } + if ok { + pw.Limits = append(pw.Limits, Limit{Scope: l.Scope, Subject: l.Subject, QOS: l.QOS, + Resource: l.Resource, Max: l.Max}) + } + } + out.Items = append(out.Items, pw) + } + writeJSON(w, http.StatusOK, out) + return nil +} diff --git a/internal/api/logs.go b/internal/api/logs.go new file mode 100644 index 0000000..e574a9d --- /dev/null +++ b/internal/api/logs.go @@ -0,0 +1,162 @@ +package api + +import ( + "context" + "errors" + "io" + "net/http" + "net/url" + "slices" + "strings" + "time" + + "golang.org/x/net/websocket" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity" + "github.com/UTMIST/Mist/internal/scheduler" +) + +// tailBytes is how much output GET /v1/jobs/{id}/log returns. +const tailBytes = 64 << 10 + +// Log is the reply of GET /v1/jobs/{id}/log. The complete output is the +// file at OutputPath; the API never reads user paths itself. +type Log struct { + State string `json:"state"` + OutputPath string `json:"output_path"` + OutputIncomplete bool `json:"output_incomplete"` + Truncated bool `json:"truncated"` // only the last 64 KiB is shown + Text string `json:"text"` +} + +// jobLog handles GET /v1/jobs/{id}/log, and ?follow=1 over a WebSocket. +func (s *Server) jobLog(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + id, err := jobID(r) + if err != nil { + return err + } + j, err := s.viewableJob(r.Context(), p, id) + if err != nil { + return err + } + if s.Logs == nil { + return unavailable("job output is not available from this API server") + } + if r.URL.Query().Get("follow") == "1" { + return s.followLog(w, r, p, j) + } + out := Log{State: string(j.State), OutputPath: j.OutputPath, OutputIncomplete: j.OutputIncomplete} + if j.Node != nil { // never admitted: no output anywhere yet + rc, err := s.Logs.Logs(r.Context(), id, false) + switch { + case errors.Is(err, executor.ErrNotFound): + case err != nil: + return err + default: + defer rc.Close() + text, truncated, err := lastBytes(rc, tailBytes) + if err != nil { + return err + } + out.Text, out.Truncated = text, truncated + } + } + writeJSON(w, http.StatusOK, out) + return nil +} + +// lastBytes reads r to the end, keeping only the last n bytes. +func lastBytes(r io.Reader, n int) (string, bool, error) { + buf := make([]byte, 0, 2*n) + chunk := make([]byte, 32<<10) + truncated := false + for { + k, err := r.Read(chunk) + buf = append(buf, chunk[:k]...) + if len(buf) > n { + buf = append(buf[:0], buf[len(buf)-n:]...) + truncated = true + } + if errors.Is(err, io.EOF) { + return strings.ToValidUTF8(string(buf), "�"), truncated, nil + } + if err != nil { + return "", false, err + } + } +} + +// followLog streams output as binary WebSocket frames until the job's +// output ends, the client goes away or the caller loses access, which is +// re-checked every StreamRecheck. +func (s *Server) followLog(w http.ResponseWriter, r *http.Request, p identity.Principal, j *scheduler.JobInfo) error { + if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { + return invalid("follow", "follow=1 needs a WebSocket upgrade") + } + if !s.originAllowed(r.Header.Get("Origin")) { + return &Error{http.StatusForbidden, CodeForbidden, "origin not allowed", nil} + } + ws := websocket.Server{ + Handshake: func(*websocket.Config, *http.Request) error { return nil }, // origin checked above + Handler: func(conn *websocket.Conn) { + defer conn.Close() + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + go func() { // a client frame or a closed connection ends the stream + io.Copy(io.Discard, conn) + cancel() + }() + go func() { + t := time.NewTicker(s.StreamRecheck) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + if err := s.Authz.Authorize(ctx, p, identity.ActionViewJob, target(j)); err != nil { + cancel() + return + } + } + } + }() + if j.Node == nil { + return // not admitted: nothing to follow yet; the client retries + } + rc, err := s.Logs.Logs(ctx, j.ID, true) + if err != nil { + return + } + go func() { <-ctx.Done(); rc.Close() }() + buf := make([]byte, 32<<10) + for { + n, err := rc.Read(buf) + if n > 0 { + if websocket.Message.Send(conn, buf[:n]) != nil { + return + } + } + if err != nil { + return + } + } + }, + } + ws.ServeHTTP(w, r) + return nil +} + +// originAllowed: requests without an Origin header are not from a browser +// and are authenticated by other means; browsers must be allowlisted. +func (s *Server) originAllowed(origin string) bool { + if origin == "" { + return true + } + u, err := url.Parse(origin) + if err != nil { + return false + } + return slices.Contains(s.AllowedOrigins, u.Scheme+"://"+u.Host) +} diff --git a/internal/api/submit.go b/internal/api/submit.go new file mode 100644 index 0000000..a015322 --- /dev/null +++ b/internal/api/submit.go @@ -0,0 +1,439 @@ +package api + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "path" + "regexp" + "strings" + "time" + "unicode" + "unicode/utf8" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/identity" + "github.com/UTMIST/Mist/internal/jobentry" + "github.com/UTMIST/Mist/internal/scheduler" +) + +// Submission bounds (api.md#submitting-post-v1jobs). +const ( + maxBody = 4 << 20 + maxScript = 1 << 20 + maxEnvTotal = 64 << 10 + maxArgsTotal = 128 << 10 + maxNameRunes = 128 + maxPath = 4096 + maxKey = 200 + maxCPUs = 1 << 20 + teamsRoot = "/mist/teams/" + defaultOutput = "mist-%j.out" +) + +var ( + slugPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,62}$`) + envPattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`) +) + +// submit handles POST /v1/jobs. +func (s *Server) submit(w http.ResponseWriter, r *http.Request, p identity.Principal) error { + key := r.Header.Get("Idempotency-Key") + if key == "" || len(key) > maxKey || strings.IndexFunc(key, func(c rune) bool { return c < 0x21 || c > 0x7e }) >= 0 { + return invalid("Idempotency-Key", "an Idempotency-Key header of 1–200 visible ASCII characters is required") + } + in, err := decodeSubmit(w, r) + if err != nil { + return err + } + canon, _ := json.Marshal(in) + sum := sha256.Sum256(canon) + idem := scheduler.Idempotency{Key: key, RequestHash: hex.EncodeToString(sum[:])} + + // A retry returns the original job, even if defaults have changed or + // the request would no longer be valid. + if id, hash, found, err := s.Jobs.LookupSubmission(r.Context(), p.Username, key); err != nil { + return err + } else if found { + if hash != idem.RequestHash { + return conflictKey() + } + return s.replySubmitted(w, r, p, id, true) + } + + cat, err := s.Jobs.Catalog(r.Context()) + if err != nil { + return err + } + req, err := s.normalize(r.Context(), p, in, cat) + if err != nil { + return err + } + id, replayed, err := s.Jobs.SubmitOnce(r.Context(), req, idem) + switch { + case errors.Is(err, scheduler.ErrIdempotencyConflict): + return conflictKey() + case errors.Is(err, scheduler.ErrLimitExceeded): + return limitExceeded("", strings.TrimPrefix(err.Error(), scheduler.ErrLimitExceeded.Error()+": ")) + case err != nil: + return err + } + return s.replySubmitted(w, r, p, id, replayed) +} + +func conflictKey() error { + return &Error{http.StatusConflict, CodeConflict, + "this Idempotency-Key was already used for a different request; use a new key to submit another job", nil} +} + +func (s *Server) replySubmitted(w http.ResponseWriter, r *http.Request, p identity.Principal, id int64, replayed bool) error { + j, err := s.viewableJob(r.Context(), p, id) + if err != nil { + return err + } + out, err := s.jobObject(r.Context(), p, j) + if err != nil { + return err + } + status := http.StatusCreated + if replayed { + w.Header().Set("Idempotent-Replayed", "true") + status = http.StatusOK + } + writeJSON(w, status, out) + return nil +} + +func decodeSubmit(w http.ResponseWriter, r *http.Request) (SubmitRequest, error) { + var in SubmitRequest + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxBody)) + if err != nil { + var tooBig *http.MaxBytesError + if errors.As(err, &tooBig) { + return in, invalid("", fmt.Sprintf("request body is larger than %d bytes", maxBody)) + } + return in, err + } + dec := json.NewDecoder(bytes.NewReader(body)) + dec.DisallowUnknownFields() + if err := dec.Decode(&in); err != nil { + return in, invalid("", "request body: "+err.Error()) + } + if dec.More() { + return in, invalid("", "request body has trailing data") + } + return in, nil +} + +// normalize validates a submission and fills every default from the +// catalog, giving the complete request the scheduler stores. Checks that +// depend on current state (membership, capacity) are repeated at +// admission; these ones make bad requests fail early and clearly. +func (s *Server) normalize(ctx context.Context, p identity.Principal, in SubmitRequest, cat *scheduler.Catalog) (scheduler.Request, error) { + var req scheduler.Request + + // What runs. + switch executor.Kind(in.Kind) { + case executor.KindBatch: + if in.Script == "" { + return req, invalid("script", "a batch job needs a script") + } + if len(in.Argv) > 0 { + return req, invalid("argv", "a batch job takes script and args, not argv") + } + case executor.KindCommand: + if len(in.Argv) == 0 || in.Argv[0] == "" { + return req, invalid("argv", "a command job needs argv, starting with the program to run") + } + if in.Script != "" || len(in.Args) > 0 { + return req, invalid("script", "a command job takes argv, not script or args") + } + case executor.KindInteractive: + return req, invalid("kind", "interactive jobs are not available yet (milestone M2)") + default: + return req, invalid("kind", `kind must be "batch" or "command"`) + } + if len(in.Script) > maxScript { + return req, invalid("script", fmt.Sprintf("script is larger than %d bytes", maxScript)) + } + if strings.ContainsRune(in.Script, 0) { + return req, invalid("script", "script contains a NUL byte") + } + total := 0 + for _, a := range append(append([]string{}, in.Args...), in.Argv...) { + if strings.ContainsRune(a, 0) { + return req, invalid("args", "arguments cannot contain NUL bytes") + } + total += len(a) + 1 + } + if total > maxArgsTotal { + return req, invalid("args", fmt.Sprintf("arguments total more than %d bytes", maxArgsTotal)) + } + + // Environment. Protected names are rejected, not silently dropped. + total = 0 + for k, v := range in.Env { + if !envPattern.MatchString(k) { + return req, invalid("env", fmt.Sprintf("%q is not a valid environment variable name", k)) + } + if jobentry.ProtectedEnv(k) { + return req, invalid("env", k+" is set by Mist and cannot be overridden") + } + if strings.ContainsRune(v, 0) { + return req, invalid("env", k+" contains a NUL byte") + } + total += len(k) + len(v) + 2 + } + if total > maxEnvTotal { + return req, invalid("env", fmt.Sprintf("environment totals more than %d bytes", maxEnvTotal)) + } + + // Name. + name := in.Name + if name == "" { + name = string(executor.KindBatch) + if len(in.Argv) > 0 { + name = path.Base(in.Argv[0]) + } + } + if !utf8.ValidString(name) || utf8.RuneCountInString(name) > maxNameRunes || + strings.IndexFunc(name, unicode.IsControl) >= 0 { + return req, invalid("name", fmt.Sprintf("name must be at most %d printable characters", maxNameRunes)) + } + + // Charged team, working directory and output path. + if in.Workdir != "" { + if err := checkPath("workdir", in.Workdir); err != nil { + return req, err + } + } + team := in.Team + if team == "" { + rest, ok := strings.CutPrefix(in.Workdir, teamsRoot) + if !ok { + if in.Workdir == "" { + return req, invalid("team", "give a team, or a workdir inside a team folder") + } + return req, invalid("workdir", "workdir must be inside a team folder (/mist/teams//...); pass a team to use your personal folder") + } + team, _, _ = strings.Cut(rest, "/") + } + if !slugPattern.MatchString(team) { + return req, invalid("team", fmt.Sprintf("%q is not a team name", team)) + } + if err := s.Authz.Authorize(ctx, p, identity.ActionSubmitJob, identity.Target{Team: team}); err != nil { + if errors.Is(err, identity.ErrForbidden) || errors.Is(err, identity.ErrRevoked) { + return req, &Error{http.StatusForbidden, CodeForbidden, + "you are not a current member of team " + team, map[string]any{"field": "team"}} + } + return req, authzError(err, false, "") + } + teamRoot := teamsRoot + team + workdir := in.Workdir + if workdir == "" { + workdir = teamRoot + "/" + p.Username + } + if workdir != teamRoot && !strings.HasPrefix(workdir, teamRoot+"/") { + return req, invalid("workdir", "workdir must be inside "+teamRoot+", the folder of the charged team") + } + output := in.Output + if output == "" { + output = defaultOutput + } + if !path.IsAbs(output) { + output = workdir + "/" + output + } + if err := checkPath("output", output); err != nil { + return req, err + } + if !strings.HasPrefix(output, teamRoot+"/") { + return req, invalid("output", "output must be inside "+teamRoot) + } + + // Partition and QoS. + pname := in.Partition + if pname == "" { + for _, c := range cat.Partitions { + if c.IsDefault { + pname = c.Name + } + } + } + part := cat.Partition(pname) + if part == nil { + return req, invalid("partition", fmt.Sprintf("no partition %q", pname)) + } + qname := in.QOS + if qname == "" { + qname = "normal" + } + qos := cat.QoS(qname) + if qos == nil { + return req, invalid("qos", fmt.Sprintf("no QoS %q", qname)) + } + if qname == "interactive" { + return req, invalid("qos", "the interactive QoS is only for interactive jobs") + } + if in.Runtime != "" && in.Runtime != part.Runtime { + return req, invalid("runtime", "partition "+part.Name+" offers only its current runtime "+part.Runtime) + } + + res, err := resolveResources(in, part) + if err != nil { + return req, err + } + timeLimit := part.DefaultTimeLimit + if in.TimeLimitSeconds != nil { + if *in.TimeLimitSeconds < 1 { + return req, invalid("time_limit_seconds", "time limit must be at least 1 second") + } + timeLimit = time.Duration(min(*in.TimeLimitSeconds, 1<<32)) * time.Second + } + maxTime := part.MaxTimeLimit + if qos.MaxTimeLimit != nil && *qos.MaxTimeLimit < maxTime { + maxTime = *qos.MaxTimeLimit + } + if timeLimit > maxTime { + return req, limitExceeded("time_limit_seconds", + fmt.Sprintf("time limit is more than the %s maximum of %s", part.Name, maxTime)) + } + grace := part.DefaultSignalGrace + if in.SignalGraceSeconds != nil { + g := *in.SignalGraceSeconds + if g < int64(executor.MinSignalGrace/time.Second) || g > int64(executor.MaxSignalGrace/time.Second) { + return req, invalid("signal_grace_seconds", fmt.Sprintf("signal grace must be %d–%d seconds", + executor.MinSignalGrace/time.Second, executor.MaxSignalGrace/time.Second)) + } + grace = time.Duration(g) * time.Second + } + + return scheduler.Request{ + Username: p.Username, Team: team, Partition: part.Name, QOS: qos.Name, Name: name, + Kind: executor.Kind(in.Kind), Resources: res, TimeLimit: timeLimit, SignalGrace: grace, + Workdir: workdir, OutputPath: output, + Script: in.Script, Args: in.Args, Argv: in.Argv, Env: in.Env, + }, nil +} + +// checkPath accepts only clean absolute paths: no "..", ".", empty or +// trailing components. The executor still opens paths beneath the team +// root without following symlinks out of it; this check only keeps +// obviously hostile or ambiguous strings out of the database. +func checkPath(field, p string) error { + if len(p) > maxPath || !utf8.ValidString(p) || strings.IndexFunc(p, unicode.IsControl) >= 0 { + return invalid(field, field+" must be a printable path of at most 4096 bytes") + } + if !path.IsAbs(p) || path.Clean(p) != p { + return invalid(field, field+" must be a clean absolute path, without .. or . components") + } + return nil +} + +func resolveResources(in SubmitRequest, p *scheduler.Partition) (executor.Resources, error) { + var res executor.Resources + get := func(v *int64, def int64) int64 { + if v == nil { + return def + } + return *v + } + res.GPUs, res.TTBoards = get(in.Resources.GPUs, 0), get(in.Resources.TTBoards, 0) + if res.GPUs < 0 || res.TTBoards < 0 { + return res, invalid("resources", "accelerator counts cannot be negative") + } + n := int64(1) // CPU and memory defaults scale with accelerators + switch p.Accelerator { + case "gpu": + res.GPUs = get(in.Resources.GPUs, p.MinAccelerators) + if res.TTBoards != 0 { + return res, invalid("resources.tt_boards", "partition "+p.Name+" has no Tenstorrent boards") + } + n = res.GPUs + case "tt": + res.TTBoards = get(in.Resources.TTBoards, p.MinAccelerators) + if res.GPUs != 0 { + return res, invalid("resources.gpus", "partition "+p.Name+" has no GPUs") + } + n = res.TTBoards + default: + if res.GPUs != 0 || res.TTBoards != 0 { + return res, invalid("resources", "partition "+p.Name+" has no accelerators; use the gpu partition for GPUs") + } + } + if p.Accelerator != "none" { + if n < p.MinAccelerators { + return res, invalid("resources", fmt.Sprintf("partition %s needs at least %d %s", p.Name, p.MinAccelerators, accelName(p))) + } + if n > p.MaxAccelerators { + return res, limitExceeded("resources", fmt.Sprintf("partition %s allows at most %d %s per job", p.Name, p.MaxAccelerators, accelName(p))) + } + } + per := "" + if p.Accelerator != "none" { + per = " per " + strings.TrimSuffix(accelName(p), "s") + } + + res.MilliCPU = p.DefaultMilliCPU * n + if in.Resources.CPUs != nil { + c := *in.Resources.CPUs + if c < 1 || c > maxCPUs { + return res, invalid("resources.cpus", "cpus must be a whole number of cores, at least 1") + } + res.MilliCPU = c * 1000 + } + if res.MilliCPU > p.MaxMilliCPU*n { + return res, limitExceeded("resources.cpus", fmt.Sprintf("partition %s allows at most %d CPUs%s", p.Name, p.MaxMilliCPU/1000, per)) + } + res.MemoryBytes = get(in.Resources.MemoryBytes, p.DefaultMemory*n) + if res.MemoryBytes < 1 { + return res, invalid("resources.memory_bytes", "memory must be positive") + } + if res.MemoryBytes > p.MaxMemory*n { + return res, limitExceeded("resources.memory_bytes", fmt.Sprintf("partition %s allows at most %s of memory%s", p.Name, gib(p.MaxMemory), per)) + } + res.ScratchBytes = get(in.Resources.ScratchBytes, p.DefaultScratch) + if res.ScratchBytes < 1 { + return res, invalid("resources.scratch_bytes", "scratch must be positive") + } + if res.ScratchBytes > p.MaxScratch { + return res, limitExceeded("resources.scratch_bytes", fmt.Sprintf("partition %s allows at most %s of scratch", p.Name, gib(p.MaxScratch))) + } + if (p.CapMilliCPU != nil && res.MilliCPU > *p.CapMilliCPU) || (p.CapMemory != nil && res.MemoryBytes > *p.CapMemory) { + return res, limitExceeded("resources", "the job is larger than partition "+p.Name+"'s total capacity") + } + + // It must fit at least one node, or it would wait forever. + reported := false + for _, node := range p.Nodes { + if node.Reported { + reported = true + if node.Total.Fits(res) { + return res, nil + } + } + } + if !reported { + return res, unavailable("partition " + p.Name + " has no nodes reported yet") + } + return res, limitExceeded("resources", "no node in partition "+p.Name+" is large enough for this job") +} + +func accelName(p *scheduler.Partition) string { + if p.Accelerator == "tt" { + return "boards" + } + return "GPUs" +} + +func gib(b int64) string { + if b%(1<<30) == 0 { + return fmt.Sprintf("%d GiB", b>>30) + } + return fmt.Sprintf("%d bytes", b) +} diff --git a/internal/api/wire.go b/internal/api/wire.go new file mode 100644 index 0000000..2633f5e --- /dev/null +++ b/internal/api/wire.go @@ -0,0 +1,207 @@ +package api + +import ( + "time" + + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/scheduler" +) + +// Wire types (api.md#job-wire-types). Times are RFC 3339 UTC, durations +// integer seconds, sizes integer bytes. + +// SubmitResources is the requested size. Omitted fields take partition +// defaults. cpus is a whole-core count. +type SubmitResources struct { + CPUs *int64 `json:"cpus,omitempty"` + MemoryBytes *int64 `json:"memory_bytes,omitempty"` + GPUs *int64 `json:"gpus,omitempty"` + TTBoards *int64 `json:"tt_boards,omitempty"` + ScratchBytes *int64 `json:"scratch_bytes,omitempty"` +} + +// SubmitRequest is the body of POST /v1/jobs. +type SubmitRequest struct { + Team string `json:"team,omitempty"` + Partition string `json:"partition,omitempty"` + QOS string `json:"qos,omitempty"` + Name string `json:"name,omitempty"` + Kind string `json:"kind"` + Script string `json:"script,omitempty"` + Args []string `json:"args,omitempty"` + Argv []string `json:"argv,omitempty"` + Workdir string `json:"workdir,omitempty"` + Output string `json:"output,omitempty"` + Env map[string]string `json:"env,omitempty"` + Resources SubmitResources `json:"resources"` + TimeLimitSeconds *int64 `json:"time_limit_seconds,omitempty"` + SignalGraceSeconds *int64 `json:"signal_grace_seconds,omitempty"` + Runtime string `json:"runtime,omitempty"` +} + +// Resources is a job's resolved size. +type Resources struct { + CPUs int64 `json:"cpus"` + MemoryBytes int64 `json:"memory_bytes"` + GPUs int64 `json:"gpus"` + TTBoards int64 `json:"tt_boards"` + ScratchBytes int64 `json:"scratch_bytes"` +} + +func wireResources(r executor.Resources) Resources { + return Resources{CPUs: r.MilliCPU / 1000, MemoryBytes: r.MemoryBytes, GPUs: r.GPUs, + TTBoards: r.TTBoards, ScratchBytes: r.ScratchBytes} +} + +// Job is the job object, for callers allowed to view the job. +type Job struct { + ID int64 `json:"id"` + Name string `json:"name"` + User string `json:"user"` + Team string `json:"team"` + Partition string `json:"partition"` + QOS string `json:"qos"` + Kind string `json:"kind"` + State string `json:"state"` + PendingReason *string `json:"pending_reason"` + PendingDetail *string `json:"pending_detail"` + ExitReason *string `json:"exit_reason"` + ExitCode *int `json:"exit_code"` + ExitDetail *string `json:"exit_detail"` + Resources Resources `json:"resources"` + TimeLimitSeconds int64 `json:"time_limit_seconds"` + SignalGraceSeconds int64 `json:"signal_grace_seconds"` + Node *string `json:"node"` + Workdir string `json:"workdir"` + OutputPath string `json:"output_path"` + OutputIncomplete bool `json:"output_incomplete"` + SubmittedAt time.Time `json:"submitted_at"` + StartedAt *time.Time `json:"started_at"` + EndedAt *time.Time `json:"ended_at"` + CleanupPending bool `json:"cleanup_pending"` + CleanupDetail *string `json:"cleanup_detail"` + AuthorizationGeneration *int64 `json:"authorization_generation"` + CanCancel bool `json:"can_cancel"` + CanAttach bool `json:"can_attach"` + + // GET /v1/jobs/{id} only. + Events []Event `json:"events,omitempty"` + Script *string `json:"script,omitempty"` // ?script=1 + Args []string `json:"args,omitempty"` + Argv []string `json:"argv,omitempty"` + Env map[string]string `json:"env,omitempty"` +} + +// Event is one state change in a job's history. +type Event struct { + At time.Time `json:"at"` + From *string `json:"from"` + To string `json:"to"` + Reason string `json:"reason"` + Actor string `json:"actor"` +} + +func nonEmpty(s string) *string { + if s == "" { + return nil + } + return &s +} + +func utc(t *time.Time) *time.Time { + if t == nil { + return nil + } + u := t.UTC() + return &u +} + +func wireJob(j *scheduler.JobInfo, canCancel bool) Job { + out := Job{ + ID: j.ID, Name: j.Name, User: j.Username, Team: j.Team, Partition: j.Partition, QOS: j.QOS, + Kind: string(j.Kind), State: string(j.State), PendingReason: j.PendingReason, + ExitReason: j.ExitReason, ExitCode: j.ExitCode, ExitDetail: nonEmpty(j.ExitDetail), + Resources: wireResources(j.Resources), + TimeLimitSeconds: int64(j.TimeLimit / time.Second), + SignalGraceSeconds: int64(j.SignalGrace / time.Second), + Node: j.Node, Workdir: j.Workdir, OutputPath: j.OutputPath, + OutputIncomplete: j.OutputIncomplete, SubmittedAt: j.SubmittedAt.UTC(), + StartedAt: utc(j.StartedAt), EndedAt: utc(j.EndedAt), + CleanupPending: j.CleanupPending, CleanupDetail: nonEmpty(j.CleanupDetail), + AuthorizationGeneration: j.AuthGeneration, + CanCancel: canCancel, + } + if j.State == scheduler.StatePending { + out.PendingDetail = nonEmpty(j.PendingDetail) + } + return out +} + +// QueueEntry is the allowlisted summary of a queued or running job. +// User, team and name are filled only for jobs the caller may view; +// workdir and output path are never shown here. +type QueueEntry struct { + ID int64 `json:"id"` + Partition string `json:"partition"` + QOS string `json:"qos"` + State string `json:"state"` + Resources Resources `json:"resources"` + AgeSeconds int64 `json:"age_seconds"` + PendingReason *string `json:"pending_reason"` + User *string `json:"user"` + Team *string `json:"team"` + Name *string `json:"name"` + Workdir *string `json:"workdir"` + OutputPath *string `json:"output_path"` +} + +// Partition is one entry of GET /v1/partitions. CPU and memory defaults +// and maximums are per accelerator when PerAccelerator is set. +type Partition struct { + Name string `json:"name"` + State string `json:"state"` + Default bool `json:"default"` + Accelerator string `json:"accelerator"` + MinAccelerators int64 `json:"min_accelerators"` + MaxAccelerators int64 `json:"max_accelerators"` + PerAccelerator bool `json:"per_accelerator"` + DefaultCPUs int64 `json:"default_cpus"` + MaxCPUs int64 `json:"max_cpus"` + DefaultMemoryBytes int64 `json:"default_memory_bytes"` + MaxMemoryBytes int64 `json:"max_memory_bytes"` + DefaultScratchBytes int64 `json:"default_scratch_bytes"` + MaxScratchBytes int64 `json:"max_scratch_bytes"` + DefaultTimeLimitSeconds int64 `json:"default_time_limit_seconds"` + MaxTimeLimitSeconds int64 `json:"max_time_limit_seconds"` + DefaultSignalGraceSeconds int64 `json:"default_signal_grace_seconds"` + CapCPUs *int64 `json:"cap_cpus"` + CapMemoryBytes *int64 `json:"cap_memory_bytes"` + Free Resources `json:"free"` // on schedulable nodes + Nodes []Node `json:"nodes"` + Limits []Limit `json:"limits"` +} + +// Node is one node of a partition. +type Node struct { + Name string `json:"name"` + State string `json:"state"` // up, drain, down + Healthy bool `json:"healthy"` + Quarantined bool `json:"quarantined"` + Total Resources `json:"total"` + Free Resources `json:"free"` +} + +// Limit is one limit that applies to the caller in a partition. +type Limit struct { + Scope string `json:"scope"` + Subject string `json:"subject"` + QOS string `json:"qos"` + Resource string `json:"resource"` + Max int64 `json:"max"` +} + +// List is the shape of every list reply. +type List[T any] struct { + Items []T `json:"items"` + NextCursor *string `json:"next_cursor"` +} diff --git a/internal/db/migrations/scheduler/0003_job_api.sql b/internal/db/migrations/scheduler/0003_job_api.sql new file mode 100644 index 0000000..66e8531 --- /dev/null +++ b/internal/db/migrations/scheduler/0003_job_api.sql @@ -0,0 +1,21 @@ +-- Job API (M1 step 6). See docs/agents/architecture/api.md. + +-- Submission idempotency. The hash is of the client's request as sent, not +-- the normalized spec, so a retry gets the original job even if defaults +-- changed in between. Kept as long as the job row. +CREATE TABLE submissions ( + username text NOT NULL, + idempotency_key text NOT NULL, + request_hash text NOT NULL, + job_id bigint NOT NULL REFERENCES jobs, + created_at timestamptz NOT NULL, + PRIMARY KEY (username, idempotency_key) +); + +-- Defaults the API fills in. Defaults live only here, never in clients. +ALTER TABLE partitions + ADD COLUMN is_default boolean NOT NULL DEFAULT false, + ADD COLUMN default_signal_grace_seconds bigint NOT NULL DEFAULT 60 + CHECK (default_signal_grace_seconds BETWEEN 1 AND 120); +CREATE UNIQUE INDEX partitions_one_default ON partitions (is_default) WHERE is_default; +UPDATE partitions SET is_default = true WHERE name = 'gpu'; diff --git a/internal/e2e/k3d_test.go b/internal/e2e/k3d_test.go new file mode 100644 index 0000000..228df8e --- /dev/null +++ b/internal/e2e/k3d_test.go @@ -0,0 +1,252 @@ +//go:build k3d + +// Package e2e runs the M1 path through the public API on a k3d cluster: +// run scripts/k3d-test.sh. HTTP submit with an idempotency key, scheduler +// admission, the K3s executor, the job's log through the API, and the +// node's cleanup reporter posting its receipt to /internal with a host +// token. Identity is the identitytest fixture. +package e2e + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/tools/clientcmd" + + "github.com/UTMIST/Mist/internal/api" + "github.com/UTMIST/Mist/internal/clock" + "github.com/UTMIST/Mist/internal/db" + "github.com/UTMIST/Mist/internal/db/dbtest" + "github.com/UTMIST/Mist/internal/executor" + "github.com/UTMIST/Mist/internal/executor/k3s" + "github.com/UTMIST/Mist/internal/identity" + "github.com/UTMIST/Mist/internal/identity/identitytest" + "github.com/UTMIST/Mist/internal/scheduler" +) + +func TestMain(m *testing.M) { dbtest.Main(m) } + +type user struct{} + +func (user) Authenticate(r *http.Request) (identity.Principal, error) { + if u := r.Header.Get("X-Test-User"); u != "" { + return identity.Principal{Kind: identity.PrincipalAccount, Username: u}, nil + } + return identity.Principal{}, api.ErrNoCredentials +} + +func TestK3dAPIEndToEnd(t *testing.T) { + kctx, node := os.Getenv("MIST_K3D_CONTEXT"), os.Getenv("MIST_K3D_NODE") + if kctx == "" || node == "" { + t.Skip("run through scripts/k3d-test.sh") + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + cfg, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( + clientcmd.NewDefaultClientConfigLoadingRules(), + &clientcmd.ConfigOverrides{CurrentContext: kctx}).ClientConfig() + if err != nil { + t.Fatal(err) + } + cfg.Impersonate.UserName = "system:serviceaccount:mist-system:mistd" + client, err := kubernetes.NewForConfig(cfg) + if err != nil { + t.Fatal(err) + } + pool, err := db.Open(ctx, dbtest.URL(t)) + if err != nil { + t.Fatal(err) + } + defer pool.Close() + if err := db.Migrate(ctx, pool); err != nil { + t.Fatal(err) + } + for _, q := range []string{ + "INSERT INTO nodes (name, executor) VALUES ($1, 'k3s')", + "INSERT INTO partition_nodes VALUES ('cpu', $1)", + "UPDATE partitions SET runtime = 'mist-test-runtime:dev'", + "SELECT setval('jobs_id_seq', (extract(epoch FROM now())::bigint % 1000000) * 100 + 90)", + } { + args := []any{node} + if !strings.Contains(q, "$1") { + args = nil + } + if _, err := pool.Exec(ctx, q, args...); err != nil { + t.Fatal(err) + } + } + + ex := k3s.New(k3s.Options{ + Config: k3s.Config{ + Namespace: "mist-jobs", EntryPath: "/usr/local/bin/mist-job-entry", + LaunchTimeout: 10 * time.Minute, BackstopSlack: time.Minute, Nameservers: []string{"1.1.1.1"}, + Volumes: map[string]k3s.Volume{"vol1": {HostNode: node, HostPath: "/mist"}}, + }, + Client: client, Pool: pool, Resync: 5 * time.Second, + }) + if err := ex.Start(ctx); err != nil { + t.Fatal(err) + } + fixture := identitytest.New() + sched, err := scheduler.New(ctx, scheduler.Config{ + Pool: pool, Identity: fixture, Executors: map[string]executor.Executor{"k3s": ex}, + }, nil) + if err != nil { + t.Fatal(err) + } + if err := sched.Cycle(ctx); err != nil { // reports the node, so submit can check fit + t.Fatal(err) + } + + tokens := filepath.Join(t.TempDir(), "tokens") + os.WriteFile(tokens, []byte(node+" cleanup "+api.HashToken("node-secret")+"\n"), 0o600) + hosts, err := api.LoadHostTokens(tokens, "cleanup") + if err != nil { + t.Fatal(err) + } + // The API's own executor, like `mistd api`: not started, logs and + // cleanup only. The scheduler learns of the cleanup through List. + apiEx := k3s.New(k3s.Options{Config: k3s.Config{Namespace: "mist-jobs"}, Client: client, Pool: pool}) + go func() { + for range apiEx.Events() { + } + }() + srv := httptest.NewUnstartedServer((&api.Server{ + Jobs: scheduler.Jobs{Pool: pool, Clock: clock.Real{}}, Authz: fixture, + Accounts: []api.Authenticator{user{}}, Hosts: hosts, Logs: apiEx, Cleanup: apiEx, + }).Handler()) + l, err := net.Listen("tcp", "0.0.0.0:0") // reachable from the node as host.k3d.internal + if err != nil { + t.Fatal(err) + } + srv.Listener = l + srv.Start() + defer srv.Close() + port := l.Addr().(*net.TCPAddr).Port + + call := func(method, path, who string, body any, into any) int { + t.Helper() + var rd io.Reader + if body != nil { + b, _ := json.Marshal(body) + rd = bytes.NewReader(b) + } + req, _ := http.NewRequestWithContext(ctx, method, srv.URL+path, rd) + req.Header.Set("X-Test-User", who) + req.Header.Set("Idempotency-Key", "e2e-1") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + if into != nil { + json.Unmarshal(b, into) + } + if resp.StatusCode >= 300 { + t.Logf("%s %s: %d %s", method, path, resp.StatusCode, b) + } + return resp.StatusCode + } + + submit := map[string]any{ + "partition": "cpu", "kind": "batch", "workdir": "/mist/teams/team-a/alice", + "script": "#!/bin/bash\necho \"hello from $MIST_JOB_ID\"\n", + "resources": map[string]any{"cpus": 1, "memory_bytes": 256 << 20, "scratch_bytes": 1 << 30}, + "time_limit_seconds": 300, "signal_grace_seconds": 5, + } + var job api.Job + if code := call("POST", "/v1/jobs", "alice", submit, &job); code != http.StatusCreated { + t.Fatalf("submit: %d", code) + } + var again api.Job + if code := call("POST", "/v1/jobs", "alice", submit, &again); code != http.StatusOK || again.ID != job.ID { + t.Fatalf("retry made job %d (%d)", again.ID, code) + } + jobPath := fmt.Sprintf("/v1/jobs/%d", job.ID) + + handled, cleaned := 0, false + events := ex.Events() + var pending []executor.Event + deadline := time.Now().Add(3 * time.Minute) + for time.Now().Before(deadline) { + if err := sched.Cycle(ctx); err != nil { + t.Fatal(err) + } + drain: + for { + select { + case ev := <-events: + pending = append(pending, ev) + default: + break drain + } + } + for ; handled < len(pending); handled++ { + if err := sched.HandleEvent(ctx, pending[handled]); err != nil { + t.Fatal(err) + } + } + call("GET", jobPath, "bob", nil, &job) // a teammate watches it + if job.State == "CLEANING" && !cleaned { + if code := call("GET", jobPath, "carol", nil, nil); code != http.StatusNotFound { + t.Fatalf("other team sees the job: %d", code) + } + var execID, uid string + pool.QueryRow(ctx, "SELECT execution_id::text FROM jobs WHERE id = $1", job.ID).Scan(&execID) + pool.QueryRow(ctx, "SELECT object_uid FROM k3s_executions WHERE execution_id = $1", execID).Scan(&uid) + cleaned = postCleanup(t, node, execID, uid, port) + } + if job.State == "COMPLETED" { + var released *time.Time + pool.QueryRow(ctx, "SELECT released_at FROM allocations WHERE job_id = $1", job.ID).Scan(&released) + if !cleaned || released == nil { + t.Fatalf("COMPLETED with cleaned=%v released=%v", cleaned, released) + } + var log api.Log + call("GET", jobPath+"/log", "alice", nil, &log) + if !strings.Contains(log.Text, fmt.Sprintf("hello from %d", job.ID)) { + t.Fatalf("log through the API: %+v", log) + } + return + } + time.Sleep(500 * time.Millisecond) + } + t.Fatalf("job ended in %s", job.State) +} + +// postCleanup runs the reporter on the node, which posts its receipt to +// the API itself, retrying while kubelet is still tearing the Pod down. +func postCleanup(t *testing.T, node, execID, uid string, port int) bool { + t.Helper() + if err := exec.Command("docker", "exec", node, "sh", "-c", + "printf node-secret > /tmp/cleanup-token").Run(); err != nil { + t.Fatal(err) + } + var last []byte + for range 60 { + out, err := exec.Command("docker", "exec", node, "/mistd", "cleanup-report", + "-pod-uid", uid, "-execution-id", execID, "-node", node, + "-post", fmt.Sprintf("http://host.k3d.internal:%d", port), "-token-file", "/tmp/cleanup-token").CombinedOutput() + if err == nil { + return true + } + last = out + time.Sleep(time.Second) + } + t.Fatalf("cleanup receipt never accepted: %s", last) + return false +} diff --git a/internal/executor/k3s/executor.go b/internal/executor/k3s/executor.go index d3ee588..7ec7bbf 100644 --- a/internal/executor/k3s/executor.go +++ b/internal/executor/k3s/executor.go @@ -588,6 +588,10 @@ type CleanupReceipt struct { Evidence string `json:"evidence"` } +// ErrCleanupRejected: a cleanup receipt does not confirm cleanup of a +// recorded execution (wrong UID or node, no outcome yet, Pod still there). +var ErrCleanupRejected = errors.New("k3s: cleanup receipt rejected") + // RecordCleanup accepts a node cleanup receipt. Only a receipt matching // the recorded execution, Pod UID and node, for a Pod the API no longer // has, confirms cleanup and lets the scheduler release resources. @@ -603,15 +607,15 @@ func (e *Executor) RecordCleanup(ctx context.Context, r CleanupReceipt) error { } } if rec == nil || rec.ObjectUID == "" || rec.ObjectUID != r.PodUID || rec.Node != r.Node { - return fmt.Errorf("k3s: cleanup receipt does not match a recorded execution") + return fmt.Errorf("%w: it does not match a recorded execution", ErrCleanupRejected) } defer e.lock(rec.JobID)() if rec.FinishedAt == nil && rec.LostAt == nil { - return fmt.Errorf("k3s: execution %s has no recorded outcome yet", rec.ExecutionID) + return fmt.Errorf("%w: execution %s has no recorded outcome yet", ErrCleanupRejected, rec.ExecutionID) } live, err := e.client.CoreV1().Pods(e.opt.Namespace).Get(ctx, Name(rec.JobID), metav1.GetOptions{}) if err == nil && string(live.UID) == rec.ObjectUID { - return fmt.Errorf("k3s: Pod %s still exists", live.Name) + return fmt.Errorf("%w: Pod %s still exists", ErrCleanupRejected, live.Name) } else if err != nil && !apierrors.IsNotFound(err) { return err } diff --git a/internal/jobentry/env.go b/internal/jobentry/env.go index 25dd9ce..8c406cf 100644 --- a/internal/jobentry/env.go +++ b/internal/jobentry/env.go @@ -11,8 +11,9 @@ import ( // runtime profile's baseline and what the executor/device runtime sets. var passthrough = []string{"PATH", "LANG", "LC_ALL", "TERM", "CUDA_VISIBLE_DEVICES", "NVIDIA_VISIBLE_DEVICES"} -// protected reports whether a user --env value may not set name. -func protected(name string) bool { +// ProtectedEnv reports whether a user --env value may not set name. The API +// rejects these at submission; the wrapper drops them again here. +func ProtectedEnv(name string) bool { switch name { case "HOME", "USER", "LOGNAME", "PWD", "SHELL", "TMPDIR", "CUDA_VISIBLE_DEVICES", "NVIDIA_VISIBLE_DEVICES", "PATH": @@ -49,7 +50,7 @@ func childEnv(s Spec, own []string, started time.Time) []string { env[k] = v } for k, v := range s.Env { - if !protected(k) && k != "" && !strings.ContainsAny(k, "=\x00") && !strings.Contains(v, "\x00") { + if !ProtectedEnv(k) && k != "" && !strings.ContainsAny(k, "=\x00") && !strings.Contains(v, "\x00") { env[k] = v } } diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 3bfc334..6d66ebe 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -2,6 +2,7 @@ package scheduler import ( "errors" + "fmt" "testing" "time" @@ -462,3 +463,36 @@ func TestSubmitQueueLimit(t *testing.T) { t.Fatalf("want ErrLimitExceeded, got %v", err) } } + +// The idempotency check inside the submit transaction, which decides +// races the API's earlier lookup cannot. +func TestSubmitOnce(t *testing.T) { + h := newHarness(t) + r := Request{ + Username: "alice", Team: "team-a", Partition: "cpu", QOS: "normal", Name: "x", Kind: executor.KindBatch, + Resources: executor.Resources{MilliCPU: 1000, MemoryBytes: 1 << 30, ScratchBytes: 1 << 30}, + TimeLimit: time.Hour, SignalGrace: time.Minute, Script: "true", + Workdir: "/mist/teams/team-a/alice", OutputPath: "/mist/teams/team-a/alice/mist-%j.out", + } + idem := Idempotency{Key: "k", RequestHash: "h1"} + id, replayed, err := h.jobs.SubmitOnce(h.ctx, r, idem) + if err != nil || replayed { + t.Fatal(err, replayed) + } + again, replayed, err := h.jobs.SubmitOnce(h.ctx, r, idem) + if err != nil || !replayed || again != id { + t.Fatalf("retry: id %d replayed %v err %v", again, replayed, err) + } + if _, _, err := h.jobs.SubmitOnce(h.ctx, r, Idempotency{Key: "k", RequestHash: "h2"}); !errors.Is(err, ErrIdempotencyConflict) { + t.Fatalf("different request under the same key: %v", err) + } + spec, err := h.jobs.Spec(h.ctx, id) + if err != nil || spec.OutputPath != fmt.Sprintf("/mist/teams/team-a/alice/mist-%d.out", id) { + t.Fatalf("output %%j: %q %v", spec.OutputPath, err) + } + var n int + h.pool.QueryRow(h.ctx, "SELECT count(*) FROM jobs").Scan(&n) + if n != 1 { + t.Fatalf("%d jobs", n) + } +} diff --git a/internal/scheduler/submit.go b/internal/scheduler/submit.go index 690fcbe..5b81113 100644 --- a/internal/scheduler/submit.go +++ b/internal/scheduler/submit.go @@ -5,6 +5,8 @@ import ( "encoding/json" "errors" "fmt" + "strconv" + "strings" "time" "github.com/jackc/pgx/v5" @@ -15,8 +17,9 @@ import ( ) // Request is a normalized submission. Filling defaults and validating -// paths, sizes and authorization happen in the API (M1 step 6); the +// paths, sizes and authorization happen in the API (internal/api); the // scheduler stores what it is given and freezes it at admission. +// OutputPath may contain %j, which Submit replaces with the job ID. type Request struct { Username string `json:"username"` Team string `json:"team"` @@ -41,6 +44,17 @@ var ErrLimitExceeded = errors.New("scheduler: limit exceeded") // ErrNotFound: no such job. var ErrNotFound = errors.New("scheduler: job not found") +// ErrIdempotencyConflict: the idempotency key was already used for a +// different request. +var ErrIdempotencyConflict = errors.New("scheduler: idempotency key reused with a different request") + +// Idempotency identifies one client submission: the key the client chose +// and a hash of the request exactly as the client sent it. +type Idempotency struct { + Key string + RequestHash string +} + // Jobs is the API-side handle: submit and cancel. It does not need // leadership; row locks serialize it with the scheduler. type Jobs struct { @@ -53,21 +67,53 @@ const notifyChannel = "mist_scheduler" // Submit inserts a PENDING job and returns its ID. func (js Jobs) Submit(ctx context.Context, r Request) (int64, error) { + id, _, err := js.SubmitOnce(ctx, r, Idempotency{}) + return id, err +} + +// LookupSubmission returns the job a previous submission with this key +// created, before the caller spends effort normalizing a retry. +func (js Jobs) LookupSubmission(ctx context.Context, username, key string) (jobID int64, requestHash string, found bool, err error) { + err = js.Pool.QueryRow(ctx, + "SELECT job_id, request_hash FROM submissions WHERE username = $1 AND idempotency_key = $2", + username, key).Scan(&jobID, &requestHash) + if errors.Is(err, pgx.ErrNoRows) { + return 0, "", false, nil + } + return jobID, requestHash, err == nil, err +} + +// SubmitOnce is Submit with an idempotency key. If the key was already +// used with the same request hash it returns the original job and +// replayed = true; with a different hash it returns ErrIdempotencyConflict. +// An empty key disables the check. +func (js Jobs) SubmitOnce(ctx context.Context, r Request, idem Idempotency) (id int64, replayed bool, err error) { if r.Resources.MilliCPU <= 0 || r.Resources.MemoryBytes <= 0 || r.Resources.ScratchBytes <= 0 || r.TimeLimit < time.Second || r.SignalGrace < executor.MinSignalGrace || r.SignalGrace > executor.MaxSignalGrace { - return 0, fmt.Errorf("scheduler: request is not normalized") - } - body, err := json.Marshal(r) - if err != nil { - return 0, err + return 0, false, fmt.Errorf("scheduler: request is not normalized") } now := js.Clock.Now() - var id int64 err = pgx.BeginFunc(ctx, js.Pool, func(tx pgx.Tx) error { - // Serialize one user's submissions so the queue-length check holds. + // Serialize one user's submissions so the queue-length check and + // the idempotency lookup hold. if _, err := tx.Exec(ctx, "SELECT pg_advisory_xact_lock(hashtext('submit:' || $1))", r.Username); err != nil { return err } + if idem.Key != "" { + var hash string + err := tx.QueryRow(ctx, + "SELECT job_id, request_hash FROM submissions WHERE username = $1 AND idempotency_key = $2", + r.Username, idem.Key).Scan(&id, &hash) + switch { + case err == nil && hash == idem.RequestHash: + replayed = true + return nil + case err == nil: + return ErrIdempotencyConflict + case !errors.Is(err, pgx.ErrNoRows): + return err + } + } var max *int64 err := tx.QueryRow(ctx, ` SELECT max FROM limits @@ -101,9 +147,21 @@ func (js Jobs) Submit(ctx context.Context, r Request) (int64, error) { int64(r.TimeLimit/time.Second), int64(r.SignalGrace/time.Second), now).Scan(&id); err != nil { return err } + r.OutputPath = strings.ReplaceAll(r.OutputPath, "%j", strconv.FormatInt(id, 10)) + body, err := json.Marshal(r) + if err != nil { + return err + } if _, err := tx.Exec(ctx, "INSERT INTO job_specs (job_id, request) VALUES ($1, $2)", id, body); err != nil { return err } + if idem.Key != "" { + if _, err := tx.Exec(ctx, ` + INSERT INTO submissions (username, idempotency_key, request_hash, job_id, created_at) + VALUES ($1, $2, $3, $4, $5)`, r.Username, idem.Key, idem.RequestHash, id, now); err != nil { + return err + } + } if _, err := tx.Exec(ctx, "INSERT INTO job_events (job_id, at, to_state, reason, actor) VALUES ($1, $2, 'PENDING', 'submitted', $3)", id, now, r.Username); err != nil { @@ -112,7 +170,14 @@ func (js Jobs) Submit(ctx context.Context, r Request) (int64, error) { _, err = tx.Exec(ctx, "SELECT pg_notify($1, '')", notifyChannel) return err }) - return id, err + return id, replayed, err +} + +// Wake asks the leader for an early cycle, for example after a cleanup +// receipt was recorded. It is only a hint; the next tick does the same. +func (js Jobs) Wake(ctx context.Context) error { + _, err := js.Pool.Exec(ctx, "SELECT pg_notify($1, '')", notifyChannel) + return err } // Cancel records a stop request and returns the job's state afterwards. diff --git a/internal/scheduler/views.go b/internal/scheduler/views.go new file mode 100644 index 0000000..8fdcc1f --- /dev/null +++ b/internal/scheduler/views.go @@ -0,0 +1,306 @@ +package scheduler + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/jackc/pgx/v5" + + "github.com/UTMIST/Mist/internal/executor" +) + +// Read-only views for the API. They need no leadership and take no row +// locks; the API decides who may see what, and redacts. + +// JobInfo is one job as the API shows it. +type JobInfo struct { + ID int64 + Username string + Team string + Partition string + QOS string + Name string + Kind executor.Kind + State State + PendingReason *string + PendingDetail string + Held bool + Resources executor.Resources + TimeLimit time.Duration + SignalGrace time.Duration + Node *string + AuthGeneration *int64 + StopCause *string + ExitCode *int + ExitReason *string + ExitDetail string + OutputIncomplete bool + SubmittedAt time.Time + StartedAt *time.Time + EndedAt *time.Time + CleanupPending bool + CleanupDetail string + Workdir string + OutputPath string +} + +const infoColumns = `j.id, j.username, j.team, j.partition, j.qos, j.name, j.kind, j.state, j.pending_reason, + j.pending_detail, j.held, j.milli_cpu, j.memory_bytes, j.gpus, j.tt_boards, j.scratch_bytes, + j.time_limit_seconds, j.signal_grace_seconds, j.node, j.auth_generation, j.stop_cause, j.exit_code, + j.exit_reason, j.exit_detail, j.output_incomplete, j.submitted_at, j.started_at, j.ended_at, + j.cleanup_pending, j.cleanup_detail, s.request->>'workdir', s.request->>'output_path' + FROM jobs j JOIN job_specs s ON s.job_id = j.id JOIN qos q ON q.name = j.qos ` + +func scanInfo(row pgx.Row) (*JobInfo, error) { + var j JobInfo + var kind, state string + var limit, grace int64 + err := row.Scan(&j.ID, &j.Username, &j.Team, &j.Partition, &j.QOS, &j.Name, &kind, &state, &j.PendingReason, + &j.PendingDetail, &j.Held, &j.Resources.MilliCPU, &j.Resources.MemoryBytes, &j.Resources.GPUs, + &j.Resources.TTBoards, &j.Resources.ScratchBytes, &limit, &grace, &j.Node, &j.AuthGeneration, + &j.StopCause, &j.ExitCode, &j.ExitReason, &j.ExitDetail, &j.OutputIncomplete, &j.SubmittedAt, + &j.StartedAt, &j.EndedAt, &j.CleanupPending, &j.CleanupDetail, &j.Workdir, &j.OutputPath) + if err != nil { + return nil, err + } + j.Kind, j.State = executor.Kind(kind), State(state) + j.TimeLimit, j.SignalGrace = time.Duration(limit)*time.Second, time.Duration(grace)*time.Second + return &j, nil +} + +func collectInfos(rows pgx.Rows, err error) ([]*JobInfo, error) { + if err != nil { + return nil, err + } + return pgx.CollectRows(rows, func(r pgx.CollectableRow) (*JobInfo, error) { return scanInfo(r) }) +} + +// Get returns one job, or ErrNotFound. +func (js Jobs) Get(ctx context.Context, id int64) (*JobInfo, error) { + j, err := scanInfo(js.Pool.QueryRow(ctx, "SELECT "+infoColumns+"WHERE j.id = $1", id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, ErrNotFound + } + return j, err +} + +// ListFilter selects jobs, newest first. Empty fields match everything. +type ListFilter struct { + Username string + Team string + States []State + Before int64 // only IDs below this (the pagination cursor) + Limit int +} + +// List returns jobs matching f, newest first. +func (js Jobs) List(ctx context.Context, f ListFilter) ([]*JobInfo, error) { + states := make([]string, len(f.States)) + for i, s := range f.States { + states[i] = string(s) + } + return collectInfos(js.Pool.Query(ctx, "SELECT "+infoColumns+` + WHERE ($1 = '' OR j.username = $1) AND ($2 = '' OR j.team = $2) + AND (cardinality($3::text[]) = 0 OR j.state = ANY($3)) AND ($4 = 0 OR j.id < $4) + ORDER BY j.id DESC LIMIT $5`, f.Username, f.Team, states, f.Before, f.Limit)) +} + +// Queue returns every nonterminal job in admission order (QoS priority, +// then FIFO), optionally for one partition. +func (js Jobs) Queue(ctx context.Context, partition string) ([]*JobInfo, error) { + return collectInfos(js.Pool.Query(ctx, "SELECT "+infoColumns+` + WHERE j.state IN ('PENDING', 'LAUNCHING', 'RUNNING', 'STOPPING', 'CLEANING') + AND ($1 = '' OR j.partition = $1) + ORDER BY q.priority DESC, j.submitted_at, j.id`, partition)) +} + +// Spec returns the normalized submission a job was created from. +func (js Jobs) Spec(ctx context.Context, id int64) (Request, error) { + var r Request + var body []byte + err := js.Pool.QueryRow(ctx, "SELECT request FROM job_specs WHERE job_id = $1", id).Scan(&body) + if errors.Is(err, pgx.ErrNoRows) { + return r, ErrNotFound + } + if err != nil { + return r, err + } + return r, json.Unmarshal(body, &r) +} + +// JobEvent is one recorded state change. +type JobEvent struct { + At time.Time + From *string + To string + Reason string + Actor string +} + +// Events returns a job's state history, oldest first. +func (js Jobs) Events(ctx context.Context, id int64) ([]JobEvent, error) { + rows, err := js.Pool.Query(ctx, + "SELECT at, from_state, to_state, reason, actor FROM job_events WHERE job_id = $1 ORDER BY id", id) + if err != nil { + return nil, err + } + return pgx.CollectRows(rows, func(r pgx.CollectableRow) (JobEvent, error) { + var e JobEvent + return e, r.Scan(&e.At, &e.From, &e.To, &e.Reason, &e.Actor) + }) +} + +// Partition is a partition's configuration, defaults and live capacity. +// CPU and memory defaults and maximums are per accelerator when the +// partition has accelerators, otherwise per job. +type Partition struct { + Name, State, Accelerator, Runtime string + IsDefault bool + MinAccelerators, MaxAccelerators int64 + DefaultMilliCPU, MaxMilliCPU int64 + DefaultMemory, MaxMemory int64 + DefaultScratch, MaxScratch int64 + DefaultTimeLimit, MaxTimeLimit time.Duration + DefaultSignalGrace time.Duration + CapMilliCPU, CapMemory *int64 + Nodes []NodeStatus +} + +// NodeStatus is one node of a partition. Total is the last reported +// capacity net of the Mist host reserve; Free also subtracts unreleased +// allocations from every partition sharing the node. +type NodeStatus struct { + Name, AdminState, Detail string + Healthy, Quarantined bool + Reported bool // the executor has reported this node at least once + Total, Free executor.Resources +} + +// QOS is one quality-of-service class. +type QOS struct { + Name string + Priority int + MaxTimeLimit *time.Duration +} + +// Limit is one limits row. +type Limit struct { + Scope, Subject, Partition, QOS, Resource string + Max int64 +} + +// Catalog is everything the API needs to normalize a submission and show +// partitions, read in one transaction so it is consistent. +type Catalog struct { + Partitions []Partition + QOS []QOS + Limits []Limit +} + +func (c *Catalog) Partition(name string) *Partition { + for i := range c.Partitions { + if c.Partitions[i].Name == name { + return &c.Partitions[i] + } + } + return nil +} + +func (c *Catalog) QoS(name string) *QOS { + for i := range c.QOS { + if c.QOS[i].Name == name { + return &c.QOS[i] + } + } + return nil +} + +// Catalog reads partitions, nodes, QoS and limits. +func (js Jobs) Catalog(ctx context.Context) (*Catalog, error) { + c := &Catalog{} + err := pgx.BeginTxFunc(ctx, js.Pool, pgx.TxOptions{IsoLevel: pgx.RepeatableRead, AccessMode: pgx.ReadOnly}, func(tx pgx.Tx) error { + rows, err := tx.Query(ctx, ` + SELECT name, state, accelerator, runtime, is_default, min_accelerators, max_accelerators, + default_milli_cpu, max_milli_cpu, default_memory_bytes, max_memory_bytes, + default_scratch_bytes, max_scratch_bytes, default_time_limit_seconds, max_time_limit_seconds, + default_signal_grace_seconds, cap_milli_cpu, cap_memory_bytes + FROM partitions ORDER BY name`) + if err != nil { + return err + } + c.Partitions, err = pgx.CollectRows(rows, func(r pgx.CollectableRow) (Partition, error) { + var p Partition + var dt, mt, dg int64 + err := r.Scan(&p.Name, &p.State, &p.Accelerator, &p.Runtime, &p.IsDefault, &p.MinAccelerators, + &p.MaxAccelerators, &p.DefaultMilliCPU, &p.MaxMilliCPU, &p.DefaultMemory, &p.MaxMemory, + &p.DefaultScratch, &p.MaxScratch, &dt, &mt, &dg, &p.CapMilliCPU, &p.CapMemory) + p.DefaultTimeLimit, p.MaxTimeLimit = time.Duration(dt)*time.Second, time.Duration(mt)*time.Second + p.DefaultSignalGrace = time.Duration(dg) * time.Second + return p, err + }) + if err != nil { + return err + } + + rows, err = tx.Query(ctx, ` + SELECT pn.partition, n.name, n.admin_state, n.health_detail, n.healthy, n.quarantined, + n.observed_at IS NOT NULL, + n.milli_cpu - n.reserve_milli_cpu, n.memory_bytes - n.reserve_memory_bytes, + n.gpus, n.tt_boards, n.scratch_bytes, n.hugepages_bytes, + coalesce(sum(a.milli_cpu), 0), coalesce(sum(a.memory_bytes), 0), coalesce(sum(a.gpus), 0), + coalesce(sum(a.tt_boards), 0), coalesce(sum(a.scratch_bytes), 0), coalesce(sum(a.hugepages_bytes), 0) + FROM partition_nodes pn JOIN nodes n ON n.name = pn.node + LEFT JOIN allocations a ON a.node = n.name AND a.released_at IS NULL + GROUP BY pn.partition, n.name ORDER BY n.name`) + if err != nil { + return err + } + for rows.Next() { + var part string + var n NodeStatus + var used executor.Resources + t := &n.Total + if err := rows.Scan(&part, &n.Name, &n.AdminState, &n.Detail, &n.Healthy, &n.Quarantined, &n.Reported, + &t.MilliCPU, &t.MemoryBytes, &t.GPUs, &t.TTBoards, &t.ScratchBytes, &t.HugepagesBytes, + &used.MilliCPU, &used.MemoryBytes, &used.GPUs, &used.TTBoards, &used.ScratchBytes, &used.HugepagesBytes); err != nil { + return err + } + n.Free = n.Total.Sub(used) + if p := c.Partition(part); p != nil { + p.Nodes = append(p.Nodes, n) + } + } + if err := rows.Err(); err != nil { + return err + } + + rows, err = tx.Query(ctx, "SELECT name, priority, max_time_limit_seconds FROM qos ORDER BY priority, name") + if err != nil { + return err + } + c.QOS, err = pgx.CollectRows(rows, func(r pgx.CollectableRow) (QOS, error) { + var q QOS + var max *int64 + err := r.Scan(&q.Name, &q.Priority, &max) + if max != nil { + d := time.Duration(*max) * time.Second + q.MaxTimeLimit = &d + } + return q, err + }) + if err != nil { + return err + } + + rows, err = tx.Query(ctx, + "SELECT scope, subject, partition, qos, resource, max FROM limits ORDER BY scope, subject, partition, qos, resource") + if err != nil { + return err + } + c.Limits, err = pgx.CollectRows(rows, pgx.RowToStructByPos[Limit]) + return err + }) + return c, err +} diff --git a/scripts/k3d-test.sh b/scripts/k3d-test.sh index d04238e..0d5b2b0 100755 --- a/scripts/k3d-test.sh +++ b/scripts/k3d-test.sh @@ -35,4 +35,4 @@ docker exec "$node" sh -c ' chown 0:40002 /mist/teams/team-b && chmod 2750 /mist/teams/team-b' MIST_K3D_CONTEXT=$ctx MIST_K3D_NODE=$node \ - go test -tags k3d -count=1 -v -run K3d ./internal/executor/k3s "$@" + go test -tags k3d -count=1 -v -run K3d ./internal/executor/k3s ./internal/e2e "$@"