Skip to content

feat(gym): steer the tail of a rollout step to a faster server group #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aoshen02 wants to merge 1 commit into
base: integration/group-routing
Choose a base branch
from
Open

feat(gym): steer the tail of a rollout step to a faster server group #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions nemo_rl/algorithms/grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -621,8 +621,12 @@ def init_train_dataloader(dataset, suffix: str = ""):
"policy.generation.server_groups does not support cluster.segment_size"
)

def _spinup_nemo_gym(base_urls, model_name):
"""Spin up the NeMo Gym actor against the given generation server URLs."""
def _spinup_nemo_gym(base_urls, model_name, generation=None):
"""Spin up the NeMo Gym actor against the given generation server URLs.

``generation`` only supplies the URL->group mapping; backends without
groups register their replicas unlabelled.
"""
t0 = time.perf_counter()
enable_router_replay = router_replay_enabled(policy_config)
routed_experts_dtype = (
Expand All @@ -633,6 +637,9 @@ def _spinup_nemo_gym(base_urls, model_name):
actor = spinup_nemo_gym_actor(
env_configs=env_configs,
base_urls=base_urls,
base_url_groups=getattr(
generation, "dp_openai_server_base_urls_by_group", None
),
model_name=model_name,
enable_router_replay=enable_router_replay,
routed_experts_dtype=routed_experts_dtype,
Expand Down Expand Up @@ -1184,6 +1191,7 @@ def initialize_generation_with_policy(
nemo_gym_actor, nemo_gym_time = _spinup_nemo_gym(
policy_generation.dp_openai_server_base_urls,
generation_config["model_name"],
generation=policy_generation,
)
worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time

Expand Down Expand Up @@ -1281,6 +1289,7 @@ def init_nemo_gym():
return _spinup_nemo_gym(
deferred_vllm.dp_openai_server_base_urls,
generation_config["model_name"],
generation=deferred_vllm,
)

# Colocated: vLLM + policy share GPUs -> sequential; otherwise parallel.
Expand Down Expand Up @@ -1381,6 +1390,7 @@ def init_trtllm():
nemo_gym_actor, nemo_gym_time = _spinup_nemo_gym(
policy_generation.dp_openai_server_base_urls,
generation_config["model_name"],
generation=policy_generation,
)
worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time

Expand Down
146 changes: 144 additions & 2 deletions nemo_rl/environments/nemo_gym.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,14 @@
# 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.
import asyncio
import contextlib
import json
import math
import os
import subprocess
import sys
import tempfile
import urllib.request
from collections import Counter
from collections.abc import AsyncGenerator
Expand Down Expand Up @@ -63,6 +66,8 @@
"int32": torch.int32,
}

DEFAULT_TAIL_ROUTE_GROUP = "low_latency"

DEFAULT_INVALID_TOOL_CALL_PATTERNS = [
"<tool_call>",
"</tool_call>",
Expand Down Expand Up @@ -141,6 +146,9 @@ def get_nemo_gym_venv_dir() -> str | None:
class NemoGymConfig(TypedDict):
model_name: str
base_urls: List[str]
# {group_name: [base_url, ...]}, used to label each replica at router
# registration. Absent means one unnamed pool.
base_url_groups: NotRequired[dict[str, List[str]] | None]
initial_global_config_dict: Dict[str, Any]
# Port range for Gym HTTP servers (head server + subprocess servers).
# Defaults to DEFAULT_GYM_PORT_RANGE_LOW/HIGH (5000-5999) from
Expand Down Expand Up @@ -169,6 +177,17 @@ class NemoGymConfig(TypedDict):
pad_dynamic_image_shapes: NotRequired[
bool
] # Normalize heterogeneous image tensors while retaining exact imgs_sizes
# Cap on rollouts in flight. Unset means the whole batch is dispatched at once,
# which with long prompts can stall a replica's HTTP server long enough that new
# connections time out and a router takes the replica out of rotation. Steering
# only reaches requests not yet dispatched, so this cap is what makes the tail
# steerable at all.
max_concurrent_rollouts: NotRequired[int | None]
# Fraction of a step's rollouts that must finish before the rest are steered to
# tail_route_group. Unset disables steering.
tail_route_threshold: NotRequired[float | None]
# Server group the tail is steered to; must match a server_groups entry.
tail_route_group: NotRequired[str]


def _detect_invalid_tool_call_and_malformed_thinking(
Expand Down Expand Up @@ -438,6 +457,9 @@ def __init__(self, cfg: NemoGymConfig):
"_attach_multimodal_data_to_user_message before enabling."
)

self._tail_sentinel: str | None = None
self._tail_group: str = cfg.get("tail_route_group", DEFAULT_TAIL_ROUTE_GROUP)

def _require_spinup(self) -> None:
"""Raise a diagnosable error if this instance never ran :meth:`_spinup`."""
if self.rh is None:
Expand All @@ -459,6 +481,17 @@ def health_check(self) -> None:
self._require_spinup()
self.rh.poll()

def _tail_route(self, on: bool) -> None:
"""Publish the tail-steering signal. Off at the start of every step."""
if not self._tail_sentinel:
return
if on:
with open(self._tail_sentinel, "w") as f:
f.write("1")
else:
with contextlib.suppress(FileNotFoundError):
os.remove(self._tail_sentinel)

def _spinup(self) -> None:
"""Start the NeMo-Gym head server and rollout collection helper.

Expand All @@ -467,6 +500,18 @@ def _spinup(self) -> None:
server URLs are available, overlapping with vLLM model loading.
"""
self.node_ip = _get_node_ip_local()
# A sentinel file rather than server state: vllm_model runs several uvicorn
# workers that share no memory, so a server-side switch would flip in only
# one of them. Gym servers are spawned from here and inherit this environment.
if self.cfg.get("tail_route_threshold") is not None:
log_dir = (self.cfg.get("initial_global_config_dict") or {}).get(
"nemo_gym_log_dir"
) or tempfile.gettempdir()
os.makedirs(log_dir, exist_ok=True)
self._tail_sentinel = os.path.join(log_dir, "tail_route.on")
self._tail_route(False)
os.environ["NEMO_GYM_TAIL_ROUTE_SENTINEL"] = self._tail_sentinel
os.environ["NEMO_GYM_TAIL_ROUTE_GROUP"] = self._tail_group
_gym_port_low = self.cfg.get("port_range_low", DEFAULT_GYM_PORT_RANGE_LOW)
_gym_port_high = self.cfg.get("port_range_high", DEFAULT_GYM_PORT_RANGE_HIGH)
self.head_server_port = _get_free_port_local(_gym_port_low, _gym_port_high)
Expand Down Expand Up @@ -494,13 +539,36 @@ def _spinup(self) -> None:
"dummy_key" # No key necessary for training.
)
router_url = initial_global_config_dict.pop("router_url", None)
# Warm every replica before it takes traffic: the first request spends
# seconds compiling Triton kernels, long enough to time out a rollout.
# Best effort -- a backend without /warmup just pays the cost later.
for base_url in filter(None, self.cfg["base_urls"]):
with contextlib.suppress(OSError):
urllib.request.urlopen(
urllib.request.Request(
f"{base_url.removesuffix('/v1')}/warmup", method="POST"
),
timeout=600,
).close()
if router_url:
# Which group each replica belongs to, so the router can honour a
# request that asks for one by name. Empty when the backend runs a
# single unnamed pool.
group_of: dict[str, str] = {
url: name
for name, urls in (self.cfg.get("base_url_groups") or {}).items()
for url in urls
if url
}
# POST /workers blocks until the router's own health probe of the
# replica succeeds, so a 200 means registered and routable.
for base_url in filter(None, self.cfg["base_urls"]):
payload: dict[str, Any] = {"url": base_url.removesuffix("/v1")}
if base_url in group_of:
payload["labels"] = {"group": group_of[base_url]}
request = urllib.request.Request(
f"{router_url.removesuffix('/v1')}/workers",
data=json.dumps({"url": base_url.removesuffix("/v1")}).encode(),
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
Expand Down Expand Up @@ -594,8 +662,20 @@ async def run_rollouts(
encode_images_in_examples(nemo_gym_examples)

timer.start("_run_rollouts_total")
limit = self.cfg.get("max_concurrent_rollouts")
nemo_gym_result_iterator = self.rch.run_examples(
examples=nemo_gym_examples, head_server_config=self.head_server_config
examples=nemo_gym_examples,
head_server_config=self.head_server_config,
semaphore=asyncio.Semaphore(limit) if limit else None,
)

# Steering is within-step state: every step starts off and flips once.
self._tail_route(False)
tail_threshold = self.cfg.get("tail_route_threshold")
tail_at = (
math.ceil(tail_threshold * len(nemo_gym_examples))
if tail_threshold is not None
else None
)

num_results = 0
Expand Down Expand Up @@ -629,6 +709,13 @@ async def run_rollouts(
raise RuntimeError("Generation logprobs contain NaN")

num_results += 1
if tail_at is not None and num_results == tail_at:
self._tail_route(True)
print(
f"[tail-route] {num_results}/{len(nemo_gym_examples)} rollouts done; "
f"steering the remainder to server group '{self._tail_group}'",
flush=True,
)
timing_metrics = None
if num_results == len(nemo_gym_examples):
timer.stop("_run_rollouts_total")
Expand Down Expand Up @@ -1035,11 +1122,51 @@ def setup_nemo_gym_config(config, tokenizer) -> None:
env_cfg.setdefault("tokenizer_config", dict(config.policy["tokenizer"]))


def _validate_tail_route(
threshold: float,
group: str,
base_url_groups: dict[str, list[str]] | None,
*,
router_url: str | None,
policy_model: Any,
) -> None:
"""Fail at setup when tail steering cannot possibly take effect.

Every one of these misconfigurations is otherwise silent: the rollout runs,
the flip is logged, and the requests go exactly where they always went.
"""
if not 0 < threshold <= 1:
raise ValueError(f"tail_route_threshold must be in (0, 1], got {threshold!r}")
if not router_url:
raise ValueError("tail_route_threshold requires env.nemo_gym.router_url")
if not base_url_groups:
raise ValueError(
"tail_route_threshold requires policy.generation.server_groups; "
"without groups there is nowhere to steer to"
)
if group not in base_url_groups:
raise ValueError(
f"tail_route_group '{group}' is not a server group: "
f"{sorted(base_url_groups)}"
)
header = (
((policy_model or {}).get("responses_api_models") or {})
.get("vllm_model", {})
.get("worker_group_header")
)
if not header:
raise ValueError(
"tail_route_threshold requires policy_model.responses_api_models."
"vllm_model.worker_group_header; without it no request is tagged"
)


def spinup_nemo_gym_actor(
env_configs: dict[str, Any],
base_urls: list[str],
model_name: str,
*,
base_url_groups: dict[str, list[str]] | None = None,
enable_router_replay: bool,
routed_experts_dtype: str,
use_fastokens: bool,
Expand Down Expand Up @@ -1080,6 +1207,17 @@ def spinup_nemo_gym_actor(
_value = nemo_gym_dict.pop(_flag, None)
if _value is not None:
multimodal_flags[_flag] = bool(_value)
max_concurrent_rollouts = nemo_gym_dict.pop("max_concurrent_rollouts", None)
tail_route_threshold = nemo_gym_dict.pop("tail_route_threshold", None)
tail_route_group = nemo_gym_dict.pop("tail_route_group", DEFAULT_TAIL_ROUTE_GROUP)
if tail_route_threshold is not None:
_validate_tail_route(
tail_route_threshold,
tail_route_group,
base_url_groups,
router_url=nemo_gym_dict.get("router_url"),
policy_model=nemo_gym_dict.get("policy_model"),
)

# Pass prebuilt cache + venv dirs through the global config so the gym reuses
# image-baked venvs instead of rebuilding them.
Expand All @@ -1093,12 +1231,16 @@ def spinup_nemo_gym_actor(
nemo_gym_cfg = NemoGymConfig(
model_name=model_name,
base_urls=base_urls,
base_url_groups=base_url_groups,
invalid_tool_call_patterns=invalid_tool_call_patterns,
thinking_tags=thinking_tags,
tokenizer_config=tokenizer_config,
require_routed_experts=enable_router_replay,
routed_experts_dtype=routed_experts_dtype,
use_fastokens=use_fastokens,
max_concurrent_rollouts=max_concurrent_rollouts,
tail_route_threshold=tail_route_threshold,
tail_route_group=tail_route_group,
initial_global_config_dict=nemo_gym_dict,
**multimodal_flags,
)
Expand Down
Loading