File size: 2,194 Bytes
e791b16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # Ray actor / standalone worker for rollout workers
"""RolloutWorker pulls prompts from a shared PromptQueue, generates responses
asynchronously using AsyncEngine, and writes experiences to a ReplayBuffer.
"""
from typing import Any, List, Optional
try:
import ray
ray_remote = ray.remote(num_gpus=1)
except ImportError:
ray = None
def ray_remote(cls):
return cls
from .async_engine import AsyncEngine
from .version_manager import VersionManager
@ray_remote
class RolloutWorker:
def __init__(
self,
model_path: Optional[str] = None,
prompt_queue: Optional[Any] = None,
replay_buffer: Optional[Any] = None,
buffer: Optional[Any] = None,
version_manager: Optional[VersionManager] = None,
engine: Optional[Any] = None,
**engine_kwargs,
):
self.prompt_queue = prompt_queue
self.replay_buffer = replay_buffer if replay_buffer is not None else buffer
self.version_manager = version_manager or VersionManager()
self.engine = AsyncEngine(
model_path=model_path,
buffer=self.replay_buffer,
version_manager=self.version_manager,
engine=engine,
**engine_kwargs,
)
async def step(self, prompts: List[dict]):
"""Run rollout on an explicit list of prompts."""
await self.engine.rollout(prompts)
async def run(self, batch_size: int = 8):
"""Continuously fetch prompts and perform rollout."""
while True:
prompts: List[dict] = []
for _ in range(batch_size):
if self.prompt_queue is not None:
if hasattr(self.prompt_queue, "get_prompt"):
if hasattr(self.prompt_queue.get_prompt, "remote"):
prompt = await self.prompt_queue.get_prompt.remote()
else:
prompt = await self.prompt_queue.get_prompt()
else:
prompt = await self.prompt_queue.get()
prompts.append(prompt)
if prompts:
await self.engine.rollout(prompts)
|