feat(backend): Add endpoint to list workloads (#83)

The endpoint can either return workload identifiers, or optionally the
workload metadata.

Additionally, the path prefix has been removed from the endpoints in the
`jobs` router, instead it is now mounted with an appropriate path prefix.

Author: AdrianoKF

backend/src/jobs_server/__main__.py

@@ -20,7 +20,7 @@ async def lifespan(app: FastAPI):
     lifespan=lifespan,
 )
 
-app.include_router(jobs.router)
+app.include_router(jobs.router, prefix="/jobs")
 
 
 @app.get("/health", include_in_schema=False)

backend/src/jobs_server/models.py

@@ -1,16 +1,13 @@
 import json
 import re
 from enum import StrEnum
-from typing import TYPE_CHECKING, Annotated, Any, Self, TypeAlias
+from typing import Annotated, Any, Self, TypeAlias
 
 from annotated_types import Ge
 from jobs import JobOptions
 from pydantic import AfterValidator, BaseModel, Field, StrictStr
 
-from jobs_server.utils.kueue import JobId, WorkloadSpec, WorkloadStatus
-
-if TYPE_CHECKING:
-    from jobs_server.dependencies import ManagedWorkload
+from jobs_server.utils.kueue import JobId, KueueWorkload, WorkloadSpec, WorkloadStatus
 
 
 def validate_image_ref(ref: str) -> str:
@@ -75,6 +72,21 @@ class WorkloadIdentifier(BaseModel):
     namespace: StrictStr
     uid: StrictStr
 
+    @classmethod
+    def from_kueue_workload(cls, workload: KueueWorkload) -> Self:
+        if len(workload.metadata.owner_references) != 1:
+            raise ValueError(
+                f"Workload {workload.metadata.uid} has multiple owner references: {workload.metadata.owner_references}"
+            )
+        owner_ref = workload.metadata.owner_references[0]
+        return cls(
+            group=owner_ref.api_version.split("/")[0],
+            version=owner_ref.api_version.split("/")[1],
+            kind=owner_ref.kind,
+            uid=owner_ref.uid,
+            namespace=workload.metadata.namespace,
+        )
+
 
 class JobStatus(StrEnum):
     PENDING = "pending"
@@ -94,7 +106,7 @@ class WorkloadMetadata(BaseModel):
     kueue_status: WorkloadStatus
 
     @classmethod
-    def from_managed_workload(cls, workload: "ManagedWorkload") -> Self:
+    def from_kueue_workload(cls, workload: KueueWorkload) -> Self:
         if workload.owner_uid is None:
             raise ValueError("Workload has no owner UID")
         return WorkloadMetadata(
@@ -111,3 +123,9 @@ class LogOptions(BaseModel):
         default=-1,
         description="Number of tail lines of logs, -1 for all",
     )
+
+
+class ListWorkloadModel(BaseModel):
+    name: str
+    id: WorkloadIdentifier
+    metadata: WorkloadMetadata | None = None

backend/src/jobs_server/routers/jobs.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, HTTPException, Response
+from fastapi import APIRouter, Depends, HTTPException, Query, Response
 from fastapi import status as http_status
 from fastapi.responses import StreamingResponse
 from jobs import Image, Job
