squarepeg 0.1.0__py3-none-any.whl

squarepeg/env_interpolation.py

@@ -0,0 +1,89 @@
+"""Substitute ${VAR} / ${VAR:-default} references in config YAML values.
+
+Runs per-file, before validation and before merging -- so error messages can
+name the exact source file, and everything downstream (validation, merging,
+manifest building) only ever sees fully-resolved strings. Interpolation is
+eager: every string leaf in the parsed document is walked, including inside
+profile bodies that are not selected on this invocation (this matches
+validate_document, which already validates every profile regardless of
+selection) and inside the free-form 'kubernetes'/'job' passthrough sections
+(interpolation is structural and doesn't care that those sections are
+otherwise unvalidated).
+
+Supported forms:
+  ${VAR}             -- substitutes $VAR; a ConfigError if VAR is unset
+  ${VAR:-default}    -- substitutes $VAR if set AND non-empty, else the
+                         literal 'default' (bash ':-' semantics: an
+                         empty-string value counts as unset)
+  $$                 -- a literal '$'
+
+Only ${...} references are recognised -- bare $VAR (no braces) is not a
+reference and passes through untouched. A default is a literal: it is never
+itself re-scanned for further references. Anything else starting with '${'
+that isn't a well-formed reference (${}, ${1BAD}, ${A B}, an unterminated
+${VAR) is a hard error, consistent with this tool's fail-fast posture
+elsewhere (missing config files, unknown config keys, unsupported flags).
+
+Mapping keys are never interpolated, only values.
+"""
+
+import os
+import re
+from collections.abc import Mapping
+from typing import Any
+
+from squarepeg.errors import ConfigError
+
+_TOKEN_RE = re.compile(
+    r"""
+      (?P<escape>\$\$)
+    | \$\{(?P<name>[A-Za-z_][A-Za-z0-9_]*)(?::-(?P<default>[^}]*))?\}
+    | (?P<bad>\$\{[^}]*\}?)
+    """,
+    re.VERBOSE,
+)
+
+
+def interpolate_string(value: str, *, source: str, path: str, env: Mapping[str, str]) -> str:
+    def replace(match: re.Match) -> str:
+        if match.group("escape") is not None:
+            return "$"
+        if match.group("bad") is not None:
+            raise ConfigError(
+                f"malformed reference {match.group('bad')!r} in {path!r} in {source}; "
+                "expected ${VAR} or ${VAR:-default}"
+            )
+        name = match.group("name")
+        default = match.group("default")
+        if default is not None:
+            resolved = env.get(name)
+            return resolved if resolved else default
+        try:
+            return env[name]
+        except KeyError:
+            raise ConfigError(
+                f"environment variable {name!r} is not set, referenced by {path!r} in {source}. "
+                f"Use ${{{name}:-default}} to supply a fallback."
+            ) from None
+
+    return _TOKEN_RE.sub(replace, value)
+
+
+def interpolate_value(value: Any, *, source: str, path: str, env: Mapping[str, str]) -> Any:
+    if isinstance(value, str):
+        return interpolate_string(value, source=source, path=path, env=env)
+    if isinstance(value, Mapping):
+        return {
+            key: interpolate_value(child, source=source, path=f"{path}.{key}" if path else str(key), env=env)
+            for key, child in value.items()
+        }
+    if isinstance(value, list):
+        return [
+            interpolate_value(child, source=source, path=f"{path}[{i}]", env=env) for i, child in enumerate(value)
+        ]
+    return value
+
+
+def interpolate_document(doc: dict, source: str, env: Mapping[str, str] | None = None) -> dict:
+    """Walk a parsed config document, substituting ${VAR}/${VAR:-default} in every string leaf."""
+    return interpolate_value(doc, source=source, path="", env=os.environ if env is None else env)

squarepeg/errors.py

