Code Execution using Morph Cloud #614

abetlen · 2025-04-17T21:21:29Z

Code Execution using Morph Cloud

Overview

This PR introduces Morph Cloud as an optional backend for code execution within Open-R1, providing users with an alternative provider for Codeforces as well as IOI code reward calculation.

Key Changes & Design

The integration supports two distinct execution pathways, chosen via configuration:

1. General Code Execution (`code_reward`)

Mechanism: Integrated via src/open_r1/utils/code_providers.py::MorphProvider
This provider supports two operational modes:

Router Mode: If morph_router_url is provided in the configuration, the provider uses a client (src/open_r1/utils/routed_morph.py::RoutedMorphSandbox) to communicate with a separate router service (scripts/morph_router.py). This pattern aims to improve performance and manage costs via sandbox reuse, similar to the E2B router.
Direct Mode: If morph_router_url is not provided, MorphProvider falls back to interacting directly with the morphcloud SDK (MorphCloudClient, Sandbox). It manages concurrent sandbox creation and execution using asyncio up to the limit specified by num_parallel.
Configuration: Enabled by setting provider_type: “morph”. The morph_router_url field is optional; omitting it selects Direct Mode (src/open_r1/configs.py). A Slurm script (slurm/morph_router.slurm) is provided for launching the router service if Router Mode is chosen.

2. IOI Problem Execution (`ioi_code_reward`)

Mechanism: Integrated via a custom client implementation (src/open_r1/utils/ioi/morph_client.py::MorphCloudExecutionClient). This client uses the lower-level APIs from the morphcloud library (morphcloud.api.MorphCloudClient) to manage Morph snapshots and instances directly. It orchestrates the complex setup required for IOI problems (multi-file uploads, custom compile/run scripts) for fine-grained control. This pathway always interacts directly with Morph Cloud and does not use the router.
Configuration: Enabled by setting ioi_provider: “morph” in the training configuration (src/open_r1/configs.py).

Supporting Changes

Updated core reward logic in src/open_r1/rewards.py to dispatch to the appropriate Morph provider/mode based on configuration settings.
Added necessary configuration fields to dataclasses in src/open_r1/configs.py.
Included comprehensive setup instructions and usage guidance in README.md and a dedicated technical overview in MORPH_INTEGRATION.md.
Updated relevant example configuration files in recipes/.

Dependencies & Usage

Requires the morphcloud package (install via uv pip install morphcloud). [Note: Verify addition to project dependencies, e.g., in setup.py]
Users must obtain a Morph API key and set it as the MORPH_API_KEY environment variable. This key is used by the Morph router service (if used) and the direct clients (MorphProvider in Direct Mode, MorphCloudExecutionClient for IOI).
Activation is opt-in via configuration flags.

Backward Compatibility

This change is fully backward compatible.
E2B remains the default provider for code_reward (provider_type=“e2b”).
Piston remains the default provider for ioi_code_reward (ioi_provider=“piston”).
Existing workflows using the default providers are unaffected.

lewtun

Thanks a lot for the nice feature @abetlen ! Overall the changes LGTM with some small suggestions to align the function args

lewtun · 2025-04-23T07:29:33Z

README.md

+To use Morph, first install the morphcloud package:
+
+```shell
+pip install morphcloud


WDYT about adding this to the code extra in setup.py so users can just run uv pip install -e '.[code]'?

lewtun · 2025-04-23T07:40:51Z

src/open_r1/rewards.py

-    if not all(v["language"] == language for v in verification_info):
-        raise ValueError("All verification_info must have the same language", verification_info)
+
+    if enforce_same_language:


Nice check!

lewtun · 2025-04-23T07:41:41Z

src/open_r1/rewards.py

@@ -398,19 +407,32 @@ def binary_code_reward(completions, num_parallel: int = 2, e2b_router_url=None,
    return output


-def code_reward(completions, num_parallel: int = 2, e2b_router_url=None, **kwargs) -> list[float]:
-    """Reward function that evaluates code snippets using the E2B code interpreter.
+# This function has been removed as we now use a unified approach in code_reward


I think we can remove these comments and code_reward_morph since it doesn't seem to be used anywhere?

lewtun · 2025-04-23T07:49:01Z

src/open_r1/rewards.py

@@ -384,8 +386,15 @@ def extract_code(completion: str, language: str = "python") -> str:
    return extracted_answer


-def binary_code_reward(completions, num_parallel: int = 2, e2b_router_url=None, **kwargs) -> list[float]:
-    rewards = code_reward(completions, num_parallel=num_parallel, e2b_router_url=e2b_router_url, **kwargs)
+def binary_code_reward(completions, num_parallel: int = 2, e2b_router_url=None, provider_type: str = None, enforce_same_language: bool = False, **kwargs) -> list[float]:


I realise this is a breaking change, but maybe it's better to do this than have a router URL arg for E2B and then have users remember they need kwargs for Morph:

Suggested change

def binary_code_reward(completions, num_parallel: int = 2, e2b_router_url=None, provider_type: str = None, enforce_same_language: bool = False, **kwargs) -> list[float]:

def binary_code_reward(completions, num_parallel: int = 2, provider_type: str = None, enforce_same_language: bool = False, **kwargs) -> list[float]:

The alternative would be to expose

lewtun · 2025-04-23T07:49:26Z

src/open_r1/rewards.py

+    rewards = code_reward(
+        completions, 
+        num_parallel=num_parallel, 
+        e2b_router_url=e2b_router_url, 


Ditto here if you agree with the above change:

Suggested change

e2b_router_url=e2b_router_url,

lewtun · 2025-04-23T07:49:45Z

src/open_r1/rewards.py

+    return code_reward(completions, num_parallel=num_parallel, provider_type="morph", **kwargs)
+
+
+def code_reward(completions, num_parallel: int = 2, e2b_router_url=None, provider_type: str = "e2b", enforce_same_language: bool = False, **kwargs) -> list[float]:


Suggested change

def code_reward(completions, num_parallel: int = 2, e2b_router_url=None, provider_type: str = "e2b", enforce_same_language: bool = False, **kwargs) -> list[float]:

def code_reward(completions, num_parallel: int = 2, provider_type: str = "e2b", enforce_same_language: bool = False, **kwargs) -> list[float]:

lewtun · 2025-04-23T07:50:32Z

src/open_r1/rewards.py

+    execution_provider = get_provider(
+        provider_type=provider_type,
+        num_parallel=num_parallel,
+        e2b_router_url=e2b_router_url,


We should propagate kwargs here to use Morph right?

Suggested change

e2b_router_url=e2b_router_url,

**kwargs,

edbeeching · 2025-04-29T11:21:12Z

scripts/morph_router.py

+        request_timeout (int): The maximum allowed time for the entire batch request in seconds.
+    """
+    scripts: List[str]
+    language: str


One assumption I made when I wrote the E2B router was that all scripts were in the same language, should we relax this assumption here?

edbeeching · 2025-04-29T11:25:50Z

slurm/morph_router.slurm

+#SBATCH --error=logs/morph_router_%j.err
+
+# Load necessary modules
+module load anaconda


I understand that you were testing this with your own conda env, but for the PR can you copy paste the e2b_router.slurm file and include your python scripts/morph_router.py --port 8001 --max_num_sandboxes 20. This will only work on our cluster, but at least the slurm files will be consistent across the project.

edbeeching · 2025-04-29T11:31:02Z

src/open_r1/utils/routed_morph.py

+        self.request_timeout = request_timeout
+
+    def run_code(self, scripts: List[str], 
+                language: str = "python", 


As mentioned above, perhaps we can relax the language here and have a list of languages?

edbeeching

LGTM in general, we have some slow_tests for the code rewards / routers that are not part of the CI in https://github.com/huggingface/open-r1/blob/main/tests/slow/test_code_reward.py

If you could update / add an example there it would be appreciated so we can validate it all works on the cluster.

orph tests

arb8020 · 2025-04-30T07:44:17Z

src/open_r1/utils/routed_sandbox.py

@@ -41,7 +41,7 @@ def __init__(self, router_url: str):
    def run_code(
        self,
        scripts: list[str],
-        language: str = "python",
+        languages: Optional[List[str]] = None,


i made the assumption that i should change the e2b router to support languages as well to keep the API consistent between the two but i can also leave the old language parameter for backwards compatibility if that's preferred

arb8020 · 2025-04-30T07:44:55Z

tests/slow/test_code_reward.py

    def test_python_code_reward(self):
        # requires E2B, see the README.md file
-        code_dataset = load_dataset("open-r1/verifiable-coding-problems-python_decontaminated-tested")
+        code_dataset = load_dataset("open-r1/verifiable-coding-problems-python_decontaminated-tested-shuffled")


couldn't find this dataset without the -shuffled prefix so i added this

arb8020 · 2025-04-30T07:45:59Z

tests/slow/test_code_reward.py

            assert result.error is None
-            assert "hello" in result.stdout or "4" in result.stdout
+            assert "hello" in result.logs['stdout'][0] or "4" in result.logs['stdout'][0]


also changed a bit of the handling here: maybe something has changed in the Execution object but i wasn't able to access the .exit_code attribute, nor .stdout without the changes i made here

arb8020 added 12 commits April 13, 2025 11:40

initial commit for morphcloud sandbox support

9f4029f

initial

c94f334

fixed prints in morph client for ioi

b008ec3

updated import

ac1d32c

context manager

e0b394d

removed unnecessary comments

638afb0

more intelligent instance/snapshot management

746172e

update

a1565d0

Merge remote-tracking branch 'upstream/main' into morph-sandbox

843d834

Add documentation for Morph integration

cdc9ecf

Merge branch 'huggingface:main' into morph-sandbox-org

0a534ff

Delete MORPH_INTEGRATION.md

07d5fa9

lewtun approved these changes Apr 23, 2025

View reviewed changes

arb8020 added 3 commits April 23, 2025 19:12

added retry and modularity to morph client

a117aad

updates to kwargs and setup.py

e2ccd7a

Update setup.py

3839681

lewtun requested a review from edbeeching April 29, 2025 10:08

edbeeching reviewed Apr 29, 2025

View reviewed changes

edbeeching requested changes Apr 29, 2025

View reviewed changes

added languages codepath + fixed slurm + added m

087d9e6

orph tests

arb8020 reviewed Apr 30, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Code Execution using Morph Cloud #614

Code Execution using Morph Cloud #614

abetlen commented Apr 17, 2025

lewtun left a comment

lewtun Apr 23, 2025

lewtun Apr 23, 2025

lewtun Apr 23, 2025

lewtun Apr 23, 2025

lewtun Apr 23, 2025

lewtun Apr 23, 2025

lewtun Apr 23, 2025

edbeeching Apr 29, 2025

edbeeching Apr 29, 2025 •

edited

Loading

edbeeching Apr 29, 2025

edbeeching left a comment

arb8020 Apr 30, 2025

arb8020 Apr 30, 2025

arb8020 Apr 30, 2025

	def binary_code_reward(completions, num_parallel: int = 2, e2b_router_url=None, provider_type: str = None, enforce_same_language: bool = False, **kwargs) -> list[float]:
	def binary_code_reward(completions, num_parallel: int = 2, provider_type: str = None, enforce_same_language: bool = False, **kwargs) -> list[float]:

		return code_reward(completions, num_parallel=num_parallel, provider_type="morph", **kwargs)


		def code_reward(completions, num_parallel: int = 2, e2b_router_url=None, provider_type: str = "e2b", enforce_same_language: bool = False, **kwargs) -> list[float]:

Code Execution using Morph Cloud #614

Are you sure you want to change the base?

Code Execution using Morph Cloud #614

Conversation

abetlen commented Apr 17, 2025

Code Execution using Morph Cloud

Overview

Key Changes & Design

1. General Code Execution (code_reward)

2. IOI Problem Execution (ioi_code_reward)

Supporting Changes

Dependencies & Usage

Backward Compatibility

lewtun left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

edbeeching Apr 29, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

edbeeching left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

1. General Code Execution (`code_reward`)

2. IOI Problem Execution (`ioi_code_reward`)

edbeeching Apr 29, 2025 •

edited

Loading