From 9cb9c787ce610eda75496d584f0d78fb7e8d45cf Mon Sep 17 00:00:00 2001 From: Jonas Yang Date: Thu, 23 Jul 2026 13:54:52 +0800 Subject: [PATCH 1/2] test: pass required NeMo Gym rollout arguments Signed-off-by: Jonas Yang --- tests/unit/experience/test_rollout_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index b43cbb7157..18bf8cd03c 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -745,6 +745,8 @@ def test_async_nemo_gym_rollout_manager_matches_original( tokenizer=nemo_gym_tokenizer, task_to_env={"nemo_gym": nemo_gym}, generation_config=nemo_gym_vllm_generation.cfg, + num_generations=num_generations, + log_full_result_tables=False, max_seq_len=nemo_gym_vllm_generation.cfg["vllm_cfg"]["max_model_len"], max_rollout_turns=None, ) From f5cc5e5ca9223cb6b2640a8ec3778cffd2cb277c Mon Sep 17 00:00:00 2001 From: ruit Date: Wed, 22 Jul 2026 23:55:35 -0700 Subject: [PATCH 2/2] test: consume async NeMo Gym rollout results Signed-off-by: ruit --- tests/unit/experience/test_rollout_manager.py | 30 ++++++++++++------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 18bf8cd03c..c4dfc5fe78 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -739,17 +739,25 @@ def test_async_nemo_gym_rollout_manager_matches_original( } ) - original_result = run_async_nemo_gym_rollout( - policy_generation=nemo_gym_vllm_generation, - input_batch=repeated_batch, - tokenizer=nemo_gym_tokenizer, - task_to_env={"nemo_gym": nemo_gym}, - generation_config=nemo_gym_vllm_generation.cfg, - num_generations=num_generations, - log_full_result_tables=False, - max_seq_len=nemo_gym_vllm_generation.cfg["vllm_cfg"]["max_model_len"], - max_rollout_turns=None, - ) + async def _collect_original_results(): + return [ + result + async for result in run_async_nemo_gym_rollout( + policy_generation=nemo_gym_vllm_generation, + input_batch=repeated_batch, + tokenizer=nemo_gym_tokenizer, + task_to_env={"nemo_gym": nemo_gym}, + generation_config=nemo_gym_vllm_generation.cfg, + num_generations=num_generations, + log_full_result_tables=False, + max_seq_len=nemo_gym_vllm_generation.cfg["vllm_cfg"]["max_model_len"], + max_rollout_turns=None, + ) + ] + + original_results = asyncio.run(_collect_original_results()) + assert len(original_results) == 1 + original_result = original_results[0] manager = RolloutManager( use_nemo_gym=True,