Kueue job priorities (#13)

* fix: Don't build container images when submitting to Ray cluster

* deps: Upgrade Ray to 2.11.0

* feat(kueue): Implement scheduling options, including priority classes

Jobs can now specify `SchedulingOptions`, which include the Kueue
local queue and priority class names.

Adds example priority classes (background, development, production)
and preemption options to the cluster queue manifest.

A sample "production" job has been added to simplify the demo of
preemption and priority based workload ordering in Kueue.

* feat(kueue): Validate existence of local queues and workload priority classes

* fix: Make all job options optional by default

* chore: Split up the Runner module

The `runner.py` module was growing to a point where readability was hampered.

This commit splits it up into the `jobs.runner` package, with submodules
grouped by execution platform.

Also, introduces a `KubernetesNamespaceMixin` to simplify determining the
current or desired namespace for a workload.

* fix: Canonicalize all paths in Docker image options

Author: AdrianoKF

example_hello.py

@@ -1,6 +1,7 @@
 import logging
+import time
 
-from jobs import ImageOptions, JobOptions, ResourceOptions, job
+from jobs import ImageOptions, JobOptions, ResourceOptions, SchedulingOptions, job
 from jobs.cli import submit_job
 
 
@@ -9,14 +10,18 @@
         # A job with explicit Dockerfile
         image=ImageOptions(
             dockerfile="Dockerfile",
-            name="hello-world-yaml",
+            name="hello-world-dev",
             tag="latest",
         ),
-        resources=ResourceOptions(memory="2Gi", cpu="1", gpu=None),
+        resources=ResourceOptions(memory="256Mi", cpu="1"),
+        scheduling=SchedulingOptions(
+            priority_class="background",
+        ),
     )
 )
 def hello_world():
     print("Hello, World!")
+    time.sleep(60)
 
 
 if __name__ == "__main__":

example_hello_prod.py

@@ -0,0 +1,27 @@
+import logging
+import time
+
+from jobs import ImageOptions, JobOptions, ResourceOptions, SchedulingOptions, job
+from jobs.cli import submit_job
+
+
+@job(
+    options=JobOptions(
+        # A job with explicit Dockerfile
+        image=ImageOptions(dockerfile="Dockerfile"),
+        resources=ResourceOptions(memory="256Mi", cpu="4"),
+        scheduling=SchedulingOptions(
+            priority_class="production",
+        ),
+    )
+)
+def prod_training():
+    print("Hello, World!")
+    time.sleep(60)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    logging.getLogger("urllib3.connectionpool").setLevel(logging.INFO)
+
+    submit_job(prod_training)

raycluster-manifest.yaml

@@ -1,5 +1,5 @@
 image:
-  tag: 2.9.0-py310
+  tag: 2.11.0-py311
 worker:
   maxReplicas: 1
   resources:

raycluster-values.yaml

@@ -171,7 +171,7 @@ pyyaml==6.0.1
     #   kubernetes
     #   pre-commit
     #   ray
-ray==2.10.0
+ray==2.11.0
 referencing==0.34.0
     # via
     #   jsonschema

requirements.txt

@@ -9,6 +9,13 @@ metadata:
   name: "cluster-queue"
 spec:
   namespaceSelector: {} # match all.
+  # Allow preemption of workloads by higher priority ones
+  preemption:
+    reclaimWithinCohort: Any
+    borrowWithinCohort:
+      policy: LowerPriority
+      maxPriorityThreshold: 100
+    withinClusterQueue: LowerPriority
   resourceGroups:
     - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
       flavors:
@@ -28,3 +35,24 @@ metadata:
   name: "user-queue"
 spec:
   clusterQueue: "cluster-queue"
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: WorkloadPriorityClass
+metadata:
+  name: background
+value: 1
+description: "Background (=lowest) priority"
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: WorkloadPriorityClass
+metadata:
+  name: development
+value: 100
+description: "Development priority"
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: WorkloadPriorityClass
+metadata:
+  name: production
+value: 1000
+description: "Production priority"

single-clusterqueue-setup.yaml

@@ -1,13 +1,21 @@
 from jobs import assembler
 from jobs.image import Image
-from jobs.job import ImageOptions, Job, JobOptions, ResourceOptions, job
+from jobs.job import (
+    ImageOptions,
+    Job,
+    JobOptions,
+    ResourceOptions,
+    SchedulingOptions,
+    job,
+)
 
 __all__ = [
     Image,
     Job,
     JobOptions,
     ImageOptions,
     ResourceOptions,
+    SchedulingOptions,
     job,
     assembler,
 ]

src/jobs/__init__.py

@@ -53,7 +53,7 @@ def submit_job(job: Job) -> None:
     logging.debug(f"Execution mode: {mode}")
 
     image: Image | None = None
-    if mode in [ExecutionMode.DOCKER, ExecutionMode.KUEUE, ExecutionMode.RAYCLUSTER]:
+    if mode in [ExecutionMode.DOCKER, ExecutionMode.KUEUE]:
         image = job.build_image()
         if image is None:
             raise RuntimeError("Could not build container image")

src/jobs/cli.py

@@ -92,6 +92,15 @@ def to_ray(self) -> RayResourceOptions:
         return remove_none_values(options)
 
 
+@dataclass(frozen=True)
+class SchedulingOptions:
+    priority_class: str | None = None
+    """Kueue priority class name"""
+
+    queue_name: str | None = None
+    """Kueue local queue name"""
+
+
 @dataclass(frozen=True)
 class ImageOptions:
     name: str | None = None
@@ -112,7 +121,7 @@ def build_mode(self) -> BuildMode:
     def _to_pathlib(self, attr: str) -> None:
         val = self.__getattribute__(attr)
         if isinstance(val, str):
-            object.__setattr__(self, attr, Path(val).absolute())
+            object.__setattr__(self, attr, Path(val).resolve())
 
     def __post_init__(self) -> None:
         def _is_yaml(path: AnyPath) -> bool:
@@ -136,14 +145,17 @@ def _is_yaml(path: AnyPath) -> bool:
             raise ValueError(f"Build context must be a directory: {self.build_context}")
 
         if self.dockerfile and not self.dockerfile.is_relative_to(self.build_context):
-            raise ValueError("Dockerfile must be relative to build context")
+            raise ValueError(
+                f"Dockerfile must be relative to build context {self.build_context}"
+            )
 
 
 @dataclass(frozen=True)
 class JobOptions:
-    resources: ResourceOptions | None
+    resources: ResourceOptions | None = None
     """Resource requests for this job in Kubernetes format (see https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/#resource-units-in-kubernetes)"""
-    image: ImageOptions | None
+    image: ImageOptions | None = None
+    scheduling: SchedulingOptions | None = None
 
 
 class Job: