From c8e8a13b30e7c8844521c8c0966edbdad2693c0e Mon Sep 17 00:00:00 2001 From: Jens Horn Date: Wed, 26 Aug 2026 15:02:54 +0200 Subject: [PATCH] feat: add Kubernetes workspace scheduling and ephemeral storage. --- .changeset/proud-lies-yawn.md | 5 + docs/deployment.md | 24 ++- docs/templates.md | 5 +- .../contracts/src/template-resources.test.ts | 22 +++ packages/contracts/src/template-schema.ts | 4 + packages/drivers/src/index.ts | 5 + .../drivers/src/kubernetes-manifests.test.ts | 106 ++++++++++++ packages/drivers/src/kubernetes-manifests.ts | 16 +- .../kubernetes-scheduling-conformance.test.ts | 163 ++++++++++++++++++ .../src/kubernetes-scheduling-driver.test.ts | 87 ++++++++++ .../drivers/src/kubernetes-scheduling.test.ts | 100 +++++++++++ packages/drivers/src/kubernetes-scheduling.ts | 47 +++++ packages/drivers/src/kubernetes.test.ts | 8 +- packages/drivers/src/kubernetes.ts | 16 +- packages/server/src/config.test.ts | 4 + packages/server/src/config.ts | 7 + .../src/kubernetes-scheduling-config.test.ts | 99 +++++++++++ .../src/kubernetes-scheduling-config.ts | 68 ++++++++ packages/server/src/lifecycle.ts | 2 + 19 files changed, 778 insertions(+), 10 deletions(-) create mode 100644 .changeset/proud-lies-yawn.md create mode 100644 packages/contracts/src/template-resources.test.ts create mode 100644 packages/drivers/src/kubernetes-manifests.test.ts create mode 100644 packages/drivers/src/kubernetes-scheduling-conformance.test.ts create mode 100644 packages/drivers/src/kubernetes-scheduling-driver.test.ts create mode 100644 packages/drivers/src/kubernetes-scheduling.test.ts create mode 100644 packages/drivers/src/kubernetes-scheduling.ts create mode 100644 packages/server/src/kubernetes-scheduling-config.test.ts create mode 100644 packages/server/src/kubernetes-scheduling-config.ts diff --git a/.changeset/proud-lies-yawn.md b/.changeset/proud-lies-yawn.md new file mode 100644 index 0000000..2b17777 --- /dev/null +++ b/.changeset/proud-lies-yawn.md @@ -0,0 +1,5 @@ +--- +"@pstdio/pocketcoder-cli": minor +--- + +Add Kubernetes workspace node scheduling and ephemeral storage limits. diff --git a/docs/deployment.md b/docs/deployment.md index f8f0ca2..175270f 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -176,6 +176,17 @@ server replica—the connection hub and scheduler are intentionally single-active. Workspace Jobs use the unprivileged `pocketcoder-workspace` service account, not the controller account. +Use driver-level scheduling settings to keep all workspace and warm-pool Jobs +on a labeled, tainted node pool: + +```sh +export POCKETCODER_KUBERNETES_NODE_SELECTOR='{"onefin.com/workload":"agent-workspace"}' +export POCKETCODER_KUBERNETES_TOLERATIONS='[{"key":"onefin.com/workload","operator":"Equal","value":"agent-workspace","effect":"NoSchedule"}]' +``` + +The selector must be a JSON object with string values. Tolerations must be a +JSON array. This release supports only the `NoSchedule` effect. + Memory-backed writable paths are mounted with the template uid/gid. Kubernetes Jobs set pod `fsGroup` to the template gid with `fsGroupChangePolicy: OnRootMismatch`; Docker tmpfs mounts set `uid`, `gid`, @@ -188,7 +199,16 @@ context before rollout: ```sh POCKETCODER_KUBERNETES_CONFORMANCE=1 \ POCKETCODER_KUBERNETES_NAMESPACE=pocketcoder \ -bun test packages/drivers/src/kubernetes.test.ts +bun test packages/drivers/src/kubernetes-conformance.test.ts +``` + +To test a labeled, tainted workspace node pool, also set +`POCKETCODER_KUBERNETES_CONFORMANCE_IMAGE` to a digest-pinned image and set the +two scheduling variables above. Then run: + +```sh +POCKETCODER_KUBERNETES_CONFORMANCE=1 \ +bun test packages/drivers/src/kubernetes-scheduling-conformance.test.ts ``` With `POCKETCODER_SECRET_PROVIDER=kubernetes`, a template value @@ -239,6 +259,8 @@ an application role with connect/usage/DML only. | `POCKETCODER_SECRET_ROOT` | required for file secrets | Deployment-owned local secret root | | `POCKETCODER_KUBERNETES_NAMESPACE` | `default` | Namespace for Jobs and input Secrets | | `POCKETCODER_KUBERNETES_SERVICE_ACCOUNT` | none | Service account assigned to workspace Jobs | +| `POCKETCODER_KUBERNETES_NODE_SELECTOR` | none | JSON object that selects nodes for workspace and warm-pool Jobs | +| `POCKETCODER_KUBERNETES_TOLERATIONS` | `[]` | JSON array of `NoSchedule` tolerations for workspace and warm-pool Jobs | | `POCKETCODER_KUBERNETES_WORKSPACE_CLAIM` | required for PVC | Claim mounted by server and workspace Jobs | | `POCKETCODER_KUBERNETES_WORKSPACE_SUBPATH` | `workspaces` | Opaque allocation prefix in the claim | | `POCKETCODER_MAX_RETAINED_BYTES` | `500Gi` | Global checkpoint quota | diff --git a/docs/templates.md b/docs/templates.md index ca53d02..fb8552f 100644 --- a/docs/templates.md +++ b/docs/templates.md @@ -40,7 +40,7 @@ command, mount, network, privilege, or driver. "idleTimeout": "10m" }, "env": { "HOME": "/home/agent" }, - "resources": { "cpu": "2", "memory": "2Gi" }, + "resources": { "cpu": "2", "memory": "2Gi", "ephemeralStorage": "10Gi" }, "timeouts": { "start": "2m", "maxAge": "2h", "idle": "20m", "disconnectGrace": "5m", "terminateGrace": "15s" }, "security": { "uid": 10001, "gid": 10001, "readOnlyRoot": true, @@ -87,6 +87,9 @@ command, mount, network, privilege, or driver. - **`timeouts`** — `start` (registration + first health), `maxAge` (hard lifetime), `idle` (no relay activity and agent not running), `disconnectGrace` (supervisor reconnect window), `terminateGrace` (TERM→KILL). +- **`resources`** — required CPU and memory sizing, plus optional scratch disk + sizing in `ephemeralStorage`. Ephemeral storage is Kubernetes-only in this + release. The Docker driver accepts the field but does not enforce it. - **`security`** — non-root uid/gid (≥1000), read-only root, memory-backed writable paths, dropped capabilities, no privilege escalation. Values can only be stricter than the defaults, never weaker. diff --git a/packages/contracts/src/template-resources.test.ts b/packages/contracts/src/template-resources.test.ts new file mode 100644 index 0000000..c3b4b7b --- /dev/null +++ b/packages/contracts/src/template-resources.test.ts @@ -0,0 +1,22 @@ +import { expect, test } from "bun:test"; +import { parseTemplateManifest } from "./index"; + +test("accepts an optional ephemeral storage resource", () => { + const parsed = parseTemplateManifest({ + apiVersion: "pocketcoder.dev/v1alpha1", + kind: "Template", + metadata: { name: "resources" }, + spec: { + version: "1.0.0", + image: `registry.test/agent@sha256:${"a".repeat(64)}`, + harness: { command: ["agent"] }, + resources: { cpu: "2", memory: "2Gi", ephemeralStorage: "10Gi" }, + }, + }); + + expect(parsed.manifest.spec.resources).toEqual({ + cpu: "2", + memory: "2Gi", + ephemeralStorage: "10Gi", + }); +}); diff --git a/packages/contracts/src/template-schema.ts b/packages/contracts/src/template-schema.ts index 2a8093f..c955e5d 100644 --- a/packages/contracts/src/template-schema.ts +++ b/packages/contracts/src/template-schema.ts @@ -123,6 +123,10 @@ export const TimeoutsSchema = z.object({ export const ResourcesSchema = z.object({ cpu: z.string().regex(/^\d+(\.\d+)?m?$/), memory: z.string().regex(/^\d+(Mi|Gi)$/), + ephemeralStorage: z + .string() + .regex(/^\d+(Mi|Gi)$/) + .optional(), }); const RepositorySchema = z.object({ diff --git a/packages/drivers/src/index.ts b/packages/drivers/src/index.ts index 4199c5f..fa5046f 100644 --- a/packages/drivers/src/index.ts +++ b/packages/drivers/src/index.ts @@ -25,6 +25,11 @@ export { KubernetesDriver, type KubernetesDriverOptions, } from "./kubernetes"; +export { + type KubernetesSchedulingOptions, + type KubernetesToleration, + validateToleration, +} from "./kubernetes-scheduling"; export { KubernetesSecretResolver, type KubernetesSecretResolverOptions, diff --git a/packages/drivers/src/kubernetes-manifests.test.ts b/packages/drivers/src/kubernetes-manifests.test.ts new file mode 100644 index 0000000..af14aa8 --- /dev/null +++ b/packages/drivers/src/kubernetes-manifests.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, test } from "bun:test"; +import type { WarmRuntimeLaunch, WorkspaceLaunch } from "@pstdio/pocketcoder-runtime-core"; +import { warmJobManifest, workspaceJobManifest } from "./kubernetes-manifests"; + +function templateSpec(ephemeralStorage?: string) { + return { + image: "registry.example/workspace@sha256:fixture", + command: ["/bin/sleep", "3600"], + env: {}, + resources: { + cpu: "1", + memory: "512Mi", + ...(ephemeralStorage ? { ephemeralStorage } : {}), + }, + security: { + uid: 10_001, + gid: 10_001, + writableMemoryPaths: ["/tmp"], + readOnlyRoot: true, + allowPrivilegeEscalation: false, + dropCapabilities: ["ALL"], + seccomp: "RuntimeDefault", + }, + network: { mode: "unrestricted" }, + }; +} + +function workspaceLaunch(ephemeralStorage?: string) { + return { + workspace: { + id: "workspace-id", + templateDigest: "sha256:template", + templateSnapshot: { spec: templateSpec(ephemeralStorage) }, + }, + mounts: [], + secrets: [], + } as unknown as WorkspaceLaunch; +} + +function warmLaunch(ephemeralStorage?: string) { + return { + runtimeId: "runtime-id", + template: { digest: "sha256:template", spec: templateSpec(ephemeralStorage) }, + } as unknown as WarmRuntimeLaunch; +} + +const options = { + imagePullPolicy: "IfNotPresent" as const, + nodeSelector: { "onefin.com/workload": "agent-workspace" }, + tolerations: [ + { + key: "onefin.com/workload", + operator: "Equal" as const, + value: "agent-workspace", + effect: "NoSchedule" as const, + }, + ], +}; + +function podSpec( + manifest: ReturnType | ReturnType, +) { + return manifest.spec.template.spec; +} + +describe("Kubernetes Job manifests", () => { + test("adds scheduling fields to workspace and warm Jobs", () => { + const manifests = [ + workspaceJobManifest(workspaceLaunch(), "workspace", "input", "egress", options), + warmJobManifest(warmLaunch(), "warm", "input", "egress", options), + ]; + + for (const manifest of manifests) { + expect(podSpec(manifest).nodeSelector).toEqual(options.nodeSelector); + expect(podSpec(manifest).tolerations).toEqual(options.tolerations); + } + }); + + test("omits scheduling fields from workspace and warm Jobs by default", () => { + const defaults = { imagePullPolicy: "IfNotPresent" as const }; + const manifests = [ + workspaceJobManifest(workspaceLaunch(), "workspace", "input", "egress", defaults), + warmJobManifest(warmLaunch(), "warm", "input", "egress", defaults), + ]; + + for (const manifest of manifests) { + expect("nodeSelector" in podSpec(manifest)).toBe(false); + expect("tolerations" in podSpec(manifest)).toBe(false); + } + }); + + test("maps ephemeral storage for workspace and warm Jobs", () => { + const manifests = [ + workspaceJobManifest(workspaceLaunch("10Gi"), "workspace", "input", "egress", options), + warmJobManifest(warmLaunch("10Gi"), "warm", "input", "egress", options), + ]; + const expected = { + requests: { cpu: "1", memory: "512Mi", "ephemeral-storage": "10Gi" }, + limits: { cpu: "1", memory: "512Mi", "ephemeral-storage": "10Gi" }, + }; + + for (const manifest of manifests) { + expect(podSpec(manifest).containers[0]?.resources).toEqual(expected); + } + }); +}); diff --git a/packages/drivers/src/kubernetes-manifests.ts b/packages/drivers/src/kubernetes-manifests.ts index 7861cc9..1a2c297 100644 --- a/packages/drivers/src/kubernetes-manifests.ts +++ b/packages/drivers/src/kubernetes-manifests.ts @@ -9,11 +9,18 @@ import { KUBERNETES_POOL_LABEL, KUBERNETES_WORKSPACE_LABEL, } from "./kubernetes-labels"; +import { + type KubernetesToleration, + resourceRequirements, + schedulingFields, +} from "./kubernetes-scheduling"; interface ManifestOptions { serviceAccountName?: string; imagePullPolicy: "Always" | "IfNotPresent" | "Never"; egressImage?: string; + nodeSelector?: Record; + tolerations?: KubernetesToleration[]; } function volumeForMount(mount: RuntimeMountRef, index: number) { @@ -132,6 +139,7 @@ export function workspaceJobManifest( spec: { restartPolicy: "Never", automountServiceAccountToken: false, + ...schedulingFields(options), ...(options.serviceAccountName ? { serviceAccountName: options.serviceAccountName } : {}), securityContext: { fsGroup: spec.security.gid, @@ -148,10 +156,7 @@ export function workspaceJobManifest( env: Object.entries(spec.env) .filter(([, value]) => !value.startsWith("secretRef:")) .map(([name, value]) => ({ name, value })), - resources: { - requests: { cpu: spec.resources.cpu, memory: spec.resources.memory }, - limits: { cpu: spec.resources.cpu, memory: spec.resources.memory }, - }, + resources: resourceRequirements(spec.resources), securityContext: { runAsUser: spec.security.uid, runAsGroup: spec.security.gid, @@ -210,6 +215,7 @@ export function warmJobManifest( spec: { restartPolicy: "Never", automountServiceAccountToken: false, + ...schedulingFields(options), ...(options.serviceAccountName ? { serviceAccountName: options.serviceAccountName } : {}), securityContext: { fsGroup: spec.security.gid, @@ -224,7 +230,7 @@ export function warmJobManifest( imagePullPolicy: options.imagePullPolicy, command: spec.command, env: Object.entries(spec.env).map(([name, value]) => ({ name, value })), - resources: { requests: spec.resources, limits: spec.resources }, + resources: resourceRequirements(spec.resources), securityContext: { runAsUser: spec.security.uid, runAsGroup: spec.security.gid, diff --git a/packages/drivers/src/kubernetes-scheduling-conformance.test.ts b/packages/drivers/src/kubernetes-scheduling-conformance.test.ts new file mode 100644 index 0000000..fd5bacc --- /dev/null +++ b/packages/drivers/src/kubernetes-scheduling-conformance.test.ts @@ -0,0 +1,163 @@ +import { describe, expect, test } from "bun:test"; +import { randomUUID } from "node:crypto"; +import type { KubernetesToleration } from "./kubernetes-scheduling"; + +const enabled = + process.env.POCKETCODER_KUBERNETES_CONFORMANCE === "1" && + Boolean(process.env.POCKETCODER_KUBERNETES_NODE_SELECTOR) && + Boolean(process.env.POCKETCODER_KUBERNETES_TOLERATIONS); + +function namespace() { + return process.env.POCKETCODER_KUBERNETES_NAMESPACE ?? "default"; +} + +function conformanceImage() { + const image = process.env.POCKETCODER_KUBERNETES_CONFORMANCE_IMAGE; + if (!image || !/^[^\s@]+@sha256:[0-9a-f]{64}$/.test(image)) { + throw new Error("POCKETCODER_KUBERNETES_CONFORMANCE_IMAGE must be pinned by digest"); + } + return image; +} + +function schedulingConfig() { + return { + nodeSelector: JSON.parse(process.env.POCKETCODER_KUBERNETES_NODE_SELECTOR as string) as Record< + string, + string + >, + tolerations: JSON.parse( + process.env.POCKETCODER_KUBERNETES_TOLERATIONS as string, + ) as KubernetesToleration[], + }; +} + +async function kubectl(args: string[], input?: string) { + const child = Bun.spawn(["kubectl", ...args], { + ...(input ? { stdin: "pipe" } : {}), + stdout: "pipe", + stderr: "pipe", + }); + if (input && child.stdin) { + child.stdin.write(input); + child.stdin.end(); + } + const [stdout, stderr, exitCode] = await Promise.all([ + new Response(child.stdout).text(), + new Response(child.stderr).text(), + child.exited, + ]); + if (exitCode !== 0) throw new Error(stderr || stdout || `kubectl exited ${exitCode}`); + return stdout; +} + +function jobManifest(name: string, includeTolerations: boolean) { + const { nodeSelector, tolerations } = schedulingConfig(); + return { + apiVersion: "batch/v1", + kind: "Job", + metadata: { name, namespace: namespace() }, + spec: { + backoffLimit: 0, + template: { + spec: { + restartPolicy: "Never", + automountServiceAccountToken: false, + nodeSelector, + ...(includeTolerations ? { tolerations } : {}), + containers: [ + { + name: "probe", + image: conformanceImage(), + command: ["sh", "-c", "echo scheduled"], + resources: { + requests: { "ephemeral-storage": "64Mi" }, + limits: { "ephemeral-storage": "64Mi" }, + }, + }, + ], + }, + }, + }, + }; +} + +async function apply(manifest: unknown) { + await kubectl(["-n", namespace(), "apply", "-f", "-"], JSON.stringify(manifest)); +} + +async function podForJob(name: string) { + const output = await kubectl([ + "-n", + namespace(), + "get", + "pods", + "-l", + `job-name=${name}`, + "-o", + "json", + ]); + const list = JSON.parse(output) as { items: Array> }; + return list.items[0]; +} + +async function waitForUnschedulable(name: string) { + const deadline = Date.now() + 30_000; + while (Date.now() < deadline) { + const pod = (await podForJob(name)) as { + spec?: { nodeName?: string }; + status?: { conditions?: Array<{ reason?: string; status?: string; type?: string }> }; + }; + const condition = pod?.status?.conditions?.find((item) => item.type === "PodScheduled"); + if ( + !pod?.spec?.nodeName && + condition?.status === "False" && + condition.reason === "Unschedulable" + ) { + return; + } + await Bun.sleep(1000); + } + throw new Error("the Job without a toleration did not become unschedulable"); +} + +async function deleteJob(name: string) { + await kubectl(["-n", namespace(), "delete", "job", name, "--ignore-not-found", "--wait=false"]); +} + +describe.skipIf(!enabled)("Kubernetes workspace scheduling conformance", () => { + test("requires the taint toleration and preserves ephemeral storage", async () => { + const suffix = randomUUID().slice(0, 8); + const blockedName = `pocketcoder-scheduling-blocked-${suffix}`; + const scheduledName = `pocketcoder-scheduling-${suffix}`; + try { + await apply(jobManifest(blockedName, false)); + await waitForUnschedulable(blockedName); + + await apply(jobManifest(scheduledName, true)); + await kubectl([ + "-n", + namespace(), + "wait", + `job/${scheduledName}`, + "--for=condition=complete", + "--timeout=120s", + ]); + const pod = (await podForJob(scheduledName)) as { + spec: { + nodeName: string; + containers: Array<{ resources: Record> }>; + }; + }; + const node = JSON.parse(await kubectl(["get", "node", pod.spec.nodeName, "-o", "json"])) as { + metadata: { labels: Record }; + }; + expect(node.metadata.labels).toMatchObject(schedulingConfig().nodeSelector); + expect(pod.spec.containers[0]?.resources).toMatchObject({ + requests: { "ephemeral-storage": "64Mi" }, + limits: { "ephemeral-storage": "64Mi" }, + }); + } finally { + await Promise.all([deleteJob(blockedName), deleteJob(scheduledName)]); + } + }, 180_000); +}); diff --git a/packages/drivers/src/kubernetes-scheduling-driver.test.ts b/packages/drivers/src/kubernetes-scheduling-driver.test.ts new file mode 100644 index 0000000..24d4e59 --- /dev/null +++ b/packages/drivers/src/kubernetes-scheduling-driver.test.ts @@ -0,0 +1,87 @@ +import { afterEach, expect, test } from "bun:test"; +import { randomUUID } from "node:crypto"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { WorkspaceLaunch } from "@pstdio/pocketcoder-runtime-core"; +import { KubernetesDriver } from "./kubernetes"; + +let temporaryDirectory: string | undefined; + +afterEach(async () => { + if (temporaryDirectory) await rm(temporaryDirectory, { recursive: true, force: true }); + temporaryDirectory = undefined; +}); + +async function fakeKubectl() { + temporaryDirectory = await mkdtemp(join(tmpdir(), "pocketcoder-kubectl-scheduling-")); + const log = join(temporaryDirectory, "calls.ndjson"); + const script = join(temporaryDirectory, "kubectl.ts"); + const bin = process.platform === "win32" ? join(temporaryDirectory, "kubectl.cmd") : script; + await writeFile( + script, + `#!/usr/bin/env bun +import { appendFileSync } from "node:fs"; +const input = await Bun.stdin.text(); +appendFileSync(${JSON.stringify(log)}, input + "\\n"); +`, + { mode: 0o755 }, + ); + if (process.platform === "win32") { + await writeFile(bin, `@${JSON.stringify(process.execPath)} ${JSON.stringify(script)} %*\r\n`); + } + return { bin, log }; +} + +function launch(): WorkspaceLaunch { + return { + workspace: { + id: randomUUID(), + templateDigest: "sha256:template", + templateSnapshot: { + spec: { + image: "registry.example/workspace@sha256:fixture", + command: ["sleep", "3600"], + env: {}, + resources: { cpu: "1", memory: "512Mi" }, + security: { + uid: 10_001, + gid: 10_001, + writableMemoryPaths: [], + readOnlyRoot: true, + allowPrivilegeEscalation: false, + dropCapabilities: ["ALL"], + seccomp: "RuntimeDefault", + }, + network: { mode: "unrestricted" }, + }, + }, + }, + input: {}, + mounts: [], + secrets: [], + } as unknown as WorkspaceLaunch; +} + +test("forwards scheduling options to the applied Job", async () => { + const fake = await fakeKubectl(); + const driver = new KubernetesDriver({ + kubectlBin: fake.bin, + nodeSelector: { dedicated: "workspace" }, + tolerations: [{ operator: "Exists", effect: "NoSchedule" }], + }); + + await driver.create(launch()); + const manifests = (await readFile(fake.log, "utf8")) + .trim() + .split("\n") + .map((line) => JSON.parse(line) as Record); + const job = manifests.find((manifest) => manifest.kind === "Job") as { + spec: { template: { spec: Record } }; + }; + + expect(job.spec.template.spec.nodeSelector).toEqual({ dedicated: "workspace" }); + expect(job.spec.template.spec.tolerations).toEqual([ + { operator: "Exists", effect: "NoSchedule" }, + ]); +}); diff --git a/packages/drivers/src/kubernetes-scheduling.test.ts b/packages/drivers/src/kubernetes-scheduling.test.ts new file mode 100644 index 0000000..698dd92 --- /dev/null +++ b/packages/drivers/src/kubernetes-scheduling.test.ts @@ -0,0 +1,100 @@ +import { describe, expect, test } from "bun:test"; +import { KubernetesDriver } from "./kubernetes"; +import { + resourceRequirements, + schedulingFields, + validateToleration, +} from "./kubernetes-scheduling"; + +describe("Kubernetes scheduling manifest fields", () => { + test("omits scheduling fields when none are configured", () => { + expect(schedulingFields({})).toEqual({}); + }); + + test("includes configured selectors and tolerations", () => { + const nodeSelector = { "onefin.com/workload": "agent-workspace" }; + const tolerations = [ + { + key: "onefin.com/workload", + operator: "Equal" as const, + value: "agent-workspace", + effect: "NoSchedule" as const, + }, + ]; + + expect(schedulingFields({ nodeSelector })).toEqual({ nodeSelector }); + expect(schedulingFields({ tolerations })).toEqual({ tolerations }); + expect(schedulingFields({ nodeSelector, tolerations })).toEqual({ nodeSelector, tolerations }); + }); +}); + +describe("Kubernetes resource requirements", () => { + test("maps CPU and memory into requests and limits", () => { + expect(resourceRequirements({ cpu: "1", memory: "512Mi" })).toEqual({ + requests: { cpu: "1", memory: "512Mi" }, + limits: { cpu: "1", memory: "512Mi" }, + }); + }); + + test("maps ephemeral storage to the Kubernetes resource name", () => { + expect(resourceRequirements({ cpu: "1", memory: "512Mi", ephemeralStorage: "10Gi" })).toEqual({ + requests: { cpu: "1", memory: "512Mi", "ephemeral-storage": "10Gi" }, + limits: { cpu: "1", memory: "512Mi", "ephemeral-storage": "10Gi" }, + }); + }); +}); + +describe("Kubernetes toleration validation", () => { + test("accepts the supported NoSchedule shapes", () => { + expect(() => + validateToleration({ + key: "dedicated", + operator: "Equal", + value: "workspace", + effect: "NoSchedule", + }), + ).not.toThrow(); + expect(() => validateToleration({ operator: "Exists", effect: "NoSchedule" })).not.toThrow(); + }); + + test("rejects a value with the Exists operator", () => { + expect(() => + validateToleration({ operator: "Exists", value: "workspace", effect: "NoSchedule" }), + ).toThrow("operator Exists must not set a value"); + }); + + test("rejects an empty key without the Exists operator", () => { + expect(() => validateToleration({ operator: "Equal", effect: "NoSchedule" })).toThrow( + "empty or omitted toleration key requires operator Exists", + ); + }); + + test("rejects unknown fields", () => { + expect(() => validateToleration({ effect: "NoSchedule", unexpected: true } as never)).toThrow( + "unknown toleration field: unexpected", + ); + }); + + test("rejects other effects", () => { + expect(() => validateToleration({ effect: "NoExecute" } as never)).toThrow( + "toleration effect must be NoSchedule", + ); + }); +}); + +describe("Kubernetes driver scheduling validation", () => { + test("rejects invalid tolerations without the server config layer", () => { + expect( + () => + new KubernetesDriver({ + tolerations: [ + { + operator: "Exists", + value: "workspace", + effect: "NoSchedule", + }, + ], + }), + ).toThrow("operator Exists must not set a value"); + }); +}); diff --git a/packages/drivers/src/kubernetes-scheduling.ts b/packages/drivers/src/kubernetes-scheduling.ts new file mode 100644 index 0000000..05c4214 --- /dev/null +++ b/packages/drivers/src/kubernetes-scheduling.ts @@ -0,0 +1,47 @@ +export interface KubernetesToleration { + key?: string; + operator?: "Exists" | "Equal"; + value?: string; + effect: "NoSchedule"; +} + +export interface KubernetesSchedulingOptions { + nodeSelector?: Record; + tolerations?: KubernetesToleration[]; +} + +export function schedulingFields(options: KubernetesSchedulingOptions) { + return { + ...(options.nodeSelector ? { nodeSelector: options.nodeSelector } : {}), + ...(options.tolerations?.length ? { tolerations: options.tolerations } : {}), + }; +} + +export function resourceRequirements(resources: { + cpu: string; + memory: string; + ephemeralStorage?: string; +}) { + const values = { + cpu: resources.cpu, + memory: resources.memory, + ...(resources.ephemeralStorage ? { "ephemeral-storage": resources.ephemeralStorage } : {}), + }; + return { requests: values, limits: values }; +} + +export function validateToleration(toleration: KubernetesToleration): void { + const allowedKeys = new Set(["key", "operator", "value", "effect"]); + for (const key of Object.keys(toleration)) { + if (!allowedKeys.has(key)) throw new Error(`unknown toleration field: ${key}`); + } + if (toleration.effect !== "NoSchedule") { + throw new Error("toleration effect must be NoSchedule in this release"); + } + if (toleration.operator === "Exists" && toleration.value !== undefined) { + throw new Error("toleration operator Exists must not set a value"); + } + if (!toleration.key && toleration.operator !== "Exists") { + throw new Error("an empty or omitted toleration key requires operator Exists"); + } +} diff --git a/packages/drivers/src/kubernetes.test.ts b/packages/drivers/src/kubernetes.test.ts index c4bc410..606689f 100644 --- a/packages/drivers/src/kubernetes.test.ts +++ b/packages/drivers/src/kubernetes.test.ts @@ -26,9 +26,10 @@ async function fakeKubectl(): Promise<{ bin: string; log: string }> { const directory = await mkdtemp(join(tmpdir(), "pocketcoder-kubectl-test-")); temporaryDirectories.push(directory); const log = join(directory, "calls.ndjson"); - const bin = join(directory, "kubectl"); + const script = join(directory, "kubectl.ts"); + const bin = process.platform === "win32" ? join(directory, "kubectl.cmd") : script; await writeFile( - bin, + script, `#!/usr/bin/env bun import { appendFileSync } from "node:fs"; const args = process.argv.slice(2); @@ -40,6 +41,9 @@ if (args.includes("version")) console.log(JSON.stringify({ serverVersion: { majo `, { mode: 0o755 }, ); + if (process.platform === "win32") { + await writeFile(bin, `@${JSON.stringify(process.execPath)} ${JSON.stringify(script)} %*\r\n`); + } return { bin, log }; } diff --git a/packages/drivers/src/kubernetes.ts b/packages/drivers/src/kubernetes.ts index d122a0b..7445692 100644 --- a/packages/drivers/src/kubernetes.ts +++ b/packages/drivers/src/kubernetes.ts @@ -12,6 +12,11 @@ import { isKubernetesName, kubectl, resourceName } from "./kubernetes-command"; import { discoveredWarmRuntimes, discoveredWorkspaces } from "./kubernetes-discovery"; import { KUBERNETES_POOL_LABEL, KUBERNETES_WORKSPACE_LABEL } from "./kubernetes-labels"; import { warmJobManifest, workspaceJobManifest } from "./kubernetes-manifests"; +import { + type KubernetesSchedulingOptions, + type KubernetesToleration, + validateToleration, +} from "./kubernetes-scheduling"; export { KUBERNETES_DIGEST_ANNOTATION, @@ -19,7 +24,7 @@ export { KUBERNETES_WORKSPACE_LABEL, } from "./kubernetes-labels"; -export interface KubernetesDriverOptions extends EgressDriverOptions { +export interface KubernetesDriverOptions extends EgressDriverOptions, KubernetesSchedulingOptions { namespace?: string; kubectlBin?: string; serviceAccountName?: string; @@ -31,6 +36,8 @@ export class KubernetesDriver implements WorkspaceDriver { private readonly namespace: string; private readonly kubectlBin: string; private readonly serviceAccountName: string | undefined; + private readonly nodeSelector: Record | undefined; + private readonly tolerations: KubernetesToleration[] | undefined; private readonly imagePullPolicy: "Always" | "IfNotPresent" | "Never"; private readonly egress: EgressDriverOptions; private sidecarsSupported = false; @@ -45,6 +52,9 @@ export class KubernetesDriver implements WorkspaceDriver { if (this.serviceAccountName && !isKubernetesName(this.serviceAccountName)) { throw new Error("serviceAccountName must be a Kubernetes resource name"); } + this.nodeSelector = options.nodeSelector; + this.tolerations = options.tolerations; + for (const toleration of this.tolerations ?? []) validateToleration(toleration); this.imagePullPolicy = options.imagePullPolicy ?? "IfNotPresent"; this.egress = { ...(options.egressImage ? { egressImage: options.egressImage } : {}), @@ -109,6 +119,8 @@ export class KubernetesDriver implements WorkspaceDriver { const manifest = workspaceJobManifest(launch, name, inputSecret, egressSecret, { serviceAccountName: this.serviceAccountName, + nodeSelector: this.nodeSelector, + tolerations: this.tolerations, imagePullPolicy: this.imagePullPolicy, egressImage: this.egress.egressImage, }); @@ -187,6 +199,8 @@ export class KubernetesDriver implements WorkspaceDriver { } const manifest = warmJobManifest(launch, name, inputSecret, egressSecret, { serviceAccountName: this.serviceAccountName, + nodeSelector: this.nodeSelector, + tolerations: this.tolerations, imagePullPolicy: this.imagePullPolicy, egressImage: this.egress.egressImage, }); diff --git a/packages/server/src/config.test.ts b/packages/server/src/config.test.ts index aa25cf7..6efb9b1 100644 --- a/packages/server/src/config.test.ts +++ b/packages/server/src/config.test.ts @@ -40,6 +40,8 @@ describe("portable persistence configuration", () => { POCKETCODER_DRIVER: "kubernetes", POCKETCODER_KUBERNETES_NAMESPACE: "agents", POCKETCODER_KUBERNETES_SERVICE_ACCOUNT: "workspace", + POCKETCODER_KUBERNETES_NODE_SELECTOR: '{"dedicated":"workspace"}', + POCKETCODER_KUBERNETES_TOLERATIONS: '[{"operator":"Exists","effect":"NoSchedule"}]', POCKETCODER_STORAGE_BACKEND: "kubernetes-pvc", POCKETCODER_KUBERNETES_WORKSPACE_CLAIM: "workspace-data", POCKETCODER_WORKSPACE_DATA_DIR: "/data/workspaces", @@ -49,6 +51,8 @@ describe("portable persistence configuration", () => { expect(config.driverKind).toBe("kubernetes"); expect(config.storageBackend).toBe("kubernetes-pvc"); expect(config.secretProvider).toBe("kubernetes"); + expect(config.kubernetesNodeSelector).toEqual({ dedicated: "workspace" }); + expect(config.kubernetesTolerations).toEqual([{ operator: "Exists", effect: "NoSchedule" }]); expect(config.workspaceServerUrl).toBe("http://pocketcoder-server.agents.svc:7080"); }); diff --git a/packages/server/src/config.ts b/packages/server/src/config.ts index b3f67f0..2b6076d 100644 --- a/packages/server/src/config.ts +++ b/packages/server/src/config.ts @@ -1,10 +1,12 @@ import { randomBytes } from "node:crypto"; import { parseDurationMs } from "@pstdio/pocketcoder-contracts"; +import type { KubernetesToleration } from "@pstdio/pocketcoder-drivers"; import { type AdmissionLimits, DEFAULT_LIMITS, type WarmPoolConfigEntry, } from "@pstdio/pocketcoder-runtime-core"; +import { resolveKubernetesScheduling } from "./kubernetes-scheduling-config"; import { DEFAULT_PERSISTENCE_LIMITS, type PersistenceLimits } from "./persistence"; type Environment = Record; @@ -33,6 +35,8 @@ export interface ServerConfig { secretRoot: string | null; kubernetesNamespace: string; kubernetesServiceAccount: string | null; + kubernetesNodeSelector: Record | null; + kubernetesTolerations: KubernetesToleration[]; kubernetesWorkspaceClaim: string | null; kubernetesWorkspaceSubPath: string; // URL workspaces use to reach this server; with the Docker driver on a @@ -317,6 +321,7 @@ export function loadConfig(env: Environment = process.env): ServerConfig { "docker", ); const kubernetesNamespace = env.POCKETCODER_KUBERNETES_NAMESPACE ?? "default"; + const scheduling = resolveKubernetesScheduling(env); const storage = resolveStorage(env); const secrets = resolveSecrets(env); assertCompatibleBackends(driverKind, storage, secrets); @@ -341,6 +346,8 @@ export function loadConfig(env: Environment = process.env): ServerConfig { ...secrets, kubernetesNamespace, kubernetesServiceAccount: env.POCKETCODER_KUBERNETES_SERVICE_ACCOUNT ?? null, + kubernetesNodeSelector: scheduling.nodeSelector ?? null, + kubernetesTolerations: scheduling.tolerations ?? [], kubernetesWorkspaceSubPath: env.POCKETCODER_KUBERNETES_WORKSPACE_SUBPATH ?? "workspaces", workspaceServerUrl: httpUrlEnv( env, diff --git a/packages/server/src/kubernetes-scheduling-config.test.ts b/packages/server/src/kubernetes-scheduling-config.test.ts new file mode 100644 index 0000000..f534751 --- /dev/null +++ b/packages/server/src/kubernetes-scheduling-config.test.ts @@ -0,0 +1,99 @@ +import { describe, expect, test } from "bun:test"; +import { resolveKubernetesScheduling } from "./kubernetes-scheduling-config"; + +describe("Kubernetes scheduling configuration", () => { + test("defaults to no scheduling fields", () => { + expect(resolveKubernetesScheduling({})).toEqual({}); + }); + + test("parses a node selector and tolerations", () => { + expect( + resolveKubernetesScheduling({ + POCKETCODER_KUBERNETES_NODE_SELECTOR: JSON.stringify({ + "onefin.com/workload": "agent-workspace", + }), + POCKETCODER_KUBERNETES_TOLERATIONS: JSON.stringify([ + { + key: "onefin.com/workload", + operator: "Equal", + value: "agent-workspace", + effect: "NoSchedule", + }, + ]), + }), + ).toEqual({ + nodeSelector: { "onefin.com/workload": "agent-workspace" }, + tolerations: [ + { + key: "onefin.com/workload", + operator: "Equal", + value: "agent-workspace", + effect: "NoSchedule", + }, + ], + }); + }); + + test("rejects malformed JSON", () => { + expect(() => + resolveKubernetesScheduling({ POCKETCODER_KUBERNETES_NODE_SELECTOR: "{" }), + ).toThrow("POCKETCODER_KUBERNETES_NODE_SELECTOR must be valid JSON"); + expect(() => resolveKubernetesScheduling({ POCKETCODER_KUBERNETES_TOLERATIONS: "[" })).toThrow( + "POCKETCODER_KUBERNETES_TOLERATIONS must be valid JSON", + ); + }); + + test("rejects the wrong top-level JSON shapes", () => { + expect(() => + resolveKubernetesScheduling({ POCKETCODER_KUBERNETES_NODE_SELECTOR: "[]" }), + ).toThrow("POCKETCODER_KUBERNETES_NODE_SELECTOR must be a JSON object"); + expect(() => resolveKubernetesScheduling({ POCKETCODER_KUBERNETES_TOLERATIONS: "{}" })).toThrow( + "POCKETCODER_KUBERNETES_TOLERATIONS must be a JSON array", + ); + }); + + test("rejects non-string selector values", () => { + expect(() => + resolveKubernetesScheduling({ + POCKETCODER_KUBERNETES_NODE_SELECTOR: JSON.stringify({ dedicated: true }), + }), + ).toThrow("POCKETCODER_KUBERNETES_NODE_SELECTOR.dedicated must be a string"); + }); + + test("rejects invalid toleration field types", () => { + expect(() => + resolveKubernetesScheduling({ + POCKETCODER_KUBERNETES_TOLERATIONS: JSON.stringify(["workspace"]), + }), + ).toThrow("POCKETCODER_KUBERNETES_TOLERATIONS[0] must be an object"); + expect(() => + resolveKubernetesScheduling({ + POCKETCODER_KUBERNETES_TOLERATIONS: JSON.stringify([{ key: true, effect: "NoSchedule" }]), + }), + ).toThrow("POCKETCODER_KUBERNETES_TOLERATIONS[0].key must be a string"); + expect(() => + resolveKubernetesScheduling({ + POCKETCODER_KUBERNETES_TOLERATIONS: JSON.stringify([ + { operator: "Unknown", effect: "NoSchedule" }, + ]), + }), + ).toThrow("POCKETCODER_KUBERNETES_TOLERATIONS[0].operator must be Exists or Equal"); + }); + + test("applies every semantic toleration check", () => { + const invalid = [ + [{ operator: "Exists", value: "workspace", effect: "NoSchedule" }, "must not set a value"], + [{ operator: "Equal", effect: "NoSchedule" }, "requires operator Exists"], + [{ effect: "NoSchedule", unknown: true }, "unknown toleration field"], + [{ operator: "Exists", effect: "NoExecute" }, "effect must be NoSchedule"], + ] as const; + + for (const [toleration, message] of invalid) { + expect(() => + resolveKubernetesScheduling({ + POCKETCODER_KUBERNETES_TOLERATIONS: JSON.stringify([toleration]), + }), + ).toThrow(message); + } + }); +}); diff --git a/packages/server/src/kubernetes-scheduling-config.ts b/packages/server/src/kubernetes-scheduling-config.ts new file mode 100644 index 0000000..ff623f5 --- /dev/null +++ b/packages/server/src/kubernetes-scheduling-config.ts @@ -0,0 +1,68 @@ +import { + type KubernetesSchedulingOptions, + type KubernetesToleration, + validateToleration, +} from "@pstdio/pocketcoder-drivers"; + +type Environment = Record; + +function parseJson(raw: string, name: string): unknown { + try { + return JSON.parse(raw); + } catch { + throw new Error(`${name} must be valid JSON`); + } +} + +function parseNodeSelector(raw: string | undefined): Record | undefined { + if (!raw) return undefined; + const name = "POCKETCODER_KUBERNETES_NODE_SELECTOR"; + const value = parseJson(raw, name); + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error(`${name} must be a JSON object`); + } + const selector = value as Record; + for (const [key, item] of Object.entries(selector)) { + if (typeof item !== "string") throw new Error(`${name}.${key} must be a string`); + } + return selector as Record; +} + +function optionalString(entry: Record, key: string, name: string) { + const value = entry[key]; + if (value !== undefined && typeof value !== "string") { + throw new Error(`${name}.${key} must be a string`); + } +} + +function parseTolerations(raw: string | undefined): KubernetesToleration[] { + if (!raw) return []; + const name = "POCKETCODER_KUBERNETES_TOLERATIONS"; + const value = parseJson(raw, name); + if (!Array.isArray(value)) throw new Error(`${name} must be a JSON array`); + return value.map((item, index) => { + const itemName = `${name}[${index}]`; + if (!item || typeof item !== "object" || Array.isArray(item)) { + throw new Error(`${itemName} must be an object`); + } + const entry = item as Record; + for (const key of ["key", "operator", "value", "effect"]) { + optionalString(entry, key, itemName); + } + if (entry.operator !== undefined && entry.operator !== "Exists" && entry.operator !== "Equal") { + throw new Error(`${itemName}.operator must be Exists or Equal`); + } + const toleration = entry as unknown as KubernetesToleration; + validateToleration(toleration); + return toleration; + }); +} + +export function resolveKubernetesScheduling(env: Environment): KubernetesSchedulingOptions { + const nodeSelector = parseNodeSelector(env.POCKETCODER_KUBERNETES_NODE_SELECTOR); + const tolerations = parseTolerations(env.POCKETCODER_KUBERNETES_TOLERATIONS); + return { + ...(nodeSelector ? { nodeSelector } : {}), + ...(tolerations.length ? { tolerations } : {}), + }; +} diff --git a/packages/server/src/lifecycle.ts b/packages/server/src/lifecycle.ts index ea26300..bcae396 100644 --- a/packages/server/src/lifecycle.ts +++ b/packages/server/src/lifecycle.ts @@ -84,6 +84,8 @@ function createWorkspaceDriver(config: ServerConfig) { return new KubernetesDriver({ ...egress, namespace: config.kubernetesNamespace, + nodeSelector: config.kubernetesNodeSelector ?? undefined, + tolerations: config.kubernetesTolerations, ...(config.kubernetesServiceAccount ? { serviceAccountName: config.kubernetesServiceAccount } : {}),