diff --git a/examples/config.yaml b/examples/config.yaml index dc9d20e..3dac320 100644 --- a/examples/config.yaml +++ b/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}} diff --git a/pyproject.toml b/pyproject.toml index 35b4823..b2a363a 100644 --- a/pyproject.toml +++ b/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" diff --git a/squarepeg/cli.py b/squarepeg/cli.py index 5425dea..a48517a 100644 --- a/squarepeg/cli.py +++ b/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) diff --git a/squarepeg/config.py b/squarepeg/config.py index 9f0769a..4d759da 100644 --- a/squarepeg/config.py +++ b/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]]: diff --git a/squarepeg/dockerargs.py b/squarepeg/dockerargs.py index 4584cdd..886e73a 100644 --- a/squarepeg/dockerargs.py +++ b/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", diff --git a/squarepeg/k8s/events.py b/squarepeg/k8s/events.py index ec71273..a58aa3b 100644 --- a/squarepeg/k8s/events.py +++ b/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", + ) diff --git a/squarepeg/k8s/orphans.py b/squarepeg/k8s/orphans.py index 777e3f6..15cb8d7 100644 --- a/squarepeg/k8s/orphans.py +++ b/squarepeg/k8s/orphans.py @@ -19,8 +19,16 @@ Never anything in Running/Pending-without-a-fail-fast-reason/Unknown, at any age -- that's the load-bearing safety property of the whole feature. + +`sweep_orphans` (used by runner.py's own per-run cleanup) and `clean_orphans` +(used by the standalone `squarepeg clean` command) share the same discovery +function but are deliberately separate orchestration loops: they differ in +output verbosity, error-hint text, and failure-reporting contract, and forcing +them into one function would need behaviour flags for what is otherwise a +ten-line loop. """ +from dataclasses import dataclass, field from datetime import datetime, timezone from kubernetes.client.rest import ApiException @@ -43,11 +51,27 @@ current_created_by, ) from squarepeg.log import chatter -from squarepeg.runspec import RunSpec DEFAULT_MIN_AGE_SECONDS = 300 +@dataclass(frozen=True) +class OrphanCandidate: + kind: str # "pod" | "job" + name: str + category: str # "terminal" | "stuck" + age_seconds: float + reason: str | None = None # the fail-fast waiting reason, for category == "stuck" + + +@dataclass +class CleanResult: + candidates: list[OrphanCandidate] = field(default_factory=list) + deleted: list[OrphanCandidate] = field(default_factory=list) + already_gone: list[OrphanCandidate] = field(default_factory=list) + failed: list[tuple[OrphanCandidate, str]] = field(default_factory=list) + + def _label_selector(created_by: str) -> str: return f"{MANAGED_BY_LABEL}={MANAGED_BY_VALUE},{CREATED_BY_LABEL}={created_by},!{KEEP_LABEL}" @@ -75,21 +99,25 @@ def _is_own_and_unexcluded(labels: dict, created_by: str, exclude_run_id: str | return True -def _sweepable_pod_age(pod, min_age_seconds: int, now: datetime) -> bool: - """True if this pod (already confirmed own-user/not-excluded) is old enough to sweep, - under either the terminal or the stuck-before-starting category.""" +def _pod_classification(pod, min_age_seconds: int, now: datetime) -> tuple[str, float, str | None] | None: + """(category, age_seconds, reason) for a pod old enough to sweep under either category, + or None if it doesn't qualify at all.""" phase = pod.status.phase if phase in TERMINAL_PHASES: age = terminal_age_seconds(pod, now) - elif phase == "Pending" and stuck_waiting_reason(pod) is not None: + if age is not None and age >= min_age_seconds: + return "terminal", age, None + return None + reason = stuck_waiting_reason(pod) if phase == "Pending" else None + if reason is not None: age = pod_age_seconds(pod, now) - else: - return False - return age is not None and age >= min_age_seconds + if age is not None and age >= min_age_seconds: + return "stuck", age, reason + return None def _find_pods(pods, created_by: str, min_age_seconds: int, exclude_run_id: str | None, now: datetime): - """Returns (bare_pod_candidates, stuck_child_job_names). + """Returns (bare_pod_candidates, stuck_child_job_info). Job-owned pods are never deleted directly -- deleting a Job's child pod while the Job still has a non-zero backoffLimit would make the controller spawn a replacement, i.e. @@ -97,27 +125,34 @@ def _find_pods(pods, created_by: str, min_age_seconds: int, exclude_run_id: str otherwise qualify is reported back so the caller can delete the parent Job instead, which cascades to the pod via Background propagation. """ - bare_candidates = [] - stuck_child_job_names = set() + bare_candidates: list[OrphanCandidate] = [] + stuck_child_job_info: dict[str, tuple[float, str | None]] = {} for pod in pods: labels = pod.metadata.labels or {} if not _is_own_and_unexcluded(labels, created_by, exclude_run_id): continue - if not _sweepable_pod_age(pod, min_age_seconds, now): + classification = _pod_classification(pod, min_age_seconds, now) + if classification is None: continue + category, age, reason = classification job_name = _job_name_for_pod(pod) if job_name is not None: - stuck_child_job_names.add(job_name) + stuck_child_job_info[job_name] = (age, reason) else: - bare_candidates.append(("pod", pod.metadata.name)) + bare_candidates.append(OrphanCandidate("pod", pod.metadata.name, category, age, reason)) - return bare_candidates, stuck_child_job_names + return bare_candidates, stuck_child_job_info def _find_jobs( - jobs, created_by: str, min_age_seconds: int, exclude_run_id: str | None, now: datetime, stuck_child_job_names: set + jobs, + created_by: str, + min_age_seconds: int, + exclude_run_id: str | None, + now: datetime, + stuck_child_job_info: dict, ): candidates = [] for job in jobs: @@ -128,30 +163,30 @@ def _find_jobs( name = job.metadata.name age = job_terminal_age_seconds(job, now) if age is not None and age >= min_age_seconds: - candidates.append(("job", name)) - elif name in stuck_child_job_names: + candidates.append(OrphanCandidate("job", name, "terminal", age)) + elif name in stuck_child_job_info: # no Complete/Failed condition yet (e.g. backoffLimit: 0 means a stuck pull # never produces a failure condition at all), but its child pod has been # stuck on a fail-fast reason long enough -- same verdict, different signal. - candidates.append(("job", name)) + child_age, reason = stuck_child_job_info[name] + candidates.append(OrphanCandidate("job", name, "stuck", child_age, reason)) return candidates -def find_orphaned_resources( +def find_orphan_candidates( session, created_by: str, min_age_seconds: int, *, exclude_run_id: str | None = None, now: datetime | None = None, -) -> list[tuple[str, str]]: - """Discover this user's own orphaned Pods/Jobs. +) -> list[OrphanCandidate]: + """Discover this user's own orphaned Pods/Jobs, with classification detail. Server-side scoped to app.kubernetes.io/managed-by=squarepeg, squarepeg.io/created-by matching `created_by`, and NOT carrying squarepeg.io/keep -- so a teammate's resource or a deliberately --keep'd one is never even returned by the apiserver. Client-side filtered - by phase/waiting-reason and age. Returns [(kind, name), ...]; deletion is identical for - every entry regardless of which category matched. + by phase/waiting-reason and age. `now` defaults to the real current time; tests pass a fixed value. """ @@ -161,15 +196,38 @@ def find_orphaned_resources( pods = session.core.list_namespaced_pod(session.namespace, label_selector=selector).items jobs = session.batch.list_namespaced_job(session.namespace, label_selector=selector).items - bare_pod_candidates, stuck_child_job_names = _find_pods(pods, created_by, min_age_seconds, exclude_run_id, now) - job_candidates = _find_jobs(jobs, created_by, min_age_seconds, exclude_run_id, now, stuck_child_job_names) + bare_pod_candidates, stuck_child_job_info = _find_pods(pods, created_by, min_age_seconds, exclude_run_id, now) + job_candidates = _find_jobs(jobs, created_by, min_age_seconds, exclude_run_id, now, stuck_child_job_info) return bare_pod_candidates + job_candidates +def find_orphaned_resources( + session, + created_by: str, + min_age_seconds: int, + *, + exclude_run_id: str | None = None, + now: datetime | None = None, +) -> list[tuple[str, str]]: + """(kind, name) pairs only -- a thin wrapper over find_orphan_candidates for callers + that don't need the classification detail. Deletion is identical for every entry + regardless of which category matched.""" + candidates = find_orphan_candidates( + session, created_by, min_age_seconds, exclude_run_id=exclude_run_id, now=now + ) + return [(c.kind, c.name) for c in candidates] + + def sweep_orphans( - session, spec: RunSpec, *, exclude_run_id: str | None = None, now: datetime | None = None + session, + min_age_seconds: int, + *, + quiet: bool = False, + exclude_run_id: str | None = None, + now: datetime | None = None, ) -> int: - """Discover and delete this user's own orphaned resources. + """Discover and delete this user's own orphaned resources, as a side effect of a run's + own cleanup (see runner.py's finally block). Best-effort throughout: a per-resource delete failure is logged and skipped rather than raised, so one undeletable orphan never aborts the rest of the sweep. The caller @@ -179,7 +237,7 @@ def sweep_orphans( created_by = current_created_by() try: candidates = find_orphaned_resources( - session, created_by, spec.orphan_sweep_min_age, exclude_run_id=exclude_run_id, now=now + session, created_by, min_age_seconds, exclude_run_id=exclude_run_id, now=now ) except ApiException as exc: wrapped = wrap_api_exception(exc, "failed to list resources for orphan sweep") @@ -191,8 +249,76 @@ def sweep_orphans( if delete_resource(session, kind, name): deleted += 1 except Exception as exc: # one bad delete must not abort the rest of the sweep - chatter(f"orphan sweep: failed to delete {kind} {name!r}: {exc}", quiet=spec.quiet, level="warn") + chatter(f"orphan sweep: failed to delete {kind} {name!r}: {exc}", quiet=quiet, level="warn") if deleted: - chatter(f"swept {deleted} orphaned resource(s) from previous runs", quiet=spec.quiet, level="success") + chatter(f"swept {deleted} orphaned resource(s) from previous runs", quiet=quiet, level="success") return deleted + + +def clean_orphans( + session, + min_age_seconds: int, + *, + quiet: bool = False, + dry_run: bool = False, + now: datetime | None = None, +) -> CleanResult: + """Discover and (unless dry_run) delete this user's own orphaned resources, for the + standalone `squarepeg clean` command. + + Unlike sweep_orphans, per-resource failures are recorded in the returned CleanResult + rather than only logged -- clean has no run whose own exit code must stay independent + of sweep failures, so the caller (cli.py) uses the failure list to decide its own exit + code. Never excludes a "current run", since there is no current run to exclude. + """ + created_by = current_created_by() + try: + candidates = find_orphan_candidates(session, created_by, min_age_seconds, now=now) + except ApiException as exc: + wrapped = wrap_api_exception(exc, "failed to list resources for cleanup") + raise ApiError(f"{wrapped} (the clean command needs 'list' on pods and jobs in the namespace)") from exc + + result = CleanResult(candidates=candidates) + if not candidates: + chatter("nothing to clean", quiet=quiet, level="info") + return result + if dry_run: + chatter( + f"would delete {len(candidates)} orphaned resource(s) (dry run; nothing was deleted)", + quiet=quiet, + level="info", + ) + return result + + chatter(f"found {len(candidates)} orphaned resource(s)", quiet=quiet, level="info") + for candidate in candidates: + try: + if delete_resource(session, candidate.kind, candidate.name): + result.deleted.append(candidate) + chatter( + f"deleted {candidate.kind} {candidate.name!r} " + f"({candidate.category}, {candidate.age_seconds:.0f}s)", + quiet=quiet, + level="success", + ) + else: + result.already_gone.append(candidate) + chatter(f"{candidate.kind} {candidate.name!r} was already gone", quiet=quiet, level="detail") + except Exception as exc: # one bad delete must not abort the rest of the cleanup + # delete_resource's own ApiError message already reads "failed to delete + # {kind} {name!r}: ..." -- don't wrap it a second time. + result.failed.append((candidate, str(exc))) + chatter(str(exc), quiet=quiet, level="warn") + + summary = f"cleaned {len(result.deleted)} of {len(candidates)} orphaned resource(s)" + extras = [] + if result.already_gone: + extras.append(f"{len(result.already_gone)} already gone") + if result.failed: + extras.append(f"{len(result.failed)} failed") + if extras: + summary += f" ({', '.join(extras)})" + chatter(summary, quiet=quiet, level="warn" if result.failed else "success") + + return result diff --git a/squarepeg/k8s/runner.py b/squarepeg/k8s/runner.py index 07291bf..aeffaf9 100644 --- a/squarepeg/k8s/runner.py +++ b/squarepeg/k8s/runner.py @@ -2,8 +2,10 @@ extract its exit code, and clean up. """ +import math import signal import threading +import time from kubernetes import watch from kubernetes.client.rest import ApiException @@ -11,7 +13,7 @@ from squarepeg import ui from squarepeg.errors import InterruptError, RunnerError from squarepeg.k8s.apierrors import delete_resource, wrap_api_exception -from squarepeg.k8s.events import print_pod_events +from squarepeg.k8s.events import WarningEventWatcher, print_pod_events from squarepeg.k8s.logs import stream_logs from squarepeg.k8s.orphans import sweep_orphans from squarepeg.k8s.podstate import FAIL_FAST_REASONS @@ -26,6 +28,11 @@ # has already started; the loop simply re-watches if this elapses. TERMINAL_WATCH_POLL_SECONDS = 5 +# the startup waits (job creating its pod, pod scheduling/pulling) are bounded by --timeout, but +# still watch in slices this long: Ctrl+C only sets stop_event, so a single --timeout-long watch +# would just resume its blocked read and not notice the interrupt until the timeout expired. +STARTUP_WATCH_SLICE_SECONDS = 2 + def create_resource(session, spec: RunSpec, manifest: dict) -> None: name = manifest["metadata"]["name"] @@ -38,48 +45,83 @@ def create_resource(session, spec: RunSpec, manifest: dict) -> None: raise wrap_api_exception(exc, f"failed to create {spec.mode} {name!r}") from exc -def discover_job_pod(session, job_name: str, timeout: int) -> str: - watcher = watch.Watch() +def _watch_pods_until(session, timeout: int, stop_event: threading.Event | None, what: str, **selector): + """Yield pod watch events matching `selector` until `timeout` seconds have passed, in + STARTUP_WATCH_SLICE_SECONDS slices, raising InterruptError as soon as stop_event is set.""" + deadline = time.monotonic() + timeout + while True: + if stop_event is not None and stop_event.is_set(): + raise InterruptError(f"interrupted while waiting for {what}") + remaining = deadline - time.monotonic() + if remaining <= 0: + return + watcher = watch.Watch() + try: + for event in watcher.stream( + session.core.list_namespaced_pod, + namespace=session.namespace, + timeout_seconds=max(1, min(STARTUP_WATCH_SLICE_SECONDS, math.ceil(remaining))), + **selector, + ): + if stop_event is not None and stop_event.is_set(): + raise InterruptError(f"interrupted while waiting for {what}") + yield event + finally: + watcher.stop() + + +def discover_job_pod( + session, + job_name: str, + timeout: int, + *, + quiet: bool = False, + stop_event: threading.Event | None = None, +) -> str: + # job-level warnings (e.g. FailedCreate from an exceeded quota) are why a job never makes a pod + with WarningEventWatcher(session, job_name, timeout=timeout, quiet=quiet): + return _watch_for_job_pod(session, job_name, timeout, stop_event) + + +def _watch_for_job_pod(session, job_name: str, timeout: int, stop_event: threading.Event | None) -> str: try: - for event in watcher.stream( - session.core.list_namespaced_pod, - namespace=session.namespace, - label_selector=f"job-name={job_name}", - timeout_seconds=timeout, + for event in _watch_pods_until( + session, timeout, stop_event, f"job {job_name!r} to create a pod", label_selector=f"job-name={job_name}" ): return event["object"].metadata.name except ApiException as exc: raise _wrap_api_exception( exc, f"failed to watch for the pod created by job {job_name!r}" ) from exc - finally: - watcher.stop() raise RunnerError(f"timed out waiting for job {job_name!r} to create a pod") def wait_until_running_or_terminal( - session, pod_name: str, timeout: int, *, quiet: bool = False + session, + pod_name: str, + timeout: int, + *, + quiet: bool = False, + stop_event: threading.Event | None = None, ) -> str: """Wait until the pod is Running/Succeeded/Failed, or fail fast on a known-bad waiting reason. This bounds only the *startup* window (image pull, scheduling, container create); once the container is actually running, log streaming plus wait_for_terminal take over with no timeout. """ + events = WarningEventWatcher(session, pod_name, timeout=timeout, quiet=quiet) try: - return _watch_for_start(session, pod_name, timeout) + with events: + return _watch_for_start(session, pod_name, timeout, stop_event) except RunnerError: - print_pod_events(session, pod_name, quiet=quiet) + print_pod_events(session, pod_name, quiet=quiet, skip=frozenset(events.shown)) raise -def _watch_for_start(session, pod_name: str, timeout: int) -> str: - watcher = watch.Watch() +def _watch_for_start(session, pod_name: str, timeout: int, stop_event: threading.Event | None) -> str: try: - for event in watcher.stream( - session.core.list_namespaced_pod, - namespace=session.namespace, - field_selector=f"metadata.name={pod_name}", - timeout_seconds=timeout, + for event in _watch_pods_until( + session, timeout, stop_event, f"pod {pod_name!r} to start", field_selector=f"metadata.name={pod_name}" ): pod = event["object"] phase = pod.status.phase @@ -93,8 +135,6 @@ def _watch_for_start(session, pod_name: str, timeout: int) -> str: return phase except ApiException as exc: raise _wrap_api_exception(exc, f"failed to watch pod {pod_name!r}") from exc - finally: - watcher.stop() raise RunnerError( f"timed out after {timeout}s waiting for pod {pod_name!r} to start" ) @@ -213,7 +253,9 @@ def run_manifest(session, spec: RunSpec, manifest: dict) -> int: slow_hint="the job controller has not produced a pod yet", slow_after=15, ): - pod_name = discover_job_pod(session, name, spec.timeout) + pod_name = discover_job_pod( + session, name, spec.timeout, quiet=spec.quiet, stop_event=handler.stop_event + ) else: pod_name = name @@ -225,7 +267,7 @@ def run_manifest(session, spec: RunSpec, manifest: dict) -> int: spinner_delay=0.0, ): phase = wait_until_running_or_terminal( - session, pod_name, spec.timeout, quiet=spec.quiet + session, pod_name, spec.timeout, quiet=spec.quiet, stop_event=handler.stop_event ) chatter( @@ -294,7 +336,9 @@ def run_manifest(session, spec: RunSpec, manifest: dict) -> int: "sweeping orphaned resources from previous runs", quiet=spec.quiet, ): - sweep_orphans(session, spec, exclude_run_id=own_run_id) + sweep_orphans( + session, spec.orphan_sweep_min_age, quiet=spec.quiet, exclude_run_id=own_run_id + ) except Exception as exc: chatter( f"orphan sweep skipped: {exc}", quiet=spec.quiet, level="warn" diff --git a/squarepeg/manifest.py b/squarepeg/manifest.py index 1eac2f5..20698b8 100644 --- a/squarepeg/manifest.py +++ b/squarepeg/manifest.py @@ -70,6 +70,9 @@ def _build_resources(spec: RunSpec) -> dict: limits["cpu"] = spec.resources.limit_cpu if spec.resources.limit_memory is not None: limits["memory"] = spec.resources.limit_memory + if spec.resources.gpus is not None: + requests[spec.resources.gpu_resource] = spec.resources.gpus + limits[spec.resources.gpu_resource] = spec.resources.gpus resources = {} if requests: resources["requests"] = requests @@ -105,6 +108,8 @@ def _build_container(spec: RunSpec) -> dict: def _build_pod_spec(spec: RunSpec) -> dict: pod_spec = {"restartPolicy": "Never", "containers": [_build_container(spec)]} + if spec.resources.gpus and spec.resources.gpu_runtime_class: + pod_spec["runtimeClassName"] = spec.resources.gpu_runtime_class volumes = _build_volumes(spec) if volumes: pod_spec["volumes"] = volumes diff --git a/squarepeg/quantities.py b/squarepeg/quantities.py index 4f872da..833b1f4 100644 --- a/squarepeg/quantities.py +++ b/squarepeg/quantities.py @@ -56,3 +56,19 @@ def docker_cpus_to_k8s(value: str) -> str: return str(int(cpus)) millicores = math.ceil(cpus * 1000) return f"{millicores}m" + + +def docker_gpus_to_k8s(value: str | int) -> int: + """Convert a docker-style --gpus value ('2' or 'count=2') to a k8s GPU count. + + Docker's 'all' and 'device=...' have no k8s equivalent: the scheduler needs an + explicit count, and picks the devices itself. + """ + text = str(value).strip() + if text.lower().startswith("count="): + text = text[len("count="):] + if not text.isdigit(): + raise UsageError( + f"invalid gpus value {value!r}: Kubernetes needs an explicit GPU count; pass --gpus N" + ) + return int(text) diff --git a/squarepeg/runspec.py b/squarepeg/runspec.py index 5b44068..3c3827b 100644 --- a/squarepeg/runspec.py +++ b/squarepeg/runspec.py @@ -1,7 +1,11 @@ from dataclasses import dataclass, field +from squarepeg.k8s.orphans import DEFAULT_MIN_AGE_SECONDS from squarepeg.volumes import VolumeMount +DEFAULT_GPU_RESOURCE = "nvidia.com/gpu" +DEFAULT_GPU_RUNTIME_CLASS = "nvidia" + @dataclass class ResourceSpec: @@ -9,6 +13,11 @@ class ResourceSpec: limit_cpu: str | None = None request_memory: str | None = None limit_memory: str | None = None + # extended resources can't be overcommitted, so one value is both the request and limit + gpus: int | None = None + gpu_resource: str = DEFAULT_GPU_RESOURCE + # some distros (e.g. k3s/RKE2 with the NVIDIA toolkit) need this set on GPU pods; None disables + gpu_runtime_class: str | None = DEFAULT_GPU_RUNTIME_CLASS @dataclass @@ -34,7 +43,7 @@ class RunSpec: timeout: int = 300 quiet: bool = False orphan_sweep: bool = True - orphan_sweep_min_age: int = 300 + orphan_sweep_min_age: int = DEFAULT_MIN_AGE_SECONDS # JSON-pointer-ish paths the CLI explicitly set, used to make config # passthrough (Phase B, see merge.py) lose to explicit CLI flags. diff --git a/squarepeg/ui.py b/squarepeg/ui.py index 7e993a7..0b21870 100644 --- a/squarepeg/ui.py +++ b/squarepeg/ui.py @@ -41,7 +41,7 @@ _SPINNER_FRAMES_UNICODE = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"] _SPINNER_FRAMES_ASCII = ["|", "/", "-", "\\"] -_SPINNER_INTERVAL = 0.1 +_SPINNER_INTERVAL = 0.25 _LOCK = threading.RLock() _ACTIVE: "Status | None" = None @@ -97,17 +97,44 @@ def animation_enabled() -> bool: return getattr(sys.stderr, "isatty", lambda: False)() -def _spinner_frames() -> list[str]: +def _stderr_is_utf() -> bool: encoding = getattr(sys.stderr, "encoding", None) or "" - if "utf" in encoding.lower(): - return _SPINNER_FRAMES_UNICODE - return _SPINNER_FRAMES_ASCII + return "utf" in encoding.lower() -def _styled_line(message: str, level: Level) -> str: +def _spinner_frames() -> list[str]: + return _SPINNER_FRAMES_UNICODE if _stderr_is_utf() else _SPINNER_FRAMES_ASCII + + +def _term_width() -> int: + # stderr, not shutil.get_terminal_size(): that reads stdout, which may be piped + try: + return os.get_terminal_size(sys.stderr.fileno()).columns + except (OSError, ValueError, AttributeError): + return 80 + + +def _label() -> str: tag = current_run_tag() - label = f"[squarepeg:{tag}] " if tag else "[squarepeg] " - prefix = click.style(label, dim=True) + return f"[squarepeg:{tag}] " if tag else "[squarepeg] " + + +def _fit(message: str, width: int) -> str: + """Truncate `message` so the prefixed line fits on one row of a `width`-column terminal. + + A redrawn line must never wrap: "\r\x1b[2K" only returns to and clears the last + wrapped row, so every redraw would otherwise leave the earlier rows behind. Everything + squarepeg puts on a spinner line is single-width, so len() is the display width. + """ + room = width - 1 - len(_label()) + if len(message) <= room: + return message + ellipsis = "…" if _stderr_is_utf() else "..." + return message[: max(room - len(ellipsis), 0)] + ellipsis + + +def _styled_line(message: str, level: Level) -> str: + prefix = click.style(_label(), dim=True) body = click.style(message, **_STYLES[level]) return prefix + body @@ -127,8 +154,9 @@ def emit(message: str, *, level: Level = "info", quiet: bool = False) -> None: class Status: """Context manager for a stage that may take a while. - On a real terminal, shows an animated spinner line (stderr only) with an elapsed - timer and, past `slow_after` seconds, a `slow_hint`. Everywhere else (non-TTY, + On a real terminal, shows an animated spinner line (stderr only), truncated to the + terminal width, with an elapsed timer; past `slow_after` seconds the `slow_hint` is + printed once on its own line above it. Everywhere else (non-TTY, --no-color, --quiet, TERM=dumb) it degrades to a single static line printed once on entry -- including the slow_hint text upfront, since there is no timer to add it later -- and nothing further is printed on exit besides the optional `success` line. @@ -154,6 +182,7 @@ def __init__( self._stop = threading.Event() self._start_time: float | None = None self._line_drawn = False + self._hint_shown = False def update(self, message: str) -> None: with _LOCK: @@ -172,7 +201,12 @@ def __enter__(self) -> "Status": self._start_time = time.monotonic() with _LOCK: - click.echo(_styled_line(self.message, "wait"), file=sys.stderr, nl=False, color=color_enabled()) + click.echo( + _styled_line(_fit(self.message, _term_width()), "wait"), + file=sys.stderr, + nl=False, + color=color_enabled(), + ) self._line_drawn = True _ACTIVE = self self._thread = threading.Thread(target=self._animate, daemon=True) @@ -210,9 +244,12 @@ def _animate(self) -> None: text = self.message if elapsed >= 2: text = f"{text} ({int(elapsed)}s)" - if self.slow_hint and elapsed >= self.slow_after: - text = f"{text} — {self.slow_hint}" - line = f"\r\x1b[2K{_styled_line(f'{next(frames)} {text}', 'wait')}" + if self.slow_hint and not self._hint_shown and elapsed >= self.slow_after: + self._erase() + click.echo(_styled_line(self.slow_hint, "detail"), file=sys.stderr, color=color_enabled()) + self._hint_shown = True + text = _fit(f"{next(frames)} {text}", _term_width()) + line = f"\r\x1b[2K{_styled_line(text, 'wait')}" click.echo(line, file=sys.stderr, nl=False, color=color_enabled()) self._line_drawn = True self._stop.wait(_SPINNER_INTERVAL) diff --git a/tests/conftest.py b/tests/conftest.py index 3b04422..dcbf515 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -82,3 +82,20 @@ def _run(*args, timeout=120, env_overrides=None): return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env) return _run + + +@pytest.fixture +def clean_cli(test_namespace): + """Invoke `squarepeg clean` as a subprocess, isolated from any local config. Sibling to + run_cli rather than a generalisation of it, so existing run_cli(*args) call sites (which + assume 'run' is prepended) don't need to change.""" + + def _clean(*args, timeout=120, env_overrides=None): + env = dict(os.environ) + env["SQUAREPEG_CONFIG"] = "" + cmd = [sys.executable, "-m", "squarepeg", "clean", "--no-default-config", "-n", test_namespace, *args] + if env_overrides: + env.update(env_overrides) + return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env) + + return _clean diff --git a/tests/test_cli.py b/tests/test_cli.py index b45e31e..47b2047 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,3 +1,6 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + import yaml from click.testing import CliRunner from kubernetes.client.rest import ApiException @@ -18,14 +21,20 @@ def _separated_runner(): class FakeSession: - """Stands in for squarepeg.k8s.session.Session so CLI-level dry-run-server tests never - touch a real cluster; records the args it was constructed with for assertions.""" + """Stands in for squarepeg.k8s.session.Session so CLI-level dry-run-server/clean tests + never touch a real cluster; records the args it was constructed with for assertions. + .core/.batch are plain MagicMocks returning empty lists by default, for 'clean' tests + that need a listable (but otherwise inert) session.""" def __init__(self, namespace=None, context=None, quiet=False): self.namespace = namespace or "default" self.context = context self.quiet = quiet self.context_name = context or "fake-context" + self.core = MagicMock() + self.batch = MagicMock() + self.core.list_namespaced_pod.return_value = SimpleNamespace(items=[]) + self.batch.list_namespaced_job.return_value = SimpleNamespace(items=[]) def test_version(): @@ -200,6 +209,55 @@ def test_run_negative_orphan_sweep_min_age_rejected(tmp_path): assert "orphan_sweep_min_age" in str(result.exception) +def _dry_run_resources(*args): + result = _separated_runner().invoke(cli, ["run", "--no-default-config", *args, "--dry-run", "alpine"]) + assert result.exit_code == 0, result.output + return yaml.safe_load(result.stdout)["spec"]["containers"][0]["resources"] + + +def test_run_dry_run_gpus_flag_sets_request_and_limit(): + resources = _dry_run_resources("--gpus", "2") + assert resources["requests"]["nvidia.com/gpu"] == 2 + assert resources["limits"]["nvidia.com/gpu"] == 2 + + +def test_run_dry_run_gpus_from_config_defaults_with_custom_resource(tmp_path): + config_file = tmp_path / "config.yaml" + config_file.write_text("defaults:\n gpus: 1\n gpu_resource: amd.com/gpu\n") + resources = _dry_run_resources("--config", str(config_file)) + assert resources == {"requests": {"amd.com/gpu": 1}, "limits": {"amd.com/gpu": 1}} + + +def test_run_dry_run_gpu_runtime_class_from_config(tmp_path): + config_file = tmp_path / "config.yaml" + config_file.write_text("defaults:\n gpu_runtime_class: nvidia-cdi\n") + args = ["run", "--no-default-config", "--config", str(config_file), "--gpus", "1", "--dry-run", "alpine"] + result = _separated_runner().invoke(cli, args) + assert result.exit_code == 0, result.output + assert yaml.safe_load(result.stdout)["spec"]["runtimeClassName"] == "nvidia-cdi" + + config_file.write_text("defaults:\n gpu_runtime_class: null\n") + result = _separated_runner().invoke(cli, args) + assert result.exit_code == 0, result.output + assert "runtimeClassName" not in yaml.safe_load(result.stdout)["spec"] + + +def test_run_dry_run_gpus_flag_beats_passthrough(tmp_path): + config_file = tmp_path / "config.yaml" + config_file.write_text( + "kubernetes:\n spec:\n containers:\n" + " - name: main\n resources: {limits: {nvidia.com/gpu: 4}}\n" + ) + resources = _dry_run_resources("--config", str(config_file), "--gpus", "1") + assert resources["limits"]["nvidia.com/gpu"] == 1 + + +def test_run_gpus_all_rejected(): + result = CliRunner().invoke(cli, ["run", "--no-default-config", "--gpus", "all", "--dry-run", "alpine"]) + assert result.exit_code != 0 + assert "explicit GPU count" in str(result.exception) + result.output + + # --- --dry-run-server --- @@ -463,3 +521,204 @@ def test_two_run_invocations_get_different_tags(monkeypatch): label1 = captured1["manifest"]["metadata"]["labels"]["squarepeg.io/run-id"] label2 = captured2["manifest"]["metadata"]["labels"]["squarepeg.io/run-id"] assert label1 != label2 + + +# --- clean --- + +CLEAN_CREATED_BY = "sam" + + +def _clean_pod(name, phase="Succeeded", finished_ago=600, waiting_reason=None, created_by=CLEAN_CREATED_BY): + import datetime as dt + + now = dt.datetime.now(dt.timezone.utc) + terminated = SimpleNamespace(finished_at=now - dt.timedelta(seconds=finished_ago)) if phase == "Succeeded" else None + waiting = SimpleNamespace(reason=waiting_reason) if waiting_reason else None + return SimpleNamespace( + metadata=SimpleNamespace( + name=name, + labels={"app.kubernetes.io/managed-by": "squarepeg", "squarepeg.io/created-by": created_by}, + owner_references=[], + creation_timestamp=now - dt.timedelta(seconds=finished_ago), + ), + status=SimpleNamespace( + phase=phase, + container_statuses=[ + SimpleNamespace(name="main", state=SimpleNamespace(waiting=waiting, terminated=terminated)) + ], + ), + ) + + +def _mock_clean_session(monkeypatch, pods=(), jobs=()): + monkeypatch.setattr("squarepeg.k8s.orphans.current_created_by", lambda: CLEAN_CREATED_BY) + + class RecordingFakeSession(FakeSession): + def __init__(self, namespace=None, context=None, quiet=False): + super().__init__(namespace=namespace, context=context, quiet=quiet) + self.core.list_namespaced_pod.return_value = SimpleNamespace(items=list(pods)) + self.batch.list_namespaced_job.return_value = SimpleNamespace(items=list(jobs)) + + monkeypatch.setattr(cli_module, "Session", RecordingFakeSession) + return RecordingFakeSession + + +def test_clean_dry_run_lists_without_deleting(monkeypatch): + _mock_clean_session(monkeypatch, pods=[_clean_pod("p")]) + result = _separated_runner().invoke(cli, ["clean", "--no-default-config", "--dry-run"]) + assert result.exit_code == 0, result.stderr + assert "pod/p" in result.stdout + assert "terminal" in result.stdout + + +def test_clean_dry_run_quiet_still_lists_on_stdout(monkeypatch): + _mock_clean_session(monkeypatch, pods=[_clean_pod("p")]) + result = _separated_runner().invoke(cli, ["clean", "--no-default-config", "--dry-run", "--quiet"]) + assert result.exit_code == 0, result.stderr + assert "pod/p" in result.stdout + assert result.stderr == "" + + +def test_clean_deletes_and_exits_zero(monkeypatch): + _mock_clean_session(monkeypatch, pods=[_clean_pod("p")]) + result = _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + assert result.exit_code == 0, result.stderr + assert "deleted pod 'p'" in result.stderr + + +def test_clean_empty_namespace_exits_zero(monkeypatch): + _mock_clean_session(monkeypatch) + result = _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + assert result.exit_code == 0, result.stderr + assert "nothing to clean" in result.stderr + + +def test_clean_partial_failure_exit_code(monkeypatch): + monkeypatch.setattr("squarepeg.k8s.orphans.current_created_by", lambda: CLEAN_CREATED_BY) + + class FailingFakeSession(FakeSession): + def __init__(self, namespace=None, context=None, quiet=False): + super().__init__(namespace=namespace, context=context, quiet=quiet) + self.core.list_namespaced_pod.return_value = SimpleNamespace(items=[_clean_pod("p1"), _clean_pod("p2")]) + self.core.delete_namespaced_pod.side_effect = [ApiException(status=500, reason="err"), None] + + monkeypatch.setattr(cli_module, "Session", FailingFakeSession) + result = _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + assert result.exit_code != 0 + assert result.exception.exit_code == 125 + assert "failed to delete 1 of 2" in str(result.exception) + + +def test_clean_listing_failure_exits_125(monkeypatch): + monkeypatch.setattr("squarepeg.k8s.orphans.current_created_by", lambda: CLEAN_CREATED_BY) + + class RaisingFakeSession(FakeSession): + def __init__(self, namespace=None, context=None, quiet=False): + super().__init__(namespace=namespace, context=context, quiet=quiet) + self.core.list_namespaced_pod.side_effect = ApiException(status=403, reason="Forbidden") + + monkeypatch.setattr(cli_module, "Session", RaisingFakeSession) + result = _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + assert result.exit_code != 0 + assert result.exception.exit_code == 125 + + +def test_clean_namespace_and_context_forwarded_to_session(monkeypatch): + constructed = {} + monkeypatch.setattr("squarepeg.k8s.orphans.current_created_by", lambda: CLEAN_CREATED_BY) + + class RecordingFakeSession(FakeSession): + def __init__(self, namespace=None, context=None, quiet=False): + constructed["namespace"] = namespace + constructed["context"] = context + super().__init__(namespace=namespace, context=context, quiet=quiet) + + monkeypatch.setattr(cli_module, "Session", RecordingFakeSession) + result = _separated_runner().invoke( + cli, ["clean", "--no-default-config", "-n", "my-ns", "--context", "my-ctx", "--dry-run"] + ) + assert result.exit_code == 0, result.stderr + assert constructed["namespace"] == "my-ns" + assert constructed["context"] == "my-ctx" + + +def test_clean_min_age_flag_overrides_default(monkeypatch): + captured = {} + monkeypatch.setattr("squarepeg.k8s.orphans.current_created_by", lambda: CLEAN_CREATED_BY) + + def fake_clean_orphans(session, min_age, **kwargs): + captured["min_age"] = min_age + from squarepeg.k8s.orphans import CleanResult + + return CleanResult() + + monkeypatch.setattr(cli_module, "Session", FakeSession) + monkeypatch.setattr(cli_module, "clean_orphans", fake_clean_orphans) + + _separated_runner().invoke(cli, ["clean", "--no-default-config", "--min-age", "60"]) + assert captured["min_age"] == 60 + + _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + assert captured["min_age"] == 300 + + +def test_clean_min_age_from_config(monkeypatch, tmp_path): + captured = {} + monkeypatch.setattr("squarepeg.k8s.orphans.current_created_by", lambda: CLEAN_CREATED_BY) + + def fake_clean_orphans(session, min_age, **kwargs): + captured["min_age"] = min_age + from squarepeg.k8s.orphans import CleanResult + + return CleanResult() + + monkeypatch.setattr(cli_module, "Session", FakeSession) + monkeypatch.setattr(cli_module, "clean_orphans", fake_clean_orphans) + + config_file = tmp_path / "config.yaml" + config_file.write_text("orphan_sweep_min_age: 42\n") + _separated_runner().invoke(cli, ["clean", "--no-default-config", "--config", str(config_file)]) + assert captured["min_age"] == 42 + + +def test_clean_negative_min_age_rejected(monkeypatch): + monkeypatch.setattr(cli_module, "Session", FakeSession) + result = _separated_runner().invoke(cli, ["clean", "--no-default-config", "--min-age", "-1"]) + assert result.exit_code != 0 + assert "--min-age" in str(result.exception) + + +def test_clean_runs_regardless_of_orphan_sweep_config_key(monkeypatch, tmp_path): + _mock_clean_session(monkeypatch, pods=[_clean_pod("p")]) + config_file = tmp_path / "config.yaml" + config_file.write_text("orphan_sweep: false\n") + result = _separated_runner().invoke(cli, ["clean", "--no-default-config", "--config", str(config_file)]) + assert result.exit_code == 0, result.stderr + assert "deleted pod 'p'" in result.stderr + + +def test_clean_tag_present_and_differs_between_invocations(monkeypatch): + import re + + _mock_clean_session(monkeypatch) + result1 = _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + result2 = _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + tag1 = re.search(r"\[squarepeg:([0-9a-f]{8})\]", result1.stderr).group(1) + tag2 = re.search(r"\[squarepeg:([0-9a-f]{8})\]", result2.stderr).group(1) + assert tag1 != tag2 + + +def test_clean_followed_by_config_leaves_config_untagged(monkeypatch): + _mock_clean_session(monkeypatch) + _separated_runner().invoke(cli, ["clean", "--no-default-config"]) + result = _separated_runner().invoke(cli, ["config", "--no-default-config"]) + assert "[squarepeg:" not in result.stderr + + +def test_format_age_examples(): + from squarepeg.cli import _format_age + + assert _format_age(5) == "5s" + assert _format_age(612) == "10m12s" + assert _format_age(11640) == "3h14m" + assert _format_age(188000) == "2d4h" diff --git a/tests/test_events.py b/tests/test_events.py index 631f6e1..553a56a 100644 --- a/tests/test_events.py +++ b/tests/test_events.py @@ -1,9 +1,11 @@ +import threading from types import SimpleNamespace from unittest.mock import MagicMock from kubernetes.client.rest import ApiException -from squarepeg.k8s.events import print_pod_events +from squarepeg.k8s import events as events_module +from squarepeg.k8s.events import WarningEventWatcher, print_pod_events def make_session(): @@ -13,8 +15,8 @@ def make_session(): return session -def event(reason, message): - return SimpleNamespace(reason=reason, message=message) +def event(reason, message, type="Normal"): + return SimpleNamespace(reason=reason, message=message, type=type) def test_print_pod_events_calls_chatter_for_each_event(monkeypatch, capsys): @@ -47,3 +49,75 @@ def test_print_pod_events_handles_empty_items(): session = make_session() session.core.list_namespaced_event.return_value = SimpleNamespace(items=[]) print_pod_events(session, "mypod", quiet=False) # must not raise + + +def test_print_pod_events_skips_already_shown(capsys): + session = make_session() + session.core.list_namespaced_event.return_value = SimpleNamespace( + items=[event("FailedScheduling", "0/12 nodes available"), event("Pulling", "pulling image")] + ) + print_pod_events(session, "mypod", skip=frozenset({("FailedScheduling", "0/12 nodes available")})) + err = capsys.readouterr().err + assert "0/12 nodes available" not in err + assert "pulling image" in err + + +# --- WarningEventWatcher --- + +UNSCHEDULABLE = event("FailedScheduling", "0/11 nodes are available: 3 Insufficient nvidia.com/gpu", "Warning") + + +def test_watcher_prints_each_warning_once_with_one_waiting_note(capsys): + session = make_session() + session.core.list_namespaced_event.return_value = SimpleNamespace( + items=[UNSCHEDULABLE, event("Scheduled", "assigned to node", "Normal")] + ) + watcher = WarningEventWatcher(session, "mypod", timeout=60) + watcher.report_new_warnings() + watcher.report_new_warnings() # same event again: must not repeat + err = capsys.readouterr().err + assert err.count("Insufficient nvidia.com/gpu") == 1 + assert err.count("keep waiting until the 60s --timeout") == 1 + assert "assigned to node" not in err # Normal events aren't surfaced early + assert watcher.shown == {("FailedScheduling", UNSCHEDULABLE.message)} + + +def test_watcher_note_printed_only_for_first_warning(capsys): + session = make_session() + session.core.list_namespaced_event.return_value = SimpleNamespace(items=[UNSCHEDULABLE]) + watcher = WarningEventWatcher(session, "mypod", timeout=60) + watcher.report_new_warnings() + session.core.list_namespaced_event.return_value = SimpleNamespace( + items=[UNSCHEDULABLE, event("FailedMount", "pvc not found", "Warning")] + ) + watcher.report_new_warnings() + err = capsys.readouterr().err + assert "pvc not found" in err + assert err.count("keep waiting") == 1 + + +def test_watcher_polls_in_background(monkeypatch, capsys): + monkeypatch.setattr(events_module, "EVENT_POLL_SECONDS", 0.01) + session = make_session() + polled = threading.Event() + session.core.list_namespaced_event.side_effect = lambda *a, **kw: ( + polled.set(), + SimpleNamespace(items=[UNSCHEDULABLE]), + )[1] + with WarningEventWatcher(session, "mypod", timeout=60): + assert polled.wait(2) + assert "Insufficient nvidia.com/gpu" in capsys.readouterr().err + + +def test_watcher_quiet_starts_no_thread_and_prints_nothing(capsys): + session = make_session() + session.core.list_namespaced_event.return_value = SimpleNamespace(items=[UNSCHEDULABLE]) + with WarningEventWatcher(session, "mypod", timeout=60, quiet=True) as watcher: + assert watcher._thread is None + assert capsys.readouterr().err == "" + + +def test_watcher_swallows_api_errors(): + session = make_session() + session.core.list_namespaced_event.side_effect = ApiException(status=403, reason="Forbidden") + WarningEventWatcher(session, "mypod", timeout=60).report_new_warnings() # must not raise diff --git a/tests/test_manifest.py b/tests/test_manifest.py index 7fb2dea..168e4bf 100644 --- a/tests/test_manifest.py +++ b/tests/test_manifest.py @@ -100,6 +100,43 @@ def test_pod_resources_block_from_resource_spec(): } +def test_pod_gpus_set_both_request_and_limit(): + resources = ResourceSpec(request_cpu="2", limit_cpu="2", gpus=1) + pod = build_pod(basic_spec(resources=resources), run_id=RUN_ID) + assert pod["spec"]["containers"][0]["resources"] == { + "requests": {"cpu": "2", "nvidia.com/gpu": 1}, + "limits": {"cpu": "2", "nvidia.com/gpu": 1}, + } + assert pod["spec"]["runtimeClassName"] == "nvidia" + + +def test_pod_no_runtime_class_without_gpus(): + for resources in (ResourceSpec(), ResourceSpec(gpus=0)): + pod = build_pod(basic_spec(resources=resources), run_id=RUN_ID) + assert "runtimeClassName" not in pod["spec"] + + +def test_pod_runtime_class_configurable_and_disableable(): + pod = build_pod(basic_spec(resources=ResourceSpec(gpus=1, gpu_runtime_class="nvidia-cdi")), run_id=RUN_ID) + assert pod["spec"]["runtimeClassName"] == "nvidia-cdi" + pod = build_pod(basic_spec(resources=ResourceSpec(gpus=1, gpu_runtime_class=None)), run_id=RUN_ID) + assert "runtimeClassName" not in pod["spec"] + + +def test_job_template_gets_gpu_runtime_class(): + job = build_job(basic_spec(mode="job", resources=ResourceSpec(gpus=1)), run_id=RUN_ID) + assert job["spec"]["template"]["spec"]["runtimeClassName"] == "nvidia" + + +def test_pod_gpus_use_configured_resource_name(): + resources = ResourceSpec(gpus=2, gpu_resource="amd.com/gpu") + pod = build_pod(basic_spec(resources=resources), run_id=RUN_ID) + assert pod["spec"]["containers"][0]["resources"] == { + "requests": {"amd.com/gpu": 2}, + "limits": {"amd.com/gpu": 2}, + } + + def test_pod_no_resources_block_when_unset(): pod = build_pod(basic_spec(), run_id=RUN_ID) assert "resources" not in pod["spec"]["containers"][0] diff --git a/tests/test_orphans.py b/tests/test_orphans.py index 910335b..8968dd0 100644 --- a/tests/test_orphans.py +++ b/tests/test_orphans.py @@ -7,7 +7,6 @@ from squarepeg.errors import ApiError from squarepeg.k8s import orphans -from squarepeg.runspec import RunSpec NOW = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) CREATED_BY = "sam" @@ -96,12 +95,6 @@ def condition(type_, status="True", last_transition_time=None): return SimpleNamespace(type=type_, status=status, last_transition_time=last_transition_time) -def basic_spec(**overrides): - defaults = dict(image="alpine", name="squarepeg-alpine-abc123", orphan_sweep_min_age=300) - defaults.update(overrides) - return RunSpec(**defaults) - - # --- selector construction --- @@ -312,7 +305,7 @@ def test_sweep_deletes_pod_and_job_with_correct_calls(): pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] jobs = [fake_job("j", conditions=[condition("Complete", last_transition_time=NOW - timedelta(seconds=600))])] session = make_session(pods=pods, jobs=jobs) - count = orphans.sweep_orphans(session, basic_spec(), now=NOW) + count = orphans.sweep_orphans(session, 300, now=NOW) assert count == 2 session.core.delete_namespaced_pod.assert_called_once_with("p", "default") args, kwargs = session.batch.delete_namespaced_job.call_args @@ -323,7 +316,7 @@ def test_sweep_deletes_pod_and_job_with_correct_calls(): def test_stuck_pod_deletes_via_identical_call_as_terminal_pod(): pods = [fake_pod("stuck", phase="Pending", container_statuses=waiting("ImagePullBackOff"), creation_offset=600)] session = make_session(pods=pods) - orphans.sweep_orphans(session, basic_spec(), now=NOW) + orphans.sweep_orphans(session, 300, now=NOW) session.core.delete_namespaced_pod.assert_called_once_with("stuck", "default") @@ -331,7 +324,7 @@ def test_delete_404_tolerated_and_not_counted(): pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] session = make_session(pods=pods) session.core.delete_namespaced_pod.side_effect = ApiException(status=404, reason="Not Found") - assert orphans.sweep_orphans(session, basic_spec(), now=NOW) == 0 + assert orphans.sweep_orphans(session, 300, now=NOW) == 0 def test_one_delete_failure_does_not_abort_the_rest(): @@ -344,7 +337,7 @@ def test_one_delete_failure_does_not_abort_the_rest(): ApiException(status=500, reason="Server Error"), None, ] - assert orphans.sweep_orphans(session, basic_spec(), now=NOW) == 1 + assert orphans.sweep_orphans(session, 300, now=NOW) == 1 # --- output --- @@ -353,20 +346,20 @@ def test_one_delete_failure_does_not_abort_the_rest(): def test_summary_message_on_nonzero_count(capsys): pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] session = make_session(pods=pods) - orphans.sweep_orphans(session, basic_spec(quiet=False), now=NOW) + orphans.sweep_orphans(session, 300, quiet=False, now=NOW) assert "swept 1 orphaned resource" in capsys.readouterr().err def test_no_message_when_nothing_swept(capsys): session = make_session() - orphans.sweep_orphans(session, basic_spec(quiet=False), now=NOW) + orphans.sweep_orphans(session, 300, quiet=False, now=NOW) assert capsys.readouterr().err == "" def test_quiet_suppresses_summary_message(capsys): pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] session = make_session(pods=pods) - orphans.sweep_orphans(session, basic_spec(quiet=True), now=NOW) + orphans.sweep_orphans(session, 300, quiet=True, now=NOW) assert capsys.readouterr().err == "" @@ -374,4 +367,172 @@ def test_list_failure_raises_apierror_naming_the_opt_out(): session = make_session() session.core.list_namespaced_pod.side_effect = ApiException(status=403, reason="Forbidden") with pytest.raises(ApiError, match="orphan_sweep"): - orphans.sweep_orphans(session, basic_spec(), now=NOW) + orphans.sweep_orphans(session, 300, now=NOW) + + +# --- find_orphan_candidates --- + + +def test_candidates_terminal_pod_classification(): + pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] + session = make_session(pods=pods) + candidates = orphans.find_orphan_candidates(session, CREATED_BY, 300, now=NOW) + assert len(candidates) == 1 + c = candidates[0] + assert (c.kind, c.name, c.category, c.reason) == ("pod", "p", "terminal", None) + assert c.age_seconds == 600 + + +def test_candidates_stuck_pod_classification(): + pods = [fake_pod("p", phase="Pending", container_statuses=waiting("ImagePullBackOff"), creation_offset=600)] + session = make_session(pods=pods) + candidates = orphans.find_orphan_candidates(session, CREATED_BY, 300, now=NOW) + assert len(candidates) == 1 + c = candidates[0] + assert (c.kind, c.name, c.category, c.reason) == ("pod", "p", "stuck", "ImagePullBackOff") + assert c.age_seconds == 600 + + +def test_candidates_job_terminal_by_condition(): + jobs = [fake_job("j", conditions=[condition("Complete", last_transition_time=NOW - timedelta(seconds=600))])] + session = make_session(jobs=jobs) + candidates = orphans.find_orphan_candidates(session, CREATED_BY, 300, now=NOW) + assert len(candidates) == 1 + c = candidates[0] + assert (c.kind, c.name, c.category, c.reason) == ("job", "j", "terminal", None) + + +def test_candidates_job_stuck_via_child_pod(): + pods = [ + fake_pod( + "child-pod", phase="Pending", job_name="stuck-job", + container_statuses=waiting("ImagePullBackOff"), creation_offset=600, + ) + ] + jobs = [fake_job("stuck-job", conditions=[])] + session = make_session(pods=pods, jobs=jobs) + candidates = orphans.find_orphan_candidates(session, CREATED_BY, 300, now=NOW) + assert len(candidates) == 1 + c = candidates[0] + assert (c.kind, c.name, c.category, c.reason) == ("job", "stuck-job", "stuck", "ImagePullBackOff") + + +def test_find_orphaned_resources_is_non_lossy_wrapper_over_candidates(): + pods = [ + fake_pod("terminal-pod", phase="Succeeded", container_statuses=finished(600)), + fake_pod("stuck-pod", phase="Pending", container_statuses=waiting("ImagePullBackOff"), creation_offset=600), + ] + jobs = [ + fake_job("terminal-job", conditions=[condition("Complete", last_transition_time=NOW - timedelta(seconds=600))]) + ] + session = make_session(pods=pods, jobs=jobs) + candidates = orphans.find_orphan_candidates(session, CREATED_BY, 300, now=NOW) + pairs = orphans.find_orphaned_resources(session, CREATED_BY, 300, now=NOW) + assert pairs == [(c.kind, c.name) for c in candidates] + + +# --- clean_orphans --- + + +def test_clean_dry_run_lists_but_deletes_nothing(): + pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] + session = make_session(pods=pods) + result = orphans.clean_orphans(session, 300, dry_run=True, now=NOW) + assert len(result.candidates) == 1 + assert result.deleted == result.already_gone == result.failed == [] + session.core.delete_namespaced_pod.assert_not_called() + session.batch.delete_namespaced_job.assert_not_called() + + +def test_clean_deletes_pod_and_job_with_correct_calls(): + pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] + jobs = [fake_job("j", conditions=[condition("Complete", last_transition_time=NOW - timedelta(seconds=600))])] + session = make_session(pods=pods, jobs=jobs) + result = orphans.clean_orphans(session, 300, now=NOW) + assert len(result.deleted) == 2 + session.core.delete_namespaced_pod.assert_called_once_with("p", "default") + args, kwargs = session.batch.delete_namespaced_job.call_args + assert args == ("j", "default") + assert kwargs["body"].propagation_policy == "Background" + + +def test_clean_404_counts_as_already_gone(): + pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] + session = make_session(pods=pods) + session.core.delete_namespaced_pod.side_effect = ApiException(status=404, reason="Not Found") + result = orphans.clean_orphans(session, 300, now=NOW) + assert result.deleted == [] + assert result.failed == [] + assert [c.name for c in result.already_gone] == ["p"] + + +def test_clean_one_delete_failure_does_not_abort_the_rest(): + pods = [ + fake_pod("p1", phase="Succeeded", container_statuses=finished(600)), + fake_pod("p2", phase="Succeeded", container_statuses=finished(600)), + ] + session = make_session(pods=pods) + session.core.delete_namespaced_pod.side_effect = [ + ApiException(status=500, reason="Server Error"), + None, + ] + result = orphans.clean_orphans(session, 300, now=NOW) + assert session.core.delete_namespaced_pod.call_count == 2 + assert [c.name for c, _err in result.failed] == ["p1"] + assert [c.name for c in result.deleted] == ["p2"] + + +def test_clean_list_failure_hint_does_not_mention_orphan_sweep(): + session = make_session() + session.core.list_namespaced_pod.side_effect = ApiException(status=403, reason="Forbidden") + with pytest.raises(ApiError, match="list") as excinfo: + orphans.clean_orphans(session, 300, now=NOW) + assert "orphan_sweep" not in str(excinfo.value) + + +def test_clean_scoped_to_own_resources_and_selector_wording(): + pods = [fake_pod("someone-elses", phase="Succeeded", created_by="alex", container_statuses=finished(600))] + session = make_session(pods=pods) + result = orphans.clean_orphans(session, 300, now=NOW) + assert result.candidates == [] + selector = session.core.list_namespaced_pod.call_args.kwargs["label_selector"] + assert f"squarepeg.io/created-by={CREATED_BY}" in selector + + +def test_clean_never_touches_kept_resources_even_at_zero_min_age(): + pods = [fake_pod("kept", phase="Succeeded", keep=True, container_statuses=finished(0))] + session = make_session(pods=pods) + result = orphans.clean_orphans(session, 0, now=NOW) + assert result.candidates == [] + + +def test_clean_never_touches_running_or_benign_pending(): + pods = [ + fake_pod("running", phase="Running", creation_offset=100000), + fake_pod( + "benign-pending", phase="Pending", container_statuses=waiting("ContainerCreating"), creation_offset=100000 + ), + ] + session = make_session(pods=pods) + result = orphans.clean_orphans(session, 0, now=NOW) + assert result.candidates == [] + + +def test_clean_does_not_exclude_any_run_id(): + pods = [fake_pod("mine", phase="Succeeded", run_id="whatever", container_statuses=finished(0))] + session = make_session(pods=pods) + result = orphans.clean_orphans(session, 0, now=NOW) + assert [c.name for c in result.candidates] == ["mine"] + + +def test_clean_output_deleted_line_and_quiet_suppression(capsys): + pods = [fake_pod("p", phase="Succeeded", container_statuses=finished(600))] + session = make_session(pods=pods) + orphans.clean_orphans(session, 300, now=NOW) + err = capsys.readouterr().err + assert "deleted pod 'p'" in err + assert "cleaned 1 of 1" in err + + session2 = make_session(pods=pods) + orphans.clean_orphans(session2, 300, quiet=True, now=NOW) + assert capsys.readouterr().err == "" diff --git a/tests/test_quantities.py b/tests/test_quantities.py index 299c774..cd61af5 100644 --- a/tests/test_quantities.py +++ b/tests/test_quantities.py @@ -1,7 +1,7 @@ import pytest from squarepeg.errors import UsageError -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 @pytest.mark.parametrize( @@ -53,3 +53,14 @@ def test_memory_rejects_junk(): def test_memory_rejects_negative(): with pytest.raises(UsageError): docker_memory_to_k8s("-512m") + + +@pytest.mark.parametrize("value,expected", [("2", 2), (2, 2), ("count=2", 2), ("0", 0)]) +def test_gpus(value, expected): + assert docker_gpus_to_k8s(value) == expected + + +@pytest.mark.parametrize("value", ["all", "device=0", "-1", "1.5", "abc"]) +def test_gpus_rejects_non_counts(value): + with pytest.raises(UsageError, match="explicit GPU count"): + docker_gpus_to_k8s(value) diff --git a/tests/test_runner.py b/tests/test_runner.py index 8d4b23b..f682bff 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -51,6 +51,19 @@ def stop(self): self.stopped = True +@pytest.fixture +def fake_clock(monkeypatch): + """Advance runner's monotonic clock one second per read, so a startup wait whose (fake) + watch returns nothing reaches its timeout instantly rather than in real seconds.""" + now = [0.0] + + def tick(): + now[0] += 1 + return now[0] + + monkeypatch.setattr(runner.time, "monotonic", tick) + + def basic_spec(**overrides): defaults = dict(image="alpine", name="squarepeg-alpine-abc123", timeout=5) defaults.update(overrides) @@ -147,7 +160,7 @@ def test_discover_job_pod_returns_name_from_first_event(monkeypatch): assert runner.discover_job_pod(session, "myjob", timeout=5) == "child-pod" -def test_discover_job_pod_timeout_raises(monkeypatch): +def test_discover_job_pod_timeout_raises(monkeypatch, fake_clock): session = make_session() monkeypatch.setattr(runner.watch, "Watch", lambda: FakeWatch([])) with pytest.raises(RunnerError, match="timed out"): @@ -173,13 +186,120 @@ def test_wait_fails_fast_on_image_pull_backoff(monkeypatch): runner.wait_until_running_or_terminal(session, "p", timeout=5) -def test_wait_timeout_raises(monkeypatch): +def test_wait_timeout_raises(monkeypatch, fake_clock): session = make_session() monkeypatch.setattr(runner.watch, "Watch", lambda: FakeWatch([])) with pytest.raises(RunnerError, match="timed out"): runner.wait_until_running_or_terminal(session, "p", timeout=5) +def test_wait_timeout_does_not_repeat_warnings_already_surfaced(monkeypatch, capsys, fake_clock): + """A warning printed live by WarningEventWatcher mustn't be printed again in the + post-timeout event dump; events not yet shown still are.""" + session = make_session() + warning = SimpleNamespace(reason="FailedScheduling", message="Insufficient nvidia.com/gpu", type="Warning") + normal = SimpleNamespace(reason="Scheduled", message="assigned", type="Normal") + session.core.list_namespaced_event.return_value = SimpleNamespace(items=[warning, normal]) + + class SurfacingWatch(FakeWatch): + def stream(self, *_args, **_kwargs): + runner_events.report_new_warnings() + yield from () + + real_watcher = runner.WarningEventWatcher + runner_events = None + + def capture(*args, **kwargs): + nonlocal runner_events + runner_events = real_watcher(*args, **kwargs) + return runner_events + + monkeypatch.setattr(runner, "WarningEventWatcher", capture) + monkeypatch.setattr(runner.watch, "Watch", lambda: SurfacingWatch([])) + with pytest.raises(RunnerError, match="timed out"): + runner.wait_until_running_or_terminal(session, "p", timeout=5) + err = capsys.readouterr().err + assert err.count("Insufficient nvidia.com/gpu") == 1 + assert "assigned" in err + + +def test_wait_rewatches_in_short_slices_until_timeout(monkeypatch, fake_clock): + session = make_session() + slices = [] + + class RecordingWatch(FakeWatch): + def stream(self, *_args, **kwargs): + slices.append(kwargs["timeout_seconds"]) + yield from () + + monkeypatch.setattr(runner.watch, "Watch", lambda: RecordingWatch([])) + with pytest.raises(RunnerError, match="timed out"): + runner.wait_until_running_or_terminal(session, "p", timeout=60) + assert len(slices) > 1 + assert max(slices) <= runner.STARTUP_WATCH_SLICE_SECONDS + + +def test_wait_stops_promptly_when_interrupted(monkeypatch, fake_clock): + """Ctrl+C only sets stop_event; the startup wait must notice it at the next slice rather + than carrying on until --timeout.""" + session = make_session() + stop_event = threading.Event() + calls = [] + + class InterruptingWatch(FakeWatch): + def stream(self, *_args, **_kwargs): + calls.append(1) + stop_event.set() + yield from () + + monkeypatch.setattr(runner.watch, "Watch", lambda: InterruptingWatch([])) + with pytest.raises(InterruptError): + runner.wait_until_running_or_terminal(session, "p", timeout=3600, stop_event=stop_event) + assert len(calls) == 1 + + +def test_wait_interrupted_mid_stream_stops_before_next_event(monkeypatch): + session = make_session() + stop_event = threading.Event() + stop_event.set() + monkeypatch.setattr( + runner.watch, "Watch", lambda: FakeWatch([{"object": pod_object(phase="Running")}]) + ) + with pytest.raises(InterruptError): + runner.wait_until_running_or_terminal(session, "p", timeout=60, stop_event=stop_event) + + +def test_discover_job_pod_stops_promptly_when_interrupted(monkeypatch, fake_clock): + session = make_session() + stop_event = threading.Event() + + class InterruptingWatch(FakeWatch): + def stream(self, *_args, **_kwargs): + stop_event.set() + yield from () + + monkeypatch.setattr(runner.watch, "Watch", lambda: InterruptingWatch([])) + with pytest.raises(InterruptError): + runner.discover_job_pod(session, "myjob", timeout=3600, stop_event=stop_event) + + +def test_run_manifest_ctrl_c_during_startup_wait_cleans_up_without_timing_out(monkeypatch, fake_clock): + """End to end: a SIGINT while the pod is stuck Pending ends the run with InterruptError + (not a timeout) and still deletes the pod.""" + session = make_session() + + class PendingThenSigint(FakeWatch): + def stream(self, *_args, **_kwargs): + os.kill(os.getpid(), signal.SIGINT) + yield {"object": pod_object(phase="Pending")} + + monkeypatch.setattr(runner.watch, "Watch", lambda: PendingThenSigint([])) + manifest = {"metadata": {"name": "squarepeg-alpine-abc123"}} + with pytest.raises(InterruptError): + runner.run_manifest(session, basic_spec(mode="pod", timeout=3600, orphan_sweep=False), manifest) + session.core.delete_namespaced_pod.assert_called_once() + + # --- extract_exit_code --- @@ -449,7 +569,7 @@ def test_run_manifest_passes_own_run_id_to_exclude(monkeypatch): monkeypatch.setattr(runner.watch, "Watch", lambda: FakeWatch([{"object": pod_object(phase="Succeeded")}])) captured = {} - def fake_sweep(_session, _spec, *, exclude_run_id=None): + def fake_sweep(_session, _min_age, *, quiet=False, exclude_run_id=None): captured["exclude_run_id"] = exclude_run_id return 0 diff --git a/tests/test_ui.py b/tests/test_ui.py index e6a84ff..1e554a9 100644 --- a/tests/test_ui.py +++ b/tests/test_ui.py @@ -161,3 +161,25 @@ def test_set_run_tag_none_clears_a_previously_set_tag(capsys): err = capsys.readouterr().err assert "[squarepeg:" not in err assert "[squarepeg] hello" in err + + +# --- spinner line fitting --- + + +def test_fit_leaves_short_message_alone(): + assert ui._fit("short", 80) == "short" + + +def test_fit_truncates_to_one_terminal_row(monkeypatch): + monkeypatch.setattr(ui, "_stderr_is_utf", lambda: True) + ui.set_run_tag("ca06ad35") + fitted = ui._fit("x" * 200, 80) + assert len(ui._label()) + len(fitted) == 79 + assert fitted.endswith("…") + + +def test_fit_uses_ascii_ellipsis_on_non_utf_stderr(monkeypatch): + monkeypatch.setattr(ui, "_stderr_is_utf", lambda: False) + fitted = ui._fit("x" * 200, 40) + assert len(ui._label()) + len(fitted) == 39 + assert fitted.endswith("...")