Skip to content

Commit cdda5b6

Browse files
TianyeGGBondclaude
andcommitted
refactor(dual-driver): PipelineHandle dataclass + local imports
_build_pipeline took six threaded dependencies (ray, MilesCoordinator, MilesPipelineConfig, get_coordinator_actor_name, get_pipeline_namespace, logger) and returned a 5-tuple that main() then re-packed into a 7-tuple, unpacked positionally in four places. A missed position would silently mis-wire a handle. - Drop the threaded deps: _build_pipeline imports ray / MilesCoordinator / get_coordinator_actor_name / get_pipeline_namespace locally (still after the env guard, since it only runs from main); MilesPipelineConfig and the logger are module-level. - Return a PipelineHandle dataclass; main fills train_group / rollout_manager after build and accesses everything by attribute. - Refresh the now-stale module docstring: drop the "disjoint-only / does not exercise overlap / via asyncio.gather" wording (overlap is supported via MILES_DUAL_* and the loop uses asyncio.wait(FIRST_EXCEPTION) + cancel). No behavior change; both drivers still import cleanly with the heavy imports deferred past the RLIX_CONTROL_PLANE guard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 6d23375 commit cdda5b6

1 file changed

Lines changed: 76 additions & 75 deletions

File tree

examples/rlix/run_miles_dual.py

