Fix miles multi-node NCCL and verbs library resolution

modal_training_gym/frameworks/miles/launcher.py

@@ -97,20 +97,33 @@ 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 = (
     "apt-get update && apt-get install -y --no-install-recommends "
     "--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.
 HARBOR_PKG_VERSION = "0.8.0"
 
 _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)

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 <this file>``.
+"""
+
+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)

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

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,15 +103,17 @@ 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)
 
     env_vars = build_ray_runtime_env(
         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)