add k8s_launcher for deepswe

tunix/experimental/examples/deepswe_dist/README.md

@@ -11,11 +11,41 @@ distributed GSM8K example:
    is already a generic PeftTrainer V2 worker.
 
 The first milestone is intentionally small: run one trainer+rollout pipeline
-step with `BETA=0.0` and `WEIGHT_SYNC_MODE=none`. Reference KL, Raiden weight
-sync, agent-sandbox prewarming, and production DeepSWE-scale settings can be
-layered on after the basic pipeline is stable.
+step with `BETA=0.0` and `WEIGHT_SYNC_MODE=none`. The default path uses the
+regular DeepSWE `SWEEnv` backend. Set `USE_AGENT_SANDBOX=1` to construct
+`SWEEnv` with `SandboxFleet` inside the rollout worker process.
 
 ```bash
 cd tunix/experimental/examples/deepswe_dist
 BETA=0.0 WEIGHT_SYNC_MODE=none MAX_STEPS=1 BATCH_SIZE=1 NUM_GENERATIONS=2 ./launcher.sh
-```
+```
+
+```bash
+cd tunix/experimental/examples/deepswe_dist
+USE_AGENT_SANDBOX=1 BETA=0.0 WEIGHT_SYNC_MODE=none MAX_STEPS=1 BATCH_SIZE=1 NUM_GENERATIONS=2 ./launcher.sh
+```
+
+For sandbox placement, set `SANDBOX_NAMESPACE`, `SANDBOX_NODE_SELECTOR_KEY`, and
+`SANDBOX_NODE_SELECTOR_VAL` before launching. The launcher forwards them to the
+rollout worker as the `agent_sandbox_rl` variables consumed by `SWEEnv`.
+
+To deploy on Kubernetes / GKE:
+
+```bash
+cd tunix/experimental/examples/deepswe_dist
+./k8s_launcher.sh --command start
+```
+
+Or run with sandbox enabled on GKE:
+
+```bash
+cd tunix/experimental/examples/deepswe_dist
+USE_AGENT_SANDBOX=1 ./k8s_launcher.sh --command start
+```
+
+To stop all jobsets:
+
+```bash
+./k8s_launcher.sh --command stop
+```
+

tunix/experimental/examples/deepswe_dist/deepswe.py

@@ -232,6 +232,16 @@ def __init__(
       group_id = prompt_id or None
     if pair_index is None:
       pair_index = group_index
+    if kwargs.get("use_agent_sandbox") and kwargs.get("fleet") is None:
+      logging.info(
+          "Initializing DeepSWE SandboxFleet in rollout worker "
+          "(max_concurrency=%s).",
+          group_size,
+      )
+      kwargs["fleet"] = swe_env._init_global_fleet(  # pylint: disable=protected-access
+          tasks=[entry],
+          max_concurrency=group_size,
+      )
 
     super().__init__(
         entry=entry,

tunix/experimental/examples/deepswe_dist/enter_kube_context.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Usage:
+#   source tunix/experimental/examples/deepswe_dist/enter_kube_context.sh
+
+PROJECT=${PROJECT:-cloud-tpu-multipod-dev}
+REGION=${REGION:-us-central1}
+ZONE=${ZONE:-us-central1-a}
+CLUSTER=${CLUSTER:-trellis-demo-0810}
+
+export KUBECONFIG="${KUBECONFIG:-$HOME/.kube/config.$PROJECT.$REGION.$CLUSTER}"
+CONTEXT_NAME="gke_${PROJECT}_${REGION}_${CLUSTER}"
+
+if kubectl config get-contexts "$CONTEXT_NAME" &>/dev/null; then
+  kubectl config use-context "$CONTEXT_NAME" >/dev/null || true
+else
+  gcloud container clusters get-credentials $CLUSTER --region=$REGION --project=$PROJECT --dns-endpoint || { echo "gcloud get-credentials failed" >&2; return 1 2>/dev/null || exit 1; }
+  kubectl config use-context "$CONTEXT_NAME" >/dev/null || { echo "kubectl use-context failed" >&2; return 1 2>/dev/null || exit 1; }
+fi
+kubectl config set-context --current --namespace=default >/dev/null || true