From 5cef1fa13e310a51ebb7111b4e5e68b49e99a6dd Mon Sep 17 00:00:00 2001 From: Shadi Noghabi Date: Fri, 7 Aug 2026 23:00:28 -0700 Subject: [PATCH] add Orchestrator implementation PiperOrigin-RevId: 961288462 --- .../orchestrator/algorithm_adapter_test.py | 85 ++ .../orchestrator/async_rl_program_test.py | 177 ++++ .../orchestrator/batch_assembly_test.py | 87 ++ .../distributed_rl_engine_test.py | 200 +++++ .../orchestrator/orchestrator_test.py | 113 +++ .../orchestrator/rl_program_test.py | 105 +++ tests/rl/rl_cluster_test.py | 5 - tunix/experimental/common/datatypes.py | 67 +- .../experimental/docs/orchestrator_plan_v2.md | 816 ++++++++++++++++++ .../orchestrator/algorithm_adapter.py | 251 ++++++ .../orchestrator/async_rl_program.py | 285 ++++++ .../orchestrator/batch_assembly.py | 262 ++++++ .../orchestrator/distributed_rl_engine.py | 312 +++++++ .../experimental/orchestrator/orchestrator.py | 172 ++++ .../orchestrator/rl_engine_interface.py | 135 +-- tunix/experimental/orchestrator/rl_program.py | 159 ++++ .../orchestrator/simple_orchestrator_nb.py | 156 ++++ .../orchestrator/startup_validation.py | 8 +- .../queue_manager/trajectory_queue_manager.py | 99 ++- .../queue_manager/group_queue_manager.py | 1 - 20 files changed, 3352 insertions(+), 143 deletions(-) create mode 100644 tests/experimental/orchestrator/algorithm_adapter_test.py create mode 100644 tests/experimental/orchestrator/async_rl_program_test.py create mode 100644 tests/experimental/orchestrator/batch_assembly_test.py create mode 100644 tests/experimental/orchestrator/distributed_rl_engine_test.py create mode 100644 tests/experimental/orchestrator/orchestrator_test.py create mode 100644 tests/experimental/orchestrator/rl_program_test.py create mode 100644 tunix/experimental/docs/orchestrator_plan_v2.md create mode 100644 tunix/experimental/orchestrator/algorithm_adapter.py create mode 100644 tunix/experimental/orchestrator/async_rl_program.py create mode 100644 tunix/experimental/orchestrator/batch_assembly.py create mode 100644 tunix/experimental/orchestrator/distributed_rl_engine.py create mode 100644 tunix/experimental/orchestrator/orchestrator.py create mode 100644 tunix/experimental/orchestrator/rl_program.py create mode 100644 tunix/experimental/orchestrator/simple_orchestrator_nb.py 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")