add a dedicated clean subcommand/GPU support

examples/config.yaml

@@ -79,10 +79,9 @@ job:
 # selected with --profile NAME; deep-merges over the rest of this config
 profiles:
   gpu:
-    defaults: {cpus: "8", memory: 32g}
+    # gpus sets the nvidia.com/gpu request+limit and runtimeClassName: nvidia; override those
+    # names with gpu_resource / gpu_runtime_class (null disables the runtime class)
+    defaults: {cpus: "8", memory: 32g, gpus: 1}
     kubernetes:
       spec:
         nodeSelector: {"nvidia.com/gpu.present": "true"}
-        containers:
-          - name: main
-            resources: {limits: {"nvidia.com/gpu": 1}}

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "squarepeg"
-version = "0.3.0"
+version = "0.4.0"
 description = "Run a docker-run-style command as a Kubernetes Pod or Job, streaming its output like a local process"
 readme = "README.md"
 requires-python = ">=3.11"

squarepeg/cli.py

@@ -15,15 +15,16 @@
     parse_env_entries,
     reject_unsupported,
 )
-from squarepeg.errors import SquarepegError, UsageError
+from squarepeg.errors import RunnerError, SquarepegError, UsageError
 from squarepeg.k8s.dryrun import server_dry_run
+from squarepeg.k8s.orphans import DEFAULT_MIN_AGE_SECONDS, clean_orphans
 from squarepeg.k8s.runner import run_manifest
 from squarepeg.k8s.session import Session
 from squarepeg.labels import RUN_ID_LABEL
 from squarepeg.log import chatter
 from squarepeg.manifest import build_job, build_pod
 from squarepeg.naming import generate_name, validate_rfc1123
-from squarepeg.quantities import docker_cpus_to_k8s, docker_memory_to_k8s
+from squarepeg.quantities import docker_cpus_to_k8s, docker_gpus_to_k8s, docker_memory_to_k8s
 from squarepeg.runspec import ResourceSpec, RunSpec
 from squarepeg.volumes import auto_mounts_from_config, parse_volume_flag
 
@@ -64,7 +65,6 @@ def _unsupported_callback(ctx, param, value):
 _UNSUPPORTED_GROUPS: list[tuple[list[str], bool]] = [
     (["--network"], False),
     (["--privileged"], True),
-    (["--gpus"], False),
     (["-p", "--publish"], False),
     (["-d", "--detach"], True),
     (["--restart"], False),
@@ -154,6 +154,7 @@ def cli():
 )
 @click.option("--cpus", "cpus", default=None)
 @click.option("-m", "--memory", "memory", default=None)
+@click.option("--gpus", "gpus", default=None, help="GPU count (sets both request and limit)")
 @click.option("--request-cpu", "request_cpu", default=None)
 @click.option("--limit-cpu", "limit_cpu", default=None)
 @click.option("--request-memory", "request_memory", default=None)
