refactor: in-process torchrun + DeepSpeed launchers; fix multi-GPU SFT/DPO/PPO spawn

[timzsu]

Purpose

Bandit follow-up (#6). Migrate the SFT / DPO / PPO distributed launchers to in-process torch.distributed.run.main and deepspeed.launcher.runner.main (the canonical entry points the torchrun and deepspeed console scripts themselves invoke). Swap B404 ↔ B603 in the bandit skip list so every subprocess call site has a per-line written rationale instead of a blanket skip. After this PR, no # nosec TODO remains in src/.

Once the launcher swap was in, the multi-GPU spawn path was actually exercised end-to-end for the first time. Three pre-existing bugs surfaced and are fixed in this PR — see the second bullet group under Changes.

Changes

Launcher swap (the original PR scope)

• src/worker/executors/utils/distributed.py (new) — run_torchrun(), run_deepspeed(), deepspeed_available(). Lazy DeepSpeed import; capability check tolerates the package's CUDA-op probe at init. run_torchrun passes --tee 3 so rank stdout+stderr land on the parent console and in the per-rank elastic log dir; without it the elastic agent reports rank crashes only as an opaque ChildFailedError with error_file: <N/A>.
• sft_executor.py / dpo_executor.py / ppo_executor.py — route the multi-GPU spawn through the new helpers; drop import subprocess (and import shutil from SFT).
• pyproject.toml — [tool.bandit].skips: add B404, drop B603. Add deepspeed.* to mypy follow_untyped_imports.
• checkpoints.py — drop the # nosec B404 — TODO on the import (now project-skipped); add per-line # nosec B603 — argv list, no shell=True, absolute path via shutil.which() on the two pigz/tar subprocess.run calls.
• AGENTS.md — adds B603 rule bullet.

Multi-GPU path fixes (newly tested e2e)

• sft_executor.py — SFT spawn was dumping the raw WorkerTaskMessage pydantic model with json.dump(task, fh), raising Object of type WorkerTaskMessage is not JSON serializable on every multi-GPU SFT submission. Match the DPO/PPO line: json.dump(task.model_dump(mode="json", by_alias=True), fh).
• *_dist_entry.py (all three) — deserializer was task = json.load(fh), then executor.run(task, ...) called task.spec and crashed with 'dict' object has no attribute 'spec'. WorkerTaskMessage serializes through dedup_json in mode="json" and has a model_validator(mode="before") that auto-restores the deduped form on model_validate. Route the loaded JSON through WorkerTaskMessage.model_validate so the dist entry hands executor.run a real envelope.
• tests/worker/test_distributed_launcher.py — round-trip test that locks the writer/reader contract; updated argv assertion to include --tee 3.
• templates/{sft,dpo_training,ppo_training}_llama_1b_multi_gpu.yaml — three smoke templates so the multi-GPU path stays reachable from flowmesh workflow submit. Bring up a single worker with multiple cuda_devices via flowmesh stack worker up --config <yaml> (the up gpu --targets a,b shortcut creates one worker per GPU, not one multi-GPU worker).

Test Plan

• CI.
• Multi-GPU e2e on two H200 NVL GPUs:

  uv run flowmesh stack build flowmesh_worker_gpu
  uv run flowmesh stack worker up --config <multi-gpu-worker.yaml>   # cuda_devices: [<id-a>, <id-b>]
  uv run flowmesh workflow submit templates/sft_llama_1b_multi_gpu.yaml         # run_deepspeed (auto-DS ZeRO-2)
  uv run flowmesh workflow submit templates/dpo_training_llama_1b_multi_gpu.yaml # run_torchrun
  uv run flowmesh workflow submit templates/ppo_training_llama_1b_multi_gpu.yaml # run_torchrun
  

  All three reach the rank training loop. SFT (DeepSpeed) finishes in ~40s on a 16-sample subset; DPO finishes on the inline 5-pair preference set. PPO ranks load both policy and reward models, enter PPOTrainer.train(), and then hang inside accelerate.synchronize_rng_state (see caveat below).

Caveat — PPO hangs upstream

PPO does not finish under run_torchrun with our current trl / accelerate pin: both ranks reach PPOTrainer.train(), the data loader iterator calls accelerate.synchronize_rng_state, and that collective never returns. Reproduced on two different GPU pairs ([0,2] and [0,3]) so it is not a NVLink-topology quirk on this host. Stack from py-spy (both ranks identical):

synchronize_rng_state    (accelerate/utils/random.py:146)
synchronize_rng_states   (accelerate/utils/random.py:156)
__iter__                 (accelerate/data_loader.py:560)
repeat_generator         (trl/trainer/ppo_trainer.py:397)
train                    (trl/trainer/ppo_trainer.py:450)

This is not a regression from this PR: DPO uses the same run_torchrun entry point and finishes cleanly, and SFT (DeepSpeed branch) finishes too. The bug is in TRL's PPOTrainer.repeat_generator interaction with accelerate.synchronize_rng_state and is independent of the launcher swap. PR #12's scope (launcher swap + the three multi-GPU bugs the e2e surfaced) is verified end-to-end; the PPO upstream hang is left as a known issue to be addressed in a follow-up dep bump or trainer patch.

---

Pre-submission Checklist

• I have read the contribution guidelines.
• I have run pre-commit run --all-files and fixed any issues.
• I have added or updated tests covering my changes (if applicable).
• I have verified that uv run pytest tests/ passes locally.
• If I changed shared schemas or proto definitions, I have checked downstream compatibility across Server and Worker. N/A.
• If I changed the SDK or CLI, I have verified the affected packages work. N/A.
• If this is a breaking change, I have prefixed the PR title with [BREAKING]. Not breaking.
• I have updated documentation or config examples if user-facing behavior changed. AGENTS.md gains the B603 rule bullet.

[kaiitunnz]

Minor comments. PTAL.

[kaiitunnz]

LGTM.