Support user-specified image pull policy (#132)

* (wip) Add image pull policy to CLI kwargs

* deps: Bump client dependency in backend

This is needed to access the newly added types in the client package.

* tests(backend): Don't install deps into external cluster for E2E tests

* 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.

* chore: Regenerate OpenAPI client

* fix(client): Adapt image pull policy type for API

* chore: Remove FIXME

* refactor: Move ImagePullPolicy type into backend

The `Runner` interface has been changed to accept an `ImageRef`,
instead of importing the (useless) `Image` type from the client package.

Co-authored-by: Max Mynter <maxmynter@users.noreply.github.com>

* docs: Update mention of default image pull policy

---------

Co-authored-by: Adrian Rumpold <a.rumpold@gmail.com>
Co-authored-by: Max Mynter <maxmynter@users.noreply.github.com>

Author: maxmynter

backend/pyproject.toml

@@ -16,13 +16,7 @@ maintainers = [
     { name = "Adrian Rumpold", email = "a.rumpold@appliedai-institute.de" },
 ]
 license = { text = "Apache-2.0" }
-dependencies = [
-    "fastapi",
-    "uvicorn",
-    "docker",
-    "kubernetes",
-    "aai-jobq",
-]
+dependencies = ["fastapi", "uvicorn", "docker", "kubernetes", "aai-jobq"]
 dynamic = ["version"]
 
 [project.optional-dependencies]

backend/src/jobq_server/models.py

@@ -32,6 +32,13 @@ def validate_image_ref(ref: str) -> str:
 
 ImageRef = Annotated[str, AfterValidator(validate_image_ref)]
 
+
+class ImagePullPolicy(StrEnum):
+    ALWAYS = "Always"
+    NEVER = "Never"
+    IFNOTPRESENT = "IfNotPresent"
+
+
 SubmissionContext: TypeAlias = dict[str, Any]
 
 
@@ -60,6 +67,7 @@ class CreateJobModel(BaseModel):
     image_ref: ImageRef
     mode: ExecutionMode
     options: JobOptions
+    pull_policy: ImagePullPolicy = ImagePullPolicy.ALWAYS
     submission_context: SubmissionContext = Field(default_factory=dict)
 
 

backend/src/jobq_server/routers/jobs.py

@@ -6,7 +6,7 @@
 from fastapi import APIRouter, Depends, HTTPException, Query, Response
 from fastapi import status as http_status
 from fastapi.responses import StreamingResponse
-from jobq import Image, Job
+from jobq import Job
 
 from jobq_server.dependencies import Kubernetes, ManagedWorkload
 from jobq_server.exceptions import PodNotReadyError
@@ -50,8 +50,9 @@ def job_fn(): ...
             detail=f"unsupported job execution mode: {opts.mode!r}",
         )
 
-    image = Image(opts.image_ref)
-    workload_id = runner.run(job, image, opts.submission_context)
+    workload_id = runner.run(
+        job, opts.image_ref, opts.submission_context, opts.pull_policy
+    )
     return workload_id
 
 

backend/src/jobq_server/runner/base.py

@@ -1,17 +1,27 @@
 import abc
 from typing import ClassVar, Self
 
-from jobq import Image, Job
+from jobq import Job
 
-from jobq_server.models import ExecutionMode, SubmissionContext, WorkloadIdentifier
+from jobq_server.models import (
+    ExecutionMode,
+    ImagePullPolicy,
+    ImageRef,
+    SubmissionContext,
+    WorkloadIdentifier,
+)
 
 
 class Runner(abc.ABC):
     _impls: ClassVar[dict[ExecutionMode, type[Self]]] = {}
 
     @abc.abstractmethod
     def run(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: ImageRef,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy = ImagePullPolicy.ALWAYS,
     ) -> WorkloadIdentifier | None: ...
 
     @classmethod

backend/src/jobq_server/runner/docker.py

@@ -2,10 +2,16 @@
 import textwrap
 
 import docker
-from jobq import Image, Job
+from jobq import Job
 from jobq.job import DockerResourceOptions