@@ -0,0 +1,34 @@
+class SquarepegError(Exception):
+    """Base class for all squarepeg errors that map to a process exit code."""
+
+    exit_code = 1
+
+
+class UsageError(SquarepegError):
+    """Bad CLI input: malformed flag value, unsupported flag, etc."""
+
+    exit_code = 2
+
+
+class UnsupportedFlagError(UsageError):
+    pass
+
+
+class ConfigError(SquarepegError):
+    """Config file missing, malformed, or invalid."""
+
+    exit_code = 2
+
+
+class RunnerError(SquarepegError):
+    """The tool itself failed to run the container to completion."""
+
+    exit_code = 125
+
+
+class ApiError(RunnerError):
+    """A Kubernetes API call failed; wraps kubernetes.client.ApiException."""
+
+
+class InterruptError(SquarepegError):
+    exit_code = 130

squarepeg/k8s/apierrors.py

@@ -0,0 +1,80 @@
+"""Shared Kubernetes API-error interpretation and the delete primitive.
+
+Split out of runner.py so both the per-run cleanup path and the orphan sweep
+(squarepeg/k8s/orphans.py) can share the exact same error wrapping and 404
+tolerance, rather than risking two copies drifting apart.
+"""
+
+import json
+
+from kubernetes import client
+from kubernetes.client.rest import ApiException
+
+from squarepeg.errors import ApiError
+
+API_ERROR_HINTS = {
+    400: "the apiserver rejected the manifest as malformed; run with --dry-run to inspect it",
+    401: "authentication failed; try 'kubectl get pods' to check your credentials",
+    403: "permission denied; check with 'kubectl auth can-i create pods'",
+    404: "not found; check the namespace exists",
+    409: "already exists; omit --name, or clean up a --keep'd previous run",
+    422: "the apiserver rejected the manifest; run with --dry-run to inspect it",
+}
+
+# A ResourceQuota admission rejection also arrives as HTTP 403, but it has nothing to do with
+# RBAC permissions -- the generic 403 hint above ("check with kubectl auth can-i") is actively
+# misleading here. Detected by the admission plugin's standard "failed quota: ..." message
+# text, which names the missing requests/limits per container.
+QUOTA_ERROR_MARKER = "failed quota"
+QUOTA_HINT = (
+    "the target namespace enforces a ResourceQuota requiring cpu/memory requests (and "
+    "possibly limits) on every container; set --cpus/--memory (or --request-cpu/--request-"
+    "memory/--limit-cpu/--limit-memory), or 'defaults.cpus'/'defaults.memory' in config so "
+    "you don't have to pass them on every run"
+)
+
+
+def api_exception_detail(exc: ApiException) -> str:
+    """Extract the Kubernetes Status object's 'message' field from the response body, if
+    present -- this is almost always far more specific than exc.reason (the bare HTTP
+    reason phrase, e.g. 'Bad Request'), naming the exact field the apiserver rejected."""
+    if exc.body:
+        try:
+            body = json.loads(exc.body)
+        except (TypeError, ValueError):
+            body = None
+        if isinstance(body, dict) and body.get("message"):
+            return body["message"]
+    return exc.reason
+
+
+def wrap_api_exception(exc: ApiException, action: str) -> ApiError:
+    detail = api_exception_detail(exc)
+    if exc.status == 403 and QUOTA_ERROR_MARKER in detail.lower():
+        hint = QUOTA_HINT
+    else:
+        hint = API_ERROR_HINTS.get(exc.status)
+    message = f"{action}: {detail}"
+    if hint:
+        message += f" ({hint})"
+    return ApiError(message)
+
+
+def delete_resource(session, kind: str, name: str) -> bool:
+    """Delete a Pod or Job by name, tolerating 404 (already gone).
+
+    Returns True if the delete actually removed something, False if it was already gone.
+    Raises ApiError for any other failure.
+    """
+    try:
+        if kind == "job":
+            session.batch.delete_namespaced_job(
+                name, session.namespace, body=client.V1DeleteOptions(propagation_policy="Background")
+            )
+        else:
+            session.core.delete_namespaced_pod(name, session.namespace)
+    except ApiException as exc:
+        if exc.status == 404:
+            return False
+        raise wrap_api_exception(exc, f"failed to delete {kind} {name!r}") from exc
+    return True

squarepeg/k8s/dryrun.py

