feat: Placement Group based subprocess manager - inference (3/5)

[praateekmahajan]

Reviewer tip: this is a zero-consumer infra PR. DynamoBackend still raises NotImplementedError at the end of this PR (unchanged from #1820). The two modules to review are nemo_curator/core/serve/subprocess_mgr.py (Ray-actor + subprocess lifecycle) and nemo_curator/core/serve/placement.py (PG planner + builders + bundle helpers). Everything else is tests or small constants moves.

Description

• PR 3/5 of the inference-server restack
• land the Ray-placement-group-based subprocess manager as a self-contained infra library; nothing in nemo_curator/ imports these modules yet — PR 4 (DynamoBackend) will
• split responsibilities across two files:
  • serve/subprocess_mgr.py — ManagedSubprocess + _SubprocessActor factory + SubprocessError + graceful_stop_actors
  • serve/placement.py — ReplicaBundleSpec, plan_replica_bundle_shape, build_pg / build_replica_pg, bundle-scoped helpers (get_bundle_node_ip, get_free_port_in_bundle), and remove_named_pgs_with_prefix for orphan sweep
• Dynamo-specific placement / naming helpers live in a new serve/dynamo/infra.py (build_infra_pg, build_worker_actor_name, engine_kwargs_to_cli_flags) so subprocess_mgr / placement stay backend-generic
• dedup CURATOR_IGNORE_RAY_HEAD_NODE parsing into a shared nemo_curator.core.utils.ignore_ray_head_node(); nemo_curator/backends/base.py now uses the same helper instead of parsing the env var inline
• add check_total_gpu_capacity next to get_available_cpu_gpu_resources in nemo_curator/backends/utils.py — coarse pre-check so insufficient-GPU cases fail fast instead of hanging on pg.ready()

Design callouts

• Detached, named PGs. Each replica PG is created with lifetime=\"detached\" and a stable name= so a reconnecting driver (server.start → pipeline.run → server.stop flow) can find and reap them across ray.shutdown() / ray.init() cycles. remove_named_pgs_with_prefix is the orphan sweep.
• Bundle-shape planner. Single-node replicas use one STRICT_PACK bundle; multi-node TP replicas use STRICT_SPREAD with an equal per-node split. Asymmetric splits (1+3 for TP=4) are rejected up front because vLLM's distributed executor requires identical local_world_size per node.
• Head-node exclusion. CURATOR_IGNORE_RAY_HEAD_NODE=1 is translated to a per-bundle {\"ray.io/node-type\": \"worker\"} label selector (Ray-native; matches Ray Serve's deployment-scheduler pattern). Auto-satisfied on Anyscale; OSS Ray users must start worker nodes with ray start --labels ray.io/node-type=worker.
• Subprocess-group reap. Subprocesses are launched with start_new_session=True so the child becomes a process-group leader. _SubprocessActor overrides __ray_terminate__ and registers an atexit hook so the subprocess tree is SIGTERM→SIGKILL'd on the process group whenever the actor shuts down gracefully.
• SIGKILL escalation in teardown. graceful_stop_actors runs in three steps per actor: (1) actor.stop() with bounded wait, (2) if stop didn't drain, actor.force_sigkill_subprocess() — a non-blocking os.killpg(..., SIGKILL) scheduled on the actor's node — then (3) ray.kill. Without step 2, ray.kill bypasses __ray_terminate__ / atexit and orphans the subprocess tree.
• ManagedSubprocess.spawn is a classmethod on the data class it returns — ergonomic factory that avoids the split between a free-function spawner and the handle type. The handle exposes .stop(), .stop_many(), .is_alive(), .pid(), .read_log_tail(), .wait() so callers don't hand-write ray.get(actor.X.remote()).

What's intentionally NOT in this PR

• any change to DynamoBackend.start() / .stop() — still the NotImplementedError placeholder from feat!: Typed Serve Config + Dyamo config stub - inference (2/5) #1820
• any consumer import of subprocess_mgr / placement from production code (tests are the only consumer)
• the Bundle(pg, bundle_index) wrapper idea — deferred until PR 4's DynamoBackend wiring shows how often the pair repeats

Internal usage (preview of how PR 4 will consume this)

from nemo_curator.core.serve.placement import (
    build_replica_pg, get_bundle_node_ip, plan_replica_bundle_shape,
)
from nemo_curator.core.serve.subprocess_mgr import ManagedSubprocess

spec = plan_replica_bundle_shape(tp_size=4)      # single STRICT_PACK bundle or N STRICT_SPREAD
pg   = build_replica_pg(spec, name="dyn_model_0")
proc = ManagedSubprocess.spawn(
    "Dynamo_DP0_Qwen3-0.6B",
    pg, bundle_index=0,
    num_gpus=spec.per_node_gpus,
    python_args=["-m", "dynamo.vllm", "--model", "Qwen/Qwen3-0.6B", "--tensor-parallel-size", "4"],
)
# ... later ...
proc.stop()                                      # SIGTERM -> SIGKILL escalation inside the actor
ManagedSubprocess.stop_many([proc1, proc2, ...]) # parallel teardown

Verification

• ruff check nemo_curator/core/serve/ tests/core/serve/
• pytest tests/core/serve/test_placement.py tests/core/serve/test_subprocess_mgr.py tests/core/serve/dynamo/test_infra.py tests/core/test_utils.py tests/backends/test_utils.py::TestCheckTotalGpuCapacity -m \"not gpu\" — 21 passed (~28s; setup-bound on Ray cluster autouse fixture)
• CUDA_VISIBLE_DEVICES=0,1 pytest tests/core/serve/test_placement.py tests/core/serve/test_subprocess_mgr.py -m gpu — 3 passed (TestReplicaLifecycle end-to-end, actor-death run_ref surfacing, orphan PG cleanup by prefix)

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[greptile-apps]

Greptile Summary

This PR lands the Ray-placement-group-based subprocess manager as a self-contained infra library, split across placement.py (PG planning + construction + bundle helpers) and subprocess_mgr.py (actor factory + subprocess lifecycle). It also deduplicates CURATOR_IGNORE_RAY_HEAD_NODE parsing into a shared helper and adds a check_total_gpu_capacity pre-check. Prior review concerns (private Ray API access, tp_size <= 0 guard, module-level constant mutation) are addressed.

Confidence Score: 5/5

Safe to merge; all remaining findings are P2 style/defensive-coding suggestions that do not affect correctness of the happy path or orphan-cleanup guarantees.

Prior review concerns (private Ray API, missing tp_size guard, constant mutation) are resolved. The two new findings are both P2: the TOCTOU race in _stop_subprocess causes unnecessary SIGKILL escalation on clean subprocess exits but leaves no orphans, and the check_total_gpu_capacity semantic mismatch is a docstring/naming concern rather than a runtime failure. Tests cover the critical lifecycle paths end-to-end.

nemo_curator/core/serve/subprocess_mgr.py for the _stop_subprocess signal-sending race; nemo_curator/backends/utils.py for the available-vs-total GPU capacity semantics.

Important Files Changed

Filename	Overview
nemo_curator/core/serve/subprocess_mgr.py	New file: Ray actor + subprocess lifecycle manager; well-structured with SIGTERM→SIGKILL escalation, but _stop_subprocess has a TOCTOU race where proc.terminate() can raise if the process exits naturally between poll() and signal delivery
nemo_curator/core/serve/placement.py	New file: PG planner + construction + bundle helpers; clean logic for single/multi-node TP splits with proper tp_size < 1 guard (addressed from prior review) and head-node filtering
nemo_curator/core/serve/dynamo/infra.py	New file: Dynamo-specific helpers (build_infra_pg, build_worker_actor_name, engine_kwargs_to_cli_flags); well-isolated from generic placement/subprocess modules
nemo_curator/backends/utils.py	Adds check_total_gpu_capacity; functional but checks currently-available rather than total installed GPUs, which may cause false-positive failures during cluster load
nemo_curator/backends/base.py	Replaces inline env-var parsing with shared ignore_ray_head_node() helper; semantically equivalent and correctly deduplicates the logic
nemo_curator/core/utils.py	Adds ignore_ray_head_node() helper; consistent with old inline parsing, well-tested with full truth table in tests/core/test_utils.py
nemo_curator/core/serve/constants.py	Adds subprocess/PG tunables (SIGTERM/SIGKILL waits, PG timeout, worker label, NOSET env fragment); well-documented with docstrings
tests/core/serve/test_subprocess_mgr.py	Comprehensive lifecycle tests covering PG readiness, bundle IP/port discovery, CUDA env injection, subprocess env overrides, graceful stop, and actor-death-surfacing via run_ref

Sequence Diagram

sequenceDiagram
    participant Driver
    participant placement.py
    participant subprocess_mgr.py
    participant RayActor as _SubprocessActor
    participant Subprocess

    Driver->>placement.py: plan_replica_bundle_shape(tp_size)
    placement.py-->>Driver: ReplicaBundleSpec
    Driver->>placement.py: build_replica_pg(spec, name)
    placement.py->>Ray: placement_group lifetime=detached
    placement.py->>Ray: pg.ready() timeout=180s
    placement.py-->>Driver: PlacementGroup
    Driver->>subprocess_mgr.py: ManagedSubprocess.spawn(label, pg, bundle_index)
    subprocess_mgr.py->>RayActor: actor_cls.options(num_gpus=N, pg_strategy).remote()
    subprocess_mgr.py->>RayActor: actor.initialize.remote(command, env, log_file)
    RayActor->>RayActor: inject CUDA_VISIBLE_DEVICES from accelerator IDs
    RayActor->>Subprocess: Popen(command, start_new_session=True)
    subprocess_mgr.py->>RayActor: actor.run.remote() returns run_ref
    subprocess_mgr.py-->>Driver: ManagedSubprocess(label, actor, run_ref)
    Driver->>subprocess_mgr.py: proc.stop()
    subprocess_mgr.py->>subprocess_mgr.py: graceful_stop_actors
    subprocess_mgr.py->>RayActor: actor.stop.remote() SIGTERM group
    RayActor->>Subprocess: os.killpg(pgid, SIGTERM)
    alt stop drains in time
        RayActor-->>subprocess_mgr.py: rc
    else timeout
        subprocess_mgr.py->>RayActor: actor.force_sigkill_subprocess.remote()
        RayActor->>Subprocess: os.killpg(pgid, SIGKILL)
    end
    subprocess_mgr.py->>Ray: ray.kill(actor, no_restart=True)

Loading

_{Reviews (4): Last reviewed commit: "chore: trigger DCO re-check after base r..." | Re-trigger Greptile}

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[praateekmahajan]

/claude review

[oyilmaz-nvidia]

@praateekmahajan Can you give a bit more context why we need these new classes (subprocess manager for instance) and functions?

[praateekmahajan]

@oyilmaz-nvidia fair question on subprocess manager, the answer is that we use this in PR4/5.

This is a completely backend (dynamo / serve) agnostic PR. Rather in future if we add native SGLang support (or vLLM) this all will be leveraged then as well.
In PR 4 and PR 5, we start dynamo as a fleet of plain processes (etcd, nats-server, python -m dynamo.frontend, N × python -m dynamo.vllm) and use these classes. Here is a high level description

• subprocess_mgr.py — classes / utilities to help spawn each process inside a named detached Ray actor and kill the whole process tree on teardown.
• placement.py — plan ray's placement group bundles for TP (e.g. if it fits one node vs when it doesn't i.e. STRICT_PACK vs STRICT_SPREAD) and resolve rank-0's node IP + free ports from inside the bundle. Backend-agnostic.

PR 4's DynamoBackend.start() is ~200 lines of glue on top: uses plan_replica_bundle_shape / ManagedSubprocess.spawn.
PR 5 adds disagg by calling the same primitives once per prefill/decode worker — no new lifecycle code.