Get current version

Author: joecummings

apps/grpo/main.py

@@ -49,8 +49,8 @@ async def main():
 
     async def continuous_rollouts():
         while True:
-            current_version = await policy.get_current_version()
-            episode = await generate_rollout()
+            version = await policy.get_current_version.choose()
+            episode = await generate_rollout(version)
             await replay_buffer.add.call(episode)
 
     rollout_task = asyncio.create_task(continuous_rollouts())

src/forge/actors/policy.py

@@ -43,6 +43,7 @@ class PolicyRouter(Actor):
     sampling_params: SamplingParams = None
     lora_request: LoRARequest = None
     tokenization_kwargs: dict = None
+    version: int = 0
 
     @endpoint
     async def setup(self):
@@ -78,6 +79,10 @@ async def setup(self):
             log_stats=None,
         )
 
+    @endpoint
+    async def get_current_version(self) -> int:
+        return self.version
+
     @endpoint
     async def generate(self, prompt: str, priority: int = 0):
         self.request_id += 1 % sys.maxsize