@@ -0,0 +1,47 @@
+"""Submit a manifest to the apiserver as a server-side dry run (nothing is persisted).
+
+Deliberately separate from runner.py: that module's concern is the full create/stream/wait/
+cleanup lifecycle of a resource that actually exists. Server-side dry-run shares none of
+that -- it's a single validate-and-discard apiserver call.
+"""
+
+from kubernetes import client
+from kubernetes.client.rest import ApiException
+
+from squarepeg.k8s.apierrors import wrap_api_exception
+from squarepeg.runspec import RunSpec
+
+SERVER_DRY_RUN = "All"  # the only value the Kubernetes API accepts
+
+_SERIALIZER: client.ApiClient | None = None
+
+
+def _api_client() -> client.ApiClient:
+    global _SERIALIZER
+    if _SERIALIZER is None:
+        _SERIALIZER = client.ApiClient()
+    return _SERIALIZER
+
+
+def _to_plain_dict(response):
+    if isinstance(response, dict):
+        return response
+    return _api_client().sanitize_for_serialization(response)
+
+
+def server_dry_run(session, spec: RunSpec, manifest: dict) -> dict:
+    """Submit `manifest` to the apiserver with dryRun=All and return its response as a plain dict.
+
+    Runs full validation, defaulting, and admission/mutating webhooks; nothing is persisted.
+    Raises ApiError (via wrap_api_exception) on rejection, reusing the same hints as a real
+    create failure.
+    """
+    name = manifest["metadata"]["name"]
+    try:
+        if spec.mode == "job":
+            response = session.batch.create_namespaced_job(session.namespace, manifest, dry_run=SERVER_DRY_RUN)
+        else:
+            response = session.core.create_namespaced_pod(session.namespace, manifest, dry_run=SERVER_DRY_RUN)
+    except ApiException as exc:
+        raise wrap_api_exception(exc, f"failed to validate {spec.mode} {name!r} against the apiserver") from exc
+    return _to_plain_dict(response)

squarepeg/k8s/events.py

@@ -0,0 +1,19 @@
+"""Surface pod events as diagnostics: the difference between 'it hung' and
+'0/12 nodes available: insufficient memory'.
+"""
+
+from kubernetes.client.rest import ApiException
+
+from squarepeg.log import chatter
+
+
+def print_pod_events(session, pod_name: str, *, quiet: bool = False) -> None:
+    try:
+        events = session.core.list_namespaced_event(
+            session.namespace, field_selector=f"involvedObject.name={pod_name}"
+        )
+    except ApiException:
+        return  # best-effort diagnostics; never fail the run because event listing failed
+
+    for event in events.items or []:
+        chatter(f"{pod_name}: {event.reason}: {event.message}", quiet=quiet)

squarepeg/k8s/logs.py

