From ef4dc19a714f652a1de223e0a64f445c941b2a8f Mon Sep 17 00:00:00 2001
From: Haoyu Gao <haoyugao@google.com>
Date: Thu, 3 Sep 2026 08:40:00 -0700
Subject: [PATCH] Add support for running DeepSWE distributed rollouts with
 SandboxFleet

PiperOrigin-RevId: 975747562
---
 .../examples/deepswe_dist/README.md             | 17 +++++++++++++----
 .../examples/deepswe_dist/deepswe.py            | 10 ++++++++++
 .../examples/deepswe_dist/launcher.sh           | 10 ++++++++++
 .../examples/deepswe_dist/run_deepswe_dist.py   |  2 +-
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/tunix/experimental/examples/deepswe_dist/README.md b/tunix/experimental/examples/deepswe_dist/README.md
index 67455f08a..80afc0a5f 100644
--- a/tunix/experimental/examples/deepswe_dist/README.md
+++ b/tunix/experimental/examples/deepswe_dist/README.md
@@ -11,11 +11,20 @@ distributed GSM8K example:
    because it 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
-```
\ No newline at end of file
+```
+
+```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`.
diff --git a/tunix/experimental/examples/deepswe_dist/deepswe.py b/tunix/experimental/examples/deepswe_dist/deepswe.py
index d90522623..f28099b0c 100644
--- a/tunix/experimental/examples/deepswe_dist/deepswe.py
+++ b/tunix/experimental/examples/deepswe_dist/deepswe.py
@@ -230,6 +230,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,
diff --git a/tunix/experimental/examples/deepswe_dist/launcher.sh b/tunix/experimental/examples/deepswe_dist/launcher.sh
index e46c1352d..378b61b28 100644
--- a/tunix/experimental/examples/deepswe_dist/launcher.sh
+++ b/tunix/experimental/examples/deepswe_dist/launcher.sh
@@ -59,6 +59,9 @@ SEED=${SEED:-42}
 ENV_BACKEND=${ENV_BACKEND:-kubernetes}
 SCAFFOLD=${SCAFFOLD:-r2egym}
 USE_AGENT_SANDBOX=${USE_AGENT_SANDBOX:-0}
+SANDBOX_NAMESPACE=${SANDBOX_NAMESPACE:-rl-tunix-swebench}
+SANDBOX_NODE_SELECTOR_KEY=${SANDBOX_NODE_SELECTOR_KEY:-}
+SANDBOX_NODE_SELECTOR_VAL=${SANDBOX_NODE_SELECTOR_VAL:-}
 STEP_TIMEOUT_SECS=${STEP_TIMEOUT_SECS:-1800}
 REWARD_TIMEOUT_SECS=${REWARD_TIMEOUT_SECS:-1800}
 ROLLOUT_MAX_CONCURRENCY=${ROLLOUT_MAX_CONCURRENCY:-64}
@@ -290,6 +293,13 @@ echo "Launching DeepSWE rollout node..."
   export TPU_CHIPS_PER_HOST_BOUNDS=${TPU_CHIPS_PER_HOST_BOUNDS}
   export TPU_HOST_BOUNDS=${TPU_HOST_BOUNDS}
   export LIBTPU_INIT_ARGS="--deepsea_chips_per_host_bounds=${TPU_CHIPS_PER_HOST_BOUNDS} --deepsea_host_bounds=${TPU_HOST_BOUNDS}"
+  if [[ "$USE_AGENT_SANDBOX" == "1" || "$USE_AGENT_SANDBOX" == "true" || "$USE_AGENT_SANDBOX" == "True" ]]; then
+    export NAMESPACE="$SANDBOX_NAMESPACE"
+    if [[ -n "$SANDBOX_NODE_SELECTOR_KEY" && -n "$SANDBOX_NODE_SELECTOR_VAL" ]]; then
+      export NODE_SELECTOR_KEY="$SANDBOX_NODE_SELECTOR_KEY"
+      export NODE_SELECTOR_VAL="$SANDBOX_NODE_SELECTOR_VAL"
+    fi
+  fi
   export PYTHONUNBUFFERED=1
   print_command "Rollout command" "${ROLLOUT_CMD[@]}"
   exec "${ROLLOUT_CMD[@]}" > "$ROLLOUT_LOG" 2>&1
diff --git a/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py b/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py
index a783531d4..ef63ef6ff 100644
--- a/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py
+++ b/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py
@@ -358,4 +358,4 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None:
 
 
 if __name__ == "__main__":
-  main(sys.argv[1:])
\ No newline at end of file
+  main(sys.argv[1:])