Implement job status API endpoint (#46)

* feat(backend): Implement job status endpoint

The `GET /jobs/{job_id}/status` API endpoint reports the high-level
execution status of a Job (identified by its UID).

The associated Kueue workload for the job is looked up and its status
is used to determine the execution status (pending, executing, failed,
succeeded).

* fix(backend): Return 404 for missing workloads

* fix(backend): Improve readability of `traverse`

* fix(backend): Rename kubernetes helper module

The rename prevents a name clash with the kubernetes client SDK package.

Author: AdrianoKF

.editorconfig

@@ -9,4 +9,4 @@ root = true
 end_of_line = lf
 # Consistency
 insert_final_newline = true
-indent_size = 2
+indent_size = 4

backend/src/jobs_server/exceptions.py

@@ -0,0 +1,12 @@
+class WorkloadNotFound(Exception):
+    """Exception indicating a non-existing Kueue workload"""
+
+    def __init__(
+        self,
+        message: str = "workload not found",
+        uid: str | None = None,
+        namespace: str | None = None,
+    ):
+        super().__init__(message)
+        self.uid = uid
+        self.namespace = namespace

backend/src/jobs_server/models.py

@@ -1,4 +1,5 @@
 import re
+from enum import StrEnum
 from typing import Annotated, Any, TypeAlias
 
 from jobs import JobOptions
@@ -49,3 +50,10 @@ class WorkloadIdentifier(BaseModel):
 
     namespace: StrictStr
     uid: StrictStr
+
+
+class JobStatus(StrEnum):
+    PENDING = "pending"
+    EXECUTING = "executing"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"

backend/src/jobs_server/routers/jobs.py

@@ -1,8 +1,15 @@
 from fastapi import APIRouter, HTTPException
 from jobs import Image, Job
 
-from jobs_server.models import CreateJobModel, ExecutionMode, WorkloadIdentifier
+from jobs_server.exceptions import WorkloadNotFound
+from jobs_server.models import (
+    CreateJobModel,
+    ExecutionMode,
+    JobId,
+    WorkloadIdentifier,
+)
 from jobs_server.runner import Runner
+from jobs_server.utils.kueue import KueueWorkload
 
 router = APIRouter(tags=["Job management"])
 
@@ -33,3 +40,12 @@ def job_fn(): ...
     image = Image(opts.image_ref)
     workload_id = runner.run(job, image, opts.submission_context)
     return workload_id
+
+
+@router.get("/jobs/{uid}/status")
+async def status(uid: JobId, namespace: str = "default"):
+    try:
+        workload = KueueWorkload.for_managed_resource(uid, namespace)
+        return workload.execution_status
+    except WorkloadNotFound as e:
+        raise HTTPException(404, "workload not found") from e

backend/src/jobs_server/runner/kueue.py

@@ -7,7 +7,7 @@
 
 from jobs_server.models import SubmissionContext, WorkloadIdentifier
 from jobs_server.runner.base import ExecutionMode, Runner, _make_executor_command
-from jobs_server.utils.kubernetes import (
+from jobs_server.utils.k8s import (
     KubernetesNamespaceMixin,
     gvk,
     k8s_annotations,

backend/src/jobs_server/runner/ray.py

@@ -20,7 +20,7 @@
 
 from jobs_server.models import SubmissionContext, WorkloadIdentifier
 from jobs_server.runner.base import ExecutionMode, Runner, _make_executor_command
-from jobs_server.utils.kubernetes import (
+from jobs_server.utils.k8s import (
     KubernetesNamespaceMixin,
     gvk,
     k8s_annotations,

backend/src/jobs_server/utils/helpers.py

@@ -10,3 +10,87 @@ def remove_none_values(d: T) -> T:
     """Remove all keys with a ``None`` value from a dict."""
     filtered_dict = {k: v for k, v in d.items() if v is not None}
     return cast(T, filtered_dict)
+
+
+def traverse(d: Any, key_path: str, sep: str = ".", strict: bool = True) -> Any:
+    """
+    Retrieve a value from a nested Mapping-like object using a key path.
+
+    If the object behaves like a Mapping (i.e., implements `__getitem__`),
+    this function will be used to access elements.
+    If it behaves like an object (i.e., `__getattr__`), the path will be
+    resolved as attributes, instead.
+
+    Parameters
+    ----------
+    d : dict
+        The object to traverse.
+    key_path : str
+        A string representing the path of keys, separated by `sep`.
+    sep : str, optional
+        The separator used to split the key path into individual keys.
+        Default is ".".
+    strict : bool, optional
+        If False, return None when a key in the path does not exist.
+        If True, raise a KeyError when a key does not exist.
+        Default is False.
+
+    Returns
+    -------
+    Any
+        The value at the specified key path, or None if a key is missing
+        and `strict` is False.
+
+    Raises
+    ------
+    KeyError
+        If `strict` is True and any key in the path does not exist.
+
+    Examples
+    --------
+    >>> d = {"foo": {"bar": {"baz": 42}}}
+    >>> traverse(d, "foo.bar.baz")
+    42
+
+    >>> traverse(d, "foo.bar.qux", strict=False) is None
+    True
+
+    >>> traverse(d, "foo.bar.qux", strict=True)
+    Traceback (most recent call last):
+    ...
+    KeyError: 'qux'
+    """
+
+    def has_item(container, key):
+        if hasattr(container, "__contains__"):
+            # Container behaves like a dict
+            return key in container
+        elif hasattr(container, key):
+            # Container behaves like an object
+            return True
+        else:
+            return False
+
+    def get_item(container, key, default=None):
+        if hasattr(container, "__getitem__"):
+            # Container behaves like a dict
+            try:
+                return container[key]
+            except (KeyError, IndexError, TypeError):
+                return default
+        elif hasattr(container, key):
+            # Container behaves like an object
+            return getattr(container, key, default)
+        else:
+            return default
+
+    keys = key_path.split(sep)
+    for key in keys:
+        # Bail out on missing entries in strict mode
+        if strict and not has_item(d, key):
+            raise KeyError(key)
+
+        d = get_item(d, key)
+        if d is None:
+            return None
+    return d

backend/src/jobs_server/utils/k8s.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import kubernetes
+from jobs.job import Job
+
+from jobs_server.models import SubmissionContext
+from jobs_server.utils.helpers import traverse
+
+
+def sanitize_rfc1123_domain_name(s: str) -> str:
+    """Sanitize a string to be compliant with RFC 1123 domain name
+
+    Note: Any invalid characters are replaced with dashes."""
+
+    # TODO: This is obviously wildly incomplete
+    return s.replace("_", "-")
+
+
+def k8s_annotations(
+    job: Job, context: SubmissionContext | None = None
+) -> dict[str, str]:
+    """Determine the Kubernetes annotations for a Job"""
+    # Store as annotations since labels have restrictive value formats
+    options = job.options.labels if job.options else {}
+    context = {"x-jobby.io/submission-context": json.dumps(context)} if context else {}
+    return options | context
+
+
+@dataclass
+class GroupVersionKind:
+    group: str
+    version: str
+    kind: str
+
+
+class KubernetesObject(Protocol):
+    @property
+    def api_version(self) -> str: ...
+
+    @property
+    def kind(self) -> str: ...
+
+
+def gvk(obj: KubernetesObject | dict[str, Any]) -> GroupVersionKind:
+    kind = obj.kind if hasattr(obj, "kind") else obj["kind"]
+    if "/" in (
+        api_version := obj.api_version
+        if hasattr(obj, "api_version")
+        else obj["apiVersion"]
+    ):
+        group, version = api_version.split("/")
+    else:
+        group, version = "", api_version
+
+    return GroupVersionKind(group, version, kind)
+
+
+class KubernetesNamespaceMixin:
+    """Determine the desired or current Kubernetes namespace."""
+
+    def __init__(self, **kwargs):
+        kubernetes.config.load_config()
+        self._namespace: str | None = kwargs.get("namespace")
+
+    @property
+    def namespace(self) -> str:
+        _, active_context = kubernetes.config.list_kube_config_contexts()
+        current_namespace = active_context["context"].get("namespace")
+        return self._namespace or current_namespace
+
+
+def filter_conditions(
+    obj: dict[str, Any],
+    typ: str | None = None,
+    reason: str | None = None,
+    message: str | None = None,
+):
+    """
+    Filters Kubernetes object conditions based on specified attributes.
+
+    This function filters the `status.conditions` field of a Kubernetes object
+    by matching conditions against the provided `type`, `reason`, and `message`
+    attributes. Only conditions that match all specified attributes are included
+    in the result.
+
+    Parameters
+    ----------
+    obj : dict[str, Any]
+        The Kubernetes object, typically a dictionary representing a Kubernetes
+        resource, containing a `status.conditions` field.
+    typ : str, optional
+        The type of condition to filter by. If `None`, this filter is not applied.
+    reason : str, optional
+        The reason attribute to filter by. If `None`, this filter is not applied.
+    message : str, optional
+        The message attribute to filter by. If `None`, this filter is not applied.
+
+    Returns
+    -------
+    list[dict[str, Any]]
+        A list of conditions that match the specified filters. Each condition
+        is represented as a dictionary.
+
+    Notes
+    -----
+    - The function assumes that the `status.conditions` field exists in the
+      provided object and that it is a list of condition dictionaries.
+    - If no conditions match the specified filters, an empty list is returned.
+
+    Examples
+    --------
+    >>> obj = {
+    ...     "status": {
+    ...         "conditions": [
+    ...             {"type": "Ready", "reason": "DeploymentCompleted", "message": "Deployment successful."},
+    ...             {"type": "Failed", "reason": "DeploymentFailed", "message": "Deployment failed due to timeout."}
+    ...         ]
+    ...     }
+    ... }
+    >>> filter_conditions(obj, typ="Ready")
+    [{'type': 'Ready', 'reason': 'DeploymentCompleted', 'message': 'Deployment successful.'}]
+
+    >>> filter_conditions(obj, reason="DeploymentFailed")
+    [{'type': 'Failed', 'reason': 'DeploymentFailed', 'message': 'Deployment failed due to timeout.'}]
+    """
+
+    def _match(cond):
+        return all([
+            typ is None or cond["type"] == typ,
+            reason is None or cond["reason"] == reason,
+            message is None or cond["message"] == message,
+        ])
+
+    return [cond for cond in traverse(obj, "status.conditions") if _match(cond)]