-
-from jobq_server.models import ExecutionMode, SubmissionContext
+from typing_extensions import override
+
+from jobq_server.models import (
+    ExecutionMode,
+    ImagePullPolicy,
+    ImageRef,
+    SubmissionContext,
+)
 from jobq_server.runner.base import Runner, _make_executor_command
 from jobq_server.utils.helpers import remove_none_values
 
@@ -15,7 +21,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: ImageRef,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy = ImagePullPolicy.ALWAYS,
+    ) -> None:
+        if pull_policy == ImagePullPolicy.ALWAYS:
+            logging.debug("Pulling container image %s", image)
+            self._client.images.pull(image)
+
         command = _make_executor_command(job)
 
         resource_kwargs: DockerResourceOptions = {
@@ -27,7 +44,7 @@ def run(self, job: Job, image: Image, context: SubmissionContext) -> None:
             resource_kwargs = res.to_docker()
 
         container: docker.api.client.ContainerApiMixin = self._client.containers.run(
-            image=image.tag,
+            image=image,
             command=command,
             detach=True,
             **remove_none_values(resource_kwargs),

backend/src/jobq_server/runner/kueue.py

@@ -1,11 +1,18 @@
 import logging
 from dataclasses import asdict
 
-from jobq import Image, Job
+from jobq import Job
 from jobq.types import K8sResourceKind
 from kubernetes import client
-
-from jobq_server.models import ExecutionMode, SubmissionContext, WorkloadIdentifier
+from typing_extensions import override
+
+from jobq_server.models import (
+    ExecutionMode,
+    ImagePullPolicy,
+    ImageRef,
+    SubmissionContext,
+    WorkloadIdentifier,
+)
 from jobq_server.runner.base import Runner, _make_executor_command
 from jobq_server.services.k8s import KubernetesService
 from jobq_server.utils.k8s import (
@@ -24,7 +31,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: ImageRef,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy,
     ) -> client.V1Job:
         if not job.options:
             raise ValueError("Job options must be specified")
@@ -39,8 +50,8 @@ def _make_job_crd(
 
         # Job container
         container = client.V1Container(
-            image=image.tag,
-            image_pull_policy="IfNotPresent",
+            image=image,
+            image_pull_policy=pull_policy.value,
             name="workload",
             command=_make_executor_command(job),
             resources=(
@@ -70,12 +81,17 @@ def _make_job_crd(
             ),
         )
 
+    @override
     def run(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: ImageRef,
+        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,11 +5,18 @@
 from dataclasses import asdict
 
 import yaml
-from jobq import Image, Job
+from jobq import Job
 from jobq.types import K8sResourceKind
 from kubernetes import client
-
-from jobq_server.models import ExecutionMode, SubmissionContext, WorkloadIdentifier
+from typing_extensions import override
+
+from jobq_server.models import (
+    ExecutionMode,
+    ImagePullPolicy,
+    ImageRef,
+    SubmissionContext,
+    WorkloadIdentifier,
+)
 from jobq_server.runner.base import Runner, _make_executor_command
 from jobq_server.services.k8s import KubernetesService
 from jobq_server.utils.k8s import (
@@ -29,7 +36,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: ImageRef,
+        context: SubmissionContext,
+        pull_policy: ImagePullPolicy,
     ) -> dict:
         """Create a ``RayJob`` Kubernetes resource for the Kuberay operator."""
 
@@ -49,9 +60,6 @@ def _create_ray_job(
         suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=4))
         job_id = f"{job.name}-{suffix}"
 
-        # FIXME: Image pull policy should be configurable
-        # It is currently hardcoded to "IfNotPresent" to support running
-        # the E2E tests in a cluster without a proper image registry.
         manifest = {
             "apiVersion": "ray.io/v1",
             "kind": "RayJob",
@@ -78,8 +86,8 @@ def _create_ray_job(
                                 "containers": [
                                     {
                                         "name": "head",
-                                        "image": image.tag,
-                                        "imagePullPolicy": "IfNotPresent",
+                                        "image": image,
+                                        "imagePullPolicy": pull_policy.value,
                                         "resources": {
                                             "requests": res_opts.to_kubernetes(
                                                 kind=K8sResourceKind.REQUESTS
@@ -100,8 +108,8 @@ def _create_ray_job(
                         "containers": [
                             {
                                 "name": "ray-submit",
-                                "image": image.tag,
-                                "imagePullPolicy": "IfNotPresent",
+                                "image": image,
+                                "imagePullPolicy": pull_policy.value,
                             }
                         ],
                     },
@@ -111,14 +119,19 @@ def _create_ray_job(
 
         return manifest
 
+    @override
     def run(
-        self, job: Job, image: Image, context: SubmissionContext
+        self,
+        job: Job,
+        image: ImageRef,
+        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/conftest.py

@@ -66,6 +66,11 @@ def cluster() -> Generator[KubernetesCluster, None, None]:
     """Create a Kubernetes cluster for testing based on environment variable"""
     cluster_type = os.getenv("E2E_CLUSTER_TYPE", "minikube").lower()
     context = os.getenv("E2E_K8S_CONTEXT")
+
+    # Don't install dependencies into an external cluster
+    install_kuberay = context is None
+    install_kueue = context is None
+
     if cluster_type == "minikube":
         cluster = MinikubeCluster(name=context)
     else:
@@ -74,8 +79,10 @@ def cluster() -> Generator[KubernetesCluster, None, None]:
     try:
         # Install Kuberay first, so that the CRDs (RayJob, RayCluster)
         # are available for Kueue to be watched.
-        setup_kuberay(cluster)
-        setup_kueue(cluster)
+        if install_kuberay:
+            setup_kuberay(cluster)
+        if install_kueue:
+            setup_kueue(cluster)
         yield cluster
     finally:
         cluster.delete()

backend/tests/e2e/test_jobs.py

@@ -45,6 +45,8 @@ def test_job_lifecycle(
             ),
             resources=ResourceOptions(cpu="1", memory="512Mi"),
         ),
+        # e2e test image is built locally and sideloaded into the cluster
+        pull_policy="IfNotPresent",
     )
 
     # Submit a job for execution

backend/uv.lock

@@ -11,6 +11,7 @@
 from jobq import Image, Job
 from jobq.submission_context import SubmissionContext
 from openapi_client import ExecutionMode
+from openapi_client.models.image_pull_policy import ImagePullPolicy
 
 
 def submit(args: argparse.Namespace, settings: Settings) -> None:
@@ -32,6 +33,7 @@ def _submit_remote_job(
     client: openapi_client.JobManagementApi,
     job: Job,
     mode: ExecutionMode,
+    pull_policy: ImagePullPolicy,
     settings: Settings,
 ) -> None:
     # Job options sent to server do not need image options
@@ -44,6 +46,7 @@ def _submit_remote_job(
         file=job.file,
         image_ref=_build_image(job, mode).tag,
         mode=mode,
+        pull_policy=openapi_client.ImagePullPolicy(pull_policy.value),
         options=openapi_client.JobOptions.model_validate(job.options.model_dump()),
         submission_context=SubmissionContext().to_dict(),
     )
@@ -57,13 +60,14 @@ def submit_job(
     settings: Settings,
 ) -> None:
     mode = args.mode
+    pull_policy = args.pull_policy
     logging.debug(f"Execution mode: {mode}")
     match mode:
         case ExecutionMode.LOCAL:
             # Run the job locally
             job()
         case _:
-            _submit_remote_job(job, mode, settings=settings)
+            _submit_remote_job(job, mode, pull_policy, settings=settings)
 
 
 def discover_job(args: argparse.Namespace) -> Job:
@@ -120,6 +124,12 @@ def add_parser(subparsers: Any, parent: argparse.ArgumentParser) -> None:
         choices=list(ExecutionMode),
         type=ExecutionMode,
     )
+    parser.add_argument(
+        "--pull-policy",
+        default=ImagePullPolicy.ALWAYS,
+        choices=list(ImagePullPolicy),
+        type=ImagePullPolicy,
+    )
 
     parser.add_argument("entrypoint")
     # TODO: Factor out into command class