From fa56d183ce98cbf974e0eac1e40f505b3c622f9c Mon Sep 17 00:00:00 2001
From: Hoda Mohammadi <hodamo@google.com>
Date: Wed, 21 Oct 2026 23:26:10 -0700
Subject: [PATCH] perf(deploy): pre-pull the worker image on every GPU node

---
 k8s/deploy/base/kustomization.yaml      |  1 +
 k8s/deploy/base/worker-prepull.yaml     | 43 +++++++++++++++++++++++++
 k8s/deploy/kind-dra/kustomization.yaml  | 10 ++++++
 k8s/deploy/lora-kind/kustomization.yaml | 12 +++++++
 4 files changed, 66 insertions(+)
 create mode 100644 k8s/deploy/base/worker-prepull.yaml

diff --git a/k8s/deploy/base/kustomization.yaml b/k8s/deploy/base/kustomization.yaml
index a5bbe7ea3..af9f4a672 100644
--- a/k8s/deploy/base/kustomization.yaml
+++ b/k8s/deploy/base/kustomization.yaml
@@ -10,6 +10,7 @@ resources:
   - sampler-router.yaml
   - redis.yaml
   - api-server.yaml
+  - worker-prepull.yaml
 
 generatorOptions:
   disableNameSuffixHash: true
diff --git a/k8s/deploy/base/worker-prepull.yaml b/k8s/deploy/base/worker-prepull.yaml
new file mode 100644
index 000000000..966ae5c8c
--- /dev/null
+++ b/k8s/deploy/base/worker-prepull.yaml
@@ -0,0 +1,43 @@
+# Pulls the worker image onto every GPU node the scheduler may use, at deploy
+# time instead of inside the first job. A worker image pull takes minutes on a
+# fresh node, and a cold job pays it for its trainer and its sampler. The
+# container only sleeps and holds no GPU. The kustomize images transformer
+# rewrites the image like every other OpenRL image, so a release that bumps
+# the tag re-pulls on every node as the DaemonSet rolls.
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: open-rl-worker-prepull
+  labels:
+    app: open-rl-worker-prepull
+spec:
+  selector:
+    matchLabels:
+      app: open-rl-worker-prepull
+  updateStrategy:
+    type: RollingUpdate
+    rollingUpdate:
+      # Nothing runs here, so every node can pull at once.
+      maxUnavailable: 100%
+  template:
+    metadata:
+      labels:
+        app: open-rl-worker-prepull
+    spec:
+      nodeSelector:
+        openrl.io/enabled: "true"
+      # GPU nodes only: TPU nodes run a different worker image and keep their taint.
+      tolerations:
+        - key: nvidia.com/gpu
+          operator: Exists
+          effect: NoSchedule
+      containers:
+        - name: prepull
+          image: ghcr.io/gke-labs/open-rl/server:latest
+          command: ["sleep", "infinity"]
+          resources:
+            requests:
+              cpu: 10m
+              memory: 16Mi
+            limits:
+              memory: 64Mi
diff --git a/k8s/deploy/kind-dra/kustomization.yaml b/k8s/deploy/kind-dra/kustomization.yaml
index 8680b6509..8b11355cc 100644
--- a/k8s/deploy/kind-dra/kustomization.yaml
+++ b/k8s/deploy/kind-dra/kustomization.yaml
@@ -32,6 +32,16 @@ configMapGenerator:
       - ENABLE_GCP_TRACE="0"
 
 patches:
+  # Workers pull from the cluster-local registry, so there is nothing to pre-pull.
+  - target:
+      kind: DaemonSet
+      name: open-rl-worker-prepull
+    patch: |
+      $patch: delete
+      apiVersion: apps/v1
+      kind: DaemonSet
+      metadata:
+        name: open-rl-worker-prepull
   - path: pvc-local.yaml
 
   - target:
diff --git a/k8s/deploy/lora-kind/kustomization.yaml b/k8s/deploy/lora-kind/kustomization.yaml
index d45e03a79..81092838e 100644
--- a/k8s/deploy/lora-kind/kustomization.yaml
+++ b/k8s/deploy/lora-kind/kustomization.yaml
@@ -4,3 +4,15 @@ namespace: openrl-system
 resources:
   - ../base
   - storage.yaml
+
+patches:
+  # Workers pull from the cluster-local registry, so there is nothing to pre-pull.
+  - target:
+      kind: DaemonSet
+      name: open-rl-worker-prepull
+    patch: |
+      $patch: delete
+      apiVersion: apps/v1
+      kind: DaemonSet
+      metadata:
+        name: open-rl-worker-prepull