diff --git a/modal_training_gym/frameworks/miles/launcher.py b/modal_training_gym/frameworks/miles/launcher.py index 1f12574c7..39a15bb77 100644 --- a/modal_training_gym/frameworks/miles/launcher.py +++ b/modal_training_gym/frameworks/miles/launcher.py @@ -97,6 +97,11 @@ def _validate_resume_checkpoint( MILES_ROOT = "/root/miles" SYSTEM_LIB_DIR = "/usr/lib/x86_64-linux-gnu" + +# Modal injects the AWS EFA/OFI stack (libfabric 1.30 + NCCL's ofi net plugin) +# on GPU hosts and prepends these dirs to the container's LD_LIBRARY_PATH. +_EFA_LIB_DIRS = ("/opt/amazon/efa/lib", "/opt/amazon/ofi-nccl/lib") + # libibverbs and the libmlx5 provider come from incompatible rdma package versions for miles multi-node training # reinstalling fixes this issue, mooncake transferengine imports successfully RDMA_RUNTIME_INSTALL_COMMAND = ( @@ -104,6 +109,8 @@ def _validate_resume_checkpoint( "--reinstall libibverbs1 ibverbs-providers && " "rm -rf /var/lib/apt/lists/*" ) + + # v0.8.0+ makes per-task CPU/memory requests configurable via enforcement # policies ("limit"/"ignore"), letting sandboxes burst on Modal and bill by # actual CPU-/RAM-second usage instead of over-provisioning a static reservation. @@ -111,6 +118,12 @@ def _validate_resume_checkpoint( _MILES_PATCHES = Path(__file__).parent / "modal_helpers" / "patches" _PATCH_SGLANG_ABORT_B64 = encode_patch("patch_sglang_abort", _MILES_PATCHES) +_PATCH_MOONCAKE_TOLERANCE_B64 = encode_patch( + "patch_mooncake_import_tolerance", _MILES_PATCHES +) +_PATCH_ROUTER_STARTUP_TIMEOUT_B64 = encode_patch( + "patch_router_startup_timeout", _MILES_PATCHES +) _PATCH_ROLLOUT_STATUS_B64 = encode_patch( "patch_rollout_status_reporting", _MILES_PATCHES ) @@ -351,6 +364,15 @@ def _build_miles_base_image(miles: MilesRecipe) -> Image: .run_commands( f"rm -rf {HF_CACHE_PATH} 2>/dev/null || true", f"echo {_PATCH_SGLANG_ABORT_B64} | base64 -d | python3", + # On EFA hosts a bind-mounted host libibverbs can break mooncake's + # TransferEngine import. Colocated sync does not use that P2P path, + # so keep mooncake out of the actor import chain and load it lazily. + f"echo {_PATCH_MOONCAKE_TOLERANCE_B64} | base64 -d | python3", + # miles allows the sglang router 30s to bind its port; the router's + # spawned child re-imports the whole stack first and overruns that + # under bring-up load. Raise the bound (it returns as soon as the + # port accepts, and still fails fast if the child dies). + f"echo {_PATCH_ROUTER_STARTUP_TIMEOUT_B64} | base64 -d | python3", f"echo {_PATCH_DIST_CKPT_QUANTIZED_B64} | base64 -d | python3", ( f"if test -f {_MEGATRON_TORCH_STRATEGY_PY}; then " @@ -369,6 +391,31 @@ def _build_miles_base_image(miles: MilesRecipe) -> Image: return image +def _overlay_local_miles(image: Image, local_miles: str) -> Image: + image = image.add_local_dir( + local_miles, + remote_path=MILES_ROOT, + copy=True, + ignore=["**/__pycache__", "**/*.pyc", "**/.git", "**/.venv"], + ) + # The local checkout just overwrote the patched miles sources; re-apply + # the built-in patches. Keep custom checkouts usable when their layout has + # intentionally diverged, matching the existing local-overlay behavior. + return image.run_commands( + f"echo {_PATCH_SGLANG_ABORT_B64} | base64 -d | python3" + " || echo 'WARNING: sglang abort patch did not apply to the" + " local_miles checkout; transient router failures during rollout" + " cleanup may crash the run'", + f"echo {_PATCH_MOONCAKE_TOLERANCE_B64} | base64 -d | python3" + " || echo 'WARNING: mooncake import tolerance patch did not apply to" + " the local_miles checkout; EFA-host actor imports may fail'", + f"echo {_PATCH_ROUTER_STARTUP_TIMEOUT_B64} | base64 -d | python3" + " || echo 'WARNING: router startup timeout patch did not apply to the" + " local_miles checkout; busy routers retain the upstream timeout'", + *_REPORTING_PATCH_COMMANDS, + ) + + def _response_parser_path(model: Any) -> str: """Import path of the model's response parser so the rollout recorder can resolve and apply it remotely. Empty when the model sets no parser.""" @@ -381,7 +428,13 @@ def _response_parser_path(model: Any) -> str: def _compose_ld_library_path() -> str: - parts = [SYSTEM_LIB_DIR] + # The injected EFA dirs must precede SYSTEM_LIB_DIR: NCCL's ofi plugin + # requires the injected libfabric 1.30 (`FABRIC_1.8`), and with the system + # dir first the loader finds the image's libfabric 1.20 instead, so the + # plugin never loads in a Ray worker — fatal at engine bring-up on hosts + # whose env pins NCCL_NET_PLUGIN=ofi. + parts = list(_EFA_LIB_DIRS) + parts.append(SYSTEM_LIB_DIR) for part in os.environ.get("LD_LIBRARY_PATH", "").split(":"): if part and part not in parts: parts.append(part) @@ -401,8 +454,9 @@ def build_ray_runtime_env( Ray workers do not pick up the container's linker path on their own, and without it the Megatron actor can resolve a libibverbs that does not match - the image's libmlx5 and die importing mooncake. The system lib dir is put - in front for that reason; the rest is read from the container, so whatever + the image's libmlx5 and die importing mooncake. The leading dirs are + ordered by ``_compose_ld_library_path``; the rest is read from the + container, so whatever the image exports — including any wheel-shipped nvidia lib dirs — is carried through. Composing it here rather than in an ``image_env`` entry keeps it independent of whether the base image exports ``LD_LIBRARY_PATH`` @@ -453,21 +507,7 @@ def build_miles_app( ) if miles.local_miles: - image = image.add_local_dir( - miles.local_miles, - remote_path=MILES_ROOT, - copy=True, - ignore=["**/__pycache__", "**/*.pyc", "**/.git", "**/.venv"], - ) - # The local checkout just overwrote the patched miles sources; - # re-apply the built-in patches. - image = image.run_commands( - f"echo {_PATCH_SGLANG_ABORT_B64} | base64 -d | python3" - " || echo 'WARNING: sglang abort patch did not apply to the" - " local_miles checkout; transient router failures during rollout" - " cleanup may crash the run'", - *_REPORTING_PATCH_COMMANDS, - ) + image = _overlay_local_miles(image, miles.local_miles) if miles.image_run_commands: image = image.run_commands(*miles.image_run_commands) diff --git a/modal_training_gym/frameworks/miles/modal_helpers/patches/patch_mooncake_import_tolerance.py b/modal_training_gym/frameworks/miles/modal_helpers/patches/patch_mooncake_import_tolerance.py new file mode 100644 index 000000000..4e6ac8074 --- /dev/null +++ b/modal_training_gym/frameworks/miles/modal_helpers/patches/patch_mooncake_import_tolerance.py @@ -0,0 +1,69 @@ +"""Import mooncake's TransferEngine only when Miles uses P2P weight transfer. + +On some Modal EFA hosts the runtime bind-mounts the host's libibverbs over +the system path, and its private ABI (IBVERBS_PRIVATE_*) does not match the +image's libmlx5 — mooncake's TransferEngine import raises. The import sits at +module level on miles' actor import chain, so it kills training runs that +never use it (colocated weight sync runs on Ray IPC; TransferEngine only +backs the p2p update_weight_from_distributed path). + +Trying and catching that import is not safe: the failed dlopen can leave +glibc's TLS bookkeeping partially mutated, and a later CUDA/NCCL load aborts +in ``_dl_allocate_tls_init``. This patch therefore removes the module-level +import entirely and performs it only inside the P2P setup function. Colocated +runs never touch mooncake; an actual P2P run gets a clear error at point of use. + +Executed at image-build time via ``python3 ``. +""" + +import pathlib + +MARKER = "PATCHED_MOONCAKE_IMPORT_TOLERANCE" + +TARGET = pathlib.Path( + "/root/miles/miles/backends/megatron_utils/update_weight/" + "update_weight_from_distributed/p2p_transfer_utils.py" +) + +OLD_IMPORT = "from mooncake.engine import TransferEngine" +NEW_IMPORT = "# " + MARKER + ": imported lazily by the P2P setup below" + +OLD_USE = " transfer_engine = TransferEngine()" +NEW_USE = ( + " try:\n" + " from mooncake.engine import TransferEngine\n" + " except (ImportError, OSError) as exc:\n" + " raise RuntimeError(\n" + ' "p2p weight transfer requires mooncake\'s TransferEngine, "\n' + ' "which failed to import on this host (verbs stack mismatch)."\n' + " ) from exc\n" + " transfer_engine = TransferEngine()" +) + + +def _patch_file(target: pathlib.Path) -> None: + if not target.exists(): + raise SystemExit( + f"{target} not found; miles layout changed — re-check the patch." + ) + + src = target.read_text() + if MARKER in src: + print("mooncake import tolerance patch already applied") + return + + if OLD_IMPORT not in src or OLD_USE not in src: + raise SystemExit( + "mooncake import tolerance patch did not match; miles' " + "p2p_transfer_utils.py has changed. Re-check the import and the " + "TransferEngine() call site before shipping." + ) + + src = src.replace(OLD_IMPORT, NEW_IMPORT, 1) + src = src.replace(OLD_USE, NEW_USE, 1) + target.write_text(src) + print("Patched mooncake TransferEngine import to load lazily for P2P only") + + +if __name__ == "__main__": + _patch_file(TARGET) diff --git a/tests/test_miles_patches.py b/tests/test_miles_patches.py index 76191999f..3d83756b6 100644 --- a/tests/test_miles_patches.py +++ b/tests/test_miles_patches.py @@ -2,12 +2,14 @@ from __future__ import annotations +import ast from pathlib import Path import pytest from modal_training_gym.frameworks.miles.modal_helpers.patches import ( patch_advantage_distribution as advantage_patcher, + patch_mooncake_import_tolerance as mooncake_patcher, patch_rollout_status_reporting as rollout_patcher, ) @@ -72,3 +74,36 @@ def test_patch_matches_golden(miles_inputs, tmp_path, request): assert actual == expected, ( f"golden mismatch for {name}; rerun with --rewrite to accept" ) + + +def test_mooncake_import_is_moved_inside_p2p_setup(tmp_path): + work = tmp_path / "p2p_transfer_utils.py" + work.write_text( + "from mooncake.engine import TransferEngine\n\n" + "def setup_transfer_engine():\n" + " transfer_engine = TransferEngine()\n" + " return transfer_engine\n" + ) + + mooncake_patcher._patch_file(work) + patched = work.read_text() + tree = ast.parse(patched) + + assert mooncake_patcher.MARKER in patched + assert not any( + isinstance(node, ast.ImportFrom) and node.module == "mooncake.engine" + for node in tree.body + ) + setup = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "setup_transfer_engine" + ) + assert any( + isinstance(node, ast.ImportFrom) and node.module == "mooncake.engine" + for node in ast.walk(setup) + ) + + # Image builds can layer the same patch more than once. + mooncake_patcher._patch_file(work) + assert work.read_text() == patched diff --git a/tests/test_miles_runtime_env.py b/tests/test_miles_runtime_env.py index eb66a068d..21c283556 100644 --- a/tests/test_miles_runtime_env.py +++ b/tests/test_miles_runtime_env.py @@ -10,6 +10,7 @@ class _FakeImage: def __init__(self): self.commands: list[str] = [] + self.operations: list[str] = [] @classmethod def from_registry(cls, _docker_image: str) -> "_FakeImage": @@ -19,9 +20,14 @@ def entrypoint(self, _entrypoint: list[str]) -> "_FakeImage": return self def run_commands(self, *commands: str) -> "_FakeImage": + self.operations.append("run_commands") self.commands.extend(commands) return self + def add_local_dir(self, *_args, **_kwargs) -> "_FakeImage": + self.operations.append("add_local_dir") + return self + def env(self, _environment: dict[str, str]) -> "_FakeImage": return self @@ -43,6 +49,28 @@ def test_single_node_image_keeps_base_rdma_runtime(monkeypatch): assert launcher.RDMA_RUNTIME_INSTALL_COMMAND not in image.commands +def test_image_applies_efa_host_patches(monkeypatch): + """The mooncake-import-tolerance and router-timeout patches, both needed + for a miles run to complete on an EFA host, are baked into the image.""" + monkeypatch.setattr(launcher, "Image", _FakeImage) + + commands = "\n".join(launcher._build_miles_base_image(MilesRecipe()).commands) + + assert launcher._PATCH_MOONCAKE_TOLERANCE_B64 in commands + assert launcher._PATCH_ROUTER_STARTUP_TIMEOUT_B64 in commands + + +def test_local_miles_overlay_reapplies_efa_host_patches(): + image = _FakeImage() + + launcher._overlay_local_miles(image, "/tmp/local-miles") + + assert image.operations == ["add_local_dir", "run_commands"] + commands = "\n".join(image.commands) + assert launcher._PATCH_MOONCAKE_TOLERANCE_B64 in commands + assert launcher._PATCH_ROUTER_STARTUP_TIMEOUT_B64 in commands + + def test_ld_library_path_comes_from_the_container(monkeypatch): """Workers get the container's linker path, behind the system lib dir.""" monkeypatch.setenv("LD_LIBRARY_PATH", "/usr/local/cuda/lib64:/wheel/nvidia/lib") @@ -52,7 +80,8 @@ def test_ld_library_path_comes_from_the_container(monkeypatch): )["env_vars"] assert env_vars["LD_LIBRARY_PATH"] == ( - "/usr/lib/x86_64-linux-gnu:/usr/local/cuda/lib64:/wheel/nvidia/lib" + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" + ":/usr/lib/x86_64-linux-gnu:/usr/local/cuda/lib64:/wheel/nvidia/lib" ) assert env_vars["MASTER_ADDR"] == "10.0.0.1" assert env_vars["no_proxy"] == "127.0.0.1,10.0.0.1" @@ -74,7 +103,7 @@ def test_recipe_environment_still_wins(monkeypatch): assert env_vars["PYTHONPATH"] == "/root/Megatron-LM/" -def test_unset_container_path_yields_only_the_system_lib_dir(monkeypatch): +def test_unset_container_path_yields_only_the_required_lib_dirs(monkeypatch): """No empty entry, which the loader would read as the working directory.""" monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) @@ -82,7 +111,9 @@ def test_unset_container_path_yields_only_the_system_lib_dir(monkeypatch): head_addr="10.0.0.1", metric_env={}, environment={} )["env_vars"] - assert env_vars["LD_LIBRARY_PATH"] == "/usr/lib/x86_64-linux-gnu" + assert env_vars["LD_LIBRARY_PATH"] == ( + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib:/usr/lib/x86_64-linux-gnu" + ) def test_system_lib_dir_is_not_duplicated(monkeypatch): @@ -93,7 +124,8 @@ def test_system_lib_dir_is_not_duplicated(monkeypatch): )["env_vars"] assert env_vars["LD_LIBRARY_PATH"] == ( - "/usr/lib/x86_64-linux-gnu:/wheel/nvidia/lib" + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" + ":/usr/lib/x86_64-linux-gnu:/wheel/nvidia/lib" ) @@ -108,3 +140,19 @@ def test_metric_env_is_preserved(monkeypatch): assert env_vars["WANDB_RUN_ID"] == "abc" assert env_vars["WANDB_RESUME"] == "allow" + + +def test_efa_dirs_lead_the_path(monkeypatch): + """The injected EFA dirs come first so NCCL's ofi plugin resolves the + injected libfabric, not the image's older one.""" + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + + env_vars = build_ray_runtime_env( + head_addr="10.0.0.1", metric_env={}, environment={} + )["env_vars"] + + path = env_vars["LD_LIBRARY_PATH"].split(":") + assert path[: len(launcher._EFA_LIB_DIRS)] == list(launcher._EFA_LIB_DIRS) + assert launcher.SYSTEM_LIB_DIR in path + # The private-prefix machinery is gone: nothing shadows the host verbs libs. + assert not any("gym-rdma" in p for p in path)