Support user-specified image pull policy

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]
@@ -64,4 +58,5 @@ markers = [
 ]
 
 [tool.uv.sources]
-aai-jobq = { git = "https://github.com/aai-institute/jobq", branch = "main", subdirectory = "client" }
+# FIXME: Revert to `main` once the `pull-policy` branch is merged
+aai-jobq = { git = "https://github.com/aai-institute/jobq", branch = "pull-policy", subdirectory = "client" }

backend/src/jobq_server/models.py

@@ -5,7 +5,7 @@
 from typing import Annotated, Any, Self, TypeAlias
 
 from annotated_types import Ge
-from jobq import JobOptions
+from jobq import ImagePullPolicy, JobOptions
 from pydantic import AfterValidator, BaseModel, Field, StrictStr
 
 from jobq_server.utils.kueue import JobId, KueueWorkload, WorkloadSpec, WorkloadStatus
@@ -60,6 +60,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

@@ -50,7 +50,7 @@ def job_fn(): ...
             detail=f"unsupported job execution mode: {opts.mode!r}",
         )
 
-    image = Image(opts.image_ref)
+    image = Image(opts.image_ref, opts.pull_policy)
     workload_id = runner.run(job, image, opts.submission_context)
     return workload_id
 

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()

client/src/cli/commands/submit.py

@@ -8,7 +8,7 @@
 import openapi_client
 from cli.types import Settings
 from cli.util import with_job_mgmt_api
-from jobq import Image, Job
+from jobq import Image, ImagePullPolicy, Job
 from jobq.submission_context import SubmissionContext
 from openapi_client import ExecutionMode
 
@@ -32,6 +32,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 +45,7 @@ def _submit_remote_job(
         file=job.file,
         image_ref=_build_image(job, mode).tag,
         mode=mode,
+        pull_policy=pull_policy,
         options=openapi_client.JobOptions.model_validate(job.options.model_dump()),
         submission_context=SubmissionContext().to_dict(),
     )
@@ -57,13 +59,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 +123,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

client/src/jobq/__init__.py

@@ -1,5 +1,5 @@
 from jobq import assembler
-from jobq.image import Image
+from jobq.image import Image, ImagePullPolicy
 from jobq.job import (
     ImageOptions,
     Job,
@@ -14,6 +14,7 @@
     "Job",
     "JobOptions",
     "ImageOptions",
+    "ImagePullPolicy",
     "ResourceOptions",
     "SchedulingOptions",
     "job",

client/src/jobq/image.py

@@ -1,3 +1,15 @@
+from enum import Enum
+
+
+class ImagePullPolicy(Enum):
+    ALWAYS = "Always"
+    NEVER = "Never"
+    IFNOTPRESENT = "IfNotPresent"
+
+
 class Image:
-    def __init__(self, tag: str) -> None:
+    def __init__(
+        self, tag: str, pull_policy: ImagePullPolicy = ImagePullPolicy.ALWAYS
+    ) -> None:
         self.tag = tag
+        self.pull_policy = pull_policy