Lines changed: 76 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,17 @@
1-
"""M11.2 dual-pipeline driver — `examples/rlix/run_miles_dual.py`.
1+
"""RLix dual-pipeline example driver — `examples/rlix/run_miles_dual.py`.
22
33
Spawns two MilesCoordinator + MilesPipeline pairs in separate Ray
4-
namespaces with disjoint ``cluster_device_mappings``. Each pipeline
5-
runs its own ``rlix_train_loop`` concurrently via ``asyncio.gather``.
6-
7-
Topology (Codex-recommended Option A — disjoint pools, no cross-pipeline
8-
GPU contention):
9-
pipeline 1: actor_train=[0,1], actor_infer=[0,1]
10-
pipeline 2: actor_train=[2,3], actor_infer=[2,3]
11-
12-
This is the minimum-viable M11.2 PASS — proves two pipelines can register
13-
+ initialize + train + sync + generate + clean up concurrently without
14-
namespace, actor-name, port, or scheduler-ledger collisions. It does NOT
15-
exercise cross-pipeline preemption (Option B/C — overlap topology — needs
16-
the deferred F22 shell-init contract per ``miles_pipeline.py:14-33``).
4+
namespaces and runs their training loops concurrently. The first loop to
5+
raise cancels its peer (see ``main``'s ``asyncio.wait`` / cancel handling)
6+
so both settle and release their scheduler allocations cleanly.
7+
8+
GPU topology is one of:
9+
- disjoint (default): the physical pool is split into two non-overlapping
10+
per-pipeline pools — see ``_split_pools_for_dual``;
11+
- explicit (overlap-capable): set all four
12+
``MILES_DUAL_P*_{TRAIN,INFER}`` env vars to map each pipeline onto
13+
specific physical GPUs, which may overlap across pipelines — see
14+
``_overlap_pools_from_env``.
1715
1816
Per-pipeline isolation that this driver enforces:
1917
- Distinct pipeline IDs from ``orchestrator.allocate_pipeline_id``
@@ -29,15 +27,18 @@
2927
- W&B / TensorBoard / Prometheus disabled (``--use-wandb`` etc must be
3028
unset); per-pipeline tracking re-enable is M11.3 follow-up
3129
32-
Per scope F13 the driver MUST NOT have a top-level ``try/except`` and
33-
MUST NOT call ``ray.shutdown()``: failure semantics = let exceptions
34-
propagate naturally → driver exits → user runs ``ray stop`` to clean up.
30+
Failure semantics: the driver has no top-level ``try/except`` and does not
31+
call ``ray.shutdown()`` — exceptions propagate, the driver exits, and the
32+
user runs ``ray stop`` to clean up.
3533
"""
3634

3735
from __future__ import annotations
3836

3937
import copy
38+
import logging
4039
import os
40+
from dataclasses import dataclass
41+
from typing import Any
4142

4243
# Support both `python -m examples.rlix.run_miles_dual` (package context) and
4344
# `python examples/rlix/run_miles_dual.py` (direct script, no parent package).
@@ -60,6 +61,22 @@
6061
"examples/rlix/run_miles_dual.py", "examples.rlix.run_miles_dual"
6162
)
6263

64+
logger = logging.getLogger("run_miles_dual")
65+
66+
67+
@dataclass
68+
class PipelineHandle:
69+
"""One pipeline's actors + handles, populated across build and setup."""
70+
71+
index: int
72+
pipeline_id: str
73+
namespace: str
74+
coordinator: Any
75+
pipeline: Any
76+
args: Any
77+
train_group: Any = None
78+
rollout_manager: Any = None
79+
6380

6481
def _split_pools_for_dual(
6582
*, num_gpus_per_node: int, infer_pool_size: int
@@ -226,13 +243,7 @@ def _build_pipeline(
226243
train_mapping: list[int],
227244
infer_mapping: list[int],
228245
orchestrator,
229-
ray,
230-
MilesCoordinator,
231-
MilesPipelineConfig,
232-
get_coordinator_actor_name,
233-
get_pipeline_namespace,
234-
logger,
235-
):
246+
) -> PipelineHandle:
236247
"""Allocate one pipeline_id, register, admit, create coordinator+pipeline.
237248
238249
``train_mapping`` / ``infer_mapping`` are the EXPLICIT physical GPU
@@ -241,8 +252,15 @@ def _build_pipeline(
241252
partial-overlap invariant is asserted; cross-pipeline overlap is
242253
asserted by ``grep_overlap_log.sh`` end-to-end.
243254
244-
Returns ``(pipeline_id, namespace, coordinator_handle, pipeline_handle, args)``.
255+
Returns a :class:`PipelineHandle`.
245256
"""
257+
import ray
258+
from rlix.pipeline.miles_coordinator import MilesCoordinator
259+
from rlix.protocol.types import (
260+
get_coordinator_actor_name,
261+
get_pipeline_namespace,
262+
)
263+
246264
pipeline_id = ray.get(orchestrator.allocate_pipeline_id.remote("miles"))
247265
pipeline_namespace = get_pipeline_namespace(pipeline_id)
248266

@@ -340,32 +358,32 @@ def _build_pipeline(
340358
pipeline_index, pipeline_id,
341359
)
342360

343-
return pipeline_id, pipeline_namespace, coordinator, pipeline, args
361+
return PipelineHandle(
362+
index=pipeline_index,
363+
pipeline_id=pipeline_id,
364+
namespace=pipeline_namespace,
365+
coordinator=coordinator,
366+
pipeline=pipeline,
367+
args=args,
368+
)
344369

345370

346371
def main():
347372
"""Dual-pipeline entry. Imports heavy modules lazily so the env-var
348373
guard above fires before transitive ``import torch`` / ``import sglang``.
349374
"""
350375
import asyncio
351-
import logging
352376

353377
import ray
354378

355379
from miles.utils.arguments import parse_args
356380
from miles.utils.logging_utils import configure_logger
357381
from miles.utils.rlix_train_loop import run_async_train_loop
358382
from miles.utils.rlix_validation import assert_rlix_topology
359-
from rlix.pipeline.miles_coordinator import MilesCoordinator
360-
from rlix.protocol.types import (
361-
get_coordinator_actor_name,
362-
get_pipeline_namespace,
363-
)
364383

365384
import rlix
366385

367386
configure_logger()
368-
logger = logging.getLogger("run_miles_dual")
369387
base_args = parse_args()
370388

371389
# F10 startup fail-fast on the BASE args. Per-pipeline arg overrides
@@ -420,75 +438,64 @@ def main():
420438
train_mapping=p1_train,
421439
infer_mapping=p1_infer,
422440
orchestrator=orchestrator,
423-
ray=ray,
424-
MilesCoordinator=MilesCoordinator,
425-
MilesPipelineConfig=MilesPipelineConfig,
426-
get_coordinator_actor_name=get_coordinator_actor_name,
427-
get_pipeline_namespace=get_pipeline_namespace,
428-
logger=logger,
429441
)
430442
p2 = _build_pipeline(
431443
base_args=base_args,
432444
pipeline_index=2,
433445
train_mapping=p2_train,
434446
infer_mapping=p2_infer,
435447
orchestrator=orchestrator,
436-
ray=ray,
437-
MilesCoordinator=MilesCoordinator,
438-
MilesPipelineConfig=MilesPipelineConfig,
439-
get_coordinator_actor_name=get_coordinator_actor_name,
440-
get_pipeline_namespace=get_pipeline_namespace,
441-
logger=logger,
442448
)
443449

444-
pipelines = [p1, p2]
450+
handles = [p1, p2]
445451

