Async code reward fixes (#546)

* expose num parallel code executions

* add e2b benchmarking script

* adds new parallel code execution with better execption handling

* style

* update default

* increase sandbox timeout

* Add pretty table and Sandbox IDs

* Add Sandbox ID

* fix merge

---------

Co-authored-by: Lewis Tunstall <[email protected]>

Author: edbeeching

scripts/benchmark_e2b.py

@@ -0,0 +1,84 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Benchmark script for the code_reward function with E2B.
+
+This script measures the performance of the code_reward function with varying numbers
+of samples and parallelization levels.
+
+Each sample is a CodeForces problem with a gold standard solution that is executed against a set of public test cases.
+"""
+
+from datasets import load_dataset
+from open_r1.rewards import code_reward
+import time
+from tqdm.auto import tqdm
+
+from dotenv import load_dotenv
+load_dotenv()
+
+def benchmark_code_reward(example):
+    start_time = time.time()
+    test_completions = [[{"content": example["gold_standard_solution"]}]]
+    reward_kwargs = {"verification_info": [example["verification_info"]]}
+    rewards = code_reward(test_completions, **reward_kwargs)
+    end_time = time.time()
+    example["test_reward"] = rewards[0]
+    example["reward_time"] = end_time - start_time
+    return example
+
+if __name__ == "__main__":
+    parallel_dict = {
+        16:[1,4,16],
+        64:[4,16, 64],
+        256:[16, 64, 96], # cap at 96 as PRO account is limited to 100
+    }
+    # Store results for table formatting
+    results = []
+    
+    for num_samples in tqdm([16, 64,256], desc="Benchmarking samples"):
+        for num_parallel in parallel_dict[num_samples]:
+            code_dataset = load_dataset("open-r1/verifiable-coding-problems-python_decontaminated")
+            code_dataset = code_dataset["train"].shuffle(seed=42).select(range(num_samples))
+
+            test_completions = [[{"content": example["gold_standard_solution"]}] for example in code_dataset]
+            reward_kwargs = {"verification_info": [example["verification_info"] for example in code_dataset]}
+
+            start_time = time.time()
+            rewards = code_reward(test_completions, num_parallel=num_parallel, **reward_kwargs)
+            execution_time = time.time() - start_time
+            
+            # Calculate some statistics about rewards
+            mean_reward = sum(rewards) / len(rewards)
+            min_reward = min(rewards)
+            max_reward = max(rewards)
+            
+            # Store results
+            results.append({
+                "num_samples": num_samples,
+                "num_parallel": num_parallel,
+                "execution_time": execution_time,
+                "mean_reward": mean_reward,
+                "min_reward": min_reward,
+                "max_reward": max_reward
+            })
+    
+    print("\n## Benchmark Results\n")
+    print("| Sample Size | Parallelization | Execution Time (s) | Mean Reward | Min Reward | Max Reward |")
+    print("|:-----------:|:---------------:|------------------:|:-----------:|:-----------:|:-----------:|")
+    
+    for result in results:
+        print(f"| {result['num_samples']:^11} | {result['num_parallel']:^15} | {result['execution_time']:17.2f} | {result['mean_reward']:^11.4f} | {result['min_reward']:^11.4f} | {result['max_reward']:^11.4f} |")
+    

setup.py

@@ -91,7 +91,7 @@ def deps_list(*pkgs):
 extras["quality"] = deps_list("ruff", "isort", "flake8")
 extras["code"] = deps_list("e2b-code-interpreter", "python-dotenv")
 extras["eval"] = deps_list("lighteval", "math-verify")
-extras["dev"] = extras["quality"] + extras["tests"] + extras["eval"]
+extras["dev"] = extras["quality"] + extras["tests"] + extras["eval"] + extras["code"]
 
 # core dependencies shared across the whole project - keep this to a bare minimum :)
 install_requires = [

src/open_r1/configs.py

@@ -154,6 +154,13 @@ class GRPOScriptArguments(trl.ScriptArguments):
             "help": "for each generation, evaluate these many test cases in parallel, then check if any of them failed (0 score): if so stop evaluating; otherwise continue with the next batch of test cases. Useful to avoid overloading the eval server + save time on wrong solutions"
         },
     )
+    parallel_code_exec_per_proc: int = field(
+        default=2,
+        metadata={
+            "help": "Number of parallel E2B code executions per process. Default of 2 is suitable for the Free Hobby tier of E2B with 8 GPUs used for training."
+        },
+    )
+
     dataset_prompt_column: str = field(
         default="prompt",
         metadata={"help": "Column to use as prompts for training."},

src/open_r1/rewards.py

@@ -1,3 +1,18 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Reward functions for GRPO training."""
 
 import asyncio
@@ -368,13 +383,13 @@ def extract_code(completion: str, language: str = "python") -> str:
     return extracted_answer
 
 
-def binary_code_reward(completions, **kwargs) -> list[float]:
-    rewards = code_reward(completions, **kwargs)
+def binary_code_reward(completions, num_parallel: int = 2, **kwargs) -> list[float]:
+    rewards = code_reward(completions, num_parallel=num_parallel, **kwargs)
     BINARY_THRESHOLD = 0.99
     return [1.0 if reward > BINARY_THRESHOLD else 0.0 for reward in rewards]
 
 
-def code_reward(completions, **kwargs) -> list[float]:
+def code_reward(completions, num_parallel: int = 2, **kwargs) -> list[float]:
     """Reward function that evaluates code snippets using the E2B code interpreter.
 
     Assumes the dataset contains a `verification_info` column with test cases.
@@ -438,7 +453,7 @@ def evaluate_code(code, test_cases):
     if not all(v["language"] == language for v in verification_info):
         raise ValueError("All verification_info must have the same language", verification_info)
     try:
-        rewards = run_async_from_sync(scripts, language)
+        rewards = run_async_from_sync(scripts, language, num_parallel)
 
     except Exception as e:
         print(f"Error from E2B executor: {e}")
@@ -463,45 +478,62 @@ def code_format_reward(completions, **kwargs):
     return code_format_reward
 
 
-def run_async_from_sync(scripts: list[str], language: str) -> list[float]:
+def run_async_from_sync(scripts: list[str], language: str, num_parallel: int) -> list[float]:
     """Function wrapping the `run_async` function."""
     # Create a new event loop and set it
     try:
         # Run the async function and get the result
-        rewards = asyncio.run(run_async(scripts, language))
+        rewards = asyncio.run(run_async(scripts, language, num_parallel))
     except Exception as e:
         print(f"Error from E2B executor async: {e}")
         raise e
 
     return rewards
 
 
-async def run_async(scripts: list[str], language: str) -> list[float]:
-    # Create the sandbox by hand, currently there's no context manager for this version
-    sbx = await AsyncSandbox.create(timeout=30, request_timeout=3)
+async def run_async(scripts: list[str], language: str, num_parallel: int) -> list[float]:
+    # Limit the number of concurrent tasks
+    semaphore = asyncio.Semaphore(num_parallel)
 
     # Create a list of tasks for running scripts concurrently
-    tasks = [run_script(sbx, script, language) for script in scripts]
+    tasks = [run_script(script, language, semaphore) for script in scripts]
 
     # Wait for all tasks to complete and gather their results as they finish
     results = await asyncio.gather(*tasks)
     rewards = list(results)  # collect results
 
-    # Kill the sandbox after all the tasks are complete
-    await sbx.kill()
-
     return rewards
 
 
-async def run_script(sbx: AsyncSandbox, script: str, language: str) -> float:
-    execution = await sbx.run_code(script, language=language)
-    try:
-        return float(execution.text)
-    except (TypeError, ValueError):
-        return 0.0
-    except Exception as e:
-        print(f"Error from E2B executor run_script: {e}")
-        return 0.0
+async def run_script(script: str, language: str, semaphore: asyncio.Semaphore) -> float:
+    # We set a timeout margin, as the AsyncSandbox timeout does not seem to work
+    # These values are based on running 256 examples with the gold solution
+    # from open-r1/verifiable-coding-problems-python_decontaminated
+    # see scripts/benchmark_e2b.py
+
+    SANDBOX_TIMEOUT = 30
+    MARGIN = 2
+    REQUEST_TIMEOUT = SANDBOX_TIMEOUT - MARGIN
+    ASYNCIO_TIMEOUT = SANDBOX_TIMEOUT + MARGIN
+
+    async with semaphore:
+        try:
+            sandbox = await AsyncSandbox.create(timeout=SANDBOX_TIMEOUT, request_timeout=REQUEST_TIMEOUT)
+            execution = await asyncio.wait_for(sandbox.run_code(script, language=language), timeout=ASYNCIO_TIMEOUT)
+            return float(execution.text)
+        except (TypeError, ValueError):
+            return 0.0
+        except asyncio.TimeoutError:
+            print("Operation timed out")
+            return 0.0
+        except Exception as e:
+            print(f"Error in `run_script` from E2B sandbox ID {sandbox.sandbox_id} : {e}")
+            return 0.0
+        finally:
+            try:
+                await sandbox.kill()
+            except Exception as e:
+                print(f"Error from E2B executor kill with sandbox ID {sandbox.sandbox_id} : {e}")
 
 
 def get_reward_funcs(script_args) -> list[Callable]:
@@ -521,8 +553,12 @@ def get_reward_funcs(script_args) -> list[Callable]:
             max_penalty=script_args.repetition_max_penalty,
         ),
         "length": len_reward,
-        "code": code_reward,
-        "binary_code": binary_code_reward,
+        "code": update_wrapper(
+            partial(code_reward, num_parallel=script_args.parallel_code_exec_per_proc), code_reward
+        ),
+        "binary_code": update_wrapper(
+            partial(binary_code_reward, num_parallel=script_args.parallel_code_exec_per_proc), binary_code_reward
+        ),
         "ioi_code": update_wrapper(
             partial(ioi_code_reward, test_batch_size=script_args.code_eval_test_batch_size), ioi_code_reward
         ),