@@ -11,6 +11,7 @@
 from jobs_server.models import (
     CreateJobModel,
     ExecutionMode,
+    ListWorkloadModel,
     LogOptions,
     WorkloadIdentifier,
     WorkloadMetadata,
@@ -22,7 +23,7 @@
 router = APIRouter(tags=["Job management"])
 
 
-@router.post("/jobs")
+@router.post("")
 async def submit_job(
     opts: CreateJobModel,
     k8s: Kubernetes,
@@ -52,20 +53,20 @@ def job_fn(): ...
     return workload_id
 
 
-@router.get("/jobs/{uid}/status")
+@router.get("/{uid}/status")
 async def status(
     workload: ManagedWorkload,
 ) -> WorkloadMetadata:
     try:
-        return WorkloadMetadata.from_managed_workload(workload)
+        return WorkloadMetadata.from_kueue_workload(workload)
     except ValueError as e:
         raise HTTPException(
             status_code=http_status.HTTP_404_NOT_FOUND,
             detail=f"Workload not found or invalid: {str(e)}",
         ) from e
 
 
-@router.get("/jobs/{uid}/logs")
+@router.get("/{uid}/logs")
 async def logs(
     workload: ManagedWorkload,
     k8s: Kubernetes,
@@ -86,7 +87,7 @@ async def logs(
         raise HTTPException(http_status.HTTP_400_BAD_REQUEST, "pod not ready") from e
 
 
-@router.post("/jobs/{uid}/stop")
+@router.post("/{uid}/stop")
 async def stop_workload(
     uid: JobId,
     workload: ManagedWorkload,
@@ -104,3 +105,28 @@ async def stop_workload(
             http_status.HTTP_500_INTERNAL_SERVER_ERROR,
             "Failed to terminate workload",
         ) from e
+
+
+@router.get("", response_model_exclude_unset=True)
+async def list_jobs(
+    k8s: Kubernetes,
+    include_metadata: Annotated[bool, Query()] = False,
+) -> list[ListWorkloadModel]:
+    workloads = k8s.list_workloads()
+    if include_metadata:
+        return [
+            ListWorkloadModel(
+                name=workload.metadata.name,
+                id=WorkloadIdentifier.from_kueue_workload(workload),
+                metadata=WorkloadMetadata.from_kueue_workload(workload),
+            )
+            for workload in workloads
+        ]
+    else:
+        return [
+            ListWorkloadModel(
+                name=workload.metadata.name,
+                id=WorkloadIdentifier.from_kueue_workload(workload),
+            )
+            for workload in workloads
+        ]

backend/src/jobs_server/services/k8s.py

@@ -113,3 +113,16 @@ def delete_resource(
             namespace=namespace,
             body=client.V1DeleteOptions(propagation_policy=propagation_policy),
         )
+
+    def list_workloads(self, namespace: str | None = None) -> list[KueueWorkload]:
+        api = client.CustomObjectsApi()
+        workloads = api.list_namespaced_custom_object(
+            group="kueue.x-k8s.io",
+            version="v1beta1",
+            namespace=namespace or self.namespace,
+            plural="workloads",
+        )
+        return [
+            KueueWorkload.model_validate(workload)
+            for workload in workloads.get("items", [])
+        ]

backend/src/jobs_server/utils/kueue.py

@@ -125,26 +125,24 @@ class KueueWorkload(BaseModel):
     spec: WorkloadSpec
     status: WorkloadStatus
 
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
     @field_validator("metadata", mode="before")
     def create_metadata(cls, metadata: client.V1ObjectMeta) -> client.V1ObjectMeta:
         return build_metadata(metadata)
 
-    owner_uid: JobId | None = None
-
-    model_config = ConfigDict(
-        arbitrary_types_allowed=True,
-    )
+    @property
+    def owner_uid(self) -> JobId:
+        return self.metadata.owner_references[0].uid
 
     @classmethod
     def for_managed_resource(cls, uid: str, namespace: str):
         workload = workload_by_managed_uid(uid, namespace)
         if workload.get("status") is None:
             raise WorkloadNotFound(uid=uid, namespace=namespace)
         result = cls.model_validate(workload)
-
-        # speed up subsequent lookups of associated resource by memoizing the managed resource UID
-        result.owner_uid = uid
-
         return result
 
     @property
@@ -175,9 +173,6 @@ def managed_resource(self):
     def pod(self) -> client.V1Pod:
         api = client.CoreV1Api()
 
-        if self.owner_uid is None:
-            self.owner_uid = self.managed_resource.metadata["uid"]
-
         if self.managed_resource.kind == "Job":
             # Jobs are simple, they directly control the pods (which we can look up by their controller UID)
             controller_uid = self.owner_uid

backend/tests/e2e/test_jobs.py

@@ -10,6 +10,7 @@
     CreateJobModel,
     ExecutionMode,
     JobStatus,
+    ListWorkloadModel,
     WorkloadIdentifier,
     WorkloadMetadata,
 )
@@ -72,6 +73,23 @@ def test_job_lifecycle(
 
         time.sleep(0.5)
 
+    # Check that the workload is listed
+    response = client.get("/jobs")
+    assert response.status_code == 200
+    workloads = [
+        ListWorkloadModel.model_validate(workload) for workload in response.json()
+    ]
+    assert managed_resource_id in [workload.id for workload in workloads]
+
+    # Detailed workload listing including metadata
+    response = client.get("/jobs?include_metadata=true")
+    assert response.status_code == 200
+    workloads = [
+        ListWorkloadModel.model_validate(workload) for workload in response.json()
+    ]
+    assert managed_resource_id in [workload.id for workload in workloads]
+    assert all(workload.metadata is not None for workload in workloads)
+
     # Check workload logs (retry if pod is not ready yet)
     while True:
         response = client.get(f"/jobs/{managed_resource_id.uid}/logs")