446452
# ---- 3. Pull handles for each pipeline. ------------------------------
447-
handles = []
448-
for pid, ns, coord, pipe, args in pipelines:
449-
train_group = ray.get(pipe.get_train_group.remote())
450-
rollout_manager = ray.get(pipe.get_rollout_manager.remote())
451-
engine_count = int(ray.get(pipe.get_declared_engine_count.remote()))
453+
for h in handles:
454+
h.train_group = ray.get(h.pipeline.get_train_group.remote())
455+
h.rollout_manager = ray.get(h.pipeline.get_rollout_manager.remote())
456+
engine_count = int(ray.get(h.pipeline.get_declared_engine_count.remote()))
452457
logger.info(
453458
"[run_miles_dual] handles ready pipeline_id=%s engines=%d",
454-
pid, engine_count,
459+
h.pipeline_id, engine_count,
455460
)
456-
handles.append((pid, ns, coord, pipe, args, train_group, rollout_manager))
457461

458-
# ---- 4. Drive 2 concurrent rlix_train_loops via asyncio.gather. -----
459-
async def _run_one_pipeline(idx, pid, pipe, args, train_group, rollout_manager):
462+
# ---- 4. Drive 2 concurrent rlix_train_loops. ------------------------
463+
async def _run_one_pipeline(h: PipelineHandle):
460464
async def _before(step: int) -> None:
461-
await pipe.before_training.remote(step)
465+
await h.pipeline.before_training.remote(step)
462466

463467
async def _after(step: int) -> None:
464-
await pipe.after_training.remote(step)
468+
await h.pipeline.after_training.remote(step)
465469

466470
async def _release_only(step: int) -> None:
467-
# R04-F1 cleanup hook: releases actor_train allocation only.
468-
await pipe.release_train_only.remote(step)
471+
# Cleanup hook: releases actor_train allocation only.
472+
await h.pipeline.release_train_only.remote(step)
469473

470474
# Per-rollout step_target = rollout_batch_size. See
471475
# MilesPipeline.signal_rollout_demand docstring for why pre-signalling
472476
# demand to the scheduler is required for 4-GPU 2-pipeline full
473477
# cross-overlap (without it, rollout 2+ hangs when both pipelines
474478
# release all DP workers between rollouts).
475-
_step_target = int(getattr(args, "rollout_batch_size", 0) or 0)
479+
_step_target = int(getattr(h.args, "rollout_batch_size", 0) or 0)
476480

477481
async def _signal_demand(rollout_id: int) -> None:
478482
if _step_target <= 0:
479483
return
480-
await pipe.signal_rollout_demand.remote(rollout_id, _step_target)
484+
await h.pipeline.signal_rollout_demand.remote(rollout_id, _step_target)
481485

482486
await run_async_train_loop(
483-
args,
484-
train_group=train_group,
485-
rollout_manager=rollout_manager,
487+
h.args,
488+
train_group=h.train_group,
489+
rollout_manager=h.rollout_manager,
486490
before_step=_before,
487491
after_step=_after,
488492
release_only=_release_only,
489493
signal_demand=_signal_demand,
490494
)
491-
logger.info("[run_miles_dual] mp%d training loop complete pipeline_id=%s", idx, pid)
495+
logger.info(
496+
"[run_miles_dual] mp%d training loop complete pipeline_id=%s",
497+
h.index, h.pipeline_id,
498+
)
492499

493500
async def _async_main():
494501
# F4 fix (m11-review.review-report.md §2): use create_task + wait(
@@ -500,13 +507,7 @@ async def _async_main():
500507
# so Phase 1's try/finally inside run_async_train_loop fires
501508
# release_only on the CancelledError path and the scheduler
502509
# ledger stays consistent.
503-
tasks = [
504-
asyncio.create_task(
505-
_run_one_pipeline(i + 1, pid, pipe, args, train_group, rollout_manager)
506-
)
507-
for i, (pid, ns, coord, pipe, args, train_group, rollout_manager)
508-
in enumerate(handles)
509-
]
510+
tasks = [asyncio.create_task(_run_one_pipeline(h)) for h in handles]
510511
try:
511512
done, pending = await asyncio.wait(
512513
tasks, return_when=asyncio.FIRST_EXCEPTION
@@ -539,13 +540,13 @@ async def _async_main():
539540
# original training exception. Codex Phase 7 review MEDIUM.
540541
try:
541542
shutdown_refs = [
542-
pipe.shutdown_hard.remote() for _, _, _, pipe, _, _, _ in handles
543+
h.pipeline.shutdown_hard.remote() for h in handles
543544
]
544545
ray.get(shutdown_refs, timeout=60.0)
545-
for pid, _, _, _, _, _, _ in handles:
546+
for h in handles:
546547
logger.info(
547548
"[run_miles_dual] shutdown_hard complete pipeline_id=%s",
548-
pid,
549+
h.pipeline_id,
549550
)
550551
except Exception as exc: # noqa: BLE001
551552
logger.warning(

0 commit comments

Comments
 (0)