@@ -0,0 +1,126 @@
+"""Follow a pod's logs and write them to stdout, reconnecting on stream drops.
+
+The kubernetes log API only accepts `since_seconds` (not `since_time`), so a
+reconnect can only ask for "everything from roughly N seconds ago" -- some
+overlap with what was already printed is unavoidable. We request timestamps
+internally, strip them before writing (so the user's stdout stays byte-clean
+unless --timestamps was passed), and suppress any re-delivered line whose
+timestamp is <= the last one we already emitted, breaking ties within the
+same timestamp with a small per-timestamp set of already-seen line hashes.
+"""
+
+import math
+import re
+import time
+from datetime import datetime, timezone
+
+from kubernetes.client.rest import ApiException
+
+from squarepeg.errors import ApiError
+from squarepeg.log import chatter
+
+_TIMESTAMP_RE = re.compile(rb"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?Z) (.*)$", re.DOTALL)
+
+MAX_RECONNECT_ATTEMPTS = 5
+MAX_BACKOFF_SECONDS = 30
+
+
+def split_timestamp(raw_line: bytes) -> tuple[str | None, bytes]:
+    match = _TIMESTAMP_RE.match(raw_line)
+    if not match:
+        return None, raw_line
+    ts, content = match.groups()
+    return ts.decode(), content
+
+
+def _parse_rfc3339(ts: str) -> datetime:
+    if ts.endswith("Z"):
+        ts = ts[:-1] + "+00:00"
+    return datetime.fromisoformat(ts)
+
+
+def since_seconds_from(last_ts: str, now: datetime | None = None) -> int:
+    now = now or datetime.now(timezone.utc)
+    elapsed = (now - _parse_rfc3339(last_ts)).total_seconds()
+    return max(1, math.ceil(elapsed) + 1)
+
+
+class _Dedupe:
+    """Suppress log lines already emitted, tolerating out-of-order redelivery on reconnect."""
+
+    def __init__(self):
+        self.last_ts: str | None = None
+        self._seen_at_last_ts: set[bytes] = set()
+
+    def should_emit(self, ts: str | None, content: bytes) -> bool:
+        if ts is None:
+            return True
+        if self.last_ts is None or ts > self.last_ts:
+            self.last_ts = ts
+            self._seen_at_last_ts = {content}
+            return True
+        if ts < self.last_ts:
+            return False
+        if content in self._seen_at_last_ts:
+            return False
+        self._seen_at_last_ts.add(content)
+        return True
+
+
+def _open_stream(session, pod_name, container_name, since_seconds):
+    try:
+        return session.core.read_namespaced_pod_log(
+            name=pod_name,
+            namespace=session.namespace,
+            container=container_name,
+            follow=True,
+            timestamps=True,
+            since_seconds=since_seconds,
+            _preload_content=False,
+        )
+    except ApiException as exc:
+        raise ApiError(f"failed to open log stream for pod {pod_name!r}: {exc.reason}") from exc
+
+
+def stream_logs(session, pod_name: str, container_name: str, stop_event, *, quiet: bool = False, sink=None) -> None:
+    """Follow the pod's logs, writing raw bytes to `sink` (default sys.stdout.buffer).
+
+    Returns when the stream ends normally (the pod is terminal) or `stop_event` is set.
+    Reconnects on a dropped stream, giving up after MAX_RECONNECT_ATTEMPTS with a warning.
+    """
+    if sink is None:
+        import sys
+
+        sink = sys.stdout.buffer
+
+    dedupe = _Dedupe()
+    since_seconds = None
+    attempt = 0
+
+    while not stop_event.is_set():
+        response = _open_stream(session, pod_name, container_name, since_seconds)
+        try:
+            for raw_line in response:
+                if stop_event.is_set():
+                    return
+                ts, content = split_timestamp(raw_line)
+                if dedupe.should_emit(ts, content):
+                    sink.write(content if content.endswith(b"\n") else content + b"\n")
+                    sink.flush()
+            return  # clean EOF: the stream ended on its own
+        except Exception as exc:  # any transport error triggers a reconnect attempt
+            attempt += 1
+            if attempt > MAX_RECONNECT_ATTEMPTS:
+                chatter(
+                    f"log stream for pod {pod_name!r} failed after {attempt - 1} reconnect attempts ({exc}); "
+                    "still waiting for the pod to finish",
+                    quiet=quiet,
+                )
+                return
+            if dedupe.last_ts is not None:
+                since_seconds = since_seconds_from(dedupe.last_ts)
+            backoff = min(2**attempt, MAX_BACKOFF_SECONDS)
+            chatter(f"log stream for pod {pod_name!r} dropped, reconnecting in {backoff}s", quiet=quiet)
+            time.sleep(backoff)
+        finally:
+            response.close()

squarepeg/k8s/orphans.py

