| |
|
|
| import asyncio |
| from typing import Any, Callable, List, Optional |
| import torch |
|
|
| from ..rollout.version_manager import VersionManager |
|
|
|
|
| class Orchestrator: |
| """Coordinates rollout, buffer, trainer, and weight synchronization. |
| |
| - Maintains a prompt queue. |
| - Tracks current policy version via VersionManager. |
| - Periodically triggers weight synchronization and training updates. |
| """ |
|
|
| def __init__( |
| self, |
| version_manager: Optional[VersionManager] = None, |
| weight_sync_fn: Optional[Callable] = None, |
| sync_interval: int = 10, |
| rollout_worker: Optional[Any] = None, |
| trainer_worker: Optional[Any] = None, |
| buffer: Optional[Any] = None, |
| ): |
| self.prompt_queue = asyncio.Queue() |
| self.version_manager = version_manager or VersionManager() |
| self.weight_sync_fn = weight_sync_fn |
| self.sync_interval = sync_interval |
| self.rollout_worker = rollout_worker |
| self.trainer_worker = trainer_worker |
| self.buffer = buffer |
| self._stop = asyncio.Event() |
|
|
| async def add_prompt(self, prompt: Any): |
| await self.prompt_queue.put(prompt) |
|
|
| async def step_training(self, batch_size: int = 16) -> Optional[float]: |
| """Execute one training step if a trainer worker is configured.""" |
| if self.trainer_worker is None: |
| return None |
| if hasattr(self.trainer_worker, "step"): |
| if hasattr(self.trainer_worker.step, "remote"): |
| return await self.trainer_worker.step.remote(batch_size) |
| return self.trainer_worker.step(batch_size) |
| return None |
|
|
| async def run(self, max_steps: Optional[int] = None): |
| """Main orchestration loop.""" |
| step = 0 |
| losses: List[float] = [] |
|
|
| while not self._stop.is_set(): |
| if max_steps is not None and step >= max_steps: |
| break |
|
|
| |
| if self.rollout_worker is not None: |
| prompts = [{"input_ids": torch.tensor([1, 2, 3]), "gt_ids": torch.tensor([2, 3]), "eos_token_id": 99}] |
| if hasattr(self.rollout_worker, "step"): |
| await self.rollout_worker.step(prompts) |
|
|
| |
| loss = await self.step_training() |
| if loss is not None: |
| losses.append(loss) |
|
|
| |
| if step % self.sync_interval == 0: |
| new_version = self.version_manager.bump() |
| if self.weight_sync_fn is not None: |
| res = self.weight_sync_fn(new_version) |
| if asyncio.iscoroutine(res): |
| await res |
|
|
| step += 1 |
| await asyncio.sleep(0.001) |
|
|
| return losses |
|
|
| def stop(self): |
| self._stop.set() |
|
|
