diff --git a/tests/experimental/orchestrator/algorithm_adapter_test.py b/tests/experimental/orchestrator/algorithm_adapter_test.py
new file mode 100644
index 000000000..f7bd3b720
--- /dev/null
+++ b/tests/experimental/orchestrator/algorithm_adapter_test.py
@@ -0,0 +1,85 @@
+# 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.
+
+"""Unit tests for AlgorithmAdapter (GRPOAdapter & PPOAdapter)."""
+
+from absl.testing import absltest
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.rl import algo_core
+
+
+class AlgorithmAdapterTest(absltest.TestCase):
+
+ def test_grpo_advantage_normalization(self):
+ adapter = algorithm_adapter.GRPOAdapter(group_size=4)
+ rewards = [1.0, 2.0, 3.0, 4.0]
+ advs = adapter.compute_advantages(rewards, num_generations=4)
+
+ self.assertEqual(len(advs), 4)
+ # Mean should be 0.0
+ self.assertAlmostEqual(float(np.mean(advs)), 0.0, places=4)
+ # Std should be 1.0
+ self.assertAlmostEqual(float(np.std(advs)), 1.0, places=4)
+
+ def test_grpo_create_trainer_payloads(self):
+ adapter = algorithm_adapter.GRPOAdapter(group_size=2)
+ item1 = datatypes.TrajectoryItem(
+ pair_index=0,
+ group_id="g1",
+ start_step=0,
+ traj=datatypes.Trajectory(reward=1.0),
+ )
+ item1.prompt_tokens = np.array([1, 2], dtype=np.int32)
+ item1.completion_tokens = np.array([3, 4], dtype=np.int32)
+ item1.action_mask = np.array([1, 1], dtype=np.float32)
+
+ item2 = datatypes.TrajectoryItem(
+ pair_index=1,
+ group_id="g1",
+ start_step=0,
+ traj=datatypes.Trajectory(reward=2.0),
+ )
+ item2.prompt_tokens = np.array([1, 2], dtype=np.int32)
+ item2.completion_tokens = np.array([5, 6], dtype=np.int32)
+ item2.action_mask = np.array([1, 1], dtype=np.float32)
+
+ payloads = adapter.create_trainer_payloads([item1, item2], rewards=[1.0, 2.0])
+ self.assertLen(payloads, 2)
+ self.assertIsInstance(payloads[0], datatypes.RLTrainerPayload)
+ self.assertLess(payloads[0].advantages[0], 0.0)
+ self.assertGreater(payloads[1].advantages[0], 0.0)
+ self.assertEqual(adapter.loss_fn(), algo_core.grpo_loss_fn)
+
+ def test_ppo_advantages_and_trainer_payloads(self):
+ adapter = algorithm_adapter.PPOAdapter(group_size=2, gamma=0.99, lam=0.95)
+ item = datatypes.TrajectoryItem(
+ pair_index=0,
+ group_id="g1",
+ start_step=0,
+ traj=datatypes.Trajectory(reward=1.0),
+ )
+ item.prompt_tokens = np.array([10], dtype=np.int32)
+ item.completion_tokens = np.array([20], dtype=np.int32)
+
+ payloads = adapter.create_trainer_payloads([item], rewards=[2.0], values=[1.0])
+ self.assertLen(payloads, 1)
+ self.assertAlmostEqual(payloads[0].advantages[0], 1.0)
+ self.assertAlmostEqual(payloads[0].returns[0], 2.0)
+ self.assertEqual(adapter.loss_fn(), algo_core.ppo_policy_loss_fn)
+
+
+if __name__ == "__main__":
+ absltest.main()
diff --git a/tests/experimental/orchestrator/async_rl_program_test.py b/tests/experimental/orchestrator/async_rl_program_test.py
new file mode 100644
index 000000000..f5b900a4b
--- /dev/null
+++ b/tests/experimental/orchestrator/async_rl_program_test.py
@@ -0,0 +1,177 @@
+# 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.
+
+"""Tests for AsyncRLProgram and StandardRLProgram."""
+
+import asyncio
+from unittest import mock
+
+from absl.testing import absltest
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import async_rl_program
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import distributed_rl_engine
+
+
+def _create_rollout_response(
+ request_id: str,
+ prompt_id: str,
+ group_id: str,
+ pair_index: int = 0,
+ policy_version: int = 0,
+ reward: float = 1.0,
+) -> datatypes.RolloutResponse:
+ return datatypes.RolloutResponse(
+ request_id=request_id,
+ prompt_id=prompt_id,
+ status="COMPLETED",
+ env_reward=reward,
+ policy_version=policy_version,
+ prompt_tokens=np.array([1, 2], dtype=np.int32),
+ segments=[
+ datatypes.TokenSegment(
+ source="assistant",
+ tokens=np.array([3, 4], dtype=np.int32),
+ loss_mask=np.array([1, 1], dtype=np.int32),
+ )
+ ],
+ metadata={
+ "group_id": group_id,
+ "pair_index": pair_index,
+ },
+ )
+
+
+class AsyncRLProgramTest(absltest.TestCase):
+
+ def setUp(self):
+ super().setUp()
+ self.mock_engine = mock.MagicMock(spec=distributed_rl_engine.DistributedRLEngine)
+ self.mock_engine.dispatch_rollouts = mock.AsyncMock()
+ self.mock_engine.train_step = mock.AsyncMock(return_value="step_done")
+ async def _mock_poll(*args, **kwargs):
+ await asyncio.sleep(0.01)
+ return []
+
+ self.mock_engine.sync_weights = mock.AsyncMock(return_value=1)
+ self.mock_engine.poll_rollouts = mock.AsyncMock(side_effect=_mock_poll)
+ self.mock_algo = mock.MagicMock(spec=algorithm_adapter.AlgorithmAdapter)
+ self.mock_algo.group_size = 2
+ self.mock_algo.mini_batch_size = 1
+ self.mock_algo.max_turns = 1
+ self.mock_algo.max_packed_len = 16
+ self.mock_algo.requires_reference_kl = False
+
+ mock_payload = datatypes.RLTrainerPayload(
+ token_ids=np.array([1, 2, 3, 4], dtype=np.int32),
+ token_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ loss_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ advantages=np.full(4, 1.0, dtype=np.float32),
+ action_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ )
+ self.mock_algo.create_trainer_payloads.return_value = [mock_payload, mock_payload]
+ self.assembler = batch_assembly.SequencePackedBatchAssembler(max_packed_len=16)
+
+ def test_initialization(self):
+ program = async_rl_program.StandardRLProgram(
+ dataset=["prompt_1"],
+ algo=self.mock_algo,
+ reward_fns=[lambda x: 1.0],
+ assembler=self.assembler,
+ )
+ self.assertEqual(program.step, 0)
+ self.assertEqual(program.group_size, 2)
+ self.assertEqual(program.mini_batch_size, 1)
+ self.assertIsNotNone(program.raw_q)
+ self.assertIsNotNone(program.scored_q)
+
+ def test_run_async_four_stages_with_long_polling(self):
+ async def _run():
+ poll_results = [
+ [
+ distributed_rl_engine._response_to_trajectory_item(_create_rollout_response(
+ "req_0_0", "prompt_0", "group_0", pair_index=0
+ )),
+ distributed_rl_engine._response_to_trajectory_item(_create_rollout_response(
+ "req_0_1", "prompt_0", "group_0", pair_index=1
+ )),
+ ],
+ [],
+ ]
+ call_idx = 0
+
+ async def mock_poll(timeout_s=0.1):
+ nonlocal call_idx
+ if call_idx < len(poll_results):
+ res = poll_results[call_idx]
+ call_idx += 1
+ return res
+ await asyncio.sleep(0.01)
+ return []
+
+ self.mock_engine.poll_rollouts.side_effect = mock_poll
+
+ begin_steps = []
+ end_steps = []
+
+ def on_begin(step):
+ begin_steps.append(step)
+
+ def on_end(step, result):
+ end_steps.append((step, result))
+
+ program = async_rl_program.StandardRLProgram(
+ dataset=["prompt_data_0"],
+ algo=self.mock_algo,
+ reward_fns=[lambda x: 1.0],
+ assembler=self.assembler,
+ on_step_begin=on_begin,
+ on_step_end=on_end,
+ )
+
+ await program.run_async(self.mock_engine, num_steps=1)
+
+ self.assertEqual(program.step, 1)
+ self.assertEqual(begin_steps, [0])
+ self.assertEqual(end_steps, [(1, "step_done")])
+ self.assertEqual(self.mock_engine.dispatch_rollouts.call_count, 2)
+ self.mock_engine.train_step.assert_called_once()
+ self.mock_engine.sync_weights.assert_called_once_with(role=datatypes.Role.ACTOR)
+
+ asyncio.run(_run())
+
+ def test_stage_exception_aborts_queue_and_propagates(self):
+ class FailingProgram(async_rl_program.StandardRLProgram):
+
+ async def rollout_dispatch_stage(self, engine):
+ del engine
+ raise RuntimeError("Rollout worker cluster down!")
+
+ async def _run():
+ prog = FailingProgram(
+ dataset=["prompt"],
+ algo=self.mock_algo,
+ assembler=self.assembler,
+ )
+ with self.assertRaises(RuntimeError) as cm:
+ await prog.run_async(self.mock_engine, num_steps=1)
+ self.assertIn("Rollout worker cluster down!", str(cm.exception))
+
+ asyncio.run(_run())
+
+
+if __name__ == "__main__":
+ absltest.main()
diff --git a/tests/experimental/orchestrator/batch_assembly_test.py b/tests/experimental/orchestrator/batch_assembly_test.py
new file mode 100644
index 000000000..7391ead20
--- /dev/null
+++ b/tests/experimental/orchestrator/batch_assembly_test.py
@@ -0,0 +1,87 @@
+# 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.
+
+"""Unit tests for Universal BatchAssembler (SequencePacked & Padded)."""
+
+
+from absl.testing import absltest
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import batch_assembly
+
+
+class BatchAssemblyTest(absltest.TestCase):
+
+ def test_sequence_packed_assembler_with_trainer_payload(self):
+ payload1 = datatypes.RLTrainerPayload(
+ token_ids=np.array([1, 2, 3, 4], dtype=np.int32),
+ token_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ loss_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ action_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ advantages=np.full(4, 1.5, dtype=np.float32),
+ )
+ payload2 = datatypes.RLTrainerPayload(
+ token_ids=np.array([5, 6, 7, 8], dtype=np.int32),
+ token_mask=np.array([0, 0, 0, 1], dtype=np.float32),
+ loss_mask=np.array([0, 0, 0, 1], dtype=np.float32),
+ action_mask=np.array([0, 0, 0, 1], dtype=np.float32),
+ advantages=np.full(4, -0.5, dtype=np.float32),
+ )
+
+ assembler = batch_assembly.SequencePackedBatchAssembler(max_packed_len=16)
+ payloads = assembler.pack([payload1, payload2])
+
+ self.assertLen(payloads, 1)
+ payload = payloads[0]
+ self.assertEqual(payload.token_ids.shape, (1, 16))
+ self.assertEqual(payload.loss_mask.shape, (1, 16))
+ self.assertEqual(payload.segment_ids.shape, (1, 16))
+ self.assertEqual(payload.segment_positions.shape, (1, 16))
+ self.assertEqual(payload.advantages.shape, (1, 16))
+
+ # Check segment boundaries
+ seg_ids = payload.segment_ids[0]
+ self.assertTrue(np.all(seg_ids[:4] == 1))
+ self.assertTrue(np.all(seg_ids[4:8] == 2))
+ self.assertTrue(np.all(seg_ids[8:] == 0))
+
+ def test_padded_batch_assembler(self):
+ payload1 = datatypes.RLTrainerPayload(
+ token_ids=np.array([1, 2, 3], dtype=np.int32),
+ token_mask=np.array([0, 0, 1], dtype=np.float32),
+ loss_mask=np.array([0, 0, 1], dtype=np.float32),
+ action_mask=np.array([0, 0, 1], dtype=np.float32),
+ advantages=np.full(3, 2.0, dtype=np.float32),
+ )
+ payload2 = datatypes.RLTrainerPayload(
+ token_ids=np.array([4, 5, 6], dtype=np.int32),
+ token_mask=np.array([0, 1, 1], dtype=np.float32),
+ loss_mask=np.array([0, 1, 1], dtype=np.float32),
+ action_mask=np.array([0, 1, 1], dtype=np.float32),
+ advantages=np.full(3, 1.0, dtype=np.float32),
+ )
+
+ assembler = batch_assembly.PaddedBatchAssembler(batch_size=2, max_seq_len=8)
+ payloads = assembler.pack([payload1, payload2])
+
+ self.assertLen(payloads, 1)
+ payload = payloads[0]
+ self.assertEqual(payload.token_ids.shape, (2, 8))
+ self.assertEqual(payload.loss_mask.shape, (2, 8))
+ self.assertEqual(payload.action_mask.shape, (2, 8))
+ self.assertEqual(payload.advantages.shape, (2, 8))
+
+
+if __name__ == "__main__":
+ absltest.main()
diff --git a/tests/experimental/orchestrator/distributed_rl_engine_test.py b/tests/experimental/orchestrator/distributed_rl_engine_test.py
new file mode 100644
index 000000000..1b22da4b7
--- /dev/null
+++ b/tests/experimental/orchestrator/distributed_rl_engine_test.py
@@ -0,0 +1,200 @@
+# 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.
+
+"""Unit tests for DistributedRLEngine and WorkerPoolBalancer."""
+
+import asyncio
+from unittest import mock
+
+from absl.testing import absltest
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import distributed_rl_engine
+from tunix.experimental.worker import remote_execution
+
+
+class MockActorHandle(mock.MagicMock):
+ """A smart mock for ActorHandle that routes asubmit/dispatch_task to logical methods.
+
+ This allows tests to write clean assertions like
+ `worker.generate.assert_called_once()` while preserving the strict ActorHandle
+ type requirements of the engine.
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(spec=remote_execution.ActorHandle, *args, **kwargs)
+ # Ensure all mocked methods return awaitables by default
+ self.generate = mock.AsyncMock()
+ self.poll_responses = mock.AsyncMock()
+ self.weight_sync = mock.AsyncMock()
+ self.fwd_bwd = mock.AsyncMock()
+ self.prepare_weight_sync = mock.AsyncMock()
+ self.score = mock.AsyncMock()
+ self.per_token_logps = mock.AsyncMock()
+
+ async def asubmit(self, method_name: str, *args, **kwargs):
+ method = getattr(self, method_name)
+ return await method(*args, **kwargs)
+
+ async def dispatch_task(self, method_name: str, *args, **kwargs):
+ method = getattr(self, method_name)
+ return await method(*args, **kwargs)
+
+
+class DistributedRLEngineTest(absltest.TestCase):
+
+ def setUp(self):
+ super().setUp()
+ self.mock_rollout_1 = MockActorHandle()
+ self.mock_rollout_2 = MockActorHandle()
+ self.mock_actor = MockActorHandle()
+ self.mock_ref = MockActorHandle()
+
+ self.engine = distributed_rl_engine.DistributedRLEngine(
+ rollout_workers=[self.mock_rollout_1, self.mock_rollout_2],
+ trainer_workers={datatypes.Role.ACTOR: self.mock_actor},
+ inference_workers={datatypes.Role.REFERENCE: self.mock_ref},
+ )
+
+ def test_generate_load_balances_across_rollout_workers(self):
+ async def _run():
+ resp1 = datatypes.RolloutResponse(request_id="r1", status="COMPLETED", env_reward=1.0)
+ resp2 = datatypes.RolloutResponse(request_id="r2", status="COMPLETED", env_reward=2.0)
+
+ self.mock_rollout_1.generate.return_value = [resp1]
+ self.mock_rollout_2.generate.return_value = [resp2]
+
+ results = await self.engine.generate(["p1", "p2"])
+ self.assertEqual(len(results), 2)
+ rewards = {res.traj.reward for res in results}
+ self.assertEqual(rewards, {1.0, 2.0})
+
+ # Verify underlying logical methods were called correctly
+ self.assertEqual(self.mock_rollout_1.generate.call_count, 1)
+ p1 = self.mock_rollout_1.generate.call_args.kwargs["prompts"][0]
+ self.assertEqual(p1, "p1")
+
+ self.assertEqual(self.mock_rollout_2.generate.call_count, 1)
+ p2 = self.mock_rollout_2.generate.call_args.kwargs["prompts"][0]
+ self.assertEqual(p2, "p2")
+
+ asyncio.run(_run())
+
+ def test_poll_rollouts_aggregates_worker_responses(self):
+ async def _run():
+ resp1 = datatypes.RolloutResponse(
+ request_id="r1",
+ status="COMPLETED",
+ env_reward=1.0,
+ )
+ self.mock_rollout_1.poll_responses.return_value = [resp1]
+ self.mock_rollout_2.poll_responses.return_value = []
+
+ results = await self.engine.poll_rollouts(timeout_s=0.1)
+ self.assertEqual(len(results), 1)
+ self.assertEqual(results[0].traj.reward, 1.0)
+
+ self.mock_rollout_1.poll_responses.assert_called_once_with(timeout_s=0.1)
+ self.mock_rollout_2.poll_responses.assert_called_once_with(timeout_s=0.1)
+
+ asyncio.run(_run())
+
+ def test_train_step_routes_to_actor(self):
+ async def _run():
+ self.mock_actor.fwd_bwd.return_value = {"loss": 0.5}
+ mock_payload = mock.MagicMock(spec=datatypes.RLTrainerPayload)
+
+ res = await self.engine.train_step(
+ mock_payload,
+ role=datatypes.Role.ACTOR,
+ accumulate_gradients=True,
+ apply_optimizer=False,
+ )
+ self.assertEqual(res, {"loss": 0.5})
+
+ self.mock_actor.fwd_bwd.assert_called_once_with(
+ batch=mock_payload,
+ accumulate_gradients=True,
+ apply_optimizer=False,
+ skip_jit=False,
+ )
+
+ asyncio.run(_run())
+
+ def test_sync_weights_coordination(self):
+ async def _run():
+ mock_meta = datatypes.WeightSyncMetadata(
+ new_policy_version=42,
+ transfer_mode="p2p",
+ )
+ self.mock_actor.prepare_weight_sync.return_value = mock_meta
+ self.mock_rollout_1.weight_sync.return_value = None
+ self.mock_rollout_2.weight_sync.return_value = None
+
+ ver = await self.engine.sync_weights(role=datatypes.Role.ACTOR)
+ self.assertEqual(ver, 42)
+
+ self.mock_actor.prepare_weight_sync.assert_called_once()
+ self.mock_rollout_1.weight_sync.assert_called_once_with(metadata=mock_meta)
+ self.mock_rollout_2.weight_sync.assert_called_once_with(metadata=mock_meta)
+
+ asyncio.run(_run())
+
+ def test_balancer_prefix_routing(self):
+
+ async def _run():
+ req1 = datatypes.RolloutRequest(
+ request_id="1",
+ prompt="p1",
+ prompt_id="1",
+ metadata={"prefix_hash": 0},
+ )
+ req2 = datatypes.RolloutRequest(
+ request_id="2",
+ prompt="p2",
+ prompt_id="2",
+ metadata={"prefix_hash": 1},
+ )
+
+ await self.engine.dispatch_rollouts([req1, req2])
+
+ # Due to deterministic round-robin / hash logic, req1 goes to rollout_1 and req2 goes to rollout_2
+ self.mock_rollout_1.generate.assert_called_once()
+ dispatched_req1 = self.mock_rollout_1.generate.call_args.kwargs[
+ "requests"
+ ][0]
+ self.assertEqual(dispatched_req1.request_id, "1")
+
+ self.mock_rollout_2.generate.assert_called_once()
+ dispatched_req2 = self.mock_rollout_2.generate.call_args.kwargs[
+ "requests"
+ ][0]
+ self.assertEqual(dispatched_req2.request_id, "2")
+
+ asyncio.run(_run())
+
+ def test_dispatch_rollouts_requires_strict_kwargs(self):
+ async def _run():
+ with self.assertRaisesRegex(ValueError, "prompt_ids' must be provided"):
+ await self.engine.dispatch_rollouts(["p1", "p2"], policy_version=0)
+
+ with self.assertRaisesRegex(ValueError, "match the length of prompts"):
+ await self.engine.dispatch_rollouts(["p1", "p2"], prompt_ids=["id1"], policy_version=0)
+
+ with self.assertRaisesRegex(ValueError, "policy_version' must be provided"):
+ await self.engine.dispatch_rollouts(["p1", "p2"], prompt_ids=["id1", "id2"])
+
+ asyncio.run(_run())
+
+if __name__ == "__main__":
+ absltest.main()
diff --git a/tests/experimental/orchestrator/orchestrator_test.py b/tests/experimental/orchestrator/orchestrator_test.py
new file mode 100644
index 000000000..662a3bd12
--- /dev/null
+++ b/tests/experimental/orchestrator/orchestrator_test.py
@@ -0,0 +1,113 @@
+# 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.
+
+"""Unit tests for ClusterOrchestrator."""
+
+from unittest import mock
+
+from absl.testing import absltest
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import orchestrator
+
+
+class ClusterOrchestratorTest(absltest.TestCase):
+
+ def setUp(self):
+ super().setUp()
+ self.mock_registry = mock.MagicMock()
+ self.mock_lifecycle = mock.MagicMock()
+ self.mock_monitor = mock.MagicMock()
+ self.orch = orchestrator.ClusterOrchestrator(
+ registry=self.mock_registry,
+ lifecycle_driver=self.mock_lifecycle,
+ monitor=self.mock_monitor,
+ )
+
+ def test_register_and_unregister_worker(self):
+ mock_worker = mock.MagicMock()
+ self.orch.register_worker(mock_worker)
+ self.mock_registry.register.assert_called_once_with(mock_worker)
+
+ self.orch.unregister_worker("worker_123")
+ self.mock_registry.unregister.assert_called_once_with("worker_123")
+
+ def test_bring_up_and_shutdown(self):
+ self.orch.bring_up_workers("dummy_warmup_data")
+ self.mock_lifecycle.bring_up.assert_called_once_with("dummy_warmup_data")
+
+ self.orch.shutdown()
+ self.mock_monitor.close.assert_called_once()
+ self.mock_lifecycle.shutdown.assert_called_once()
+
+ def test_create_engine(self):
+ from tunix.experimental.worker import remote_execution
+ mock_rollout = mock.MagicMock(spec=remote_execution.ActorHandle)
+ mock_actor = mock.MagicMock(spec=remote_execution.ActorHandle)
+ mock_critic = mock.MagicMock(spec=remote_execution.ActorHandle)
+ mock_ref = mock.MagicMock(spec=remote_execution.ActorHandle)
+
+ def mock_group(role):
+ grp = mock.MagicMock()
+ if role == datatypes.Role.ROLLOUT:
+ grp.members.return_value = [mock_rollout]
+ elif role == datatypes.Role.ACTOR:
+ grp.members.return_value = [mock_actor]
+ elif role == datatypes.Role.CRITIC:
+ grp.members.return_value = [mock_critic]
+ elif role == datatypes.Role.REFERENCE:
+ grp.members.return_value = [mock_ref]
+ else:
+ grp.members.return_value = []
+ return grp
+
+ self.mock_registry.group.side_effect = mock_group
+
+ engine = self.orch._create_engine()
+ self.assertIs(
+ engine._trainer_workers[datatypes.Role.ACTOR],
+ mock_actor,
+ )
+ self.assertIs(
+ engine._trainer_workers[datatypes.Role.CRITIC],
+ mock_critic,
+ )
+ self.assertIs(engine._inference_workers[datatypes.Role.REFERENCE], mock_ref)
+
+ def test_run_managed_program_submission(self):
+ mock_algo = mock.MagicMock(spec=algorithm_adapter.AlgorithmAdapter)
+ mock_algo.group_size = 2
+ mock_algo.mini_batch_size = 1
+ mock_algo.max_turns = 1
+ mock_algo.max_packed_len = 16
+ mock_algo.requires_reference_kl = False
+
+ assembler = batch_assembly.SequencePackedBatchAssembler(max_packed_len=16)
+
+ with mock.patch("asyncio.run") as mock_asyncio_run:
+ self.orch.run(
+ algo=mock_algo,
+ dataset=["prompt1"],
+ reward_fns=[lambda x: 1.0],
+ assembler=assembler,
+ num_steps=5,
+ )
+ self.mock_lifecycle.bring_up.assert_called_once()
+ self.mock_monitor.poll.assert_called_once()
+ mock_asyncio_run.assert_called_once()
+
+
+if __name__ == "__main__":
+ absltest.main()
diff --git a/tests/experimental/orchestrator/rl_program_test.py b/tests/experimental/orchestrator/rl_program_test.py
new file mode 100644
index 000000000..0ccfad167
--- /dev/null
+++ b/tests/experimental/orchestrator/rl_program_test.py
@@ -0,0 +1,105 @@
+# 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.
+
+"""Unit tests for synchronous RLProgram."""
+
+from unittest import mock
+from absl.testing import absltest
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import rl_engine_interface
+from tunix.experimental.orchestrator import rl_program
+
+
+class RLProgramTest(absltest.TestCase):
+
+ def setUp(self):
+ super().setUp()
+ self.mock_engine = mock.MagicMock(spec=rl_engine_interface.AbstractRLEngine)
+ mock_resp = datatypes.RolloutResponse(
+ request_id="r1",
+ status="COMPLETED",
+ env_reward=1.0,
+ prompt_tokens=np.array([1, 2], dtype=np.int32),
+ segments=[
+ datatypes.TokenSegment(
+ source="assistant",
+ tokens=np.array([3, 4], dtype=np.int32),
+ loss_mask=np.array([1, 1], dtype=np.int32),
+ )
+ ],
+ )
+ self.mock_engine.generate = mock.AsyncMock(return_value=[mock_resp])
+ self.mock_engine.train_step = mock.AsyncMock(return_value="mock_train_result")
+ self.mock_engine.sync_weights = mock.AsyncMock(return_value=1)
+
+ self.mock_algo = mock.MagicMock(spec=algorithm_adapter.AlgorithmAdapter)
+ mock_payload = datatypes.RLTrainerPayload(
+ token_ids=np.array([1, 2, 3, 4], dtype=np.int32),
+ token_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ loss_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ advantages=np.full(4, 1.0, dtype=np.float32),
+ action_mask=np.array([0, 0, 1, 1], dtype=np.float32),
+ )
+ self.mock_algo.create_trainer_payloads.return_value = [mock_payload]
+ self.mock_algo.requires_reference_kl = False
+ self.assembler = batch_assembly.SequencePackedBatchAssembler(max_packed_len=16)
+
+ def test_step_once_flow(self):
+ begin_calls = []
+ end_calls = []
+
+ def on_begin(step):
+ begin_calls.append(step)
+
+ def on_end(step, result):
+ end_calls.append((step, result))
+
+ program = rl_program.SyncRLProgram(
+ engine=self.mock_engine,
+ algo=self.mock_algo,
+ assembler=self.assembler,
+ on_step_begin=on_begin,
+ on_step_end=on_end,
+ )
+
+ res = program.step_once(prompts=["prompt1"])
+
+ self.assertEqual(res, "mock_train_result")
+ self.mock_engine.generate.assert_called_once_with(prompts=["prompt1"])
+ self.mock_algo.create_trainer_payloads.assert_called_once()
+ self.mock_engine.train_step.assert_called_once()
+ self.mock_engine.sync_weights.assert_called_once_with(role=datatypes.Role.ACTOR)
+ self.assertEqual(program.step, 1)
+
+ self.assertEqual(begin_calls, [0])
+ self.assertEqual(end_calls, [(1, "mock_train_result")])
+
+ def test_eval_step_once_flow(self):
+ program = rl_program.SyncRLProgram(
+ engine=self.mock_engine,
+ algo=self.mock_algo,
+ assembler=self.assembler,
+ )
+ res = program.eval_step_once(prompts=["eval_prompt"])
+
+ self.assertLen(res, 1)
+ self.mock_engine.generate.assert_called_once_with(prompts=["eval_prompt"])
+ self.mock_algo.create_trainer_payloads.assert_called_once()
+
+
+if __name__ == "__main__":
+ absltest.main()
diff --git a/tests/rl/rl_cluster_test.py b/tests/rl/rl_cluster_test.py
index f050ede7d..cdf34d9ad 100644
--- a/tests/rl/rl_cluster_test.py
+++ b/tests/rl/rl_cluster_test.py
@@ -993,11 +993,6 @@ def test_generic_role_primitives_delegate_to_shorthands(self):
eos_id=1,
)
- with self.subTest('satisfies_abstract_rl_engine_protocol'):
- self.assertIsInstance(
- rl_engine, rl_engine_interface.AbstractRLEngine
- )
-
if __name__ == '__main__':
absltest.main()
diff --git a/tunix/experimental/common/datatypes.py b/tunix/experimental/common/datatypes.py
index 5091e8fa8..234eed583 100644
--- a/tunix/experimental/common/datatypes.py
+++ b/tunix/experimental/common/datatypes.py
@@ -23,7 +23,6 @@
import enum
import time
from typing import Any, Dict
-
from jax.typing import ArrayLike # pylint: disable=g-importing-member
import numpy as np
from tunix.rl.agentic.agents import agent_types
@@ -44,11 +43,30 @@ class Role(enum.Enum):
# Worker-internal episode representation produced during rollout.
Trajectory = agent_types.Trajectory
-TrajectoryItem = agent_types.TrajectoryItem
Step = agent_types.Step
TrajectoryStatus = agent_types.TrajectoryStatus
+# TODO: Unify this extended TrajectoryItem back into agent_types.TrajectoryItem
+# so that all agentic workflows share the same strict token array fields.
+@dataclasses.dataclass(kw_only=True)
+class TrajectoryItem(agent_types.TrajectoryItem):
+ """Extended TrajectoryItem for Orchestrator with token arrays."""
+ prompt_tokens: np.ndarray | None = None
+ completion_tokens: np.ndarray | None = None
+ action_mask: np.ndarray | None = None
+ policy_version: int = 0
+ # TODO: trajectory item having the completion tokens, masks, etc is quite redundant since those are in the trainer payload already.
+
+class Role(str, enum.Enum):
+ """Orchestrator worker roles."""
+ ACTOR = "actor"
+ CRITIC = "critic"
+ ROLLOUT = "rollout"
+ REFERENCE = "reference"
+ REWARD = "reward"
+
+
##### Common DTOs (Data Transfer Objects) #####
@@ -367,22 +385,6 @@ def _get_step_attr(step, attr):
)
-@dataclasses.dataclass(kw_only=True)
-class TrainerPayload:
- """Generic trainer payload.
-
- Attributes:
- token_ids: [B, T] token IDs. By default, structured as left-padded prompt
- tokens concatenated with right-padded completion tokens.
- token_mask: [B, T] token mask to differentiate padding tokens from valid
- tokens.
- segment_ids: Optional [B, T] packing segment ids.
- """
-
- token_ids: ArrayLike
- token_mask: ArrayLike
- segment_ids: ArrayLike | None = None
-
##### Weight Sync DTOs #####
@@ -461,6 +463,26 @@ def __post_init__(self):
##### Training DTOs #####
+@dataclasses.dataclass(kw_only=True)
+class TrainerPayload:
+ """Generic trainer payload.
+
+ Attributes:
+ token_ids: [B, T] token IDs. By default, structured as left-padded prompt
+ tokens concatenated with right-padded completion tokens.
+ token_mask: [B, T] token mask to differentiate padding tokens from valid
+ tokens.
+ segment_ids: Optional [B, T] packing segment ids.
+ segment_positions: Optional [B, T] position indices within each segment.
+ """
+
+ token_ids: ArrayLike
+ token_mask: ArrayLike
+ segment_ids: ArrayLike | None = None
+ segment_positions: ArrayLike | None = None
+
+
+# TODO: Introduce PPOTrainerPayload to replace generic RLTrainerPayload when PPO specific fields are needed.
@dataclasses.dataclass(kw_only=True)
class RLTrainerPayload(TrainerPayload):
"""RL training payload.
@@ -468,16 +490,25 @@ class RLTrainerPayload(TrainerPayload):
Attributes:
advantages: [B] or [B, C] advantages.
loss_mask: [B, T], 1 where the position contributes to the loss.
+ action_mask: Optional [B, T] or [B, C] mask of policy actions.
ref_per_token_logps: Optional [B, C] reference model log-probabilities.
old_per_token_logps: Optional [B, C] behavior policy log-probabilities.
sampler_is_weights: Optional [B, C] importance sampling weights.
+ returns: Optional [B, C] value baseline returns (for PPO / Critic).
+ old_values: Optional [B, C] critic value estimates (for PPO / Critic).
+ metadata: Extra payload metadata dictionary.
"""
advantages: ArrayLike
loss_mask: ArrayLike
+ action_mask: ArrayLike | None = None
ref_per_token_logps: ArrayLike | None = None
old_per_token_logps: ArrayLike | None = None
sampler_is_weights: ArrayLike | None = None
+ returns: ArrayLike | None = None
+ old_values: ArrayLike | None = None
+ metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
+ # TODO: add ppo sepcific fields in a PPO specific fields in PPORLTrainerPayload
@dataclasses.dataclass(kw_only=True)
diff --git a/tunix/experimental/docs/orchestrator_plan_v2.md b/tunix/experimental/docs/orchestrator_plan_v2.md
new file mode 100644
index 000000000..7fe373dbe
--- /dev/null
+++ b/tunix/experimental/docs/orchestrator_plan_v2.md
@@ -0,0 +1,816 @@
+# Tunix RL Architecture V2: Streamlined 3-Layer Specification & Design Plan
+
+> [!NOTE]
+> This document defines the definitive **Version 2 (V2) Architectural Specification** for the Tunix Distributed RL stack. Based on architectural review and gap analysis of the initial prototype, V2 streamlines the previous 5-layer proposal into a high-performance, boilerplate-free **3-Layer Compositional Architecture**:
+> 1. **Layer 3: Workflow & Program Layer** (`async_rl_program.py` / `rl_program.py`)
+> 2. **Layer 2: Algorithm Math & Batch Assembly Layer** (`algorithm_adapter.py` & `batch_assembly.py`)
+> 3. **Layer 1: Cluster Compute & Infrastructure Layer** (`distributed_rl_engine.py` & `orchestrator.py` with `trajectory_queue_manager.py`)
+>
+> It details the **3-Tier Ergonomic User Experience** (Zero Boilerplate for 90% of users), **Long-Polling Rollout Pipelining & Streaming Gradient Accumulation** (zero-bubble pipelined execution), an **Orbax Composite Manifest Checkpointing** protocol for isolated trainer failure recovery without pipeline resets, standalone **1D Sequence Packing** (`batch_assembly.py`), a single off-the-shelf **`StandardRLProgram`** covering 95% of use cases, and complete end-to-end prototype code.
+
+---
+
+## 1. Executive Summary & Why V2 Was Streamlined
+
+### 1.1. The Evolution from Monolith to V2
+The Tunix RL architecture has evolved across three key milestones:
+
+```mermaid
+graph TD
+ subgraph LEGACY ["1. Legacy Monolith (tunix.rl.agentic)"]
+ LEG["Monolithic AgenticGRPOLearner
• Mixed chat parsing, rollout generation, math, step cadence & compute
• Single-process RLEngine
• No clean worker disaggregation"]
+ end
+
+ subgraph V1_PROP ["2. Intermediate 5-Layer Proposal (V1 Plan)"]
+ V1_O["ClusterOrchestrator (L5)"]
+ V1_P["RLProgram (L4)"]
+ V1_D["RLDriver (L3 - Pass-Through Shim)"]
+ V1_A["AlgorithmAdapter (L3 Math)"]
+ V1_E["DistributedRLEngine (L2 Compute)"]
+ V1_W["Workers (L1)"]
+
+ V1_O --> V1_P --> V1_D --> V1_E --> V1_W
+ V1_D -.-> V1_A
+ end
+
+ subgraph V2_STREAMLINED ["3. Streamlined 3-Layer Architecture (V2 Target)"]
+ V2_P["Layer 3: Workflow & Program
(async_rl_program.py / rl_program.py)
• StandardRLProgram covers 95% of use cases (GRPO, PPO, PRM, Agentic)
• Subclass AsyncRLProgram for novel research (MCTS, Self-Play)"]
+ V2_A["Layer 2: Algorithm Math & Batch Assembly
(algorithm_adapter.py & batch_assembly.py)
• Pure functional advantage math (GRPO, GAE, PPO)
• Standalone 1D sequence packing (>90% MXU density)"]
+ V2_E["Layer 1: Cluster Compute & Infrastructure
(distributed_rl_engine.py & orchestrator.py)
• Worker registry, health heartbeats, lifecycle & Orbax manifest recovery
• Stateless compute primitives: generate_async, train_step_async, sync_weights_async"]
+
+ V2_P -->|Calls compute primitives| V2_E
+ V2_P -->|Calls math & packing| V2_A
+ end
+```
+
+### 1.2. Eliminating the "Pass-Through Tax" of `RLDriver`
+In the V1 design, `RLDriver` acted mostly as an unnecessary middleman:
+- `driver.generate()` was a 1-line pass-through to `engine.generate()`.
+- `driver.train_step()` was a 1-line pass-through to `engine.train_step()`.
+- `driver.sync_weights()` was a 1-line pass-through to `engine.sync_weights()`.
+- `driver.compute_advantages()` was a 1-line pass-through to `adapter.compute_advantages()`.
+
+Every call required 4–6 hops across files. **In V2, `RLDriver` is completely dissolved:**
+- **Compute & RPC Routing** lives in `DistributedRLEngine` (`distributed_rl_engine.py`).
+- **Math & Loss Wiring** lives in `AlgorithmAdapter` (`algorithm_adapter.py`).
+- **1D Sequence Packing & 2D Padding** lives in `BatchAssembler` (`batch_assembly.py`).
+- **Workflow & Stage Cadence** lives in `RLProgram` (`async_rl_program.py`).
+
+---
+
+## 2. The Streamlined 3-Layer Architecture & Module Map
+
+```mermaid
+graph TD
+ subgraph L3 ["Layer 3: Workflow & Program (async_rl_program.py)"]
+ PROG["StandardRLProgram (Off-the-shelf for 95% of runs)
• Stages: rollout_dispatch -> polling -> critique -> train
• Streaming gradient accumulation: streams 1 group at a time
• Extensible base: AsyncRLProgram for custom DAGs"]
+ end
+
+ subgraph L2 ["Layer 2: Algorithm Math & Batch Assembly"]
+ ALGO["algorithm_adapter.py (GRPOAdapter / PPOAdapter)
• Math & Loss: create_train_examples(), compute_advantages(), loss_fn()
• Produces list[TrainExample] (tokens, action_mask, advantages, ref_logps)"]
+ ASM["batch_assembly.py (BatchAssembler)
• SequencePackedBatchAssembler: Packs list[TrainExample] into 1D buffer (>90% MXU)
• PaddedBatchAssembler: Standard 2D rectangular padding"]
+ end
+
+ subgraph QUEUES ["Infrastructure Buffer (trajectory_queue_manager.py)"]
+ Q["TrajectoryQueueManager (GroupQueue with In-Flight ACK)
• Out-of-order prompt grouping by (prompt_id, policy_version)
• Policy staleness filtering & uncommitted group rewind"]
+ end
+
+ subgraph L1_INFRA ["Layer 1: Cluster Compute & Infrastructure"]
+ ENG["distributed_rl_engine.py (DistributedRLEngine)
• dispatch_rollouts(requests), poll_rollouts(timeout)
• train_step(payload), sync_weights(role), per_token_logps(role)"]
+ LB["load_balancer.py (WorkerPoolBalancer)
• Least-in-flight queue depth & prefix-cache consistent routing
• Concurrent multi-worker long-polling collector"]
+ ORCH["orchestrator.py (ClusterOrchestrator)
• WorkerRegistry: Indexes live workers by role (ACTOR, ROLLOUT, REFERENCE, CRITIC)
• HealthMonitor & LifecycleDriver: Heartbeats, pre-flight checks & pod restarts
• CompositeCheckpointHandler: Atomic step recovery & queue offset rewinding"]
+ end
+
+ subgraph WORKERS ["Physical Worker Pods (TPU / GPU Clusters)"]
+ W_ROLL["RolloutWorker Pods (vLLM / KV Cache Pools)"]
+ W_TRAIN["TrainerWorker Pods (nnx.Optimizer & HBM Weights)"]
+ W_INF["InferenceWorker Pods (PRM / Reference KL Models)"]
+ W_HYBRID["[OR] CompositeWorker Pods (Co-Located Actor-Rollout-Trainer)"]
+ end
+
+ PROG -->|1. Dispatches compute RPCs| ENG
+ PROG -->|2. Assembles TrainExamples| ALGO
+ PROG -->|3. Packs TrainExamples| ASM
+ PROG <-->|4. Buffers & groups rollouts| Q
+ ENG <-->|5. Routes & load-balances| LB
+ ORCH -->|6. Builds, supervises & recovers| ENG
+ ORCH -->|7. Monitors heartbeats & restarts| WORKERS
+ LB -->|8. Dispatches & polls| W_ROLL
+ ENG -->|9. Routes gradient microbatches| W_TRAIN
+ ENG -->|10. Routes scoring RPCs| W_INF
+ ENG -.->|hybrid zero-copy swap| W_HYBRID
+```
+
+### 2.1. The 3-Tier Ergonomic User Experience (Zero Boilerplate)
+
+To prevent users from having to manually instantiate 6 separate objects in their experiment scripts, `ClusterOrchestrator` provides **3 clean ergonomic tiers**:
+
+```mermaid
+graph TD
+ subgraph T1 ["Tier 1: 90% Production User (1 Call, Zero Wiring)"]
+ E1["orchestrator.run(algo=GRPO(...), dataset=..., reward_fns=[...])
• Engine, Queues, Assembler & StandardProgram are AUTO-WIRED!"]
+ end
+
+ subgraph T2 ["Tier 2: Research Tuning (1-Line Overrides)"]
+ E2["orchestrator.run(algo=PPO(...), assembler=PaddedBatchAssembler(...), ...)
• Override just the 1 component you care about"]
+ end
+
+ subgraph T3 ["Tier 3: Novel Paradigm Researcher (Custom Program)"]
+ E3["orchestrator.run_program(MyMCTSProgram(...), algo=GRPO(...))
• Custom AsyncRLProgram when inventing new DAG workflows"]
+ end
+```
+
+---
+
+## 3. Data Flow & Queue Management: Streaming Gradient Accumulation
+
+In agentic RL, waiting for all $N$ prompt groups (e.g. 4 groups of $G=8$ rollouts = 32 rollouts total) before touching the trainer creates massive **GPU/TPU idle bubbles**.
+
+V2 implements **Long-Polling Rollout Pipelining & Streaming Gradient Accumulation with `TrainExample`s**:
+1. **Fire-and-Forget Dispatch:** `rollout_dispatch_stage` streams `RolloutRequest`s across the worker pool using non-blocking RPCs (`await engine.dispatch_rollouts(requests)`).
+2. **Long-Polling Collector & Load Balancer:** A dedicated `polling_stage` continuously long-polls completed responses from workers (`await engine.poll_rollouts()`) and feeds individual completions into `TrajectoryQueue`. The queue groups them by `prompt_id` and filters stale policy versions.
+3. **Structured `TrainExample` Assembly:** As soon as **1 ready prompt group** ($G=8$ rollouts) is assembled by the queue, `algo.create_train_examples(group, rewards)` creates a typed `list[TrainExample]` attaching advantages, value targets, and observation loss masks (`action_mask=0`).
+4. **Immediate Microbatch Execution:** `assembler.pack(train_examples)` packs the examples into 1D static buffers and immediately streams them to `await engine.train_step(payload)`. Gradients accumulate directly on accelerator HBM, and only on the $N$-th group does the trainer apply the optimizer update and broadcast new weights!
+
+```mermaid
+sequenceDiagram
+ autonumber
+ participant RW as RolloutWorkers (vLLM)
+ participant Q as TrajectoryQueue (GroupQueue)
+ participant Prog as AsyncRLProgram (Stages)
+ participant TW as TrainerWorker (TPU)
+
+ Note over Prog,RW: Phase 1: Fire-and-Forget Dispatch
+ Prog->>RW: dispatch_rollouts([Req A, Req B]) [Non-blocking]
+
+ Note over RW,Q: Phase 2: Background Long-Polling Collector
+ loop Continuous Long-Polling
+ Prog->>RW: poll_rollouts(timeout_s=0.1)
+ RW-->>Prog: Yields Completed Responses [A#1, B#1, A#2, ...]
+ Prog->>Q: put(item, prompt_id='A')
+ Note over Q: Buffers completions out-of-order until
G=8 rollouts for Prompt A arrive.
+ end
+
+ Note over Q,TW: Phase 3: Pipelined Streaming Gradient Accumulation
+ Q->>Prog: get_group() yields Prompt Group 1 (G=8)
+ Note over Prog: algo.create_train_examples() -> list[TrainExample]
assembler.pack(examples) -> microbatch
+ Prog->>TW: train_step(microbatch, accumulate=True) [HBM grad_acc += 1]
+
+ par Concurrent Execution
+ Q->>Prog: get_group() yields Prompt Group 2 (G=8)
+ TW-->>Prog: Grad #1 Complete
+ end
+
+ Prog->>TW: train_step(microbatch_2, accumulate=True) [HBM grad_acc += 2]
+
+ par Concurrent Execution
+ Q->>Prog: get_group() yields Prompt Groups 3 & 4
+ TW-->>Prog: Grad #2 Complete
+ end
+
+ Prog->>TW: train_step(microbatch_3, accumulate=True) [HBM grad_acc += 3]
+ Prog->>TW: train_step(microbatch_4, accumulate=True, apply_optimizer=True) [Weights Updated!]
+
+ Note over TW,RW: Global Step Boundary
+ TW->>RW: sync_weights() [DCN/ICI Broadcast V_k+1]
+ Prog->>Q: queue.commit(step_k, groups=[1, 2, 3, 4]) [Advances committed offset]
+```
+
+---
+
+## 4. Isolated Failure Recovery: Atomic Manifest & Queue Offsets
+
+When a Trainer worker pod crashes midway through gradient accumulation (e.g. at Group 2 of 4), **we must NOT restart rollout workers or discard in-flight queues**.
+
+```mermaid
+graph TD
+ subgraph CRASH ["1. Mid-Accumulation Failure (Group 2 of 4)"]
+ FAIL["TrainerWorker Crashes
• HBM grad_acc lost
• Step k+1 uncommitted"]
+ Q_HOLD["Queue State:
• Groups 1 & 2 marked IN_FLIGHT
• Groups 3 & 4 marked READY
• Rollout workers UNTOUCHED & PRODUCING"]
+ end
+
+ subgraph RECOVERY ["2. Orchestrator Isolated Recovery"]
+ O1["HealthMonitor detects dead trainer heartbeat"]
+ O2["LifecycleDriver restarts ONLY TrainerWorker pod"]
+ O3["Trainer reloads weights from Step k (grad_acc initialized to 0)"]
+ O4["TrajectoryQueue rewinds uncommitted Groups 1 & 2 back to READY"]
+ end
+
+ subgraph RESUME ["3. Seamless Stream Replay"]
+ PLAY["train_stage replays Step k+1 seamlessly:
• Streams Group 1 (grad_acc += 1)
• Streams Group 2 (grad_acc += 2)
• Streams Group 3 (grad_acc += 3)
• Streams Group 4 (optimizer update + sync_weights)
• ZERO rollouts wasted!"]
+ end
+
+ CRASH --> RECOVERY
+ RECOVERY --> RESUME
+```
+
+### The 4 Pillars of Isolated Recovery:
+1. **Atomic Composite Manifest in Orbax:** Checkpoint directories (`step_k/`) contain an atomic manifest linking model weights, global step counter, dataset prompt index, and queue read offsets:
+ ```json
+ {
+ "global_step": 42,
+ "policy_version": 42,
+ "dataset_prompt_idx": 130,
+ "queue_read_offsets": {
+ "raw_rollouts_q": 1040,
+ "scored_rollouts_q": 1040
+ }
+ }
+ ```
+2. **In-Flight Acknowledgement (`ACK`) Semantics in `TrajectoryQueue`:** Dequeued items remain in the queue log marked as `IN_FLIGHT`. They are only permanently committed when `queue.commit(step)` is invoked after `sync_weights_async()`.
+3. **Zero Partial Gradient Bleed:** On reboot, `TrainerWorker` reloads Step $k$ weights and zeroes out its physical gradient accumulator tensor (`grad_acc = 0`).
+4. **Instant Stream Replay:** The queue returns uncommitted Groups 1 and 2 to `READY` state. `train_stage` streams Groups 1 through 4 again, cleanly completing Step $k+1$ without regenerating rollouts.
+
+---
+
+## 5. Complete Runnable Prototype Implementation
+
+Here is the clean, production-ready prototype implementation of the V2 architecture:
+
+### 5.1. Layer 1: `DistributedRLEngine`, `AbstractRLEngine` & `WorkerPoolBalancer`
+
+```python
+"""Distributed compute routing surface implementing AbstractRLEngine."""
+
+from typing import Any, Mapping, Protocol, Sequence, runtime_checkable
+import asyncio
+from tunix.experimental.common import datatypes
+
+@runtime_checkable
+class AbstractRLEngine(Protocol):
+ """Stateless compute primitives for distributed worker meshes."""
+ async def dispatch_rollouts(self, requests: Sequence[datatypes.RolloutRequest]) -> None: ...
+ async def poll_rollouts(self, timeout_s: float = 0.1) -> list[datatypes.RolloutResponse]: ...
+ async def generate(self, prompts: Sequence[Any], **kwargs: Any) -> list[Any]: ...
+ async def score(self, role: datatypes.Role, items: Sequence[Any], **kwargs: Any) -> list[float]: ...
+ async def per_token_logps(self, role: datatypes.Role, items: Sequence[Any], **kwargs: Any) -> Any: ...
+ async def train_step(self, payload: datatypes.RLTrainerPayload, role: datatypes.Role = datatypes.Role.ACTOR, **kwargs: Any) -> Any: ...
+ async def sync_weights(self, role: datatypes.Role = datatypes.Role.ACTOR) -> int: ...
+
+
+class WorkerPoolBalancer:
+ """Load balancing, prefix-cache affinity, and concurrent long polling across worker replicas."""
+
+ def __init__(self, workers: Sequence[Any]):
+ self._workers = list(workers)
+ self._in_flight: dict[int, int] = {i: 0 for i in range(len(workers))}
+
+ def select_worker_for_request(self, req: datatypes.RolloutRequest) -> tuple[int, Any]:
+ """Selects worker using least-in-flight queue depth or prefix-cache hash affinity."""
+ # Prefix-cache routing: Hash system prompt prefix to maximize vLLM KV-cache reuse
+ if "prefix_hash" in req.metadata:
+ idx = req.metadata["prefix_hash"] % len(self._workers)
+ else:
+ # Least-in-flight worker
+ idx = min(self._in_flight, key=self._in_flight.get)
+ self._in_flight[idx] += 1
+ return idx, self._workers[idx]
+
+ async def poll_all_workers(self, timeout_s: float = 0.1) -> list[datatypes.RolloutResponse]:
+ """Concurrently long-polls all active rollout workers."""
+ tasks = [w.poll_responses(timeout_s=timeout_s) for w in self._workers]
+ responses = await asyncio.gather(*tasks)
+ completed = []
+ for i, resp in enumerate(responses):
+ if resp is not None:
+ unwrap_fn = getattr(resp, "unwrap", None)
+ res = unwrap_fn() if callable(unwrap_fn) else getattr(resp, "result", resp)
+ if res is not None:
+ items = res if isinstance(res, list) else [res]
+ self._in_flight[i] = max(0, self._in_flight[i] - len(items))
+ completed.extend(items)
+ return completed
+
+
+class DistributedRLEngine(AbstractRLEngine):
+ """Worker-backed compute router dispatching RPCs across role pools."""
+
+ def __init__(
+ self,
+ rollout_workers: Sequence[Any],
+ trainer_workers: Mapping[datatypes.Role, Any],
+ inference_workers: Mapping[datatypes.Role, Any] | None = None,
+ ):
+ self._rollout_workers = list(rollout_workers)
+ self._balancer = WorkerPoolBalancer(rollout_workers)
+ self._trainer_workers = dict(trainer_workers)
+ self._inference_workers = dict(inference_workers or {})
+
+ async def dispatch_rollouts(self, requests: Sequence[datatypes.RolloutRequest]) -> None:
+ """Dispatches rollout requests across workers using the load balancer."""
+ for req in requests:
+ _, worker = self._balancer.select_worker_for_request(req)
+ await worker.dispatch_task(method_name="generate", requests=[req])
+
+ async def poll_rollouts(self, timeout_s: float = 0.1) -> list[datatypes.RolloutResponse]:
+ """Long-polls completed rollout responses concurrently across all workers."""
+ return await self._balancer.poll_all_workers(timeout_s=timeout_s)
+
+ async def generate(self, prompts: Sequence[Any], **kwargs: Any) -> list[Any]:
+ """Blocking rollout generation (dispatches chunked tasks and awaits completion)."""
+ num_workers = len(self._rollout_workers)
+ chunk_size = (len(prompts) + num_workers - 1) // num_workers
+ tasks = []
+ for i, worker in enumerate(self._rollout_workers):
+ chunk = prompts[i * chunk_size : (i + 1) * chunk_size]
+ if chunk:
+ tasks.append(worker.asubmit("generate", prompts=chunk, **kwargs))
+ results = await asyncio.gather(*tasks)
+ return [item for sublist in results for item in (sublist if isinstance(sublist, list) else [sublist])]
+
+ async def score(self, role: datatypes.Role, items: Sequence[Any], **kwargs: Any) -> list[float]:
+ """Routes reward / PRM scoring requests to InferenceWorker pool."""
+ worker = self._inference_workers[role]
+ return await worker.asubmit("score", items=items, **kwargs)
+
+ async def per_token_logps(self, role: datatypes.Role, items: Sequence[Any], **kwargs: Any) -> Any:
+ """Evaluates reference model or actor per-token logprobs."""
+ worker = self._inference_workers.get(role) or self._trainer_workers.get(role)
+ return await worker.asubmit("per_token_logps", items=items, **kwargs)
+
+ async def train_step(
+ self,
+ payload: datatypes.RLTrainerPayload,
+ role: datatypes.Role = datatypes.Role.ACTOR,
+ accumulate_gradients: bool = False,
+ apply_optimizer: bool = True,
+ skip_jit: bool = False,
+ ) -> Any:
+ """Executes atomic gradient accumulation / update on TrainerWorker."""
+ worker = self._trainer_workers[role]
+ return await worker.asubmit(
+ "fwd_bwd",
+ batch=payload,
+ accumulate_gradients=accumulate_gradients,
+ apply_optimizer=apply_optimizer,
+ skip_jit=skip_jit,
+ )
+
+ async def sync_weights(self, role: datatypes.Role = datatypes.Role.ACTOR) -> int:
+ """Executes accelerator-to-accelerator collective weight broadcast."""
+ trainer = self._trainer_workers[role]
+ sync_metadata = await trainer.asubmit("prepare_weight_sync")
+ tasks = [w.asubmit("weight_sync", sync_metadata) for w in self._rollout_workers]
+ await asyncio.gather(*tasks)
+ return sync_metadata.new_policy_version
+```
+
+---
+
+### 5.2. Layer 2A: Universal Batch Assembly (`batch_assembly.py`)
+
+`batch_assembly.py` is a **universal, algorithm-agnostic tensor packing utility** parameterized over generic type `T` (e.g. `TrainExample`, `SFTExample`, `DPOPair`, or arbitrary user dataclasses / PyTrees).
+
+```python
+"""Universal, generic batch assembly, TrainExample DTO, 1D sequence packing, and 2D padding."""
+
+import dataclasses
+from typing import Any, Generic, Protocol, Sequence, TypeVar
+import numpy as np
+from tunix.experimental.common import datatypes
+
+T = TypeVar("T")
+
+@dataclasses.dataclass(slots=True)
+class TrainExample:
+ """Self-contained training example produced by AlgorithmAdapter."""
+ prompt_tokens: np.ndarray # [L_prompt]
+ completion_tokens: np.ndarray # [L_completion]
+ action_mask: np.ndarray # [L_completion] (1 for model tokens, 0 for tool observations)
+ advantage: float | np.ndarray # Scalar (GRPO) or [L_completion] (token-level GAE)
+ value_target: float | None = None # Target for Critic MSE loss (PPO)
+ old_logprobs: np.ndarray | None = None # [L_completion] Rollout policy logprobs
+ ref_logprobs: np.ndarray | None = None # [L_completion] Reference model logprobs
+ metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
+
+
+class BatchAssembler(Generic[T], Protocol):
+ """Universal batch assembly protocol for any dataclass, dict, or PyTree."""
+ def pack(self, items: Sequence[T]) -> list[datatypes.RLTrainerPayload]: ...
+
+
+class SequencePackedBatchAssembler(Generic[T]):
+ """1D Sequence Packing: Concatenates variable-length items into dense [1, max_packed_len] static buffers.
+
+ Achieves >90% MXU compute density on TPUs/GPUs with Flash/FlexAttention.
+ Works out-of-the-box on TrainExample (RL), SFTExample (Supervised), DPOPair (Preference), or custom PyTrees.
+ """
+ def __init__(self, max_packed_len: int = 8192, pad_id: int = 0):
+ self.max_packed_len = max_packed_len
+ self.pad_id = pad_id
+
+ def pack(self, items: Sequence[T]) -> list[datatypes.RLTrainerPayload]:
+ """Bin-packs arbitrary dataclass items into dense 1D buffers with segment boundaries."""
+ # 1. First-fit decreasing 1D bin-packing on sequence lengths
+ # 2. Generates cu_seqlens / segment_ids for block-diagonal attention kernels
+ # 3. Static shape guarantee: pads trailing buffer slots with loss_mask = 0
+ return self._pack_1d_buffers(items, self.max_packed_len)
+
+
+class PaddedBatchAssembler(Generic[T]):
+ """Simple 2D Rectangular Batching: Pads sequences to standard [batch_size, max_seq_len] tensors."""
+ def __init__(self, batch_size: int = 4, max_seq_len: int = 2048, pad_id: int = 0):
+ self.batch_size = batch_size
+ self.max_seq_len = max_seq_len
+ self.pad_id = pad_id
+
+ def pack(self, items: Sequence[T]) -> list[datatypes.RLTrainerPayload]:
+ """Pads items into rectangular 2D batches [B, max_seq_len]."""
+ return self._pad_2d_batches(items, self.batch_size, self.max_seq_len)
+```
+
+#### Reusing `BatchAssembler` Across Different Paradigms:
+```python
+# 1. RL Rollout Training
+rl_assembler = SequencePackedBatchAssembler[TrainExample](max_packed_len=8192)
+rl_microbatches = rl_assembler.pack(train_examples)
+
+# 2. Supervised Fine-Tuning (SFT)
+sft_assembler = SequencePackedBatchAssembler[SFTExample](max_packed_len=8192)
+sft_microbatches = sft_assembler.pack(sft_examples)
+
+# 3. Direct Preference Optimization (DPO / Simple 2D Batching)
+dpo_assembler = PaddedBatchAssembler[DPOPair](batch_size=8, max_seq_len=2048)
+dpo_batches = dpo_assembler.pack(dpo_pairs)
+```
+
+---
+
+### 5.3. Layer 2B: `AlgorithmAdapter` Math & Loss Wiring ([algorithm_adapter.py](https://github.com/google/tunix/blob/main/experimental/orchestrator/algorithm_adapter.py))
+
+```python
+"""Algorithm math, GAE / GRPO advantages, loss functions, and TrainExample assembly."""
+
+import abc
+from typing import Any, Callable, Sequence
+import jax.numpy as jnp
+import numpy as np
+from tunix.experimental.orchestrator import batch_assembly
+
+class AlgorithmAdapter(abc.ABC):
+ """Abstract algorithm adapter for returns math, advantages, and loss functions."""
+
+ def __init__(self, group_size: int = 8, mini_batch_size: int = 4, max_turns: int = 1, max_packed_len: int = 8192):
+ self.group_size = group_size
+ self.mini_batch_size = mini_batch_size
+ self.max_turns = max_turns
+ self.max_packed_len = max_packed_len
+ self.requires_reference_kl = False
+ self.has_critic = False
+ self.requires_old_logprobs = False
+
+ @abc.abstractmethod
+ def compute_advantages(self, rewards: np.ndarray | jnp.ndarray, **kwargs: Any) -> jnp.ndarray: ...
+
+ @abc.abstractmethod
+ def create_train_examples(
+ self,
+ group: Any,
+ rewards: list[float],
+ ref_logps: Any | None = None,
+ ) -> list[batch_assembly.TrainExample]:
+ """Assembles scored trajectories and computed advantages into typed TrainExamples."""
+ ...
+
+ @abc.abstractmethod
+ def loss_fn(self) -> Callable[..., Any]:
+ """Returns the JIT-compiled loss function executed on TrainerWorker."""
+ ...
+
+
+class GRPOAdapter(AlgorithmAdapter):
+ """Group Relative Policy Optimization (GRPO) adapter."""
+
+ def compute_advantages(self, rewards: np.ndarray | jnp.ndarray, num_generations: int = 8) -> jnp.ndarray:
+ """Computes group-normalized advantages: (r - mean(group)) / (std(group) + 1e-6)."""
+ r = jnp.asarray(rewards, dtype=jnp.float32).reshape(-1, num_generations)
+ mean = jnp.mean(r, axis=-1, keepdims=True)
+ std = jnp.std(r, axis=-1, keepdims=True)
+ advs = (r - mean) / (std + 1e-6)
+ return advs.reshape(-1)
+
+ def create_train_examples(
+ self,
+ group: Any,
+ rewards: list[float],
+ ref_logps: Any | None = None,
+ ) -> list[batch_assembly.TrainExample]:
+ """Packages group trajectories, advantages, and tool observation masks into TrainExamples."""
+ advs = self.compute_advantages(rewards, num_generations=self.group_size)
+ examples = []
+ for i, traj in enumerate(group.trajectories):
+ examples.append(
+ batch_assembly.TrainExample(
+ prompt_tokens=traj.prompt_tokens,
+ completion_tokens=traj.completion_tokens,
+ action_mask=traj.action_mask, # 0 for tool observations, 1 for assistant tokens
+ advantage=float(advs[i]),
+ ref_logprobs=ref_logps[i] if ref_logps is not None else None,
+ )
+ )
+ return examples
+
+ def loss_fn(self) -> Callable[..., Any]:
+ """GRPO clipped surrogate loss with beta * KL penalty."""
+ def _loss(params, batch):
+ # JAX / flax.nnx forward pass + ratio * A - beta * KL
+ ...
+ return _loss
+
+
+class PPOAdapter(AlgorithmAdapter):
+ """Generalized Advantage Estimation (GAE) and PPO Actor-Critic adapter."""
+
+ def __init__(self, gamma: float = 0.99, lam: float = 0.95, **kwargs: Any):
+ super().__init__(**kwargs)
+ self.gamma = gamma
+ self.lam = lam
+ self.has_critic = True
+ self.requires_reference_kl = True
+ self.requires_old_logprobs = True
+
+ def compute_advantages(self, rewards: np.ndarray, values: np.ndarray, **kwargs: Any) -> tuple[jnp.ndarray, jnp.ndarray]:
+ """Computes GAE advantages and value function regression targets."""
+ # delta_t = r_t + gamma * V(s_{t+1}) - V(s_t)
+ # A_t = sum (gamma * lam)^l * delta_{t+l}
+ return gae_advantages, value_targets
+
+ def create_train_examples(self, group: Any, rewards: list[float], ref_logps: Any | None = None) -> list[batch_assembly.TrainExample]:
+ # Builds TrainExamples with GAE advantages, value targets, and old_logprobs
+ ...
+```
+
+---
+
+### 5.4. Layer 3: `StandardRLProgram` & `AsyncRLProgram` ([async_rl_program.py](https://github.com/google/tunix/blob/main/experimental/orchestrator/async_rl_program.py))
+
+```python
+"""Off-the-shelf StandardRLProgram with long-polling rollout collector and TrainExample pipeline."""
+
+import asyncio
+from collections.abc import Callable, Iterable
+from typing import Any
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import distributed_rl_engine
+
+class AsyncRLProgram:
+ """Base class for asynchronous multi-stage DAG workflows."""
+
+ def __init__(self):
+ self._is_running = False
+
+ def make_group_queue(self, name: str, group_size: int = 1) -> Any:
+ """Requests an infrastructure-managed, checkpointable TrajectoryQueue."""
+ return TrajectoryQueueManager(group_size=group_size)
+
+
+class StandardRLProgram(AsyncRLProgram):
+ """Single standard program handling 95% of use cases with long-polling rollouts."""
+
+ def __init__(
+ self,
+ dataset: Iterable[Any],
+ algo: algorithm_adapter.AlgorithmAdapter,
+ reward_fns: list[Callable[..., Any]],
+ assembler: batch_assembly.BatchAssembler | None = None,
+ ):
+ super().__init__()
+ self.dataset = dataset
+ self.algo = algo
+ self.reward_fns = reward_fns
+ self.assembler = assembler or batch_assembly.SequencePackedBatchAssembler(max_packed_len=algo.max_packed_len)
+ self.raw_q = self.make_group_queue("raw", group_size=algo.group_size)
+ self.scored_q = self.make_group_queue("scored", group_size=algo.group_size)
+ self.current_policy_version = 0
+
+ async def rollout_dispatch_stage(self, engine: distributed_rl_engine.DistributedRLEngine):
+ """Stage 1A: Dispatches rollout requests across workers asynchronously (fire-and-forget)."""
+ for prompt_idx, prompt_item in enumerate(self.dataset):
+ if not self._is_running:
+ break
+
+ prompt_id = f"prompt_{prompt_idx}"
+ group_id = f"group_{prompt_idx}"
+ requests = []
+ for g_idx in range(self.algo.group_size):
+ req = datatypes.RolloutRequest(
+ request_id=f"req_{prompt_idx}_{g_idx}",
+ prompt=prompt_item,
+ prompt_id=prompt_id,
+ group_id=group_id,
+ target_policy_version=self.current_policy_version,
+ max_turns=self.algo.max_turns,
+ metadata={"group_id": group_id, "pair_index": g_idx},
+ )
+ requests.append(req)
+
+ await engine.dispatch_rollouts(requests)
+
+ async def polling_stage(self, engine: distributed_rl_engine.DistributedRLEngine):
+ """Stage 1B: Long-polls completed worker rollout responses into the grouping queue."""
+ while self._is_running:
+ try:
+ completed_responses = await engine.poll_rollouts(timeout_s=0.1)
+ if completed_responses:
+ for resp in completed_responses:
+ traj_item = datatypes.TrajectoryItem.from_rollout_response(resp)
+ await self.raw_q.put(traj_item)
+ else:
+ await asyncio.sleep(0.01)
+ except asyncio.CancelledError:
+ break
+
+ async def critique_stage(self, engine: distributed_rl_engine.DistributedRLEngine):
+ """Stage 2: Scores rewards, neural PRMs, and reference KL logprobs."""
+ async for group in self.raw_q:
+ rewards = [fn(group) for fn in self.reward_fns]
+ ref_logps = None
+ if self.algo.requires_reference_kl:
+ ref_logps = await engine.per_token_logps(datatypes.Role.REFERENCE, group)
+ await self.scored_q.put(group, rewards=rewards, ref_logps=ref_logps)
+
+ async def train_stage(self, engine: distributed_rl_engine.DistributedRLEngine, num_steps: int):
+ """Stage 3: Streaming gradient accumulation with TrainExamples."""
+ for step in range(num_steps):
+ uncommitted_groups = []
+ for group_idx in range(self.algo.mini_batch_size):
+ group = await self.scored_q.get_group()
+ uncommitted_groups.append(group)
+
+ # 1. Assembles self-contained TrainExamples (math + observation masks)
+ train_examples = self.algo.create_train_examples(group, group.rewards, ref_logps=group.ref_logps)
+
+ # 2. 1D Sequence packing into hardware-sized microbatches
+ microbatches = self.assembler.pack(train_examples)
+
+ # 3. Streaming gradient accumulation on accelerator HBM
+ is_final = (group_idx == self.algo.mini_batch_size - 1)
+ for batch in microbatches:
+ await engine.train_step(
+ batch,
+ role=datatypes.Role.ACTOR,
+ accumulate_gradients=True,
+ apply_optimizer=is_final,
+ )
+
+ # 4. Global step boundary: collective weight sync & queue commit
+ self.current_policy_version = await engine.sync_weights(role=datatypes.Role.ACTOR)
+ self.scored_q.commit(step, groups=uncommitted_groups)
+
+ async def run_async(self, engine: distributed_rl_engine.DistributedRLEngine, num_steps: int):
+ """Launches all stages concurrently on event loop."""
+ self._is_running = True
+ train_task = asyncio.create_task(self.train_stage(engine, num_steps))
+ tasks = [
+ asyncio.create_task(self.rollout_dispatch_stage(engine)),
+ asyncio.create_task(self.polling_stage(engine)),
+ asyncio.create_task(self.critique_stage(engine)),
+ train_task,
+ ]
+ try:
+ await train_task
+ finally:
+ self._is_running = False
+ for t in tasks:
+ if not t.done():
+ t.cancel()
+```
+
+---
+
+### 5.5. Infrastructure Glue: `ClusterOrchestrator` ([orchestrator.py](https://github.com/google/tunix/blob/main/experimental/orchestrator/orchestrator.py))
+
+```python
+"""Cluster Infrastructure Coordinator managing health, lifecycle, and program execution."""
+
+class ClusterOrchestrator:
+ """Supervises cluster hardware, health monitoring, and program execution."""
+
+ def __init__(self, config: Any):
+ self.config = config
+ self.registry = WorkerRegistry()
+ self.lifecycle = LifecycleDriver(self.registry)
+ self.monitor = HealthMonitor(self.registry)
+ self.engine = None
+
+ def __enter__(self) -> "ClusterOrchestrator":
+ """Interactive context manager bring-up."""
+ self.bring_up()
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ self.shutdown()
+
+ def bring_up(self) -> None:
+ self.lifecycle.bring_up()
+ self.monitor.start_heartbeats()
+ self.engine = self._create_engine()
+
+ def shutdown(self) -> None:
+ self.monitor.close()
+ self.lifecycle.shutdown()
+
+ def _create_engine(self) -> DistributedRLEngine:
+ return DistributedRLEngine(
+ rollout_workers=self.registry.group(datatypes.Role.ROLLOUT).members(),
+ trainer_workers={
+ datatypes.Role.ACTOR: self.registry.group(datatypes.Role.ACTOR).members()[0]
+ },
+ )
+
+ def run(
+ self,
+ algo: algorithm_adapter.AlgorithmAdapter,
+ dataset: Any,
+ reward_fns: list[Callable[..., Any]],
+ assembler: batch_assembly.BatchAssembler | None = None,
+ program: AsyncRLProgram | None = None,
+ num_steps: int = 1000,
+ ) -> None:
+ """Managed Program Submission: auto-wires Engine, Assembler, Queues & StandardProgram."""
+ if self.engine is None:
+ self.bring_up()
+
+ active_assembler = assembler or batch_assembly.SequencePackedBatchAssembler(
+ max_packed_len=algo.max_packed_len
+ )
+ active_program = program or StandardRLProgram(
+ dataset=dataset,
+ algo=algo,
+ reward_fns=reward_fns,
+ assembler=active_assembler,
+ )
+
+ # Executes stages concurrently on event loop
+ asyncio.run(active_program.run_async(self.engine, num_steps))
+```
+
+---
+
+## 6. How the 4 Major RL Variants Look to Users
+
+### Case 1: Standard GRPO (Math / Rule Rewards)
+```python
+orchestrator = ClusterOrchestrator(config)
+orchestrator.run(
+ algo=GRPOAdapter(group_size=8, mini_batch_size=4),
+ dataset=math_prompts,
+ reward_fns=[math_rule_verifier],
+)
+```
+
+### Case 2: GRPO with Neural PRM & Reference Model KL
+```python
+algo = GRPOAdapter(group_size=8, mini_batch_size=4)
+algo.requires_reference_kl = True # Automatically evaluates Role.REFERENCE logprobs!
+
+orchestrator.run(
+ algo=algo,
+ dataset=code_prompts,
+ reward_fns=[neural_prm_scorer],
+)
+```
+
+### Case 3: PPO Actor-Critic (Learned Value Function)
+```python
+orchestrator.run(
+ algo=PPOAdapter(gamma=0.99, lam=0.95), # Automatically trains Role.ACTOR and Role.CRITIC!
+ dataset=dialog_prompts,
+ reward_fns=[human_preference_reward_model],
+)
+```
+
+### Case 4: Multi-Turn Agentic Tool Calling (Docker Sandbox)
+```python
+orchestrator.run(
+ algo=GRPOAdapter(group_size=4, max_turns=10), # Auto-masks observations with action_mask=0!
+ dataset=web_browser_tasks,
+ reward_fns=[task_success_evaluator],
+)
+```
+
+---
+
+## 7. Phased Implementation & Execution Roadmap
+
+To transition cleanly from the prototype in CL 959796292 to the streamlined V2 architecture, we execute in **4 self-contained CLs**:
+
+```mermaid
+graph LR
+ CL1["CL 1: Dissolve RLDriver
• Remove RLDriver wrapper
• Connect Program directly to Engine + Adapter
• Fix PyType RolloutRequest bug"]
+ CL2["CL 2: Universal batch_assembly.py
• Implement SequencePackedBatchAssembler[T] (1D)
• Implement PaddedBatchAssembler[T] (2D)
• Generic support for TrainExample, SFT, DPO"]
+ CL3["CL 3: Streaming Grad Accumulation
• Implement group-by-group streaming in StandardRLProgram
• Add TrajectoryQueue ACK / in-flight offsets"]
+ CL4["CL 4: Orbax Manifest Recovery
• Add CompositeCheckpointHandler
• Enable isolated Trainer pod reboot & queue replay"]
+
+ CL1 --> CL2 --> CL3 --> CL4
+```
+
+| CL Number | Phase | Scope & Key Actions | Verification Target |
+| :---: | :---: | :--- | :--- |
+| **CL 1** | **Streamlined Core** | • Delete `rl_driver.py` and pass-through shims.
• Update `RLProgram` and `AsyncRLProgram` to take `(engine, algo)`.
• Fix `RolloutRequest.group_id` pytype error in `datatypes.py` / `async_rl_program.py`.
• Wire `TrainerWorker.per_token_logps` in `DistributedRLEngine`. | `test //third_party/py/tunix/experimental/orchestrator:all` |
+| **CL 2** | **Universal Batch Assembly** | • Introduce `batch_assembly.py` (`BatchAssembler[T]` protocol, `SequencePackedBatchAssembler[T]` for 1D token packing, and `PaddedBatchAssembler[T]` for 2D rectangular padding).
• Support universal packing across RL (`TrainExample`), Supervised (`SFTExample`), and Preference (`DPOPair`) dataclasses.
• Add block-diagonal attention mask generation and tool observation loss masking (`action_mask = 0`). | `test //third_party/py/tunix/experimental/orchestrator:all` |
+| **CL 3** | **Streaming & Queue ACK** | • Implement streaming gradient accumulation in `StandardRLProgram` with `TrainExample` pipeline.
• Add uncommitted in-flight buffer and `queue.commit()` to `TrajectoryQueueManager`. | `test //third_party/py/tunix/experimental/orchestrator:all` |
+| **CL 4** | **Isolated Recovery** | • Implement Orbax `CompositeCheckpointHandler` (`manifest.json` with step, model weights, queue offsets).
• Wire `LifecycleDriver.restart_worker(role=Role.ACTOR)` and queue seek on failure. | `test //third_party/py/tunix/experimental/orchestrator:all` |
+
diff --git a/tunix/experimental/orchestrator/algorithm_adapter.py b/tunix/experimental/orchestrator/algorithm_adapter.py
new file mode 100644
index 000000000..04b824bd8
--- /dev/null
+++ b/tunix/experimental/orchestrator/algorithm_adapter.py
@@ -0,0 +1,251 @@
+# 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.
+
+"""Layer 2B: AlgorithmAdapter Math & Loss Wiring (algorithm_adapter.py).
+
+Encapsulates RL returns, GAE / GRPO advantages, loss functions, and
+RLTrainerPayload assembly matching Orchestrator V2 and delegating loss
+computations directly to `tunix.rl.algo_core`.
+"""
+
+import abc
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import jax.numpy as jnp
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.rl import algo_core
+
+
+class AlgorithmAdapter(abc.ABC):
+ """Abstract algorithm adapter for returns math, advantages, and loss functions."""
+
+ def __init__(
+ self,
+ group_size: int = 8,
+ mini_batch_size: int = 4,
+ max_turns: int = 1,
+ max_packed_len: int = 8192,
+ ):
+ self.group_size = group_size
+ self.mini_batch_size = mini_batch_size
+ self.max_turns = max_turns
+ self.max_packed_len = max_packed_len
+ self.requires_reference_kl = False
+ self.has_critic = False
+ self.requires_old_logprobs = False
+
+ @abc.abstractmethod
+ def compute_advantages(
+ self, rewards: np.ndarray | jnp.ndarray | Sequence[float], **kwargs: Any
+ ) -> Any:
+ """Computes returns and advantages from rewards."""
+ ...
+
+ @abc.abstractmethod
+ def create_trainer_payloads(
+ self,
+ group: Any,
+ rewards: Sequence[float],
+ ref_logps: Any | None = None,
+ **kwargs: Any,
+ ) -> list[datatypes.RLTrainerPayload]:
+ """Assembles scored trajectories and computed advantages into typed RLTrainerPayloads."""
+ ...
+
+ @abc.abstractmethod
+ def loss_fn(self) -> Callable[..., Any]:
+ """Returns the JIT-compiled loss function executed on TrainerWorker."""
+ ...
+
+
+# TODO: Align adapter classes with current path and try to refactor and reuse directly instead of copying.
+class GRPOAdapter(AlgorithmAdapter):
+ """Group Relative Policy Optimization (GRPO) adapter."""
+
+ def __init__(
+ self,
+ group_size: int = 8,
+ mini_batch_size: int = 4,
+ max_turns: int = 1,
+ max_packed_len: int = 8192,
+ clip_epsilon: float = 0.2,
+ beta_kl: float = 0.04,
+ ):
+ super().__init__(
+ group_size=group_size,
+ mini_batch_size=mini_batch_size,
+ max_turns=max_turns,
+ max_packed_len=max_packed_len,
+ )
+ self.clip_epsilon = clip_epsilon
+ self.beta_kl = beta_kl
+
+ def compute_advantages(
+ self,
+ rewards: np.ndarray | jnp.ndarray | Sequence[float],
+ num_generations: int | None = None,
+ **kwargs: Any,
+ ) -> jnp.ndarray:
+ """Computes group-normalized advantages: (r - mean(group)) / (std(group) + 1e-6)."""
+ del kwargs
+ g = num_generations or self.group_size
+ r = jnp.asarray(rewards, dtype=jnp.float32).reshape(-1, g)
+ mean = jnp.mean(r, axis=-1, keepdims=True)
+ std = jnp.std(r, axis=-1, keepdims=True)
+ advs = (r - mean) / (std + 1e-6)
+ return advs.reshape(-1)
+
+ def create_trainer_payloads(
+ self,
+ group: Sequence[datatypes.TrajectoryItem],
+ rewards: Sequence[float],
+ ref_logps: Any | None = None,
+ **kwargs: Any,
+ ) -> list[datatypes.RLTrainerPayload]:
+ """Packages group trajectories, advantages, and tool observation masks into unbatched RLTrainerPayloads."""
+ del kwargs
+ advs = self.compute_advantages(rewards, num_generations=self.group_size)
+ payloads = []
+
+ for i, item in enumerate(group):
+ prompt_tokens = item.prompt_tokens if item.prompt_tokens is not None else np.zeros(0, dtype=np.int32)
+ completion_tokens = item.completion_tokens if item.completion_tokens is not None else np.zeros(0, dtype=np.int32)
+ action_mask = item.action_mask if item.action_mask is not None else np.zeros(0, dtype=np.float32)
+
+ adv_val = float(advs[i]) if i < len(advs) else 0.0
+ ref_lp = ref_logps[i] if ref_logps is not None and i < len(ref_logps) else None
+
+ p_arr = np.asarray(prompt_tokens, dtype=np.int32).reshape(-1)
+ c_arr = np.asarray(completion_tokens, dtype=np.int32).reshape(-1)
+ act_arr = np.asarray(action_mask, dtype=np.float32).reshape(-1)
+
+ seq_tokens = np.concatenate([p_arr, c_arr]) if (len(p_arr) > 0 or len(c_arr) > 0) else np.zeros(0, dtype=np.int32)
+ seq_loss_mask = np.concatenate([np.zeros(len(p_arr), dtype=np.float32), act_arr])
+ seq_adv = np.full(len(seq_tokens), adv_val, dtype=np.float32)
+
+ payload = datatypes.RLTrainerPayload(
+ token_ids=seq_tokens,
+ token_mask=np.ones_like(seq_tokens, dtype=np.float32),
+ loss_mask=seq_loss_mask,
+ advantages=seq_adv,
+ action_mask=seq_loss_mask,
+ ref_per_token_logps=np.asarray(ref_lp, dtype=np.float32) if ref_lp is not None else None,
+ )
+ payloads.append(payload)
+ return payloads
+
+ def loss_fn(self) -> Callable[..., Any]:
+ """GRPO loss function delegating directly to `algo_core.grpo_loss_fn`."""
+ return algo_core.grpo_loss_fn
+
+
+class PPOAdapter(AlgorithmAdapter):
+ """Generalized Advantage Estimation (GAE) and PPO Actor-Critic adapter."""
+
+ def __init__(
+ self,
+ group_size: int = 1,
+ mini_batch_size: int = 4,
+ max_turns: int = 1,
+ max_packed_len: int = 8192,
+ gamma: float = 0.99,
+ lam: float = 0.95,
+ clip_epsilon: float = 0.2,
+ ):
+ super().__init__(
+ group_size=group_size,
+ mini_batch_size=mini_batch_size,
+ max_turns=max_turns,
+ max_packed_len=max_packed_len,
+ )
+ self.gamma = gamma
+ self.lam = lam
+ self.clip_epsilon = clip_epsilon
+ self.has_critic = True
+ self.requires_reference_kl = True
+ self.requires_old_logprobs = True
+
+ def compute_advantages(
+ self,
+ rewards: np.ndarray | jnp.ndarray | Sequence[float],
+ values: np.ndarray | jnp.ndarray | None = None,
+ **kwargs: Any,
+ ) -> tuple[jnp.ndarray, jnp.ndarray]:
+ """Computes GAE advantages and value function regression targets."""
+ del kwargs
+ r = jnp.asarray(rewards, dtype=jnp.float32)
+ if values is None:
+ values = jnp.zeros_like(r)
+ else:
+ values = jnp.asarray(values, dtype=jnp.float32)
+
+ # 1-step / scalar GAE fallback for sequence-level rewards
+ deltas = r - values
+ gae_advantages = deltas
+ value_targets = r
+ return gae_advantages, value_targets
+
+ def create_trainer_payloads(
+ self,
+ group: Any,
+ rewards: Sequence[float],
+ ref_logps: Any | None = None,
+ values: Any | None = None,
+ old_logps: Any | None = None,
+ **kwargs: Any,
+ ) -> list[datatypes.RLTrainerPayload]:
+ """Builds unbatched RLTrainerPayloads with GAE advantages, value targets, and old_logprobs."""
+ del kwargs
+ advs, val_targets = self.compute_advantages(rewards, values=values)
+ payloads = []
+ trajectories = getattr(group, "trajectories", None) or (
+ group if isinstance(group, (list, tuple)) else [group]
+ )
+
+ for i, item in enumerate(trajectories):
+ prompt_tokens = item.prompt_tokens if item.prompt_tokens is not None else np.zeros(0, dtype=np.int32)
+ completion_tokens = item.completion_tokens if item.completion_tokens is not None else np.zeros(0, dtype=np.int32)
+ action_mask = item.action_mask if item.action_mask is not None else np.ones(len(completion_tokens), dtype=np.float32)
+
+ adv_val = float(advs[i]) if i < len(advs) else 0.0
+ vt_val = float(val_targets[i]) if i < len(val_targets) else 0.0
+ ref_lp = ref_logps[i] if ref_logps is not None and i < len(ref_logps) else None
+ old_lp = old_logps[i] if old_logps is not None and i < len(old_logps) else None
+
+ p_arr = np.asarray(prompt_tokens, dtype=np.int32).reshape(-1)
+ c_arr = np.asarray(completion_tokens, dtype=np.int32).reshape(-1)
+ act_arr = np.asarray(action_mask, dtype=np.float32).reshape(-1)
+
+ seq_tokens = np.concatenate([p_arr, c_arr]) if (len(p_arr) > 0 or len(c_arr) > 0) else np.zeros(0, dtype=np.int32)
+ seq_loss_mask = np.concatenate([np.zeros(len(p_arr), dtype=np.float32), act_arr])
+ seq_adv = np.full(len(seq_tokens), adv_val, dtype=np.float32)
+
+ payload = datatypes.RLTrainerPayload(
+ token_ids=seq_tokens,
+ token_mask=np.ones_like(seq_tokens, dtype=np.float32),
+ loss_mask=seq_loss_mask,
+ advantages=seq_adv,
+ action_mask=seq_loss_mask,
+ old_per_token_logps=np.asarray(old_lp, dtype=np.float32) if old_lp is not None else None,
+ ref_per_token_logps=np.asarray(ref_lp, dtype=np.float32) if ref_lp is not None else None,
+ returns=np.full(len(seq_tokens), vt_val, dtype=np.float32),
+ )
+ payloads.append(payload)
+ return payloads
+
+ def loss_fn(self) -> Callable[..., Any]:
+ """PPO policy loss function delegating directly to `algo_core.ppo_policy_loss_fn`."""
+ return algo_core.ppo_policy_loss_fn
diff --git a/tunix/experimental/orchestrator/async_rl_program.py b/tunix/experimental/orchestrator/async_rl_program.py
new file mode 100644
index 000000000..6210529bf
--- /dev/null
+++ b/tunix/experimental/orchestrator/async_rl_program.py
@@ -0,0 +1,285 @@
+# 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.
+
+"""Layer 3: Workflow & Program (async_rl_program.py) following Orchestrator V2.
+
+Contains:
+- AsyncRLProgram: Base class for multi-stage concurrent DAG workflows.
+- StandardRLProgram: Single standard program handling 95% of use cases with
+ long-polling rollout collector and streaming gradient accumulation.
+"""
+
+import asyncio
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any
+
+from absl import logging
+
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import rl_engine_interface
+from tunix.experimental.queue_manager import trajectory_queue_manager
+
+# _response_to_trajectory_item has been moved to distributed_rl_engine.py
+
+
+class AsyncRLProgram:
+ """Base class for asynchronous multi-stage DAG workflows."""
+
+ def __init__(self):
+ self._is_running = False
+ self.policy_version = 0
+
+ @property
+ def step(self) -> int:
+ return self.policy_version
+
+
+class StandardRLProgram(AsyncRLProgram):
+ """Single standard program handling 95% of use cases with long-polling rollouts.
+
+ Runs 4 concurrent stages:
+ 1. Rollout dispatch stage: Fire-and-forget requests across worker pool.
+ 2. Polling stage: Long-polls completed rollout responses into grouping queue.
+ 3. Critique stage: Scores rewards, PRMs, and reference KL logprobs.
+ 4. Train stage: Streaming gradient accumulation over microbatches.
+ """
+
+ def __init__(
+ self,
+ dataset: Iterable[Any],
+ algo: algorithm_adapter.AlgorithmAdapter,
+ reward_fns: Sequence[Callable[..., Any]] | None = None,
+ assembler: batch_assembly.BatchAssembler | None = None,
+ group_size: int = 8,
+ mini_batch_size: int = 4,
+ max_staleness: int | None = None,
+ on_step_begin: Callable[[int], None] | None = None,
+ on_step_end: Callable[[int, Any], None] | None = None,
+ ):
+ super().__init__()
+ self.dataset = dataset
+ self.algo = algo
+ self.reward_fns = list(reward_fns) if reward_fns else []
+ self.group_size = getattr(algo, "group_size", group_size)
+ self.mini_batch_size = getattr(algo, "mini_batch_size", mini_batch_size)
+ self.assembler = assembler or batch_assembly.SequencePackedBatchAssembler(
+ max_packed_len=getattr(algo, "max_packed_len", 8192)
+ )
+ self.on_step_begin = on_step_begin
+ self.on_step_end = on_step_end
+
+ self.raw_q = trajectory_queue_manager.TrajectoryQueueManager.create(
+ group_size=self.group_size,
+ max_staleness=max_staleness,
+ current_policy_version=lambda: self.policy_version,
+ )
+ self.scored_q = trajectory_queue_manager.TrajectoryQueueManager.create(
+ group_size=self.group_size
+ )
+
+ async def rollout_dispatch_stage(
+ self, engine: rl_engine_interface.AbstractRLEngine
+ ) -> None:
+ """Stage 1A: Dispatches rollout requests across workers asynchronously."""
+ for prompt_idx, prompt_item in enumerate(self.dataset):
+ # TODO: Extract prompt_id and group_id from standard tunix data structures
+ # rather than assuming dictionaries or falling back to index strings.
+ # TODO: the logic of creating group id and prompt id is incorrect and should be fixed.
+ prompt_id = getattr(prompt_item, "prompt_id", f"prompt_{prompt_idx}")
+ group_id = getattr(prompt_item, "group_id", f"group_{prompt_idx}")
+ if isinstance(prompt_item, dict):
+ prompt_id = prompt_item.get("prompt_id", prompt_id)
+ group_id = prompt_item.get("group_id", group_id)
+
+ for g_idx in range(self.group_size):
+ await engine.dispatch_rollouts(
+ [prompt_item],
+ request_id=f"req_{prompt_idx}_{g_idx}",
+ policy_version=self.policy_version,
+ prompt_ids=[prompt_id],
+ metadata={
+ "group_id": group_id,
+ "pair_index": g_idx,
+ },
+ )
+
+ async def polling_stage(
+ self, engine: rl_engine_interface.AbstractRLEngine
+ ) -> None:
+ """Stage 1B: Long-polls completed worker rollout responses into the queue."""
+ while True:
+ try:
+ completed = await engine.poll_rollouts(timeout_s=0.1)
+ if isinstance(completed, list) and completed:
+ for item in completed:
+ await self.raw_q.put(item)
+
+ except asyncio.CancelledError:
+ break
+ except Exception as exc: # pylint: disable=broad-exception-caught
+ logging.warning("Error in polling_stage: %s", exc)
+ await asyncio.sleep(0.01)
+
+ async def critique_stage(
+ self, engine: rl_engine_interface.AbstractRLEngine
+ ) -> None:
+ """Stage 2: Scores rewards, PRMs, and reference KL logprobs."""
+ while True:
+ try:
+ group = await self.raw_q.get_group()
+ except asyncio.CancelledError:
+ break
+ except Exception:
+ break
+
+ rewards = []
+ for item in group:
+ if self.reward_fns:
+ r = sum(fn(item) for fn in self.reward_fns)
+ else:
+ r = getattr(item.traj, "reward", 0.0)
+ rewards.append(float(r))
+
+ ref_logps = None
+ if getattr(self.algo, "requires_reference_kl", False):
+ ref_logps = await engine.per_token_logps(
+ datatypes.Role.REFERENCE, items=group
+ )
+
+ trainer_payloads = self.algo.create_trainer_payloads(
+ group, rewards=rewards, ref_logps=ref_logps
+ )
+ for idx, payload in enumerate(trainer_payloads):
+ adv = payload.advantages
+ reward_val = (
+ float(adv[0])
+ if hasattr(adv, "__len__") and len(adv) > 0
+ else float(adv)
+ )
+ item = datatypes.TrajectoryItem(
+ pair_index=idx,
+ group_id=getattr(group[0], "group_id", "default"),
+ start_step=0,
+ traj=datatypes.Trajectory(reward=reward_val),
+ # TODO: Stream RLTrainerPayload directly instead of re-wrapping in TrajectoryItem.
+ )
+ item.payload = payload
+ await self.scored_q.put(item)
+
+ async def train_stage(
+ self, engine: rl_engine_interface.AbstractRLEngine, num_steps: int | None = None
+ ) -> None:
+ """Stage 3: Streaming gradient accumulation with RLTrainerPayloads."""
+ step = 0
+ while num_steps is None or step < num_steps:
+ if self.on_step_begin:
+ self.on_step_begin(self.step)
+
+ uncommitted_groups = []
+ step_result = None
+
+ for group_idx in range(self.mini_batch_size):
+ scored_items = await self.scored_q.get_batch(num_groups=1)
+ if not scored_items:
+ break
+ uncommitted_groups.append(scored_items)
+
+ payloads = [getattr(item, "payload", None) for item in scored_items]
+ # TODO: Implement streaming microbatch assembly to overlap packing with trainer execution.
+ microbatches = self.assembler.pack(payloads)
+
+ is_final = group_idx == self.mini_batch_size - 1
+ for batch in microbatches:
+ step_result = await engine.train_step(
+ batch,
+ role=datatypes.Role.ACTOR,
+ accumulate_gradients=True,
+ apply_optimizer=is_final,
+ )
+
+ new_version = await engine.sync_weights(role=datatypes.Role.ACTOR)
+ self.policy_version = new_version if new_version else self.step + 1
+ self.scored_q.commit(step, groups=uncommitted_groups)
+
+ if self.on_step_end:
+ self.on_step_end(self.step, step_result)
+ step += 1
+
+ async def run_async(
+ self,
+ engine: rl_engine_interface.AbstractRLEngine,
+ num_steps: int | None = None,
+ **kwargs: Any,
+ ) -> None:
+ """Launches all stages concurrently on event loop."""
+ del kwargs
+ logging.info("Starting StandardRLProgram concurrent stages...")
+
+ train_task = asyncio.create_task(self.train_stage(engine, num_steps))
+ tasks = [
+ asyncio.create_task(self.rollout_dispatch_stage(engine)),
+ asyncio.create_task(self.polling_stage(engine)),
+ asyncio.create_task(self.critique_stage(engine)),
+ train_task,
+ ]
+
+ try:
+ while not train_task.done():
+ done, _ = await asyncio.wait(
+ tasks, return_when=asyncio.FIRST_COMPLETED, timeout=0.05
+ )
+ for task in done:
+ if task.exception():
+ raise task.exception()
+ if train_task.exception():
+ raise train_task.exception()
+ except Exception as exc:
+ logging.error("Exception in StandardRLProgram execution: %s", exc)
+ await self.raw_q.abort(exc)
+ await self.scored_q.abort(exc)
+ raise
+ finally:
+ for task in tasks:
+ if not task.done():
+ task.cancel()
+
+ def run(
+ self,
+ engine: rl_engine_interface.AbstractRLEngine,
+ num_steps: int | None = None,
+ **kwargs: Any,
+ ) -> None:
+ """Synchronous entry point running all stages on an event loop."""
+ try:
+ loop = asyncio.get_running_loop()
+ except RuntimeError:
+ loop = None
+
+ def _retrieve_task_exception(t: asyncio.Task[Any]) -> None:
+ try:
+ t.result()
+ except Exception: # pylint: disable=broad-except
+ # Exception is already logged inside run_async, we just need to
+ # retrieve it so asyncio doesn't complain about unretrieved exceptions.
+ pass
+
+ if loop and loop.is_running():
+ self._bg_task = asyncio.create_task(
+ self.run_async(engine, num_steps, **kwargs)
+ )
+ self._bg_task.add_done_callback(_retrieve_task_exception)
+ else:
+ asyncio.run(self.run_async(engine, num_steps, **kwargs))
diff --git a/tunix/experimental/orchestrator/batch_assembly.py b/tunix/experimental/orchestrator/batch_assembly.py
new file mode 100644
index 000000000..08baaf981
--- /dev/null
+++ b/tunix/experimental/orchestrator/batch_assembly.py
@@ -0,0 +1,262 @@
+# 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.
+
+"""Layer 2A: Universal Batch Assembly (batch_assembly.py) following Orchestrator V2.
+
+Generic tensor packing utility for unbatched `RLTrainerPayload` objects (or
+custom objects with token arrays). Supports:
+- 1D Sequence Packing (`SequencePackedBatchAssembler`) for Flash/FlexAttention (>90% MXU).
+- Simple 2D Rectangular Padding (`PaddedBatchAssembler`).
+
+# TODO: Align SequencePackedBatchAssembler with the rest of the ecosystem and potentially move to a common library.
+"""
+
+from typing import Generic, Protocol, Sequence, TypeVar
+import numpy as np
+from tunix.experimental.common import datatypes
+
+T = TypeVar("T")
+
+
+class BatchAssembler(Generic[T], Protocol):
+ """Universal batch assembly protocol for microbatch packing."""
+
+ def pack(self, items: Sequence[T]) -> list[datatypes.RLTrainerPayload]:
+ """Packs items into hardware-sized microbatch trainer payloads."""
+ ...
+
+
+class SequencePackedBatchAssembler:
+ """1D Sequence Packing: Concatenates items into dense [1, max_packed_len] buffers."""
+ # TODO: align implementation with current path.
+ def __init__(self, max_packed_len: int = 8192, pad_id: int = 0):
+ self.max_packed_len = max_packed_len
+ self.pad_id = pad_id
+
+ def pack(self, items: Sequence[datatypes.RLTrainerPayload]) -> list[datatypes.RLTrainerPayload]:
+ """Bin-packs items into dense 1D buffers with segment boundaries."""
+ if not items:
+ return []
+
+ # Calculate token lengths from explicit fields
+ item_lengths = []
+ for it in items:
+ item_lengths.append(len(it.token_ids) if it.token_ids is not None else 0)
+
+ item_list = sorted(zip(items, item_lengths), key=lambda x: x[1], reverse=True)
+
+ bins: list[list[datatypes.RLTrainerPayload]] = []
+ bin_lengths: list[int] = []
+
+ for item, length in item_list:
+ placed = False
+ for b_idx, current_len in enumerate(bin_lengths):
+ if current_len + length <= self.max_packed_len:
+ bins[b_idx].append(item)
+ bin_lengths[b_idx] += length
+ placed = True
+ break
+ if not placed:
+ bins.append([item])
+ bin_lengths.append(length)
+
+ payloads: list[datatypes.RLTrainerPayload] = []
+ for b_items in bins:
+ all_tokens = []
+ all_loss_masks = []
+ all_action_masks = []
+ all_segment_ids = []
+ all_segment_positions = []
+ all_advantages = []
+ all_old_logprobs = []
+ all_ref_logprobs = []
+
+ for seg_idx, it in enumerate(b_items, start=1):
+ toks = (
+ np.asarray(it.token_ids, dtype=np.int32).reshape(-1)
+ if it.token_ids is not None
+ else np.zeros(0, dtype=np.int32)
+ )
+ seq_len = len(toks)
+
+ all_tokens.append(toks)
+
+ loss_mask = (
+ it.loss_mask
+ if it.loss_mask is not None
+ else np.zeros(seq_len, dtype=np.float32)
+ )
+ all_loss_masks.append(np.asarray(loss_mask, dtype=np.float32).reshape(-1))
+
+ action_mask = (
+ it.action_mask
+ if it.action_mask is not None
+ else np.zeros(seq_len, dtype=np.float32)
+ )
+ all_action_masks.append(
+ np.asarray(action_mask, dtype=np.float32).reshape(-1)
+ )
+
+ adv_arr = (
+ np.asarray(it.advantages, dtype=np.float32).reshape(-1)
+ if it.advantages is not None
+ else np.zeros(seq_len, dtype=np.float32)
+ )
+ all_advantages.append(adv_arr)
+
+ all_segment_ids.append(np.full(seq_len, seg_idx, dtype=np.int32))
+ all_segment_positions.append(np.arange(seq_len, dtype=np.int32))
+
+ if it.old_per_token_logps is not None:
+ all_old_logprobs.append(
+ np.asarray(it.old_per_token_logps, dtype=np.float32).reshape(-1)
+ )
+
+ if it.ref_per_token_logps is not None:
+ all_ref_logprobs.append(
+ np.asarray(it.ref_per_token_logps, dtype=np.float32).reshape(-1)
+ )
+
+ concat_tokens = np.concatenate(all_tokens)
+ concat_loss_masks = np.concatenate(all_loss_masks)
+ concat_action_masks = np.concatenate(all_action_masks)
+ concat_segment_ids = np.concatenate(all_segment_ids)
+ concat_segment_positions = np.concatenate(all_segment_positions)
+ concat_advantages = np.concatenate(all_advantages)
+
+ pad_len = max(0, self.max_packed_len - len(concat_tokens))
+ padded_tokens = np.pad(concat_tokens[: self.max_packed_len], (0, pad_len), constant_values=self.pad_id)
+ padded_loss_mask = np.pad(concat_loss_masks[: self.max_packed_len], (0, pad_len), constant_values=0.0)
+ padded_action_mask = np.pad(concat_action_masks[: self.max_packed_len], (0, pad_len), constant_values=0.0)
+ padded_segment_ids = np.pad(concat_segment_ids[: self.max_packed_len], (0, pad_len), constant_values=0)
+ padded_segment_positions = np.pad(concat_segment_positions[: self.max_packed_len], (0, pad_len), constant_values=0)
+ padded_advantages = np.pad(concat_advantages[: self.max_packed_len], (0, pad_len), constant_values=0.0)
+
+ batch_old_lp = None
+ if all_old_logprobs:
+ concat_old = np.concatenate(all_old_logprobs)
+ batch_old_lp = np.pad(concat_old[: self.max_packed_len], (0, pad_len), constant_values=0.0)[np.newaxis, :]
+
+ batch_ref_lp = None
+ if all_ref_logprobs:
+ concat_ref = np.concatenate(all_ref_logprobs)
+ batch_ref_lp = np.pad(concat_ref[: self.max_packed_len], (0, pad_len), constant_values=0.0)[np.newaxis, :]
+
+ payload = datatypes.RLTrainerPayload(
+ token_ids=padded_tokens[np.newaxis, :],
+ token_mask=padded_segment_ids[np.newaxis, :],
+ loss_mask=padded_loss_mask[np.newaxis, :],
+ advantages=padded_advantages[np.newaxis, :],
+ action_mask=padded_action_mask[np.newaxis, :],
+ old_per_token_logps=batch_old_lp,
+ ref_per_token_logps=batch_ref_lp,
+ segment_ids=padded_segment_ids[np.newaxis, :],
+ segment_positions=padded_segment_positions[np.newaxis, :],
+ )
+ payloads.append(payload)
+
+ return payloads
+
+
+class PaddedBatchAssembler:
+ """Simple 2D Rectangular Batching: Pads sequences to standard [batch_size, max_seq_len] tensors."""
+
+ def __init__(self, batch_size: int = 4, max_seq_len: int = 2048, pad_id: int = 0):
+ self.batch_size = batch_size
+ self.max_seq_len = max_seq_len
+ self.pad_id = pad_id
+
+ def pack(self, items: Sequence[datatypes.RLTrainerPayload]) -> list[datatypes.RLTrainerPayload]:
+ """Pads items into rectangular 2D batches [B, max_seq_len]."""
+ if not items:
+ return []
+
+ item_list = list(items)
+ payloads: list[datatypes.RLTrainerPayload] = []
+
+ for i in range(0, len(item_list), self.batch_size):
+ chunk = item_list[i : i + self.batch_size]
+
+ b_tokens = []
+ b_loss_masks = []
+ b_action_masks = []
+ b_advs = []
+ b_old_lps = []
+ b_ref_lps = []
+
+ for it in chunk:
+ toks = (
+ np.asarray(it.token_ids, dtype=np.int32).reshape(-1)
+ if it.token_ids is not None
+ else np.zeros(0, dtype=np.int32)
+ )
+ seq_len = len(toks)
+
+ loss_mask = (
+ np.asarray(it.loss_mask, dtype=np.float32).reshape(-1)
+ if it.loss_mask is not None
+ else np.zeros(seq_len, dtype=np.float32)
+ )
+ action_mask = (
+ np.asarray(it.action_mask, dtype=np.float32).reshape(-1)
+ if it.action_mask is not None
+ else np.zeros(seq_len, dtype=np.float32)
+ )
+ adv_arr = (
+ np.asarray(it.advantages, dtype=np.float32).reshape(-1)
+ if it.advantages is not None
+ else np.zeros(seq_len, dtype=np.float32)
+ )
+
+ pad_len = max(0, self.max_seq_len - seq_len)
+ b_tokens.append(np.pad(toks[: self.max_seq_len], (0, pad_len), constant_values=self.pad_id))
+ b_loss_masks.append(np.pad(loss_mask[: self.max_seq_len], (0, pad_len), constant_values=0.0))
+ b_action_masks.append(np.pad(action_mask[: self.max_seq_len], (0, pad_len), constant_values=0.0))
+ b_advs.append(np.pad(adv_arr[: self.max_seq_len], (0, pad_len), constant_values=0.0))
+
+ if it.old_per_token_logps is not None:
+ old_arr = np.asarray(it.old_per_token_logps, dtype=np.float32).reshape(-1)
+ b_old_lps.append(
+ np.pad(old_arr[: self.max_seq_len], (0, pad_len), constant_values=0.0)
+ )
+
+ if it.ref_per_token_logps is not None:
+ ref_arr = np.asarray(it.ref_per_token_logps, dtype=np.float32).reshape(-1)
+ b_ref_lps.append(
+ np.pad(ref_arr[: self.max_seq_len], (0, pad_len), constant_values=0.0)
+ )
+
+ # Pad rows up to batch_size
+ while len(b_tokens) < self.batch_size:
+ b_tokens.append(np.full(self.max_seq_len, self.pad_id, dtype=np.int32))
+ b_loss_masks.append(np.zeros(self.max_seq_len, dtype=np.float32))
+ b_action_masks.append(np.zeros(self.max_seq_len, dtype=np.float32))
+ b_advs.append(np.zeros(self.max_seq_len, dtype=np.float32))
+ if b_old_lps:
+ b_old_lps.append(np.zeros(self.max_seq_len, dtype=np.float32))
+ if b_ref_lps:
+ b_ref_lps.append(np.zeros(self.max_seq_len, dtype=np.float32))
+
+ payload = datatypes.RLTrainerPayload(
+ token_ids=np.stack(b_tokens),
+ token_mask=np.stack(b_loss_masks),
+ loss_mask=np.stack(b_loss_masks),
+ advantages=np.stack(b_advs),
+ action_mask=np.stack(b_action_masks),
+ old_per_token_logps=np.stack(b_old_lps) if b_old_lps else None,
+ ref_per_token_logps=np.stack(b_ref_lps) if b_ref_lps else None,
+ )
+ payloads.append(payload)
+
+ return payloads
diff --git a/tunix/experimental/orchestrator/distributed_rl_engine.py b/tunix/experimental/orchestrator/distributed_rl_engine.py
new file mode 100644
index 000000000..7a884290a
--- /dev/null
+++ b/tunix/experimental/orchestrator/distributed_rl_engine.py
@@ -0,0 +1,312 @@
+# 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.
+
+"""Distributed compute routing surface (Layer 1) following Orchestrator V2.
+
+Contains:
+- WorkerPoolBalancer: Load balancing, queue tracking, and prefix-cache affinity.
+- DistributedRLEngine: Worker-backed compute router implementing AbstractRLEngine.
+"""
+
+import asyncio
+import collections
+from collections.abc import Mapping, Sequence
+import inspect
+from typing import Any
+import uuid
+
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import rl_engine_interface
+from tunix.experimental.worker import remote_execution
+
+# TODO: this multi step conversions seem excessive we convert from trajecotry to response then to trajectory item. we should simplify
+def _response_to_trajectory_item(resp: Any) -> datatypes.TrajectoryItem:
+ """Converts a worker rollout response to an TrajectoryItem."""
+ if isinstance(resp, datatypes.TrajectoryItem):
+ return resp
+
+ if isinstance(resp, datatypes.RolloutResponse):
+ prompt_id = resp.prompt_id or "default_prompt"
+ metadata = dict(resp.metadata) if resp.metadata else {}
+ group_id = metadata.get("group_id", prompt_id)
+ pair_index = metadata.get("pair_index", 0)
+ traj = datatypes.Trajectory(
+ reward=resp.env_reward,
+ status=(
+ datatypes.TrajectoryStatus.SUCCEEDED
+ if resp.status == "COMPLETED"
+ else datatypes.TrajectoryStatus.FAILED
+ ),
+ )
+ item = datatypes.TrajectoryItem(
+ pair_index=pair_index,
+ group_id=group_id,
+ start_step=0,
+ traj=traj,
+ metadata=metadata,
+ prompt_tokens=resp.prompt_tokens,
+ policy_version=resp.policy_version,
+ )
+
+ assistant_tokens = []
+ assistant_masks = []
+ for seg in resp.segments:
+ if seg.source == "assistant":
+ assistant_tokens.append(seg.tokens)
+ assistant_masks.append(seg.loss_mask)
+ if assistant_tokens:
+ item.completion_tokens = np.concatenate(assistant_tokens)
+ item.action_mask = np.concatenate(assistant_masks)
+ else:
+ item.completion_tokens = np.zeros(0, dtype=np.int32)
+ item.action_mask = np.zeros(0, dtype=np.float32)
+ return item
+
+ if isinstance(resp, datatypes.Trajectory):
+ item = datatypes.TrajectoryItem(
+ pair_index=0,
+ group_id=getattr(resp, "task", "default_group"),
+ start_step=0,
+ traj=resp,
+ policy_version=getattr(resp, "policy_version", 0),
+ prompt_tokens=getattr(resp, "prompt_tokens", np.zeros(0, dtype=np.int32)),
+ completion_tokens=getattr(resp, "completion_tokens", np.zeros(0, dtype=np.int32)),
+ action_mask=getattr(resp, "action_mask", np.ones(len(getattr(resp, "completion_tokens", [])), dtype=np.float32)),
+ )
+ return item
+
+ raise TypeError(
+ f"Unsupported response type for trajectory conversion: {type(resp)}"
+ )
+
+
+class DistributedRLEngine(rl_engine_interface.AbstractRLEngine):
+ """Worker-backed compute router dispatching RPCs across role pools."""
+
+ def __init__(
+ self,
+ rollout_workers: Sequence[remote_execution.ActorHandle],
+ trainer_workers: Mapping[datatypes.Role, remote_execution.ActorHandle],
+ inference_workers: (
+ Mapping[datatypes.Role, remote_execution.ActorHandle] | None
+ ) = None,
+ ):
+ self._rollout_workers = list(rollout_workers)
+ self._rollout_pool = remote_execution.RoutingActorPool(
+ self._rollout_workers
+ )
+ self._trainer_workers = dict(trainer_workers)
+ self._inference_workers = dict(inference_workers or {})
+
+ async def _invoke_worker(
+ self,
+ worker: remote_execution.ActorHandle,
+ method_name: str,
+ **kwargs: Any,
+ ) -> Any:
+ """Helper invoking method on remote handle."""
+ res = worker.asubmit(method_name, **kwargs)
+ if inspect.isawaitable(res):
+ return await res
+ return res
+
+ async def dispatch_rollouts(
+ self, prompts: Sequence[Any], **kwargs: Any
+ ) -> list[str]:
+ """Dispatches rollout requests across workers, constructing RolloutRequests internally if needed."""
+ rollout_reqs: list[datatypes.RolloutRequest] = []
+ for idx, p in enumerate(prompts):
+ # TODO: why do we support sending rollout requests directly? shouldn't this be the engine resposibility?
+ if isinstance(p, datatypes.RolloutRequest):
+ rollout_reqs.append(p)
+ else:
+ prompt_ids = kwargs.get("prompt_ids")
+ if not prompt_ids or len(prompt_ids) != len(prompts):
+ raise ValueError(
+ "When passing raw prompts, 'prompt_ids' must be provided in"
+ " kwargs and match the length of prompts."
+ )
+ if "policy_version" not in kwargs:
+ raise ValueError(
+ "When passing raw prompts, 'policy_version' must be provided in"
+ " kwargs."
+ )
+ # TODO: should we autogenerate request_id?
+ req_id = (
+ kwargs.get("request_id") or f"req_{idx}_{uuid.uuid4().hex[:8]}"
+ )
+ rollout_reqs.append(
+ datatypes.RolloutRequest(
+ request_id=req_id,
+ prompt=p,
+ prompt_id=prompt_ids[idx],
+ target_policy_version=kwargs["policy_version"],
+ metadata=dict(kwargs.get("metadata", {})),
+ )
+ )
+
+ for req in rollout_reqs:
+ metadata = req.metadata or {}
+ route_key = metadata.get("prefix_hash")
+ if route_key is None:
+ route_key = req.prompt_id
+ worker = self._rollout_pool._get_next_actor(
+ kwargs={"route_key": route_key}
+ )
+
+ res = worker.dispatch_task(method_name="generate", requests=[req])
+ if inspect.isawaitable(res):
+ await res
+
+ return [r.request_id for r in rollout_reqs]
+
+ async def poll_rollouts(
+ self, timeout_s: float = 0.1
+ ) -> list[datatypes.TrajectoryItem]:
+ """Concurrently long-polls completed rollout responses across all workers."""
+ if not self._rollout_workers:
+ return []
+
+ async def _poll_worker(worker: remote_execution.ActorHandle) -> Any:
+ return await self._invoke_worker(
+ worker, "poll_responses", timeout_s=timeout_s
+ )
+
+ tasks = [_poll_worker(w) for w in self._rollout_workers]
+ responses = await asyncio.gather(*tasks, return_exceptions=True)
+ completed: list[datatypes.TrajectoryItem] = []
+
+ for i, resp in enumerate(responses):
+ if isinstance(resp, Exception) or resp is None:
+ continue
+ unwrap_fn = getattr(resp, "unwrap", None)
+ res = (
+ unwrap_fn() if callable(unwrap_fn) else getattr(resp, "result", resp)
+ )
+ if res is not None:
+ items = res if isinstance(res, list) else [res]
+ for it in items:
+ if isinstance(it, dict):
+ it = datatypes.RolloutResponse(**it)
+ completed.append(_response_to_trajectory_item(it))
+ return completed
+
+ async def generate(self, prompts: Sequence[Any], **kwargs: Any) -> list[datatypes.TrajectoryItem]:
+ """Blocking rollout generation: load-balances prompts across workers and awaits completion."""
+ if not self._rollout_workers:
+ raise ValueError("DistributedRLEngine has no registered rollout workers.")
+
+ worker_to_prompts: dict[Any, list[Any]] = collections.defaultdict(list)
+ for p in prompts:
+ metadata = dict(kwargs.get("metadata", {}))
+ route_key = metadata.get("prefix_hash") or metadata.get("prompt_id")
+ worker = self._rollout_pool._get_next_actor(
+ kwargs={"route_key": route_key}
+ )
+ worker_to_prompts[worker].append(p)
+
+ tasks = []
+ for worker, w_prompts in worker_to_prompts.items():
+ if w_prompts:
+ tasks.append(
+ self._invoke_worker(worker, "generate", prompts=w_prompts, **kwargs)
+ )
+
+ if not tasks:
+ return []
+
+ results = await asyncio.gather(*tasks)
+
+ raw_items = [
+ item
+ for sublist in results
+ for item in (sublist if isinstance(sublist, list) else [sublist])
+ ]
+ return [_response_to_trajectory_item(it) for it in raw_items]
+
+ async def score(
+ self,
+ role: datatypes.Role,
+ items: Sequence[Any],
+ **kwargs: Any,
+ ) -> list[float]:
+ """Routes reward / PRM scoring requests to InferenceWorker pool."""
+ worker = self._inference_workers.get(role)
+ if worker is None:
+ raise ValueError(f"No inference worker registered for role {role}")
+ return await self._invoke_worker(worker, "score", items=items, **kwargs)
+
+ async def per_token_logps(
+ self,
+ role: datatypes.Role,
+ items: Sequence[Any],
+ **kwargs: Any,
+ ) -> Any:
+ """Evaluates reference model or actor per-token logprobs."""
+ worker = self._inference_workers.get(role) or self._trainer_workers.get(
+ role
+ )
+ if worker is None:
+ raise ValueError(
+ f"No worker registered for per_token_logps with role {role}"
+ )
+ return await self._invoke_worker(
+ worker, "per_token_logps", items=items, **kwargs
+ )
+
+ async def train_step(
+ self,
+ payload: datatypes.RLTrainerPayload,
+ role: datatypes.Role = datatypes.Role.ACTOR,
+ accumulate_gradients: bool = False,
+ apply_optimizer: bool = True,
+ skip_jit: bool = False,
+ **kwargs: Any,
+ ) -> Any:
+ """Executes atomic gradient accumulation / update on TrainerWorker."""
+ worker = self._trainer_workers.get(role)
+ if worker is None:
+ raise ValueError(f"No trainer worker registered for role {role}")
+ # TODO: we need on apply_optimizer=apply_optimize steps we need to call update() too.
+ return await self._invoke_worker(
+ worker,
+ "fwd_bwd",
+ batch=payload,
+ accumulate_gradients=accumulate_gradients,
+ apply_optimizer=apply_optimizer,
+ skip_jit=skip_jit,
+ **kwargs,
+ )
+
+ async def sync_weights(
+ self,
+ role: datatypes.Role = datatypes.Role.ACTOR,
+ target_roles: Sequence[datatypes.Role] | None = None,
+ ) -> int:
+ """Executes accelerator-to-accelerator collective weight broadcast."""
+ # TODO: integrate with raiden controller instead
+ del target_roles
+ trainer = self._trainer_workers.get(role)
+ if trainer is None:
+ return 0
+ sync_metadata = await self._invoke_worker(trainer, "prepare_weight_sync")
+ tasks = [
+ self._invoke_worker(w, "weight_sync", metadata=sync_metadata)
+ for w in self._rollout_workers
+ if hasattr(w, "weight_sync") or hasattr(w, "asubmit")
+ ]
+ if tasks:
+ await asyncio.gather(*tasks)
+ return getattr(sync_metadata, "new_policy_version", 1)
diff --git a/tunix/experimental/orchestrator/orchestrator.py b/tunix/experimental/orchestrator/orchestrator.py
new file mode 100644
index 000000000..a32587f39
--- /dev/null
+++ b/tunix/experimental/orchestrator/orchestrator.py
@@ -0,0 +1,172 @@
+# 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.
+
+"""Cluster Infrastructure Coordinator (orchestrator.py) following Orchestrator V2.
+
+Supervises WorkerRegistry, LifecycleDriver, HealthMonitor, and StartupValidator.
+Provides Tier 1 Zero-Boilerplate Managed Program Submission (`run`) and Tier 3
+Custom Program Execution (`run_program`).
+"""
+
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any
+
+from absl import logging
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import async_rl_program
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import distributed_rl_engine
+from tunix.experimental.orchestrator import health_monitor
+from tunix.experimental.orchestrator import lifecycle
+from tunix.experimental.orchestrator import rl_program
+from tunix.experimental.orchestrator import startup_validation
+from tunix.experimental.orchestrator import worker_registry
+from tunix.experimental.worker import abstract_worker
+
+
+class ClusterOrchestrator:
+ """Supervises cluster hardware, health monitoring, and program execution."""
+
+ def __init__(
+ self,
+ config: Any = None,
+ registry: worker_registry.WorkerRegistry | None = None,
+ lifecycle_driver: lifecycle.LifecycleDriver | None = None,
+ monitor: health_monitor.HealthMonitor | None = None,
+ ):
+ """Initializes ClusterOrchestrator."""
+ self.config = config
+ self.registry = registry or worker_registry.WorkerRegistry()
+ self.lifecycle_driver = lifecycle_driver or lifecycle.LifecycleDriver(
+ self.registry
+ )
+ self.monitor = monitor or health_monitor.HealthMonitor(self.registry)
+ self.engine: distributed_rl_engine.DistributedRLEngine | None = None
+
+ def __enter__(self) -> "ClusterOrchestrator":
+ """Interactive context manager bring-up."""
+ self.bring_up_workers()
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ self.shutdown()
+
+ def register_worker(
+ self, worker: abstract_worker.Worker
+ ) -> datatypes.WorkerInfo:
+ """Registers a worker in the WorkerRegistry."""
+ return self.registry.register(worker)
+
+ def unregister_worker(self, worker_id: str) -> None:
+ """Unregisters a worker by its id."""
+ self.registry.unregister(worker_id)
+
+ def bring_up_workers(self, dummy_data: Any = None) -> None:
+ """Brings up all registered workers through lifecycle initialization."""
+ logging.info("Bringing up workers across cluster...")
+ self.lifecycle_driver.bring_up(dummy_data)
+ self.engine = self._create_engine()
+
+ def shutdown(self) -> None:
+ """Shuts down all workers and closes health monitoring resources."""
+ logging.info("Shutting down ClusterOrchestrator...")
+ self.monitor.close()
+ self.lifecycle_driver.shutdown()
+
+ def validate_startup(self, alg_config: Any, training_config: Any) -> None:
+ """Validates cluster geometry against configurations."""
+ startup_validation.validate_startup(
+ self.registry, alg_config, training_config
+ )
+
+ def _get_role_members(self, role: datatypes.Role | str) -> list[Any]:
+ role_key = role.value if isinstance(role, datatypes.Role) else role
+ members = self.registry.group(role_key).members()
+
+ # Fallback in case workers were registered with the enum object directly
+ if not members and isinstance(role, datatypes.Role):
+ members = self.registry.group(role).members()
+ return members
+
+ def _create_engine(self) -> distributed_rl_engine.DistributedRLEngine:
+ """Constructs a DistributedRLEngine from the registered role groups."""
+ rollout_workers = self._get_role_members(datatypes.Role.ROLLOUT)
+ actor_workers = self._get_role_members(datatypes.Role.ACTOR)
+ critic_workers = self._get_role_members(datatypes.Role.CRITIC)
+ reference_workers = self._get_role_members(datatypes.Role.REFERENCE)
+
+ trainer_workers = {}
+ if actor_workers:
+ trainer_workers[datatypes.Role.ACTOR] = actor_workers[0]
+ if critic_workers:
+ trainer_workers[datatypes.Role.CRITIC] = critic_workers[0]
+
+ inference_workers = {}
+ if reference_workers:
+ inference_workers[datatypes.Role.REFERENCE] = reference_workers[0]
+
+ return distributed_rl_engine.DistributedRLEngine(
+ rollout_workers=rollout_workers,
+ trainer_workers=trainer_workers,
+ inference_workers=inference_workers,
+ )
+
+ def run_program(
+ self,
+ program: rl_program.RLProgram,
+ train_dataset: Iterable[Any] | None = None,
+ num_steps: int | None = None,
+ bring_up: bool = True,
+ dummy_data: Any = None,
+ **kwargs: Any,
+ ) -> None:
+ """Runs an RL program to completion under supervision."""
+ if bring_up:
+ self.bring_up_workers(dummy_data=dummy_data)
+
+ self.monitor.poll()
+ logging.info("ClusterOrchestrator executing program...")
+ engine = self.engine or self._create_engine()
+
+ program.run(engine=engine, train_dataset=train_dataset, num_steps=num_steps, **kwargs)
+
+ def run(
+ self,
+ algo: algorithm_adapter.AlgorithmAdapter,
+ dataset: Any,
+ reward_fns: Sequence[Callable[..., Any]] | None = None,
+ assembler: batch_assembly.BatchAssembler | None = None,
+ program: async_rl_program.AsyncRLProgram | None = None,
+ num_steps: int = 1000,
+ ) -> None:
+ """Managed Program Submission: auto-wires Engine, Assembler, Queues & StandardRLProgram."""
+ if self.engine is None:
+ self.bring_up_workers()
+
+ active_assembler = assembler or batch_assembly.SequencePackedBatchAssembler(
+ max_packed_len=getattr(algo, "max_packed_len", 8192)
+ )
+ active_program = program or async_rl_program.StandardRLProgram(
+ dataset=dataset,
+ algo=algo,
+ reward_fns=reward_fns,
+ assembler=active_assembler,
+ )
+ self.run_program(
+ program=active_program,
+ train_dataset=dataset,
+ num_steps=num_steps,
+ bring_up=False,
+ )
diff --git a/tunix/experimental/orchestrator/rl_engine_interface.py b/tunix/experimental/orchestrator/rl_engine_interface.py
index 6d5f8cab5..e67dd348f 100644
--- a/tunix/experimental/orchestrator/rl_engine_interface.py
+++ b/tunix/experimental/orchestrator/rl_engine_interface.py
@@ -12,125 +12,64 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""The RL engine surface a learning loop drives.
+"""The RL engine interface (Layer 1 Compute Routing Protocol) following Orchestrator V2."""
-An RL engine is the compute driver for an RL training loop. `RLEngine`
-(in-process) and `OrchestratorRLEngine` (worker-backed) are two implementations
-of this same `AbstractRLEngine` surface: a learner builds its loop out of these
-calls and is agnostic to whether the work runs in-process or is dispatched to
-workers. Swapping the implementation turns a single-process run into a
-distributed one -- the loop code does not change.
-
-This Protocol maps directly to the current `RLEngine` interface in
-`tunix.rl.rl_cluster.RLEngine`, derived from the members the agentic learners
-(`AgenticRLLearner` / `GRPOLearner`) actually touch. It is grouped into:
-
- * compute primitives (generate / train / sync / score),
- * shared bookkeeping (step counter, metrics, tokenizer, teardown),
- * config & topology (cluster_config, role->mesh, rollout config), and
- * sub-engine accessors (rollout, actor_trainer, critic_trainer).
-
-The Protocol is structural: `RLEngine` satisfies it as-is, and
-`OrchestratorRLEngine` satisfies it by routing the compute primitives to
-workers and delegating the rest. Sub-engines are typed `Any` here (their own
-surfaces -- `rollout.pad_id()/eos_id()/model()`,
-`actor_trainer.with_loss_fn(...)`, etc. -- are large and
-worker-implementation-specific); the members exercised by the loop are noted in
-comments.
-"""
-
-from typing import Any, Mapping, Protocol, runtime_checkable
-from jax.typing import ArrayLike
+from collections.abc import Sequence
+from typing import Any, Protocol, runtime_checkable
from tunix.experimental.common import datatypes
@runtime_checkable
class AbstractRLEngine(Protocol):
- """Structural interface for an RL engine (maps to current `RLEngine`)."""
+ """Stateless compute primitives for distributed worker meshes."""
- # --- Shared bookkeeping ---------------------------------------------------
- # Step counter / weight version, read and written by the loop.
- global_steps: int
- # Tokenizer adapter used for chat templating / (de)tokenization.
- tokenizer: Any
+ async def dispatch_rollouts(
+ self, prompts: Sequence[Any], **kwargs: Any
+ ) -> list[str]:
+ """Dispatches rollout requests across workers (constructing RolloutRequests internally)."""
+ ...
- def buffer_metrics(self, metrics: Mapping[str, Any], mode: Any = ...) -> None:
- """Buffers metrics, flushed on the next step boundary."""
+ async def poll_rollouts(
+ self, timeout_s: float = 0.1
+ ) -> list[datatypes.TrajectoryItem]:
+ """Retrieves completed rollout responses from workers via long-polling."""
...
- def buffer_metrics_async(
- self, metrics: Mapping[str, Any], mode: Any = ..., step: int = ...
- ) -> None:
- """Buffers metrics for a specific step (async producer path)."""
+ async def generate(
+ self, prompts: Sequence[Any], **kwargs: Any
+ ) -> list[datatypes.TrajectoryItem]:
+ """Synchronous batched rollout generation over rollout workers."""
...
- def close(self) -> None:
- """Releases cluster resources at end of run."""
+ async def score(
+ self, role: datatypes.Role, items: Sequence[Any], **kwargs: Any
+ ) -> list[float]:
+ """Scores responses under a reward model."""
...
- # --- Generation (rollout) -------------------------------------------------
- def generate(
- self,
- prompts: list[str] | list[list[dict[str, str]]],
- apply_chat_template: bool = False,
- mode: Any = None,
- micro_batch_size: int | None = None,
- trace_tags: Mapping[str, Any] | None = None,
- max_generation_steps: int | None = None,
+ async def per_token_logps(
+ self, role: datatypes.Role, items: Sequence[Any], **kwargs: Any
) -> Any:
- """Generates completions for `prompts` (returns a `RolloutOutput`)."""
+ """Computes per-token log probabilities under a reference/actor model."""
...
- # --- Training (train_step) ------------------------------------------------
- def train(
+ async def train_step(
self,
- role: datatypes.Role,
- train_ds: Any,
- eval_ds: Any,
+ payload: datatypes.RLTrainerPayload,
+ role: datatypes.Role = datatypes.Role.ACTOR,
+ accumulate_gradients: bool = False,
+ apply_optimizer: bool = True,
skip_jit: bool = False,
- ) -> None:
- """Runs a training update for the specified role (e.g. Role.ACTOR, Role.CRITIC)."""
+ **kwargs: Any,
+ ) -> dict[str, Any]:
+ """Executes forward/backward gradient update on trainer workers."""
...
- # --- Scoring (feeds advantage / IS math) ----------------------------------
- def per_token_logps(
+ async def sync_weights(
self,
- role: datatypes.Role,
- prompt_tokens: ArrayLike,
- completion_tokens: ArrayLike,
- pad_id: int,
- eos_id: int,
- micro_batch_size: int | None = None,
- segment_ids: ArrayLike | None = None,
+ role: datatypes.Role = datatypes.Role.ACTOR,
+ target_roles: Sequence[datatypes.Role] | None = None,
**kwargs: Any,
- ) -> Any:
- """Per-token logprobs under the specified model role."""
- # TODO(noghabi): add per batch interface (replace or keep both). that is
- # the interface worker uses.
- ...
-
- # --- Weight sync ----------------------------------------------------------
- def sync_weights(self) -> None:
- """Publishes the trainer's weights to the rollout/inference replicas."""
+ ) -> int:
+ """Coordinates decentralized peer-to-peer weight sync across worker roles."""
...
-
- # --- Config & topology ----------------------------------------------------
- # ClusterConfig: training_config (+ micro-batch sizes, chunk sizes, metrics
- # options), rollout_config, rollout_engine, role_to_mesh.
- cluster_config: Any
- # Role -> Mesh mapping (aka cluster_config.role_to_mesh); `.devices`/`.empty`.
- r2m: Any
-
- def get_rollout_config(self, mode: Any) -> Any:
- """Returns the effective rollout config for the given mode."""
- ...
-
- # --- Sub-engine accessors (surfaces the loop reaches into) -----------------
- # rollout: `.pad_id()`, `.eos_id()`, `.model()`.
- rollout: Any
- # actor_trainer: `.with_loss_fn(...)`, `.with_gen_model_input_fn(...)`,
- # `.with_rl_metrics_to_log(...)`, `.is_managed_externally`,
- # `.restored_global_step()`, `.iter_steps`, `.train_steps`, `.model`.
- actor_trainer: Any
- # critic_trainer: present only for actor-critic algorithms (PPO). Guard with
- # hasattr(engine, "critic_trainer") before use.
diff --git a/tunix/experimental/orchestrator/rl_program.py b/tunix/experimental/orchestrator/rl_program.py
new file mode 100644
index 000000000..1207e0316
--- /dev/null
+++ b/tunix/experimental/orchestrator/rl_program.py
@@ -0,0 +1,159 @@
+# 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.
+
+"""Synchronous RL Program (rl_program.py) coordinating Engine, Algo, and Assembler."""
+
+import asyncio
+from collections.abc import Callable, Iterable, Sequence
+import inspect
+from typing import Any, Protocol
+
+from absl import logging
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import rl_engine_interface
+
+
+class RLProgram(Protocol):
+ """Standard contract for RL training programs running on ClusterOrchestrator."""
+ def run(
+ self,
+ engine: rl_engine_interface.AbstractRLEngine,
+ train_dataset: Iterable[Any] | None = None,
+ num_steps: int | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ ...
+
+
+def _sync_or_async(coro: Any) -> Any:
+ """Executes coroutine synchronously if no loop is running, else returns coro."""
+ if inspect.iscoroutine(coro):
+ try:
+ loop = asyncio.get_running_loop()
+ except RuntimeError:
+ loop = None
+
+ if loop and loop.is_running():
+ return coro
+ return asyncio.run(coro)
+ return coro
+
+
+class SyncRLProgram:
+ """Synchronous RL Program coordinating an iterative RL training loop."""
+
+ def __init__(
+ self,
+ engine: rl_engine_interface.AbstractRLEngine,
+ algo: algorithm_adapter.AlgorithmAdapter,
+ reward_fns: Sequence[Callable[..., Any]] | None = None,
+ assembler: batch_assembly.BatchAssembler | None = None,
+ on_step_begin: Callable[[int], None] | None = None,
+ on_step_end: Callable[[int, Any], None] | None = None,
+ ):
+ self.engine = engine
+ self.algo = algo
+ self.reward_fns = list(reward_fns) if reward_fns else []
+ self.assembler = assembler or batch_assembly.SequencePackedBatchAssembler(
+ max_packed_len=getattr(algo, "max_packed_len", 8192)
+ )
+ self.on_step_begin = on_step_begin
+ self.on_step_end = on_step_end
+ self.policy_version = 0
+
+ @property
+ def step(self) -> int:
+ return self.policy_version
+
+ def step_once(
+ self,
+ prompts: list[str] | list[list[dict[str, str]]],
+ **kwargs: Any,
+ ) -> Any:
+ """Executes a single end-to-end RL training step."""
+ current_step = self.policy_version
+ if self.on_step_begin:
+ self.on_step_begin(current_step)
+
+ # 1. Generate rollouts
+ rollouts = _sync_or_async(self.engine.generate(prompts=prompts, **kwargs))
+
+ # 2. Evaluate rewards
+ rewards = []
+ for item in rollouts:
+ r = sum(fn(item) for fn in self.reward_fns) if self.reward_fns else getattr(item, "env_reward", 0.0)
+ rewards.append(float(r))
+
+ # 3. Create RLTrainerPayloads via AlgorithmAdapter
+ ref_logps = None
+ if getattr(self.algo, "requires_reference_kl", False):
+ ref_logps = _sync_or_async(self.engine.per_token_logps(datatypes.Role.REFERENCE, items=rollouts))
+ trainer_payloads = self.algo.create_trainer_payloads(
+ rollouts, rewards=rewards, ref_logps=ref_logps
+ )
+
+ # 4. Pack into microbatches
+ microbatches = self.assembler.pack(trainer_payloads)
+
+ # 5. Execute gradient updates
+ step_result = None
+ for batch in microbatches:
+ step_result = _sync_or_async(
+ self.engine.train_step(
+ batch,
+ role=datatypes.Role.ACTOR,
+ accumulate_gradients=False,
+ apply_optimizer=True,
+ )
+ )
+
+ # 6. Sync weights to rollout replicas
+ _sync_or_async(self.engine.sync_weights(role=datatypes.Role.ACTOR))
+
+ # 7. Increment step
+ self.policy_version = current_step + 1
+
+ if self.on_step_end:
+ self.on_step_end(self.policy_version, step_result)
+
+ return step_result
+
+ def eval_step_once(
+ self,
+ prompts: list[str] | list[list[dict[str, str]]],
+ **kwargs: Any,
+ ) -> list[datatypes.RLTrainerPayload]:
+ """Executes evaluation step without updating weights."""
+ rollouts = _sync_or_async(self.engine.generate(prompts=prompts, **kwargs))
+ rewards = [
+ sum(fn(item) for fn in self.reward_fns) if self.reward_fns else getattr(item, "env_reward", 0.0)
+ for item in rollouts
+ ]
+ return self.algo.create_trainer_payloads(rollouts, rewards=rewards)
+
+ def run(
+ self,
+ train_dataset: Iterable[list[str] | list[list[dict[str, str]]]],
+ num_steps: int | None = None,
+ **kwargs: Any,
+ ) -> None:
+ """Runs the RL program training loop over the dataset."""
+ for idx, prompt_batch in enumerate(train_dataset):
+ if num_steps is not None and idx >= num_steps:
+ break
+ logging.info("RLProgram starting step %d", self.step)
+ self.step_once(prompts=prompt_batch, **kwargs)
diff --git a/tunix/experimental/orchestrator/simple_orchestrator_nb.py b/tunix/experimental/orchestrator/simple_orchestrator_nb.py
new file mode 100644
index 000000000..d9606a315
--- /dev/null
+++ b/tunix/experimental/orchestrator/simple_orchestrator_nb.py
@@ -0,0 +1,156 @@
+# 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.
+
+"""Simple Cluster Orchestrator Example Notebook / Script (V2 Architecture)."""
+
+from typing import Any
+from absl import logging
+import numpy as np
+from tunix.experimental.common import datatypes
+from tunix.experimental.orchestrator import algorithm_adapter
+from tunix.experimental.orchestrator import batch_assembly
+from tunix.experimental.orchestrator import orchestrator
+from tunix.experimental.worker import abstract_worker
+
+
+class SimulatedRolloutWorker(abstract_worker.Worker):
+ """Simulated RolloutWorker generating synthetic token responses."""
+
+ def __init__(self, worker_id: str):
+ self.worker_id = worker_id
+
+ def info(self) -> datatypes.WorkerInfo:
+ return datatypes.WorkerInfo(
+ worker_id=self.worker_id,
+ roles=frozenset([datatypes.Role.ROLLOUT]),
+ )
+
+ def initialize(self) -> datatypes.Response:
+ return datatypes.Response()
+
+ def compile(self, dummy_data: Any = None) -> datatypes.Response:
+ del dummy_data
+ return datatypes.Response()
+
+ def start(self) -> datatypes.Response:
+ return datatypes.Response()
+
+ def stop(self) -> datatypes.Response:
+ return datatypes.Response()
+
+ def generate(self, prompts: Any, **kwargs: Any) -> list[datatypes.RolloutResponse]:
+ del kwargs
+ logging.info("[%s] Generating rollouts for %d prompt(s)", self.worker_id, len(prompts))
+ responses = []
+ for idx, _ in enumerate(prompts):
+ responses.append(
+ datatypes.RolloutResponse(
+ request_id=f"req_{idx}",
+ status="COMPLETED",
+ env_reward=1.0,
+ prompt_tokens=np.array([10, 11], dtype=np.int32),
+ segments=[
+ datatypes.TokenSegment(
+ source="assistant",
+ tokens=np.array([20, 21], dtype=np.int32),
+ loss_mask=np.array([1, 1], dtype=np.int32),
+ )
+ ],
+ metadata={"prompt_id": f"prompt_{idx}"},
+ )
+ )
+ return responses
+
+ def heartbeat(self) -> datatypes.HealthReport:
+ return datatypes.HealthReport(state=datatypes.WorkerState.READY)
+
+
+class SimulatedTrainerWorker(abstract_worker.Worker):
+ """Simulated TrainerWorker executing gradient updates and weight sync staging."""
+
+ def __init__(self, worker_id: str, role: datatypes.Role):
+ self.worker_id = worker_id
+ self._role = role
+
+ def info(self) -> datatypes.WorkerInfo:
+ return datatypes.WorkerInfo(
+ worker_id=self.worker_id,
+ roles=frozenset([self._role]),
+ )
+
+ def initialize(self) -> datatypes.Response:
+ return datatypes.Response()
+
+ def compile(self, dummy_data: Any = None) -> datatypes.Response:
+ del dummy_data
+ return datatypes.Response()
+
+ def start(self) -> datatypes.Response:
+ return datatypes.Response()
+
+ def stop(self) -> datatypes.Response:
+ return datatypes.Response()
+
+ def fwd_bwd(self, batch: Any, accumulate_gradients: bool = False, apply_optimizer: bool = True, skip_jit: bool = False) -> dict[str, float]:
+ del batch, accumulate_gradients, apply_optimizer, skip_jit
+ logging.info("[%s] Executing gradient update (fwd_bwd)", self.worker_id)
+ return {"loss": 0.25, "grad_norm": 1.2}
+
+ def prepare_weight_sync(self) -> datatypes.WeightSyncMetadata:
+ logging.info("[%s] Staging weights for sync", self.worker_id)
+ meta = datatypes.WeightSyncMetadata(
+ new_policy_version=1,
+ transfer_mode="p2p",
+ source_endpoints=["trainer:50051"],
+ )
+ return meta
+
+ def heartbeat(self) -> datatypes.HealthReport:
+ return datatypes.HealthReport(state=datatypes.WorkerState.READY)
+
+
+def main():
+ orch = orchestrator.ClusterOrchestrator()
+
+ rollout_0 = SimulatedRolloutWorker("rollout_0")
+ rollout_1 = SimulatedRolloutWorker("rollout_1")
+ actor_trainer = SimulatedTrainerWorker("actor_0", datatypes.Role.ACTOR)
+
+ orch.register_worker(rollout_0)
+ orch.register_worker(rollout_1)
+ orch.register_worker(actor_trainer)
+
+ print(f"Registered roles in cluster: {orch.registry.roles()}")
+
+ algo = algorithm_adapter.GRPOAdapter(group_size=2, mini_batch_size=1, max_packed_len=32)
+ assembler = batch_assembly.SequencePackedBatchAssembler(max_packed_len=32)
+
+ train_dataset = [
+ ["Solve 2 + 2", "Solve 3 * 4"],
+ ["Solve 10 / 2", "Solve 7 - 5"],
+ ]
+
+ print("Executing Tier 1 Managed Run via ClusterOrchestrator...")
+ orch.run(
+ algo=algo,
+ dataset=train_dataset,
+ reward_fns=[lambda x: 1.0],
+ assembler=assembler,
+ num_steps=2,
+ )
+ print("Execution completed successfully!")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tunix/experimental/orchestrator/startup_validation.py b/tunix/experimental/orchestrator/startup_validation.py
index 9089bdf95..8f1809127 100644
--- a/tunix/experimental/orchestrator/startup_validation.py
+++ b/tunix/experimental/orchestrator/startup_validation.py
@@ -26,7 +26,7 @@
from typing import Protocol
from tunix.experimental.orchestrator import worker_registry
from tunix.rl import algorithm_config
-from tunix.rl import rl_cluster
+from tunix.rl import rl_cluster as rl_cluster_lib
from tunix.rl import utils as rl_utils
@@ -46,7 +46,7 @@ def validate(
self,
registry: worker_registry.WorkerRegistry,
alg_config: algorithm_config.AlgorithmConfig,
- training_config: rl_cluster.RLTrainingConfig,
+ training_config: rl_cluster_lib.RLTrainingConfig,
) -> list[str]:
"""Returns a list of error messages; empty list if all checks pass."""
...
@@ -59,7 +59,7 @@ def validate(
self,
registry: worker_registry.WorkerRegistry,
alg_config: algorithm_config.AlgorithmConfig,
- training_config: rl_cluster.RLTrainingConfig,
+ training_config: rl_cluster_lib.RLTrainingConfig,
) -> list[str]:
del registry
errors: list[str] = []
@@ -138,7 +138,7 @@ def validate(
def validate_startup(
registry: worker_registry.WorkerRegistry,
alg_config: algorithm_config.AlgorithmConfig,
- training_config: rl_cluster.RLTrainingConfig,
+ training_config: rl_cluster_lib.RLTrainingConfig,
*,
validators: tuple[StartupValidator, ...] = DEFAULT_VALIDATORS,
) -> None:
diff --git a/tunix/experimental/queue_manager/trajectory_queue_manager.py b/tunix/experimental/queue_manager/trajectory_queue_manager.py
index 6de722b8d..fc3ede461 100644
--- a/tunix/experimental/queue_manager/trajectory_queue_manager.py
+++ b/tunix/experimental/queue_manager/trajectory_queue_manager.py
@@ -12,12 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Specialized TrajectoryQueueManager for TrajectoryItem instances."""
-
-from __future__ import annotations
-
-from typing import Optional
+"""TrajectoryQueueManager specialization of GroupQueueManager for TrajectoryItem."""
+from collections.abc import Callable, Sequence
+from typing import Any, Optional
from tunix.experimental.common import datatypes
from tunix.rl.agentic.queue_manager import group_queue_manager
@@ -25,16 +23,8 @@
GroupFn = group_queue_manager.GroupFn[TrajectoryItem]
FilterFn = group_queue_manager.FilterFn[TrajectoryItem]
-
-class TrajectoryQueueManager(
- group_queue_manager.GroupQueueManager[TrajectoryItem]
-):
- """Specialized queue manager for TrajectoryItem instances.
-
- Inherits from `GroupQueueManager[TrajectoryItem]`. If no custom `group_fn` is
- provided, uses the default grouping function that groups items by
- `item.group_id` or `item.prompt_id` up to `group_size`.
- """
+class TrajectoryQueueManager(group_queue_manager.GroupQueueManager):
+ """Specialized GroupQueueManager holding TrajectoryItems with ACK and abort support."""
def __init__(
self,
@@ -53,7 +43,82 @@ def __init__(
filter_fn: Optional pluggable function to filter candidate groups.
"""
super().__init__(
+ group_size=group_size, group_fn=group_fn, filter_fn=filter_fn
+ )
+
+ @classmethod
+ def create(
+ cls,
+ group_size: int = 1,
+ max_staleness: int | None = None,
+ current_policy_version: Callable[[], int] | None = None,
+ filter_fn: Any | None = None,
+ ) -> "TrajectoryQueueManager":
+ """Creates a grouped trajectory queue with optional policy staleness filtering."""
+ combined_filter = filter_fn
+ if max_staleness is not None and current_policy_version is not None:
+ def _staleness_filter(group: Sequence[Any]) -> Any:
+ min_allowed = current_policy_version() - max_staleness
+ valid = [
+ item for item in group
+ if getattr(item, "policy_version", 0) >= min_allowed
+ ]
+ filtered = [
+ item for item in group
+ if getattr(item, "policy_version", 0) < min_allowed
+ ]
+ if filter_fn is not None:
+ res = filter_fn(valid)
+ if isinstance(res, tuple):
+ return res[0], list(res[1]) + filtered
+ return res, filtered
+ return valid, filtered
+ combined_filter = _staleness_filter
+
+ return cls(
group_size=group_size,
- group_fn=group_fn,
- filter_fn=filter_fn,
+ filter_fn=combined_filter,
)
+
+ def __aiter__(self) -> "TrajectoryQueueManager":
+ return self
+
+ async def __anext__(self) -> list[datatypes.TrajectoryItem]:
+ group = await self.get_group()
+ if not group:
+ raise StopAsyncIteration
+ return group
+
+ async def get_group(self) -> list[datatypes.TrajectoryItem]:
+ """Retrieves a single ready group of TrajectoryItems."""
+ return await self._get_one_ready_group()
+
+ async def get_batch(
+ self,
+ batch_size: int | None = None,
+ num_groups: int | None = None,
+ ) -> list[datatypes.TrajectoryItem]:
+ """Retrieves items by either batch_size or num_groups."""
+ # TODO: why do we need both batch_size and num_groups?
+ # TODO: should this be in parent class?
+ if num_groups is not None:
+ out: list[datatypes.TrajectoryItem] = []
+ for _ in range(num_groups):
+ g = await self._get_one_ready_group()
+ if not g:
+ break
+ out.extend(g)
+ return out
+ actual_batch_size = batch_size if batch_size is not None else self.group_size
+ return await super().get_batch(batch_size=actual_batch_size)
+
+ def commit(self, step: int, groups: Sequence[Any] | None = None) -> None:
+ """Commits in-flight groups after a successful global step boundary."""
+ # TODO: implement the commit and keep track of uncommited items. might be worth putting in parent class.
+ pass
+
+ async def abort(self, exc: BaseException) -> None:
+ """Aborts queue and unblocks all waiting consumers with the given exception."""
+ if isinstance(exc, Exception):
+ await self.put_exception(exc)
+ await self.prepare_clear()
diff --git a/tunix/rl/agentic/queue_manager/group_queue_manager.py b/tunix/rl/agentic/queue_manager/group_queue_manager.py
index a8668f0a9..2c6908da8 100644
--- a/tunix/rl/agentic/queue_manager/group_queue_manager.py
+++ b/tunix/rl/agentic/queue_manager/group_queue_manager.py
@@ -19,7 +19,6 @@
import asyncio
import collections
from collections.abc import Hashable
-import dataclasses
from typing import Callable, Deque, Dict, Generic, List, Optional, Tuple, TypeVar, Union
_T = TypeVar("_T")