@@ -0,0 +1,198 @@
+"""Opportunistic sweep for orphaned squarepeg-managed Pods/Jobs left behind by a
+previous invocation that crashed or was killed before it could clean up after
+itself (machine died, kill -9, network partition, ...).
+
+Runs alongside -- never instead of -- the per-run cleanup in runner.py, and is
+scoped strictly to the CURRENT user's own resources (never a teammate's, even
+in a shared namespace) via the server-side label selector.
+
+Two independent categories of orphan, because a resource that never managed to
+start looks nothing like one that finished and wasn't deleted:
+
+  - terminal:  Succeeded/Failed, older than the threshold since it finished.
+  - stuck:     Pending with a container waiting on a reason squarepeg already
+               treats as fatal-on-sight (see podstate.FAIL_FAST_REASONS -- the
+               same set runner.py uses to fail fast on ITS OWN pod). A bare
+               Pod has no backoffLimit/TTL, so a pod stuck like this (e.g. a
+               typo'd image reference) never reaches Succeeded/Failed and
+               would otherwise sit in the namespace forever.
+
+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.
+"""
+
+from datetime import datetime, timezone
+
+from kubernetes.client.rest import ApiException
+
+from squarepeg.errors import ApiError
+from squarepeg.k8s.apierrors import delete_resource, wrap_api_exception
+from squarepeg.k8s.podstate import (
+    TERMINAL_PHASES,
+    job_terminal_age_seconds,
+    pod_age_seconds,
+    stuck_waiting_reason,
+    terminal_age_seconds,
+)
+from squarepeg.labels import (
+    CREATED_BY_LABEL,
+    KEEP_LABEL,
+    MANAGED_BY_LABEL,
+    MANAGED_BY_VALUE,
+    RUN_ID_LABEL,
+    current_created_by,
+)
+from squarepeg.log import chatter
+from squarepeg.runspec import RunSpec
+
+DEFAULT_MIN_AGE_SECONDS = 300
+
+
+def _label_selector(created_by: str) -> str:
+    return f"{MANAGED_BY_LABEL}={MANAGED_BY_VALUE},{CREATED_BY_LABEL}={created_by},!{KEEP_LABEL}"
+
+
+def _job_name_for_pod(pod) -> str | None:
+    labels = pod.metadata.labels or {}
+    if "job-name" in labels:
+        return labels["job-name"]
+    for ref in pod.metadata.owner_references or []:
+        if ref.kind == "Job":
+            return ref.name
+    return None
+
+
+def _is_own_and_unexcluded(labels: dict, created_by: str, exclude_run_id: str | None) -> bool:
+    if labels.get(CREATED_BY_LABEL) != created_by:
+        return False
+    if exclude_run_id is not None and labels.get(RUN_ID_LABEL) == exclude_run_id:
+        return False
+    # belt-and-braces: the server-side selector already excludes anything carrying the keep
+    # label, but re-checking client-side means a selector typo can never turn into deleting
+    # a resource the user explicitly asked to keep.
+    if labels.get(KEEP_LABEL):
+        return False
+    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."""
+    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:
+        age = pod_age_seconds(pod, now)
+    else:
+        return False
+    return age is not None and age >= min_age_seconds
+
+
+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).
+
+    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.
+    the sweep would start work rather than clean it up. Instead, a Job-owned pod that would
+    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()
+
+    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):
+            continue
+
+        job_name = _job_name_for_pod(pod)
+        if job_name is not None:
+            stuck_child_job_names.add(job_name)
+        else:
+            bare_candidates.append(("pod", pod.metadata.name))
+
+    return bare_candidates, stuck_child_job_names
+
+
+def _find_jobs(
+    jobs, created_by: str, min_age_seconds: int, exclude_run_id: str | None, now: datetime, stuck_child_job_names: set
+):
+    candidates = []
+    for job in jobs:
+        labels = job.metadata.labels or {}
+        if not _is_own_and_unexcluded(labels, created_by, exclude_run_id):
+            continue
+
+        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:
+            # 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))
+    return 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]]:
+    """Discover this user's own orphaned Pods/Jobs.
+
+    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.
+
+    `now` defaults to the real current time; tests pass a fixed value.
+    """
+    now = now or datetime.now(timezone.utc)
+    selector = _label_selector(created_by)
+
+    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)
+    return bare_pod_candidates + job_candidates
+
+
+def sweep_orphans(
+    session, spec: RunSpec, *, exclude_run_id: str | None = None, now: datetime | None = None
+) -> int:
+    """Discover and delete this user's own orphaned resources.
+
+    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
+    (run_manifest) is responsible for making sure a failure here can never change the
+    current invocation's own exit code -- see runner.py's finally block.
+    """
+    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
+        )
+    except ApiException as exc:
+        wrapped = wrap_api_exception(exc, "failed to list resources for orphan sweep")
+        raise ApiError(f"{wrapped} (set 'orphan_sweep: false' in config to disable)") from exc
+
+    deleted = 0
+    for kind, name in candidates:
+        try:
+            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)
+
+    if deleted:
+        chatter(f"swept {deleted} orphaned resource(s) from previous runs", quiet=spec.quiet)
+    return deleted