[https://nvbugs/6384589][fix] Mount a 640G /dev/shm for online-EPLB perf cases

jenkins/scripts/perf/benchmark_utils.py

@@ -14,6 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+from typing import Iterator, Optional
+
+import yaml
+
 GEN_ONLY_NO_CONTEXT_MODE = "gen_only_no_context"
 
 DISAGG_BENCHMARK_MODES = ("e2e", "gen_only")
@@ -99,3 +104,110 @@ def parse_positive_concurrency(value: object) -> int:
     if concurrency <= 0:
         raise ValueError(f"benchmark.concurrency_list must be a positive integer, got {value!r}")
     return concurrency
+
+
+# Online EPLB stages a host copy of every MoE layer's expert weights in POSIX
+# shared memory on each node, so the container's /dev/shm has to hold all of
+# them: ~532 GiB for Kimi-K2.5 NVFP4, ~380 GiB for GLM-5 NVFP4. /dev/shm
+# defaults to half of host RAM (480 GiB on a GB300 node).
+EPLB_SHM_TMPFS_MOUNT = "tmpfs:/dev/shm:size=640G"
+
+
+def _iter_worker_llm_configs(config: Optional[dict]) -> Iterator[dict]:
+    """Yield every per-server LLM API config dict in a perf-sanity YAML.
+
+    Disaggregated YAMLs keep them under ``worker_config.{ctx,gen}``; aggregated
+    YAMLs keep one per ``server_configs`` entry.
+    """
+    for role_config in ((config or {}).get("worker_config") or {}).values():
+        if isinstance(role_config, dict):
+            yield role_config
+    for server_config in (config or {}).get("server_configs") or []:
+        if isinstance(server_config, dict):
+            yield server_config
+
+
+def _load_yaml(path: str, llm_src: str) -> object:
+    """Load a YAML file; a relative ``path`` is resolved against ``llm_src``."""
+    if not os.path.isabs(path):
+        path = os.path.join(llm_src, path)
+    with open(path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def _merged_llm_api_config(server_config: dict, llm_src: str) -> dict:
+    """Return the LLM API config one server is launched with.
+
+    The file at ``extra_llm_api_config_path`` is the base and the inline fields
+    replace its top-level keys, matching ``ServerConfig.merged_llm_api_config_data``
+    in ``tests/integration/defs/perf/test_perf_sanity.py``.
+    """
+    extra_config_path = server_config.get("extra_llm_api_config_path")
+    if not extra_config_path:
+        return server_config
+    extra_config = _load_yaml(extra_config_path, llm_src) or {}
+    if not isinstance(extra_config, dict):
+        raise ValueError(
+            f"extra_llm_api_config_path {extra_config_path} does not contain a YAML mapping"
+        )
+    return {**extra_config, **server_config}
+
+
+def uses_online_eplb(config: Optional[dict], llm_src: str = "") -> bool:
+    """Return True when any server in the YAML enables online EPLB.
+
+    Online EPLB (``moe_config.load_balancer.layer_updates_per_iter > 0``) is the
+    only EPLB mode that keeps expert weights in host shared memory.
+    ``moe_config`` is read from the config each server is launched with, so a
+    load balancer set only in its ``extra_llm_api_config_path`` file counts.
+    ``load_balancer`` is either an inline dict or a path to an EPLB YAML.
+    Relative paths are resolved against ``llm_src``.
+
+    Args:
+        config: The parsed perf-sanity config YAML.
+        llm_src: TensorRT-LLM source root used to resolve relative paths.
+
+    Returns:
+        bool
+    """
+    for server_config in _iter_worker_llm_configs(config):
+        llm_config = _merged_llm_api_config(server_config, llm_src)
+        load_balancer = (llm_config.get("moe_config") or {}).get("load_balancer")
+        if isinstance(load_balancer, str) and load_balancer:
+            load_balancer = _load_yaml(load_balancer, llm_src)
+        if (
+            isinstance(load_balancer, dict)
+            and (load_balancer.get("layer_updates_per_iter") or 0) > 0
+        ):
+            return True
+    return False
+
+
+def _mount_target(entry: str) -> str:
+    """Return the container path of one ``--container-mounts`` entry.
+
+    pyxis mounts a source-only entry such as ``/dev/shm`` on the same path.
+    """
+    parts = entry.split(":")
+    target = parts[1] if len(parts) > 1 else parts[0]
+    return target.rstrip("/") or "/"
+
+
+def add_eplb_shm_mount(mounts: str, config: Optional[dict], llm_src: str = "") -> str:
+    """Append the large /dev/shm tmpfs mount when the YAML uses online EPLB.
+
+    Args:
+        mounts: Comma-separated ``--container-mounts`` value (may be empty).
+        config: The parsed perf-sanity config YAML.
+        llm_src: TensorRT-LLM source root used to resolve relative paths.
+
+    Returns:
+        str: ``mounts`` unchanged, or with ``EPLB_SHM_TMPFS_MOUNT`` appended. A
+        caller-provided mount on ``/dev/shm`` is kept as is.
+    """
+    if not uses_online_eplb(config, llm_src):
+        return mounts
+    entries = [entry for entry in mounts.split(",") if entry]
+    if any(_mount_target(entry) == "/dev/shm" for entry in entries):
+        return mounts
+    return ",".join(entries + [EPLB_SHM_TMPFS_MOUNT])

jenkins/scripts/perf/local/submit.py

@@ -32,6 +32,7 @@
     DISAGG_BENCHMARK_MODES,
     DISAGG_CONFIG_MODES,
     GEN_ONLY_NO_CONTEXT_MODE,
+    add_eplb_shm_mount,
     gen_only_no_context_server_counts,
     gen_only_no_context_world_size,
     is_gen_only_no_context,
@@ -569,7 +570,9 @@ def generate_sbatch_params(args, hardware_config, work_dir):
     return lines
 
 
-def generate_srun_args(args, runtime_mode, timestamp, llm_src="", hardware_config=None):
+def generate_srun_args(
+    args, runtime_mode, timestamp, llm_src="", hardware_config=None, config=None
+):
     """Generate srun arguments."""
     is_aggr = runtime_mode == "aggregated"
     container_name = f"{'aggr' if is_aggr else 'disagg'}_test-{timestamp}"
@@ -587,6 +590,7 @@ def generate_srun_args(args, runtime_mode, timestamp, llm_src="", hardware_confi
     if llm_src and llm_src not in mounts:
         llm_src_mount = f"{llm_src}:{llm_src}"
         mounts = f"{mounts},{llm_src_mount}" if mounts else llm_src_mount
+    mounts = add_eplb_shm_mount(mounts, config, llm_src)
     if mounts:
         lines.append(f"--container-mounts={mounts}")
 
@@ -957,7 +961,9 @@ def main():
     sbatch_lines = generate_sbatch_params(args, hardware_config, work_dir)
 
     # Generate srun args
-    srun_args_lines = generate_srun_args(args, runtime_mode, timestamp, llm_src, hardware_config)
+    srun_args_lines = generate_srun_args(
+        args, runtime_mode, timestamp, llm_src, hardware_config, config
+    )
 
     # Generate pytest command
     pytest_command, test_list_content, test_list_path = generate_pytest_command(

jenkins/scripts/perf/submit.py

@@ -53,6 +53,7 @@
     AGGREGATED_DISAGG_YAML_MODES,
     DISAGG_BENCHMARK_MODES,
     GEN_ONLY_NO_CONTEXT_MODE,
+    add_eplb_shm_mount,
     gen_only_no_context_server_counts,
     gen_only_no_context_world_size,
     is_gen_only_no_context,
@@ -907,6 +908,13 @@ def main():
     with open(args.srun_args, "r") as f:
         srun_args_content = f.read()
     srun_args_lines = srun_args_content.split()
+    mounts_prefix = "--container-mounts="
+    srun_args_lines = [
+        mounts_prefix + add_eplb_shm_mount(line[len(mounts_prefix) :], config, args.llm_src)
+        if line.startswith(mounts_prefix)
+        else line
+        for line in srun_args_lines
+    ]
 
     (
         agg_pytest_command,