Merge

tdene · tdene · commit e29f89c29f80 · 2025-11-21T15:49:57.000-08:00
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
@@ -193,9 +193,10 @@ async def base_generate(self, request: InferenceRequest):
             self._client.add_request(prompt=prompt, sampling_params=sampling_params)
             for prompt in request.prompt
         ]
-        responses = await asyncio.gather(
+        records = await asyncio.gather(
             *requests
         )
+        responses = [record[-1] for record in records]
         return [
             InferenceResponse(
                 response=r.generated_text,

Original file line number	Diff line number	Diff line change
`@@ -193,9 +193,10 @@ async def base_generate(self, request: InferenceRequest):`
`193`	`193`	`self._client.add_request(prompt=prompt, sampling_params=sampling_params)`
`194`	`194`	`for prompt in request.prompt`
`195`	`195`	`]`
`196`		`- responses = await asyncio.gather(`
	`196`	`+ records = await asyncio.gather(`
`197`	`197`	`*requests`
`198`	`198`	`)`
	`199`	`+ responses = [record[-1] for record in records]`
`199`	`200`	`return [`
`200`	`201`	`InferenceResponse(`
`201`	`202`	`response=r.generated_text,`