feat(backend): Support for image pull policies

This commit adds support for pull policies in the runners.

The job submission endpoint passes the pull policy to the runner, by
default the policy is set to `Always` to prevent issues with stale
images.

The e2e test has been updated to pass `IfNotPresent` as the pull policy,
since the job image for the test suite is built locally and cannot be
pulled from a registry.

Author: AdrianoKF

backend/src/jobq_server/routers/jobs.py

@@ -51,7 +51,7 @@ def job_fn(): ...
         )
 
     image = Image(opts.image_ref, opts.pull_policy)
-    workload_id = runner.run(job, image, opts.submission_context)
+    workload_id = runner.run(job, image, opts.submission_context, opts.pull_policy)
     return workload_id
 
 

backend/src/jobq_server/runner/base.py

@@ -1,7 +1,7 @@
 import abc
 from typing import ClassVar, Self
 
-from jobq import Image, Job
+from jobq import Image, ImagePullPolicy, Job
 
 from jobq_server.models import ExecutionMode, SubmissionContext, WorkloadIdentifier
 
@@ -11,7 +11,11 @@ class Runner(abc.ABC):
 
     @abc.abstractmethod
     def run(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: Image,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy = ImagePullPolicy.ALWAYS,
     ) -> WorkloadIdentifier | None: ...
 
     @classmethod

backend/src/jobq_server/runner/docker.py

@@ -2,8 +2,9 @@
 import textwrap
 
 import docker
-from jobq import Image, Job
+from jobq import Image, ImagePullPolicy, Job
 from jobq.job import DockerResourceOptions
+from typing_extensions import override
 
 from jobq_server.models import ExecutionMode, SubmissionContext
 from jobq_server.runner.base import Runner, _make_executor_command
@@ -15,7 +16,18 @@ def __init__(self, **kwargs):
         super().__init__()
         self._client = docker.from_env()
 
-    def run(self, job: Job, image: Image, context: SubmissionContext) -> None:
+    @override
+    def run(
+        self,
+        job: Job,
+        image: Image,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy = ImagePullPolicy.ALWAYS,
+    ) -> None:
+        if pull_policy == ImagePullPolicy.ALWAYS:
+            logging.debug("Pulling container image %s", image.tag)
+            self._client.images.pull(image.tag)
+
         command = _make_executor_command(job)
 
         resource_kwargs: DockerResourceOptions = {

backend/src/jobq_server/runner/kueue.py

@@ -1,9 +1,10 @@
 import logging
 from dataclasses import asdict
 
-from jobq import Image, Job
+from jobq import Image, ImagePullPolicy, Job
 from jobq.types import K8sResourceKind
 from kubernetes import client
+from typing_extensions import override
 
 from jobq_server.models import ExecutionMode, SubmissionContext, WorkloadIdentifier
 from jobq_server.runner.base import Runner, _make_executor_command
@@ -24,7 +25,11 @@ def __init__(self, k8s: KubernetesService, **kwargs: str) -> None:
         self._queue = kwargs.get("local_queue", "user-queue")
 
     def _make_job_crd(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: Image,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy,
     ) -> client.V1Job:
         if not job.options:
             raise ValueError("Job options must be specified")
@@ -40,7 +45,7 @@ def _make_job_crd(
         # Job container
         container = client.V1Container(
             image=image.tag,
-            image_pull_policy="IfNotPresent",
+            image_pull_policy=pull_policy.value,
             name="workload",
             command=_make_executor_command(job),
             resources=(
@@ -70,12 +75,17 @@ def _make_job_crd(
             ),
         )
 
+    @override
     def run(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: Image,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy = ImagePullPolicy.ALWAYS,
     ) -> WorkloadIdentifier:
         logging.info(f"Submitting job {job.name} to Kueue")
 
-        k8s_job = self._make_job_crd(job, image, context)
+        k8s_job = self._make_job_crd(job, image, context, pull_policy)
         batch_api = client.BatchV1Api()
         resource: client.V1Job = batch_api.create_namespaced_job(
             self._k8s.namespace, k8s_job

backend/src/jobq_server/runner/ray.py

@@ -5,9 +5,10 @@
 from dataclasses import asdict
 
 import yaml
-from jobq import Image, Job
+from jobq import Image, ImagePullPolicy, Job
 from jobq.types import K8sResourceKind
 from kubernetes import client
+from typing_extensions import override
 
 from jobq_server.models import ExecutionMode, SubmissionContext, WorkloadIdentifier
 from jobq_server.runner.base import Runner, _make_executor_command
@@ -29,7 +30,11 @@ def __init__(self, k8s: KubernetesService, **kwargs):
         self._k8s = k8s
 
     def _create_ray_job(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: Image,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy,
     ) -> dict:
         """Create a ``RayJob`` Kubernetes resource for the Kuberay operator."""
 
@@ -79,7 +84,7 @@ def _create_ray_job(
                                     {
                                         "name": "head",
                                         "image": image.tag,
-                                        "imagePullPolicy": "IfNotPresent",
+                                        "imagePullPolicy": pull_policy.value,
                                         "resources": {
                                             "requests": res_opts.to_kubernetes(
                                                 kind=K8sResourceKind.REQUESTS
@@ -101,7 +106,7 @@ def _create_ray_job(
                             {
                                 "name": "ray-submit",
                                 "image": image.tag,
-                                "imagePullPolicy": "IfNotPresent",
+                                "imagePullPolicy": pull_policy.value,
                             }
                         ],
                     },
@@ -111,14 +116,19 @@ def _create_ray_job(
 
         return manifest
 
+    @override
     def run(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: Image,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy,
     ) -> WorkloadIdentifier:
         logging.info(
             f"Submitting RayJob {job.name} to namespace {self._k8s.namespace!r}"
         )
 
-        manifest = self._create_ray_job(job, image, context)
+        manifest = self._create_ray_job(job, image, context, pull_policy)
         api = client.CustomObjectsApi()
         obj = api.create_namespaced_custom_object(
             "ray.io", "v1", self._k8s.namespace, "rayjobs", manifest

backend/tests/e2e/test_jobs.py

@@ -45,6 +45,7 @@ def test_job_lifecycle(
             ),
             resources=ResourceOptions(cpu="1", memory="512Mi"),
         ),
+        pull_policy="IfNotPresent",
     )
 
     # Submit a job for execution