diff --git a/README.md b/README.md index 63a9f22..d2033df 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,23 @@ 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. +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 +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 new file mode 100644 index 0000000..5b18545 --- /dev/null +++ b/cmd/mist-job-entry/main.go @@ -0,0 +1,43 @@ +// 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] +// mist-job-entry -await-isolation host:port[,host:port] (init container) +package main + +import ( + "flag" + "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) + 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/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 2b7d7d5..4cc49d8 100644 --- a/cmd/mistd/main.go +++ b/cmd/mistd/main.go @@ -2,51 +2,84 @@ // // mistd api HTTP API // mistd scheduler queue, admission and executor calls (one active leader) +// 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" "flag" "fmt" "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" "github.com/UTMIST/Mist/internal/scheduler" ) 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(ctx, args, stdout) + } fs := flag.NewFlagSet("mistd "+cmd, flag.ContinueOnError) fs.SetOutput(stderr) dbURL := fs.String("database-url", os.Getenv("MIST_DATABASE_URL"), "Postgres URL") 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") } @@ -64,23 +97,66 @@ func run(ctx context.Context, args []string, 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, + 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) @@ -96,3 +172,68 @@ func run(ctx context.Context, args []string, stderr io.Writer) error { return fmt.Errorf("unknown subcommand %q", cmd) } } + +// 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, 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 + } + evidence, err := k3s.VerifyCleanup(*kubelet, *cgroups, *podUID) + if err != nil { + return err + } + 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 4d41be2..5581b06 100644 --- a/cmd/mistd/main_test.go +++ b/cmd/mistd/main_test.go @@ -3,6 +3,12 @@ package main import ( "bytes" "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" "strings" "sync" "testing" @@ -11,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) } @@ -53,7 +60,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 @@ -130,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/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/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 1e13f03..fdd48f3 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 @@ -352,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 @@ -412,6 +449,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/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/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 0303345..27574a2 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 @@ -92,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 @@ -100,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. @@ -113,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). @@ -120,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/go.mod b/go.mod index 17e6521..935f1a4 100644 --- a/go.mod +++ b/go.mod @@ -5,14 +5,53 @@ 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 + 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 caad47e..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,24 +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= -golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= -golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= +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/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/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/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/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/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/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/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..c786de2 --- /dev/null +++ b/internal/executor/fake/fake.go @@ -0,0 +1,597 @@ +// 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()) +} + +// 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() + 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..6781c43 --- /dev/null +++ b/internal/executor/fake/fake_test.go @@ -0,0 +1,423 @@ +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, + Partition: "gpu", + QOS: "normal", + 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") + } + }) +} 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..7ec7bbf --- /dev/null +++ b/internal/executor/k3s/executor.go @@ -0,0 +1,771 @@ +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"` +} + +// 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. +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("%w: it does not match a recorded execution", ErrCleanupRejected) + } + defer e.lock(rec.JobID)() + if rec.FinishedAt == nil && rec.LostAt == nil { + 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("%w: Pod %s still exists", ErrCleanupRejected, 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/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/jobentry/env.go b/internal/jobentry/env.go new file mode 100644 index 0000000..8c406cf --- /dev/null +++ b/internal/jobentry/env.go @@ -0,0 +1,89 @@ +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"} + +// 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": + 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 !ProtectedEnv(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/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/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 +} diff --git a/internal/scheduler/admit.go b/internal/scheduler/admit.go new file mode 100644 index 0000000..7b25509 --- /dev/null +++ b/internal/scheduler/admit.go @@ -0,0 +1,235 @@ +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, + 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, + 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..6d66ebe --- /dev/null +++ b/internal/scheduler/scheduler_test.go @@ -0,0 +1,498 @@ +package scheduler + +import ( + "errors" + "fmt" + "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) + } +} + +// 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/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..5b81113 --- /dev/null +++ b/internal/scheduler/submit.go @@ -0,0 +1,204 @@ +package scheduler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + "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 (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"` + 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") + +// 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 { + 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) { + 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, false, fmt.Errorf("scheduler: request is not normalized") + } + now := js.Clock.Now() + err = pgx.BeginFunc(ctx, js.Pool, func(tx pgx.Tx) error { + // 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 + 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 + } + 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 { + return err + } + _, err = tx.Exec(ctx, "SELECT pg_notify($1, '')", notifyChannel) + return 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. +// 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 +} 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 new file mode 100755 index 0000000..0d5b2b0 --- /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 ./internal/e2e "$@"