From 74f3364d4887953f587b4dca05e3b9bb99b507ce Mon Sep 17 00:00:00 2001 From: zhouhelena1 Date: Tue, 1 Sep 2026 15:23:47 -0400 Subject: [PATCH 1/7] Fix miles multi-node NCCL and verbs library resolution Two root-caused library-resolution bugs in the Ray worker environment, affecting every multi-node miles run: - LD_LIBRARY_PATH put the system lib dir first, shadowing Modal's injected EFA stack (libfabric 1.30 + ofi NCCL plugin): fatal "Failed to initialize any NET plugin" on hosts pinning NCCL_NET_PLUGIN=ofi, and a silent slow-transport fallback elsewhere (~20x slower rollout generation measured). EFA dirs now precede the system dir. - On EFA hosts Modal bind-mounts the host's libibverbs over the system path, breaking the IBVERBS_PRIVATE_* coupling with the image's libmlx5, so mooncake's TransferEngine import fails. The image now snapshots its matched verbs set to /opt/gym-rdma/lib, and each node probes its system pair once: the prefix enters the worker search path only where the pair is broken. Unconditional inclusion dual-loads the verbs libs on healthy hosts and degrades NCCL IB 20-50x. Co-Authored-By: Claude Fable 5 --- .../frameworks/miles/launcher.py | 92 ++++++++++++++++++- 1 file changed, 87 insertions(+), 5 deletions(-) diff --git a/modal_training_gym/frameworks/miles/launcher.py b/modal_training_gym/frameworks/miles/launcher.py index 1f12574c7..9b6af9bfc 100644 --- a/modal_training_gym/frameworks/miles/launcher.py +++ b/modal_training_gym/frameworks/miles/launcher.py @@ -4,6 +4,7 @@ import os import shlex import subprocess +import sys import shutil import tempfile import threading @@ -30,7 +31,6 @@ timing_debug_env, ) from modal_training_gym.common.metrics import ( - apply_metric_image, metric_metadata, metric_runtime_env, metric_secrets, @@ -97,6 +97,12 @@ 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 +110,62 @@ def _validate_resume_checkpoint( "--reinstall libibverbs1 ibverbs-providers && " "rm -rf /var/lib/apt/lists/*" ) + +# On EFA hosts the Modal runtime bind-mounts the *host's* libibverbs.so.1 and +# libefa.so.1 over the system paths. libibverbs and its provider libraries are +# coupled through private symbol versions (IBVERBS_PRIVATE_*), so a host +# libibverbs older than the image's libmlx5 breaks every verbs consumer — +# mooncake's TransferEngine import died with "version 'IBVERBS_PRIVATE_34' not +# found" on exactly the attempts that drew EFA hosts. A bind mount cannot be +# replaced by apt or shadowed at its own path, so the image's matched verbs +# set is copied to a private prefix at build time and Ray workers resolve it +# from there; everything outside the set (libfabric included) links only the +# public IBVERBS_1.x ABI and is unaffected. +# +# The prefix is engaged only when the system pair is actually broken: with the +# prefix on the path, absolute-path loads (the provider loader) still map the +# system copies, so a healthy Mellanox node ends up with two live instances of +# libibverbs/libmlx5 and NCCL's IB data path degrades badly (a weight sync +# measured 20-50x slower). Where the system pair is intact — every host that +# mounts nothing — the prefix must stay out of the path. +_GYM_RDMA_DIR = "/opt/gym-rdma/lib" +GYM_RDMA_COPY_COMMAND = ( + f"mkdir -p {_GYM_RDMA_DIR} && " + f"cp -a {SYSTEM_LIB_DIR}/libibverbs.so.1* {SYSTEM_LIB_DIR}/libmlx5.so.1* " + f"{_GYM_RDMA_DIR}/ && " + f"if ls {SYSTEM_LIB_DIR}/libefa.so.1* >/dev/null 2>&1; then " + f"cp -a {SYSTEM_LIB_DIR}/libefa.so.1* {_GYM_RDMA_DIR}/; fi && " + f"if test -d {SYSTEM_LIB_DIR}/libibverbs; then " + f"cp -a {SYSTEM_LIB_DIR}/libibverbs {_GYM_RDMA_DIR}/; fi" +) + + +_system_verbs_broken_cache: bool | None = None + + +def _system_verbs_broken() -> bool: + """Whether this node's system verbs pair fails mooncake's import. + + Probed once per container under the system search order (no prefix). True + means the host bind-mounted a foreign libibverbs over the image's — the + only case where the private prefix should enter the search path. + """ + global _system_verbs_broken_cache + if _system_verbs_broken_cache is not None: + return _system_verbs_broken_cache + probe = [sys.executable, "-c", "from mooncake.engine import TransferEngine"] + check = subprocess.run(probe, capture_output=True, text=True) + broken = check.returncode != 0 and "IBVERBS_PRIVATE" in check.stderr + if broken: + print( + "WARNING: this node's system verbs pair is mismatched " + f"({check.stderr.strip().splitlines()[-1]}); resolving the image's " + f"matched set from {_GYM_RDMA_DIR} for Ray workers." + ) + _system_verbs_broken_cache = broken + return broken + + # 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. @@ -363,7 +425,9 @@ def _build_miles_base_image(miles: MilesRecipe) -> Image: ) ) if miles.total_nodes > 1: - image = image.run_commands(RDMA_RUNTIME_INSTALL_COMMAND) + # The copy must follow the reinstall so the private prefix snapshots + # the freshly matched verbs set. + image = image.run_commands(RDMA_RUNTIME_INSTALL_COMMAND, GYM_RDMA_COPY_COMMAND) if miles.image_env: image = image.env(miles.image_env) return image @@ -381,7 +445,23 @@ def _response_parser_path(model: Any) -> str: def _compose_ld_library_path() -> str: - parts = [SYSTEM_LIB_DIR] + # Ordering carries two constraints: + # - 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. + # - The private rdma prefix precedes SYSTEM_LIB_DIR only on nodes whose + # system verbs pair is broken (an EFA host's bind-mounted libibverbs + # breaking the private-ABI coupling with the image's providers — see + # GYM_RDMA_COPY_COMMAND). On healthy nodes it must stay out: it would + # dual-load libibverbs/libmlx5 (path lookups hit the prefix, the + # provider loader's absolute paths hit the system copies) and degrade + # NCCL's IB data path badly. + parts = [d for d in _EFA_LIB_DIRS if os.path.isdir(d)] + if os.path.isdir(_GYM_RDMA_DIR) and _system_verbs_broken(): + parts.append(_GYM_RDMA_DIR) + 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) @@ -479,7 +559,6 @@ def build_miles_app( if isinstance(dataset, HarborDataset): image = image.uv_pip_install(f"harbor=={HARBOR_PKG_VERSION}") - image = apply_metric_image(image, miles.metrics) image = image.add_local_python_source("modal_training_gym", copy=True) image = image.uv_pip_install("randomname") image = mount_tools_dir(image) @@ -913,7 +992,6 @@ def convert_checkpoint( image=image, gpu=gpu_spec, memory=miles.memory, - cpu=miles.cpu, ephemeral_disk=train_ephemeral_disk, cloud=miles.cloud, region=miles.region, @@ -940,6 +1018,10 @@ async def train( if framework_status_token: os.environ["TRAINING_GYM_FRAMEWORK_STATUS_TOKEN"] = framework_status_token + # Decide (and log) this node's verbs resolution before anything reads + # the composed library path. + _system_verbs_broken() + await asyncio.gather( hf_cache_volume.reload.aio(), data_volume.reload.aio(), From cd9bb6d2f0aede6e1c66075cc33fb703a7238e0f Mon Sep 17 00:00:00 2001 From: zhouhelena1 Date: Tue, 1 Sep 2026 16:29:09 -0400 Subject: [PATCH 2/7] Make the rdma prefix decision per node, and restore clobbered main lines Ray ships one env_vars mapping to every node, so the head's probe decided verbs resolution cluster-wide. The path now carries a fixed alias (/opt/gym-rdma/enabled) that each node symlinks to the copy only after its own probe fails; on healthy nodes the alias does not exist and the loader skips it. Also rebuilds the launcher change on top of current main, restoring apply_metric_image and cpu=miles.cpu, which the previous revision clobbered by starting from an older copy of the file. Co-Authored-By: Claude Fable 5 --- .../frameworks/miles/launcher.py | 91 +++++++++---------- tests/test_miles_runtime_env.py | 9 +- 2 files changed, 48 insertions(+), 52 deletions(-) diff --git a/modal_training_gym/frameworks/miles/launcher.py b/modal_training_gym/frameworks/miles/launcher.py index 9b6af9bfc..1a5edf6d9 100644 --- a/modal_training_gym/frameworks/miles/launcher.py +++ b/modal_training_gym/frameworks/miles/launcher.py @@ -31,6 +31,7 @@ timing_debug_env, ) from modal_training_gym.common.metrics import ( + apply_metric_image, metric_metadata, metric_runtime_env, metric_secrets, @@ -102,7 +103,6 @@ def _validate_resume_checkpoint( # 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 = ( @@ -115,20 +115,21 @@ def _validate_resume_checkpoint( # libefa.so.1 over the system paths. libibverbs and its provider libraries are # coupled through private symbol versions (IBVERBS_PRIVATE_*), so a host # libibverbs older than the image's libmlx5 breaks every verbs consumer — -# mooncake's TransferEngine import died with "version 'IBVERBS_PRIVATE_34' not -# found" on exactly the attempts that drew EFA hosts. A bind mount cannot be -# replaced by apt or shadowed at its own path, so the image's matched verbs -# set is copied to a private prefix at build time and Ray workers resolve it -# from there; everything outside the set (libfabric included) links only the -# public IBVERBS_1.x ABI and is unaffected. +# mooncake's TransferEngine import dies with "version 'IBVERBS_PRIVATE_34' not +# found". A bind mount cannot be replaced by apt or shadowed at its own path, +# so the image's matched verbs set is copied to a private prefix at build time +# and Ray workers resolve it from there. # -# The prefix is engaged only when the system pair is actually broken: with the -# prefix on the path, absolute-path loads (the provider loader) still map the -# system copies, so a healthy Mellanox node ends up with two live instances of -# libibverbs/libmlx5 and NCCL's IB data path degrades badly (a weight sync -# measured 20-50x slower). Where the system pair is intact — every host that -# mounts nothing — the prefix must stay out of the path. +# The prefix is engaged per node, and only where the system pair is actually +# broken: with the prefix on the path, absolute-path loads (the provider +# loader) still map the system copies, so a healthy node ends up with two live +# instances of libibverbs/libmlx5 and NCCL's IB data path degrades badly (a +# weight sync measured 20-50x slower). Ray distributes one env_vars mapping to +# every worker, so the path carries a fixed alias that each node materializes +# (as a symlink to the copy) only after its own probe fails — on healthy nodes +# the alias does not exist and the loader skips it. _GYM_RDMA_DIR = "/opt/gym-rdma/lib" +_GYM_RDMA_ENABLED_DIR = "/opt/gym-rdma/enabled" GYM_RDMA_COPY_COMMAND = ( f"mkdir -p {_GYM_RDMA_DIR} && " f"cp -a {SYSTEM_LIB_DIR}/libibverbs.so.1* {SYSTEM_LIB_DIR}/libmlx5.so.1* " @@ -140,30 +141,24 @@ def _validate_resume_checkpoint( ) -_system_verbs_broken_cache: bool | None = None - - -def _system_verbs_broken() -> bool: - """Whether this node's system verbs pair fails mooncake's import. - - Probed once per container under the system search order (no prefix). True - means the host bind-mounted a foreign libibverbs over the image's — the - only case where the private prefix should enter the search path. - """ - global _system_verbs_broken_cache - if _system_verbs_broken_cache is not None: - return _system_verbs_broken_cache +def _enable_rdma_prefix_if_broken() -> None: + """Point this node's prefix alias at the matched verbs copy when the + system pair fails mooncake's import (the EFA bind-mount case).""" + if not os.path.isdir(_GYM_RDMA_DIR) or os.path.lexists(_GYM_RDMA_ENABLED_DIR): + return probe = [sys.executable, "-c", "from mooncake.engine import TransferEngine"] check = subprocess.run(probe, capture_output=True, text=True) - broken = check.returncode != 0 and "IBVERBS_PRIVATE" in check.stderr - if broken: - print( - "WARNING: this node's system verbs pair is mismatched " - f"({check.stderr.strip().splitlines()[-1]}); resolving the image's " - f"matched set from {_GYM_RDMA_DIR} for Ray workers." - ) - _system_verbs_broken_cache = broken - return broken + if check.returncode == 0 or "IBVERBS_PRIVATE" not in check.stderr: + return + print( + "WARNING: this node's system verbs pair is mismatched " + f"({check.stderr.strip().splitlines()[-1]}); resolving the image's " + f"matched set from {_GYM_RDMA_DIR} for Ray workers." + ) + try: + os.symlink(_GYM_RDMA_DIR, _GYM_RDMA_ENABLED_DIR) + except FileExistsError: + pass # v0.8.0+ makes per-task CPU/memory requests configurable via enforcement @@ -451,16 +446,13 @@ def _compose_ld_library_path() -> str: # 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. - # - The private rdma prefix precedes SYSTEM_LIB_DIR only on nodes whose - # system verbs pair is broken (an EFA host's bind-mounted libibverbs - # breaking the private-ABI coupling with the image's providers — see - # GYM_RDMA_COPY_COMMAND). On healthy nodes it must stay out: it would - # dual-load libibverbs/libmlx5 (path lookups hit the prefix, the - # provider loader's absolute paths hit the system copies) and degrade - # NCCL's IB data path badly. + # - The private rdma prefix alias precedes SYSTEM_LIB_DIR; it resolves + # only on nodes that materialized it because their system verbs pair is + # broken (see _enable_rdma_prefix_if_broken). This path is composed once + # on the head and shipped to every node's workers, so the decision has + # to live in the per-node filesystem, not here. parts = [d for d in _EFA_LIB_DIRS if os.path.isdir(d)] - if os.path.isdir(_GYM_RDMA_DIR) and _system_verbs_broken(): - parts.append(_GYM_RDMA_DIR) + parts.append(_GYM_RDMA_ENABLED_DIR) parts.append(SYSTEM_LIB_DIR) for part in os.environ.get("LD_LIBRARY_PATH", "").split(":"): if part and part not in parts: @@ -481,8 +473,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`` @@ -559,6 +552,7 @@ def build_miles_app( if isinstance(dataset, HarborDataset): image = image.uv_pip_install(f"harbor=={HARBOR_PKG_VERSION}") + image = apply_metric_image(image, miles.metrics) image = image.add_local_python_source("modal_training_gym", copy=True) image = image.uv_pip_install("randomname") image = mount_tools_dir(image) @@ -992,6 +986,7 @@ def convert_checkpoint( image=image, gpu=gpu_spec, memory=miles.memory, + cpu=miles.cpu, ephemeral_disk=train_ephemeral_disk, cloud=miles.cloud, region=miles.region, @@ -1018,9 +1013,7 @@ async def train( if framework_status_token: os.environ["TRAINING_GYM_FRAMEWORK_STATUS_TOKEN"] = framework_status_token - # Decide (and log) this node's verbs resolution before anything reads - # the composed library path. - _system_verbs_broken() + _enable_rdma_prefix_if_broken() await asyncio.gather( hf_cache_volume.reload.aio(), diff --git a/tests/test_miles_runtime_env.py b/tests/test_miles_runtime_env.py index eb66a068d..25a05ab19 100644 --- a/tests/test_miles_runtime_env.py +++ b/tests/test_miles_runtime_env.py @@ -52,7 +52,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/gym-rdma/enabled:/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" @@ -82,7 +83,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/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" + ) def test_system_lib_dir_is_not_duplicated(monkeypatch): @@ -93,7 +96,7 @@ 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/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu:/wheel/nvidia/lib" ) From 54a3f01920bbdebebb2b6313f96cf378d85f9ec7 Mon Sep 17 00:00:00 2001 From: zhouhelena1 Date: Tue, 1 Sep 2026 16:35:52 -0400 Subject: [PATCH 3/7] Test EFA precedence and the probe-driven prefix alias Co-Authored-By: Claude Fable 5 --- tests/test_miles_runtime_env.py | 54 +++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/tests/test_miles_runtime_env.py b/tests/test_miles_runtime_env.py index 25a05ab19..001f54740 100644 --- a/tests/test_miles_runtime_env.py +++ b/tests/test_miles_runtime_env.py @@ -2,6 +2,8 @@ from __future__ import annotations +import subprocess + from modal_training_gym.frameworks.miles import launcher from modal_training_gym.frameworks.miles.launcher import build_ray_runtime_env from modal_training_gym.train_recipes.miles_recipe import MilesRecipe @@ -111,3 +113,55 @@ def test_metric_env_is_preserved(monkeypatch): assert env_vars["WANDB_RUN_ID"] == "abc" assert env_vars["WANDB_RESUME"] == "allow" + + +def test_efa_dirs_precede_the_system_lib_dir(monkeypatch): + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + monkeypatch.setattr( + launcher.os.path, "isdir", lambda d: d in launcher._EFA_LIB_DIRS + ) + + env_vars = build_ray_runtime_env( + head_addr="10.0.0.1", metric_env={}, environment={} + )["env_vars"] + + assert env_vars["LD_LIBRARY_PATH"] == ( + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" + ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" + ) + + +def _probe(monkeypatch, returncode: int, stderr: str) -> list[tuple[str, str]]: + created: list[tuple[str, str]] = [] + monkeypatch.setattr( + launcher.os.path, "isdir", lambda d: d == launcher._GYM_RDMA_DIR + ) + monkeypatch.setattr(launcher.os.path, "lexists", lambda d: False) + monkeypatch.setattr( + launcher.subprocess, + "run", + lambda *a, **k: subprocess.CompletedProcess(a, returncode, "", stderr), + ) + monkeypatch.setattr( + launcher.os, "symlink", lambda src, dst: created.append((src, dst)) + ) + launcher._enable_rdma_prefix_if_broken() + return created + + +def test_broken_verbs_pair_materializes_the_prefix_alias(monkeypatch): + created = _probe( + monkeypatch, 1, "ImportError: version `IBVERBS_PRIVATE_34' not found" + ) + + assert created == [(launcher._GYM_RDMA_DIR, launcher._GYM_RDMA_ENABLED_DIR)] + + +def test_healthy_verbs_pair_leaves_the_alias_absent(monkeypatch): + assert _probe(monkeypatch, 0, "") == [] + + +def test_unrelated_probe_failure_leaves_the_alias_absent(monkeypatch): + assert ( + _probe(monkeypatch, 1, "ModuleNotFoundError: No module named 'mooncake'") == [] + ) From 06c4c07082f1f5e47db6ae43070119db89522de1 Mon Sep 17 00:00:00 2001 From: zhouhelena1 Date: Tue, 1 Sep 2026 17:04:38 -0400 Subject: [PATCH 4/7] Keep EFA runtime paths independent of Ray head --- modal_training_gym/frameworks/miles/launcher.py | 2 +- tests/test_miles_runtime_env.py | 17 +++++++++-------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/modal_training_gym/frameworks/miles/launcher.py b/modal_training_gym/frameworks/miles/launcher.py index 1a5edf6d9..92936ac53 100644 --- a/modal_training_gym/frameworks/miles/launcher.py +++ b/modal_training_gym/frameworks/miles/launcher.py @@ -451,7 +451,7 @@ def _compose_ld_library_path() -> str: # broken (see _enable_rdma_prefix_if_broken). This path is composed once # on the head and shipped to every node's workers, so the decision has # to live in the per-node filesystem, not here. - parts = [d for d in _EFA_LIB_DIRS if os.path.isdir(d)] + parts = list(_EFA_LIB_DIRS) parts.append(_GYM_RDMA_ENABLED_DIR) parts.append(SYSTEM_LIB_DIR) for part in os.environ.get("LD_LIBRARY_PATH", "").split(":"): diff --git a/tests/test_miles_runtime_env.py b/tests/test_miles_runtime_env.py index 001f54740..8cd5e9bf7 100644 --- a/tests/test_miles_runtime_env.py +++ b/tests/test_miles_runtime_env.py @@ -54,7 +54,8 @@ def test_ld_library_path_comes_from_the_container(monkeypatch): )["env_vars"] assert env_vars["LD_LIBRARY_PATH"] == ( - "/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" + ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" ":/usr/local/cuda/lib64:/wheel/nvidia/lib" ) assert env_vars["MASTER_ADDR"] == "10.0.0.1" @@ -77,7 +78,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) @@ -86,7 +87,8 @@ def test_unset_container_path_yields_only_the_system_lib_dir(monkeypatch): )["env_vars"] assert env_vars["LD_LIBRARY_PATH"] == ( - "/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" + ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" ) @@ -98,7 +100,8 @@ def test_system_lib_dir_is_not_duplicated(monkeypatch): )["env_vars"] assert env_vars["LD_LIBRARY_PATH"] == ( - "/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu:/wheel/nvidia/lib" + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" + ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu:/wheel/nvidia/lib" ) @@ -115,11 +118,9 @@ def test_metric_env_is_preserved(monkeypatch): assert env_vars["WANDB_RESUME"] == "allow" -def test_efa_dirs_precede_the_system_lib_dir(monkeypatch): +def test_efa_dirs_precede_the_system_lib_dir_when_absent_on_head(monkeypatch): monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) - monkeypatch.setattr( - launcher.os.path, "isdir", lambda d: d in launcher._EFA_LIB_DIRS - ) + monkeypatch.setattr(launcher.os.path, "isdir", lambda _d: False) env_vars = build_ray_runtime_env( head_addr="10.0.0.1", metric_env={}, environment={} From 628741c28572c46330bf15446e39c403d8c69d17 Mon Sep 17 00:00:00 2001 From: zhouhelena1 Date: Tue, 1 Sep 2026 23:39:17 -0400 Subject: [PATCH 5/7] Keep AWS EFA libraries authoritative for Miles --- .../frameworks/miles/launcher.py | 88 +++++-------------- .../patch_mooncake_import_tolerance.py | 66 ++++++++++++++ tests/test_miles_patches.py | 35 ++++++++ tests/test_miles_runtime_env.py | 71 +++++---------- 4 files changed, 144 insertions(+), 116 deletions(-) create mode 100644 modal_training_gym/frameworks/miles/modal_helpers/patches/patch_mooncake_import_tolerance.py diff --git a/modal_training_gym/frameworks/miles/launcher.py b/modal_training_gym/frameworks/miles/launcher.py index 92936ac53..0cdecb8a1 100644 --- a/modal_training_gym/frameworks/miles/launcher.py +++ b/modal_training_gym/frameworks/miles/launcher.py @@ -4,7 +4,6 @@ import os import shlex import subprocess -import sys import shutil import tempfile import threading @@ -111,55 +110,6 @@ def _validate_resume_checkpoint( "rm -rf /var/lib/apt/lists/*" ) -# On EFA hosts the Modal runtime bind-mounts the *host's* libibverbs.so.1 and -# libefa.so.1 over the system paths. libibverbs and its provider libraries are -# coupled through private symbol versions (IBVERBS_PRIVATE_*), so a host -# libibverbs older than the image's libmlx5 breaks every verbs consumer — -# mooncake's TransferEngine import dies with "version 'IBVERBS_PRIVATE_34' not -# found". A bind mount cannot be replaced by apt or shadowed at its own path, -# so the image's matched verbs set is copied to a private prefix at build time -# and Ray workers resolve it from there. -# -# The prefix is engaged per node, and only where the system pair is actually -# broken: with the prefix on the path, absolute-path loads (the provider -# loader) still map the system copies, so a healthy node ends up with two live -# instances of libibverbs/libmlx5 and NCCL's IB data path degrades badly (a -# weight sync measured 20-50x slower). Ray distributes one env_vars mapping to -# every worker, so the path carries a fixed alias that each node materializes -# (as a symlink to the copy) only after its own probe fails — on healthy nodes -# the alias does not exist and the loader skips it. -_GYM_RDMA_DIR = "/opt/gym-rdma/lib" -_GYM_RDMA_ENABLED_DIR = "/opt/gym-rdma/enabled" -GYM_RDMA_COPY_COMMAND = ( - f"mkdir -p {_GYM_RDMA_DIR} && " - f"cp -a {SYSTEM_LIB_DIR}/libibverbs.so.1* {SYSTEM_LIB_DIR}/libmlx5.so.1* " - f"{_GYM_RDMA_DIR}/ && " - f"if ls {SYSTEM_LIB_DIR}/libefa.so.1* >/dev/null 2>&1; then " - f"cp -a {SYSTEM_LIB_DIR}/libefa.so.1* {_GYM_RDMA_DIR}/; fi && " - f"if test -d {SYSTEM_LIB_DIR}/libibverbs; then " - f"cp -a {SYSTEM_LIB_DIR}/libibverbs {_GYM_RDMA_DIR}/; fi" -) - - -def _enable_rdma_prefix_if_broken() -> None: - """Point this node's prefix alias at the matched verbs copy when the - system pair fails mooncake's import (the EFA bind-mount case).""" - if not os.path.isdir(_GYM_RDMA_DIR) or os.path.lexists(_GYM_RDMA_ENABLED_DIR): - return - probe = [sys.executable, "-c", "from mooncake.engine import TransferEngine"] - check = subprocess.run(probe, capture_output=True, text=True) - if check.returncode == 0 or "IBVERBS_PRIVATE" not in check.stderr: - return - print( - "WARNING: this node's system verbs pair is mismatched " - f"({check.stderr.strip().splitlines()[-1]}); resolving the image's " - f"matched set from {_GYM_RDMA_DIR} for Ray workers." - ) - try: - os.symlink(_GYM_RDMA_DIR, _GYM_RDMA_ENABLED_DIR) - except FileExistsError: - pass - # v0.8.0+ makes per-task CPU/memory requests configurable via enforcement # policies ("limit"/"ignore"), letting sandboxes burst on Modal and bill by @@ -168,6 +118,12 @@ def _enable_rdma_prefix_if_broken() -> None: _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 ) @@ -408,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 " @@ -420,9 +385,7 @@ def _build_miles_base_image(miles: MilesRecipe) -> Image: ) ) if miles.total_nodes > 1: - # The copy must follow the reinstall so the private prefix snapshots - # the freshly matched verbs set. - image = image.run_commands(RDMA_RUNTIME_INSTALL_COMMAND, GYM_RDMA_COPY_COMMAND) + image = image.run_commands(RDMA_RUNTIME_INSTALL_COMMAND) if miles.image_env: image = image.env(miles.image_env) return image @@ -440,19 +403,12 @@ def _response_parser_path(model: Any) -> str: def _compose_ld_library_path() -> str: - # Ordering carries two constraints: - # - 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. - # - The private rdma prefix alias precedes SYSTEM_LIB_DIR; it resolves - # only on nodes that materialized it because their system verbs pair is - # broken (see _enable_rdma_prefix_if_broken). This path is composed once - # on the head and shipped to every node's workers, so the decision has - # to live in the per-node filesystem, not here. + # 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(_GYM_RDMA_ENABLED_DIR) parts.append(SYSTEM_LIB_DIR) for part in os.environ.get("LD_LIBRARY_PATH", "").split(":"): if part and part not in parts: @@ -1013,8 +969,6 @@ async def train( if framework_status_token: os.environ["TRAINING_GYM_FRAMEWORK_STATUS_TOKEN"] = framework_status_token - _enable_rdma_prefix_if_broken() - await asyncio.gather( hf_cache_volume.reload.aio(), data_volume.reload.aio(), 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..9fd884872 --- /dev/null +++ b/modal_training_gym/frameworks/miles/modal_helpers/patches/patch_mooncake_import_tolerance.py @@ -0,0 +1,66 @@ +"""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 8cd5e9bf7..aa1953ec2 100644 --- a/tests/test_miles_runtime_env.py +++ b/tests/test_miles_runtime_env.py @@ -2,8 +2,6 @@ from __future__ import annotations -import subprocess - from modal_training_gym.frameworks.miles import launcher from modal_training_gym.frameworks.miles.launcher import build_ray_runtime_env from modal_training_gym.train_recipes.miles_recipe import MilesRecipe @@ -45,6 +43,17 @@ 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_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") @@ -55,8 +64,7 @@ def test_ld_library_path_comes_from_the_container(monkeypatch): assert env_vars["LD_LIBRARY_PATH"] == ( "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" - ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" - ":/usr/local/cuda/lib64:/wheel/nvidia/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" @@ -87,8 +95,7 @@ def test_unset_container_path_yields_only_the_required_lib_dirs(monkeypatch): )["env_vars"] assert env_vars["LD_LIBRARY_PATH"] == ( - "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" - ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" + "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib:/usr/lib/x86_64-linux-gnu" ) @@ -101,7 +108,7 @@ def test_system_lib_dir_is_not_duplicated(monkeypatch): assert env_vars["LD_LIBRARY_PATH"] == ( "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" - ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu:/wheel/nvidia/lib" + ":/usr/lib/x86_64-linux-gnu:/wheel/nvidia/lib" ) @@ -118,51 +125,17 @@ def test_metric_env_is_preserved(monkeypatch): assert env_vars["WANDB_RESUME"] == "allow" -def test_efa_dirs_precede_the_system_lib_dir_when_absent_on_head(monkeypatch): +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) - monkeypatch.setattr(launcher.os.path, "isdir", lambda _d: False) env_vars = build_ray_runtime_env( head_addr="10.0.0.1", metric_env={}, environment={} )["env_vars"] - assert env_vars["LD_LIBRARY_PATH"] == ( - "/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib" - ":/opt/gym-rdma/enabled:/usr/lib/x86_64-linux-gnu" - ) - - -def _probe(monkeypatch, returncode: int, stderr: str) -> list[tuple[str, str]]: - created: list[tuple[str, str]] = [] - monkeypatch.setattr( - launcher.os.path, "isdir", lambda d: d == launcher._GYM_RDMA_DIR - ) - monkeypatch.setattr(launcher.os.path, "lexists", lambda d: False) - monkeypatch.setattr( - launcher.subprocess, - "run", - lambda *a, **k: subprocess.CompletedProcess(a, returncode, "", stderr), - ) - monkeypatch.setattr( - launcher.os, "symlink", lambda src, dst: created.append((src, dst)) - ) - launcher._enable_rdma_prefix_if_broken() - return created - - -def test_broken_verbs_pair_materializes_the_prefix_alias(monkeypatch): - created = _probe( - monkeypatch, 1, "ImportError: version `IBVERBS_PRIVATE_34' not found" - ) - - assert created == [(launcher._GYM_RDMA_DIR, launcher._GYM_RDMA_ENABLED_DIR)] - - -def test_healthy_verbs_pair_leaves_the_alias_absent(monkeypatch): - assert _probe(monkeypatch, 0, "") == [] - - -def test_unrelated_probe_failure_leaves_the_alias_absent(monkeypatch): - assert ( - _probe(monkeypatch, 1, "ModuleNotFoundError: No module named 'mooncake'") == [] - ) + 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) From ff5d45bb39a158af193d558e00d72b9d900b1647 Mon Sep 17 00:00:00 2001 From: zhouhelena1 Date: Tue, 1 Sep 2026 23:47:04 -0400 Subject: [PATCH 6/7] Reapply EFA patches after local Miles overlay --- .../frameworks/miles/launcher.py | 41 ++++++++++++------- tests/test_miles_runtime_env.py | 17 ++++++++ 2 files changed, 43 insertions(+), 15 deletions(-) diff --git a/modal_training_gym/frameworks/miles/launcher.py b/modal_training_gym/frameworks/miles/launcher.py index 0cdecb8a1..39a15bb77 100644 --- a/modal_training_gym/frameworks/miles/launcher.py +++ b/modal_training_gym/frameworks/miles/launcher.py @@ -391,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.""" @@ -482,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/tests/test_miles_runtime_env.py b/tests/test_miles_runtime_env.py index aa1953ec2..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 @@ -54,6 +60,17 @@ def test_image_applies_efa_host_patches(monkeypatch): 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") From 2e281dea07e53a959662d1fe87536d27b948a6ac Mon Sep 17 00:00:00 2001 From: zhouhelena1 Date: Wed, 2 Sep 2026 00:07:18 -0400 Subject: [PATCH 7/7] Format Mooncake import patch --- .../modal_helpers/patches/patch_mooncake_import_tolerance.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 index 9fd884872..4e6ac8074 100644 --- 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 @@ -40,9 +40,12 @@ " 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.") + raise SystemExit( + f"{target} not found; miles layout changed — re-check the patch." + ) src = target.read_text() if MARKER in src: