Skip to content

Commit f0ebac5

Browse files
jon-barkerjon-barker
committed
Use agent registry to prevent arbitrary code execution from env configs
Signed-off-by: Jon Barker <jbarker@nvidia.com>
1 parent 9a19203 commit f0ebac5

19 files changed

Lines changed: 110 additions & 50 deletions

File tree

‎examples/rl/environment_configs/countdown.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
- agent_type: examples.rl.environments.countdown.countdown_agent.CountdownAgent
1+
- agent_type: CountdownAgent
22
agent_args:
33
hf_dataset_name: "Jiayi-Pan/Countdown-Tasks-3to4"
44
split: "train"

‎examples/rl/environment_configs/dapo.yaml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
- agent_type: examples.rl.environments.math.dapo_agent.DAPOAgent
1+
- agent_type: DAPOAgent
22
agent_args:
33
format_reward: 0.0
44
weight: 1.0
5-
- agent_type: examples.rl.environments.math.aime_agent.AIMEAgent
5+
- agent_type: AIMEAgent
66
agent_args:
77
format_reward: 0.0
88
weight: 0.0

‎examples/rl/environment_configs/default.yaml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
- agent_type: examples.rl.environments.countdown.countdown_agent.CountdownAgent
1+
- agent_type: CountdownAgent
22
agent_args:
33
hf_dataset_name: "Jiayi-Pan/Countdown-Tasks-3to4"
44
split: "train"
55
weight: 1.0
6-
- agent_type: examples.rl.environments.math.openmath_agent.OpenMathInstructAgent
6+
- agent_type: OpenMathInstructAgent
77
agent_args: {}
88
weight: 1.0

‎examples/rl/environment_configs/gsm8k.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
- agent_type: examples.rl.environments.math.gsm8k_agent.GSM8KAgent
1+
- agent_type: GSM8KAgent
22
agent_args:
33
answer_format: "boxed"
44
format_reward: 0.5

‎examples/rl/environment_configs/gsm8k_nanov3.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
- agent_type: examples.rl.environments.math.gsm8k_agent.GSM8KAgent
1+
- agent_type: GSM8KAgent
22
agent_args:
33
answer_format: "boxed"
44
format_reward: 0.5

‎examples/rl/environment_configs/math.yaml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
- agent_type: examples.rl.environments.math.openmath_agent.OpenMathInstructAgent
1+
- agent_type: OpenMathInstructAgent
22
agent_args: {}
33
weight: 1.0
4-
- agent_type: examples.rl.environments.math.bigmath_agent.BigMathAgent
4+
- agent_type: BigMathAgent
55
agent_args: {}
66
weight: 1.0
7-
- agent_type: examples.rl.environments.math.aime_agent.AIMEAgent
7+
- agent_type: AIMEAgent
88
agent_args: {}
99
weight: 0.0
1010
evaluation_only: true

‎examples/rl/environment_configs/openmathinstructv2.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
1-
- agent_type: examples.rl.environments.math.openmath_agent.OpenMathInstructAgent
1+
- agent_type: OpenMathInstructAgent
22
agent_args: {}
33
weight: 1.0

‎megatron/rl/__init__.py‎

Lines changed: 1 addition & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -2,39 +2,14 @@
22

33
import asyncio
44
import functools
5-
import importlib
6-
import os
7-
import sys
85
import time
96
import traceback
10-
from typing import Callable, Coroutine, Type
7+
from typing import Callable, Coroutine
118

129
from pydantic import BaseModel, ConfigDict, Field
1310
from typing_extensions import Self, Type
1411

1512

16-
def import_class(class_path: str) -> Type:
17-
"""Import a class from a string path.
18-
19-
Args:
20-
class_path: String path to the class (e.g. 'examples.rl.environments.countdown.countdown_agent.CountdownAgent' or '../environments.countdown.py:CountdownAgent')
21-
22-
Returns:
23-
The class object
24-
"""
25-
if '.py:' in class_path:
26-
# filepath.py:Classname branch.
27-
module_path, class_name = class_path.split(':')
28-
abs_path = os.path.abspath(module_path)
29-
spec = importlib.util.spec_from_file_location('acemath_agent', abs_path)
30-
module = importlib.util.module_from_spec(spec)
31-
spec.loader.exec_module(module)
32-
else:
33-
module_path, class_name = class_path.rsplit('.', 1)
34-
module = importlib.import_module(module_path, package=__package__)
35-
return getattr(module, class_name)
36-
37-
3813
class TypeLookupable(BaseModel, extra='allow'):
3914
"""Supports 'unwrapping' of base class into subclasses."""
4015

‎megatron/rl/agent/registry.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
3+
from typing import Type
4+
5+
from examples.rl.environments.countdown.countdown_agent import CountdownAgent
6+
from examples.rl.environments.math.aime_agent import AIMEAgent
7+
from examples.rl.environments.math.bigmath_agent import BigMathAgent
8+
from examples.rl.environments.math.dapo_agent import DAPOAgent
9+
from examples.rl.environments.math.gsm8k_agent import GSM8KAgent
10+
from examples.rl.environments.math.openmath_agent import OpenMathInstructAgent
11+
from megatron.rl.agent.remote_agent import RemoteAgent
12+
13+
AGENT_REGISTRY: dict[str, Type] = {
14+
"RemoteAgent": RemoteAgent,
15+
"CountdownAgent": CountdownAgent,
16+
"OpenMathInstructAgent": OpenMathInstructAgent,
17+
"BigMathAgent": BigMathAgent,
18+
"DAPOAgent": DAPOAgent,
19+
"GSM8KAgent": GSM8KAgent,
20+
"AIMEAgent": AIMEAgent,
21+
}
22+
23+
try:
24+
from nemogym2mrl.nemo_gym_agent import NemoGymAgent
25+
26+
AGENT_REGISTRY["NemoGymAgent"] = NemoGymAgent
27+
except ImportError:
28+
pass
29+
30+
try:
31+
from environments.acemath_agent import AceMathAgent
32+
33+
AGENT_REGISTRY["AceMathAgent"] = AceMathAgent
34+
except ImportError:
35+
pass
36+
37+
38+
def get_agent_class(agent_name: str) -> Type:
39+
"""Resolve a config agent_type string to a registered agent class.
40+
41+
Only explicitly registered agent names are allowed. This prevents arbitrary
42+
code execution from untrusted environment configuration files.
43+
"""
44+
try:
45+
return AGENT_REGISTRY[agent_name]
46+
except KeyError as exc:
47+
known = ", ".join(sorted(AGENT_REGISTRY))
48+
raise ValueError(
49+
f"Unknown agent_type {agent_name!r}. "
50+
f"Registered agent types: {known or '(none)'}"
51+
) from exc

‎megatron/rl/agent/weighted_multi_task.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
import numpy as np
77

8-
from .. import import_class
8+
from .registry import get_agent_class
99
from .api import (
1010
AgentBaseModel,
1111
ContrastiveRollout,
@@ -73,7 +73,7 @@ def from_config(
7373
7474
Args:
7575
config: List of dicts with keys:
76-
- agent_type: String path to agent class
76+
- agent_type: Registered agent name (see megatron.rl.agent.registry)
7777
- agent_args: Dict of arguments to pass to agent constructor
7878
- weight: Float weight for this agent
7979
@@ -87,8 +87,7 @@ def from_config(
8787
agent_args = entry.get('agent_args', {})
8888
agent_args['parallel_generation_tasks'] = parallel_generation_tasks
8989

90-
# Import and instantiate the agent class
91-
agent_type = import_class(entry['agent_type'])
90+
agent_type = get_agent_class(entry['agent_type'])
9291
agent_configs.append(
9392
AgentConfig(
9493
agent_type=agent_type,

0 commit comments

Comments
 (0)