fix(services): start ml jobs with the corvus-job entry point

docs/reference/control-plane-api.md

@@ -225,8 +225,8 @@ The runner is selected with `CORVUS_API_JOB_RUNNER`:
 
 | Runner | Behaviour |
 |---|---|
-| `local` | Starts `CORVUS_API_JOB_LOCAL_COMMAND` (JSON list, default `["corvus-ml"]`) plus the job type as a subprocess without a shell. The child does not inherit `CORVUS_API_*`, `CORVUS_PSEUDONYM_KEY*` or `CORVUS_KEK`. Process handles are in memory: after an API restart a running job is reported `failed` ("process handle lost"). |
-| `kubernetes` | Creates a `batch/v1` Job through the official Python client (in-cluster config, else kubeconfig) in `CORVUS_API_JOB_K8S_NAMESPACE`: image, command, `CPU`/`memory` requests and limits (`CORVUS_API_JOB_K8S_*`), labels `app.kubernetes.io/name=corvus-job` and `corvus.onelitefeather.net/job-id`/`job-type`, `ttlSecondsAfterFinished`, `activeDeadlineSeconds`, `backoffLimit` 0, `restartPolicy: Never`, non-root, no privilege escalation, all capabilities dropped. `CORVUS_API_JOB_K8S_ENV_SECRET` names a Secret exposed as environment. Status comes from the Job's `Complete` and `Failed` conditions. The API's service account needs `create` and `get` on `jobs` in that namespace. |
+| `local` | Starts `CORVUS_API_JOB_LOCAL_COMMAND` (JSON list, default `["corvus-job"]`) plus the job type as a subprocess without a shell. The child does not inherit `CORVUS_API_*`, `CORVUS_PSEUDONYM_KEY*` or `CORVUS_KEK`. Process handles are in memory: after an API restart a running job is reported `failed` ("process handle lost"). |
+| `kubernetes` | Creates a `batch/v1` Job through the official Python client (in-cluster config, else kubeconfig) in `CORVUS_API_JOB_K8S_NAMESPACE`: image, command, `CPU`/`memory` requests and limits (`CORVUS_API_JOB_K8S_*`), labels `app.kubernetes.io/name=corvus-job` and `corvus.onelitefeather.net/job-id`/`job-type`, `ttlSecondsAfterFinished`, `activeDeadlineSeconds`, `backoffLimit` 0, a sized `emptyDir` at `/scratch` (`CORVUS_API_JOB_K8S_SCRATCH_SIZE_LIMIT`, default `20Gi`; `CORVUS_JOB_SCRATCH_DIR` and `TMPDIR` point there), `restartPolicy: Never`, non-root, no privilege escalation, all capabilities dropped. `CORVUS_API_JOB_K8S_ENV_SECRET` names a Secret exposed as environment. Status comes from the Job's `Complete` and `Failed` conditions. The API's service account needs `create` and `get` on `jobs` in that namespace. |
 
 Contract with the job process, identical for both runners: arguments are the command followed by the job type; the
 environment carries `CORVUS_JOB_ID`, `CORVUS_JOB_TYPE`, `CORVUS_JOB_SNAPSHOT_ID` and `CORVUS_JOB_PARAMS` (JSON).

services/src/corvus_services/api/jobs.py

@@ -168,6 +168,10 @@ class K8sJobConfig:
     backoff_limit: int
     service_account: str | None = None
     env_secret: str | None = None
+    scratch_size_limit: str = "20Gi"
+
+
+SCRATCH_MOUNT = "/scratch"
 
 
 class KubernetesRunner:
@@ -198,7 +202,15 @@ def build_job(self, spec: JobSpec) -> dict[str, Any]:
             "image": c.image,
             "command": list(c.command),
             "args": [spec.type.value],
