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
33Spawns 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
1816Per-pipeline isolation that this driver enforces:
1917- Distinct pipeline IDs from ``orchestrator.allocate_pipeline_id``
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
3735from __future__ import annotations
3836
3937import copy
38+ import logging
4039import 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).
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
6481def _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
346371def 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