(wip) Add image pull policy to CLI kwargs

Author: maxmynter

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
 

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