-            "env": [{"name": k, "value": v} for k, v in spec.env().items()],
+            "env": [
+                {"name": k, "value": v}
+                for k, v in {
+                    **spec.env(),
+                    "CORVUS_JOB_SCRATCH_DIR": SCRATCH_MOUNT,
+                    "TMPDIR": SCRATCH_MOUNT,
+                }.items()
+            ],
+            "volumeMounts": [{"name": "scratch", "mountPath": SCRATCH_MOUNT}],
             "resources": {
                 "requests": {"cpu": c.cpu_request, "memory": c.memory_request},
                 "limits": {"cpu": c.cpu_limit, "memory": c.memory_limit},
@@ -213,6 +225,8 @@ def build_job(self, spec: JobSpec) -> dict[str, Any]:
         pod_spec: dict[str, Any] = {
             "restartPolicy": "Never",
             "containers": [container],
+            # Feature builds can be large: scratch space on a sized emptyDir, not the image layer.
+            "volumes": [{"name": "scratch", "emptyDir": {"sizeLimit": c.scratch_size_limit}}],
             "securityContext": {"runAsNonRoot": True, "seccompProfile": {"type": "RuntimeDefault"}},
         }
         if c.service_account:
@@ -275,6 +289,7 @@ def build_runner(settings: ApiSettings) -> JobRunner:
         backoff_limit=settings.job_k8s_backoff_limit,
         service_account=settings.job_k8s_service_account,
         env_secret=settings.job_k8s_env_secret,
+        scratch_size_limit=settings.job_k8s_scratch_size_limit,
     )
     return KubernetesRunner(client.BatchV1Api(), cfg)
 

services/src/corvus_services/api/settings.py

@@ -99,10 +99,10 @@ class ApiSettings(BaseSettings):
     # Jobs
     job_runner: Literal["local", "kubernetes"] = "local"
     # Local runner: argv prefix; the job type and arguments are appended.
-    job_local_command: list[str] = ["corvus-ml"]
+    job_local_command: list[str] = ["corvus-job"]
     job_k8s_namespace: str = "corvus"
     job_k8s_image: str = ""
-    job_k8s_command: list[str] = ["corvus-ml"]
+    job_k8s_command: list[str] = ["corvus-job"]
     job_k8s_service_account: str | None = None
     job_k8s_cpu_request: str = "1"
     job_k8s_cpu_limit: str = "4"
@@ -111,6 +111,8 @@ class ApiSettings(BaseSettings):
     job_k8s_ttl_after_finished_s: int = Field(default=86_400, ge=0)
     job_k8s_active_deadline_s: int = Field(default=14_400, ge=1)
     job_k8s_backoff_limit: int = Field(default=0, ge=0)
+    # Size limit of the emptyDir mounted at /scratch (CORVUS_JOB_SCRATCH_DIR) in job pods.
+    job_k8s_scratch_size_limit: str = "20Gi"
     # Name of a Secret that is exposed as environment (S3/DB credentials of the job); optional.
     job_k8s_env_secret: str | None = None
 

services/tests/api/test_api_job_command.py

@@ -0,0 +1,50 @@
+"""The default job command must be the entry point the ml image really provides."""
+
+from __future__ import annotations
+
+import re
+import tomllib
+from pathlib import Path
+
+import pytest
+
+from corvus_services.api.settings import ApiSettings
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+ML_DIR = REPO_ROOT / "ml"
+
+
+def _settings() -> ApiSettings:
+    return ApiSettings(_env_file=None)  # type: ignore[call-arg]
+
+
+def test_default_kubernetes_job_command_is_corvus_job():
+    assert _settings().job_k8s_command == ["corvus-job"]
+
+
+def test_default_local_job_command_is_corvus_job():
+    assert _settings().job_local_command == ["corvus-job"]
+
+
+@pytest.mark.skipif(
+    not (ML_DIR / "pyproject.toml").is_file(),
+    reason="ml/ is not checked out next to services/ (standalone services run)",
+)
+def test_default_job_commands_are_console_scripts_of_the_ml_package():
+    scripts = tomllib.loads((ML_DIR / "pyproject.toml").read_text())["project"]["scripts"]
+    settings = _settings()
+
+    assert settings.job_k8s_command[0] in scripts
+    assert settings.job_local_command[0] in scripts
+
+
+@pytest.mark.skipif(
+    not (ML_DIR / "Dockerfile").is_file(),
+    reason="ml/ is not checked out next to services/ (standalone services run)",
+)
+def test_default_kubernetes_command_equals_the_ml_image_entrypoint():
+    match = re.search(r"^ENTRYPOINT\s+(\[.*\])\s*$", (ML_DIR / "Dockerfile").read_text(), re.M)
+    assert match, "ml/Dockerfile has no exec-form ENTRYPOINT"
+    entrypoint = tomllib.loads(f"v = {match.group(1)}")["v"]
+
+    assert _settings().job_k8s_command == entrypoint

