Support DeepSWE on MaxText models in Tunix by niting · Pull Request #2051 · google/tunix

niting · 2026-08-31T19:13:41Z

Primarily adds a script train_maxtext_nb.py that uses MaxText models on a DeepSWE setup. The DeepSWE setup uses agentic sandboxes which maintains warmpools for faster training execution. Modifies Dockerfile accordingly to install necessary dependencies, adds a flag to swe_env.py to allow limiting warmpool replicas. Makes various robustness fixes for error logging etc.

I have added all the necessary unit tests for my change.
I have verified that my change does not break existing code and all unit tests pass.
I have added all appropriate doc-strings/documentation.
My PR is based on the latest changes of the main branch (if unsure, rebase the code).
I have signed the Contributor License Agreement.
I have followed Contribution Guidelines.

tianshub · 2026-08-31T21:59:47Z

    images = self._extract_images(batch)
    if images and self.fleet:
+      target_replicas = (
+          self.max_warmpool_size


nit: this can be self.max_warmpool_size or self.num_generations

tianshub · 2026-08-31T22:02:34Z

+NODE_SELECTOR_VAL = args.node_selector_val
+
+
+# Monkeypatch r2egym DockerRuntime to dynamically configure Kubernetes nodeSelector.


put this this in a standalone r2e_gym_helper.py file? so we can reuse it later

I also realized I didn't need the agentic sandbox patch. Only the patch for overriding the node selector in r2egym.

tianshub · 2026-08-31T22:03:48Z

-              "Fatal error in runner for pair %d: %s",
-              env.extra_kwargs["pair_index"],
+              "Fatal error in runner for pair %s: %s",
+              env.extra_kwargs.get("pair_index", "unknown"),


when will pair_index be None?

I removed this entirely. It's possible this was happening in certain conditoins that are not common.

tianshub · 2026-08-31T22:04:51Z


-RUN bash /app/scripts/install_tunix_vllm_requirement.sh
+# Install SFT/MaxText dependencies (unconditional)
+RUN pip install --upgrade flax && \


uv pip install upgrade?

wang2yn84 · 2026-09-01T19:27:29Z

        max_concurrent=effective_max_concurrent,
        window_size=batch_size,
-        max_warmpool_size=num_generations,
+        max_warmpool_size=max_warmpool_size if max_warmpool_size is not None else num_generations,


Is it better to use max_concurrency here?

That will likely not be right? max_concurrency for my script is 256, but I want two replicas per stream so I set max_warmpool_size to 2 so that I have 512 sandboxes with 256 on standby at a time.

I see, so the max_warmpool_size is per image. The naming is a bit misleading. nit, maybe max_warmpool_per_stream or per_image?

wang2yn84 · 2026-09-01T19:38:21Z

Hi Nitin, please squash the commits into one before submission.

- Add train_maxtext_nb for agentic GRPO training with MaxText and vLLM. - Add r2e_gym_helper with nodeSelector patching for Kubernetes execution on GKE. - Add max_warmpool_size to swe_env to allow limiting warm pool replicas. - Pass tokenizer_pad_id in AutoModel for MaxText causal LM loading. - Update Dockerfile to optimize layer caching and configure DeepSWE MaxText dependencies.

-- f98e5632d4fc224a1b0213709b867cceee35e165 by Nitin Gangahar <niting@google.com>: Support DeepSWE on MaxText models in Tunix - Add train_maxtext_nb for agentic GRPO training with MaxText and vLLM. - Add r2e_gym_helper with nodeSelector patching for Kubernetes execution on GKE. - Add max_warmpool_replicas to swe_env to allow limiting warm pool replicas. - Pass tokenizer_pad_id in AutoModel for MaxText causal LM loading. - Update Dockerfile to optimize layer caching and configure DeepSWE MaxText dependencies. COPYBARA_INTEGRATE_REVIEW=#2051 from niting:deepswe f98e5632d4fc224a1b0213709b867cceee35e165 PiperOrigin-RevId: 974800060

niting · 2026-09-02T21:17:22Z

This was merged separately: https://github.com/google/tunix/blob/main/examples/deepswe/train_maxtext_nb.py, due to some shenanigans with our internal (Google) importer. Closing this PR out.

niting requested review from abheesht17, hgao327, jiangyangmu, lc5211, s-noghabi, sizhit2, tianshub and wang2yn84 as code owners August 31, 2026 19:13

github-actions Bot assigned lc5211 Aug 31, 2026

niting had a problem deploying to testing August 31, 2026 19:14 — with GitHub Actions Error

niting force-pushed the deepswe branch from 4e8eac3 to b232764 Compare August 31, 2026 19:20

niting changed the title ~~Add max_warmpool_size to swe_env~~ Support DeepSWE on MaxText models in Tunix Aug 31, 2026

niting temporarily deployed to testing August 31, 2026 19:21 — with GitHub Actions Inactive

niting force-pushed the deepswe branch from b232764 to e006637 Compare August 31, 2026 20:51

niting had a problem deploying to testing August 31, 2026 20:51 — with GitHub Actions Error

niting force-pushed the deepswe branch from e006637 to 68a98cd Compare August 31, 2026 21:00

niting had a problem deploying to testing August 31, 2026 21:01 — with GitHub Actions Error

niting force-pushed the deepswe branch from 68a98cd to 6334bcf Compare August 31, 2026 21:10

niting had a problem deploying to testing August 31, 2026 21:11 — with GitHub Actions Error

niting force-pushed the deepswe branch from 6334bcf to cd16b1d Compare August 31, 2026 21:14

niting temporarily deployed to testing August 31, 2026 21:15 — with GitHub Actions Inactive

tianshub reviewed Aug 31, 2026

View reviewed changes

niting force-pushed the deepswe branch from cd16b1d to 686878e Compare August 31, 2026 23:05

niting had a problem deploying to testing August 31, 2026 23:06 — with GitHub Actions Error

niting force-pushed the deepswe branch from 686878e to da221e2 Compare August 31, 2026 23:08

niting had a problem deploying to testing August 31, 2026 23:09 — with GitHub Actions Error

niting force-pushed the deepswe branch from da221e2 to f356dba Compare August 31, 2026 23:09

niting temporarily deployed to testing August 31, 2026 23:10 — with GitHub Actions Inactive

tianshub approved these changes Aug 31, 2026

View reviewed changes

niting force-pushed the deepswe branch from f356dba to c086afe Compare September 1, 2026 19:09

wang2yn84 reviewed Sep 1, 2026

View reviewed changes

niting force-pushed the deepswe branch from c086afe to 0ec9787 Compare September 1, 2026 20:09

wang2yn84 added the ready-to-submit label Sep 1, 2026

wang2yn84 temporarily deployed to testing September 1, 2026 20:21 — with GitHub Actions Inactive

niting force-pushed the deepswe branch from 0ec9787 to 9085e32 Compare September 1, 2026 21:19

niting had a problem deploying to testing September 1, 2026 21:19 — with GitHub Actions Error

niting force-pushed the deepswe branch from 9085e32 to f98e563 Compare September 1, 2026 21:23

niting temporarily deployed to testing September 1, 2026 21:24 — with GitHub Actions Inactive

wang2yn84 added ready-to-submit and removed ready-to-submit labels Sep 2, 2026

wang2yn84 temporarily deployed to testing September 2, 2026 00:05 — with GitHub Actions Inactive

niting closed this Sep 2, 2026

		NODE_SELECTOR_VAL = args.node_selector_val


		# Monkeypatch r2egym DockerRuntime to dynamically configure Kubernetes nodeSelector.

Uh oh!

Conversation

niting commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wang2yn84 Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wang2yn84 commented Sep 1, 2026

Uh oh!

niting commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

niting commented Aug 31, 2026 •

edited

Loading

wang2yn84 Sep 1, 2026 •

edited

Loading