@@ -198,6 +199,7 @@ def run(
     pull,
     cpus,
     memory,
+    gpus,
     request_cpu,
     limit_cpu,
     request_memory,
@@ -300,6 +302,22 @@ def run(
     if cli_set_memory and limit_memory is None:
         claims.add(f"/spec/containers/[name={container_name}]/resources/limits/memory")
 
+    if "gpu_resource" in config_defaults:
+        resources.gpu_resource = str(config_defaults["gpu_resource"])
+    if "gpu_runtime_class" in config_defaults:
+        # null/"" disables it, for clusters where the NVIDIA runtime is the node default
+        resources.gpu_runtime_class = config_defaults["gpu_runtime_class"] or None
+    cli_set_gpus = gpus is not None
+    if gpus is None and "gpus" in config_defaults:
+        gpus = config_defaults["gpus"]
+    if gpus is not None:
+        resources.gpus = docker_gpus_to_k8s(gpus)
+    if cli_set_gpus:
+        for kind in ("requests", "limits"):
+            claims.add(
+                f"/spec/containers/[name={container_name}]/resources/{kind}/{resources.gpu_resource}"
+            )
+
     if name is not None:
         validate_rfc1123(name, what="--name")
         claims.add("/metadata/name")
@@ -343,7 +361,7 @@ def run(
         resolved_config.get("orphan_sweep", True), "orphan_sweep"
     )
     orphan_sweep_min_age = coerce_int(
-        resolved_config.get("orphan_sweep_min_age", 300), "orphan_sweep_min_age"
+        resolved_config.get("orphan_sweep_min_age", DEFAULT_MIN_AGE_SECONDS), "orphan_sweep_min_age"
     )
     if orphan_sweep_min_age < 0:
         raise UsageError(
@@ -432,6 +450,97 @@ def run(
     raise SystemExit(exit_code)
 
 
+def _format_age(seconds: float) -> str:
+    """kubectl-style compact duration: the two largest non-zero units, e.g. '10m12s',
+    '3h14m', '2d4h'. Used only for --dry-run's human-readable listing."""
+    seconds = int(seconds)
+    if seconds < 60:
+        return f"{seconds}s"
+    minutes, secs = divmod(seconds, 60)
+    if minutes < 60:
+        return f"{minutes}m{secs}s" if secs else f"{minutes}m"
+    hours, mins = divmod(minutes, 60)
+    if hours < 24:
+        return f"{hours}h{mins}m" if mins else f"{hours}h"
+    days, hrs = divmod(hours, 24)
+    return f"{days}d{hrs}h" if hrs else f"{days}d"
+
+
+def _candidate_detail(candidate) -> str:
+    age = _format_age(candidate.age_seconds)
+    if candidate.category == "terminal":
+        return f"finished {age} ago"
+    return f"{candidate.reason} for {age}"
+
+
+@cli.command("clean")
+@click.option("--dry-run", "dry_run", is_flag=True, help="list what would be deleted without deleting anything")
+@click.option(
+    "--min-age",
+    "min_age",
+    type=int,
+    default=None,
+    help="seconds a resource must be eligible before it's cleaned (default: config orphan_sweep_min_age, else 300)",
+)
+@click.option("-n", "--namespace", "namespace", default=None)
+@click.option("--context", "context", default=None, help="kubeconfig context to use")
+@click.option("--quiet", "quiet", is_flag=True)
+@click.option(
+    "--no-color",
+    "no_color",
+    is_flag=True,
+    help="disable colour/animation in squarepeg's own stderr output",
+)
+@_add_config_options
+def clean(dry_run, min_age, namespace, context, quiet, no_color, config_paths, no_default_config, profile):
+    """Delete this user's own orphaned Pods/Jobs left behind by crashed or killed runs.
+
+    Scoped to your own resources only (never a teammate's), and --keep'd resources are
+    always exempt -- there is no flag to override either. Runs regardless of the
+    'orphan_sweep' config key, which only affects the automatic per-run sweep. There is
+    no confirmation prompt: the candidates are, by construction, either already finished
+    or stuck and unable to start, so nothing running is ever at risk -- use --dry-run to
+    preview what would be deleted first if you want to check.
+    """
+    ui.set_color_override(False if no_color else None)
+    ui.set_run_tag(secrets.token_hex(4))
+
+    resolved_config, _sources = _resolve_config(config_paths, no_default_config, profile)
+
+    if namespace is None:
+        namespace = resolved_config.get("namespace")
+    quiet = quiet or coerce_bool(resolved_config.get("quiet", False), "quiet")
+
+    if min_age is None:
+        min_age = coerce_int(
+            resolved_config.get("orphan_sweep_min_age", DEFAULT_MIN_AGE_SECONDS), "orphan_sweep_min_age"
+        )
+        if min_age < 0:
+            raise UsageError(f"'orphan_sweep_min_age' must not be negative, got {min_age}")
+    elif min_age < 0:
+        raise UsageError(f"--min-age must not be negative, got {min_age}")
+
+    session = Session(namespace=namespace, context=context, quiet=quiet)
+
+    with ui.Status("scanning for orphaned resources", quiet=quiet):
+        result = clean_orphans(session, min_age, quiet=quiet, dry_run=dry_run)
+
+    if dry_run:
+        if result.candidates:
+            kind_name_width = max(len(f"{c.kind}/{c.name}") for c in result.candidates)
+            category_width = max(len(c.category) for c in result.candidates)
+            for c in result.candidates:
+                kind_name = f"{c.kind}/{c.name}".ljust(kind_name_width)
+                category = c.category.ljust(category_width)
+                click.echo(f"{kind_name}  {category}  {_candidate_detail(c)}")
+        return
+
+    if result.failed:
+        raise RunnerError(
+            f"failed to delete {len(result.failed)} of {len(result.candidates)} orphaned resource(s)"
+        )
+
+
 def _print_config(config_paths, no_default_config, profile):
     ui.set_run_tag(None)  # config has no run context, even if invoked right after a 'run' in-process
     resolved_config, sources = _resolve_config(config_paths, no_default_config, profile)

squarepeg/config.py

@@ -40,7 +40,16 @@
     "job",
     "profiles",
 }
-ALLOWED_DEFAULTS_KEYS = {"cpus", "memory", "image_pull_policy", "workdir", "env"}
+ALLOWED_DEFAULTS_KEYS = {
+    "cpus",
+    "memory",
+    "gpus",
+    "gpu_resource",
+    "gpu_runtime_class",
+    "image_pull_policy",
+    "workdir",
+    "env",
+}
 
 
 def _resolve_env_paths() -> list[tuple[Path, str]]:

squarepeg/dockerargs.py

@@ -7,8 +7,6 @@
     "--network": "container networking is a cluster-level concept; use the 'kubernetes' passthrough "
     "section in your config (spec.dnsPolicy, spec.hostNetwork, etc.) if you need to control it",
     "--privileged": "use 'kubernetes.spec.containers[].securityContext.privileged' in your config",
-    "--gpus": "request GPUs via 'kubernetes.spec.containers[].resources' (e.g. 'nvidia.com/gpu: 1') "
-    "in your config",
     "-p": "there is no host to publish a port to on a remote cluster; use a Kubernetes Service instead",
     "--publish": "there is no host to publish a port to on a remote cluster; use a Kubernetes Service instead",
     "-d": "squarepeg always attaches to stream output; there is no detached mode",

squarepeg/k8s/events.py

@@ -2,18 +2,89 @@
 '0/12 nodes available: insufficient memory'.
 """
 
+import threading
+
 from kubernetes.client.rest import ApiException
 
 from squarepeg.log import chatter
 
+# how often WarningEventWatcher re-lists events while waiting for a pod to start
+EVENT_POLL_SECONDS = 3
+
 
-def print_pod_events(session, pod_name: str, *, quiet: bool = False) -> None:
+def _list_events(session, name: str) -> list:
     try:
         events = session.core.list_namespaced_event(
-            session.namespace, field_selector=f"involvedObject.name={pod_name}"
+            session.namespace, field_selector=f"involvedObject.name={name}"
         )
     except ApiException:
-        return  # best-effort diagnostics; never fail the run because event listing failed
+        return []  # best-effort diagnostics; never fail the run because event listing failed
+    return list(events.items or [])
+
+
+def _event_key(event) -> tuple[str, str]:
+    return (event.reason, event.message)
+
+
+def print_pod_events(
+    session, pod_name: str, *, quiet: bool = False, skip: frozenset[tuple[str, str]] = frozenset()
+) -> None:
+    """Print every event for `pod_name`, except those already shown (by (reason, message))."""
+    for event in _list_events(session, pod_name):
+        if _event_key(event) not in skip:
+            chatter(f"{pod_name}: {event.reason}: {event.message}", quiet=quiet, level="detail")
+
+
+class WarningEventWatcher:
+    """Context manager that, while squarepeg waits on `name` (a pod, or a job waiting to
+    create one), polls its Warning events in the background and prints each new one as
+    soon as it appears.
+
+    These are usually the real answer to "why isn't it starting?" -- an unschedulable pod,
+    a missing RuntimeClass, an exceeded quota, a volume that won't mount -- and without
+    this they'd only be printed once --timeout had already expired. squarepeg still waits
+    out the timeout rather than failing on them, since some do resolve on their own (a
+    cluster autoscaler adding a node, a PVC binding), and says so once.
+    """
+
+    def __init__(self, session, name: str, *, timeout: int, quiet: bool = False):
+        self.session = session
+        self.name = name
+        self.timeout = timeout
+        self.quiet = quiet
+        self.shown: set[tuple[str, str]] = set()
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def __enter__(self) -> "WarningEventWatcher":
+        if not self.quiet:
+            self._thread = threading.Thread(target=self._poll, daemon=True)
+            self._thread.start()
+        return self
+
+    def __exit__(self, *_exc) -> bool:
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=2)
+        return False
+
+    def _poll(self) -> None:
+        while not self._stop.wait(EVENT_POLL_SECONDS):
+            self.report_new_warnings()
 
-    for event in events.items or []:
-        chatter(f"{pod_name}: {event.reason}: {event.message}", quiet=quiet, level="detail")
+    def report_new_warnings(self) -> None:
+        for event in _list_events(self.session, self.name):
+            key = _event_key(event)
+            if event.type != "Warning" or key in self.shown:
+                continue
+            first = not self.shown
+            self.shown.add(key)
+            chatter(f"{self.name}: {event.reason}: {event.message}", quiet=self.quiet, level="warn")
+            if first:
+                chatter(
+                    f"this may not resolve on its own, but squarepeg will keep waiting until the "
+                    f"{self.timeout}s --timeout in case it does (e.g. a node scaling up); "
+                    "Ctrl+C to give up now",
+                    quiet=self.quiet,
+                    level="detail",
+                )