services/tests/api/test_api_jobs.py

@@ -52,12 +52,12 @@ def launch(self, argv, env):
 
 async def test_local_runner_starts_the_configured_command_with_the_job_contract():
     launcher = FakeLauncher()
-    runner = LocalRunner(["corvus-ml", "--quiet"], launcher)
+    runner = LocalRunner(["corvus-job", "--quiet"], launcher)
 
     ref = await runner.submit(SPEC)
 
     argv, env = launcher.calls[0]
-    assert argv == ["corvus-ml", "--quiet", "train"]
+    assert argv == ["corvus-job", "--quiet", "train"]
     assert env == {
         "CORVUS_JOB_ID": str(JOB_ID),
         "CORVUS_JOB_TYPE": "train",
@@ -136,7 +136,7 @@ def read_namespaced_job_status(self, name, namespace):
 CONFIG = K8sJobConfig(
     namespace="corvus-ml",
     image="ghcr.io/onelitefeathernet/corvus-ml:1.2.3",
-    command=["corvus-ml"],
+    command=["corvus-job"],
     cpu_request="500m",
     cpu_limit="2",
     memory_request="1Gi",
@@ -175,7 +175,7 @@ async def test_kubernetes_runner_creates_a_batch_job_with_limits_labels_and_ttl(
     assert pod["securityContext"]["runAsNonRoot"] is True
     container = pod["containers"][0]
     assert container["image"] == "ghcr.io/onelitefeathernet/corvus-ml:1.2.3"
-    assert container["command"] == ["corvus-ml"]
+    assert container["command"] == ["corvus-job"]
     assert container["args"] == ["train"]
     assert container["resources"] == {
         "requests": {"cpu": "500m", "memory": "1Gi"},
@@ -466,3 +466,25 @@ def create_namespaced_job(self, namespace, body):
 
     assert seen
     assert seen[0] != threading.get_ident()
+
+
+def test_kubernetes_job_mounts_a_sized_scratch_volume_and_points_the_job_at_it():
+    cfg = K8sJobConfig(**{**CONFIG.__dict__, "scratch_size_limit": "5Gi"})
+
+    pod = KubernetesRunner(FakeBatchApi(), cfg).build_job(SPEC)["spec"]["template"]["spec"]
+
+    assert pod["volumes"] == [{"name": "scratch", "emptyDir": {"sizeLimit": "5Gi"}}]
+    container = pod["containers"][0]
+    assert container["volumeMounts"] == [{"name": "scratch", "mountPath": "/scratch"}]
+    env = {e["name"]: e["value"] for e in container["env"]}
+    assert env["CORVUS_JOB_SCRATCH_DIR"] == "/scratch"
+    assert env["TMPDIR"] == "/scratch"
+
+
+def test_kubernetes_job_passes_the_job_type_as_first_argument_and_in_the_environment():
+    container = KubernetesRunner(FakeBatchApi(), CONFIG).build_job(SPEC)["spec"]["template"][
+        "spec"
+    ]["containers"][0]
+
+    env = {e["name"]: e["value"] for e in container["env"]}
+    assert container["args"][0] == env["CORVUS_JOB_TYPE"] == "train"