diff --git a/components/src/dynamo/profiler/profile_sla.py b/components/src/dynamo/profiler/profile_sla.py index 132151886e0b..ad428d4e448d 100644 --- a/components/src/dynamo/profiler/profile_sla.py +++ b/components/src/dynamo/profiler/profile_sla.py @@ -30,6 +30,7 @@ from dynamo.profiler.utils.config_modifiers.parallelization_mapping import ( PickedParallelConfig, ) +from dynamo.profiler.utils.config_modifiers.protocol import apply_dgd_overrides from dynamo.profiler.utils.defaults import SearchStrategy from dynamo.profiler.utils.dgd_generation import generate_dgd_config_with_planner from dynamo.profiler.utils.dgdr_v1beta1_types import DynamoGraphDeploymentRequestSpec @@ -279,6 +280,18 @@ async def run_profile( else: final_config = dgd_config + # --- Apply DGD overrides (user-supplied partial DGD) --- + if final_config and dgdr.overrides and dgdr.overrides.dgd: + if isinstance(final_config, list): + # only apply to the last element (the actual DGD) + # previous elements are generated config maps + final_config[-1] = apply_dgd_overrides( + final_config[-1], dgdr.overrides.dgd + ) + elif isinstance(final_config, dict): + final_config = apply_dgd_overrides(final_config, dgdr.overrides.dgd) + logger.info("Applied DGD overrides to the final config.") + output_file = f"{ops.output_dir}/final_config.yaml" if not final_config: if ops.dry_run: diff --git a/components/src/dynamo/profiler/thorough.py b/components/src/dynamo/profiler/thorough.py index a316bf9af47e..45017a0414a8 100644 --- a/components/src/dynamo/profiler/thorough.py +++ b/components/src/dynamo/profiler/thorough.py @@ -39,6 +39,7 @@ ) from dynamo.profiler.utils.config import Config, get_service_name_by_type from dynamo.profiler.utils.config_modifiers import CONFIG_MODIFIERS +from dynamo.profiler.utils.config_modifiers.protocol import apply_dgd_overrides from dynamo.profiler.utils.dgdr_v1beta1_types import DynamoGraphDeploymentRequestSpec from dynamo.profiler.utils.profile_common import ProfilerOperationalConfig from dynamo.profiler.utils.profile_decode import get_num_request_range @@ -86,6 +87,22 @@ async def run_thorough( len(decode_candidates), ) + # --- Apply DGD overrides to each candidate --- + if dgdr.overrides and dgdr.overrides.dgd: + for candidate in prefill_candidates: + candidate.dgd_config = apply_dgd_overrides( + candidate.dgd_config, dgdr.overrides.dgd + ) + for candidate in decode_candidates: + candidate.dgd_config = apply_dgd_overrides( + candidate.dgd_config, dgdr.overrides.dgd + ) + logger.info( + "Applied DGD overrides to %d prefill + %d decode candidates.", + len(prefill_candidates), + len(decode_candidates), + ) + config_modifier = CONFIG_MODIFIERS[backend] # --- Stage 2: Benchmarking --- diff --git a/components/src/dynamo/profiler/utils/config_modifiers/protocol.py b/components/src/dynamo/profiler/utils/config_modifiers/protocol.py index 62bb7c4843ca..5730b600494f 100644 --- a/components/src/dynamo/profiler/utils/config_modifiers/protocol.py +++ b/components/src/dynamo/profiler/utils/config_modifiers/protocol.py @@ -15,6 +15,7 @@ from __future__ import annotations +import copy import logging from typing import Any, Protocol, Tuple @@ -624,3 +625,96 @@ def _apply_agg_worker( cls._apply_worker_config( cfg.spec.services[svc_name], agg_cli_args, agg_replicas, agg_gpus ) + + +# --------------------------------------------------------------------------- +# DGD override merging (module-level, backend-agnostic) +# --------------------------------------------------------------------------- + +# Services whose CLI args are fully replaced by overrides. +# For engine-worker services (everything else), the main container args +# are *appended* because they contain profiler-generated sweep results. +_OVERRIDE_NON_WORKER_SERVICES = frozenset({"Frontend", "Planner"}) + +# The exact path suffix where profiler-generated CLI args live inside a +# service dict. Only this specific location gets append semantics. +_WORKER_ARGS_SUFFIX = ("extraPodSpec", "mainContainer", "args") + + +def _is_worker_main_container_args(path: list[str]) -> bool: + """True when *path* is ``spec.services..extraPodSpec.mainContainer.args``.""" + if len(path) != 6: + return False + return ( + path[0] == "spec" + and path[1] == "services" + and path[2] not in _OVERRIDE_NON_WORKER_SERVICES + and tuple(path[3:]) == _WORKER_ARGS_SUFFIX + ) + + +def _deep_merge_overrides( + target: dict, + overrides: dict, + path: list[str], +) -> None: + """Recursively merge *overrides* into *target* (mutates *target* in-place). + + Rules: + - Dicts are merged recursively; missing intermediate keys are created. + - ``spec.services.`` that does not exist in *target* is skipped + with a warning (all nested overrides under that service are dropped). + - Only ``spec.services..extraPodSpec.mainContainer.args`` is + *appended* to the existing list (preserving profiler-generated CLI + args). ``args`` at any other path is replaced normally. + - All other leaf values replace the target value. + """ + for key, value in overrides.items(): + current_path = path + [key] + + # Guard: skip overrides for services that don't exist in the DGD + if ( + len(current_path) == 3 + and current_path[0] == "spec" + and current_path[1] == "services" + ): + services = target.get("services", target) if path == ["spec"] else target + if key not in services: + logger.warning( + "Service '%s' does not exist in the generated DGD config; " + "overrides for this service will not be applied.", + key, + ) + continue + + if isinstance(value, dict) and isinstance(target.get(key), dict): + _deep_merge_overrides(target[key], value, current_path) + elif isinstance(value, dict) and key not in target: + target[key] = copy.deepcopy(value) + elif ( + key == "args" + and isinstance(value, list) + and _is_worker_main_container_args(current_path) + ): + existing = target.get(key) or [] + target[key] = list(existing) + list(value) + else: + target[key] = ( + copy.deepcopy(value) if isinstance(value, (dict, list)) else value + ) + + +def apply_dgd_overrides(dgd_config: dict, overrides: dict) -> dict: + """Deep-merge an ``overrides.dgd`` dict onto a generated DGD config. + + Args: + dgd_config: The generated DynamoGraphDeployment config dict. + overrides: A partial DGD dict with the same structure. Leaf values + overwrite the corresponding keys in *dgd_config*. + + Returns: + A new dict with the overrides applied (the original is not mutated). + """ + result = copy.deepcopy(dgd_config) + _deep_merge_overrides(result, overrides, path=[]) + return result diff --git a/tests/profiler/configs/10_thorough_override_security_context.yaml b/tests/profiler/configs/10_thorough_override_security_context.yaml new file mode 100644 index 000000000000..82d0bab1a0d5 --- /dev/null +++ b/tests/profiler/configs/10_thorough_override_security_context.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Case 10: Thorough sweep with DGD overrides for imagePullSecrets. +# Verifies that overrides can inject new spec-level fields (imagePullSecrets) +# that do not exist in the base DGD template. +model: "Qwen/Qwen3-32B" +backend: trtllm +image: "nvcr.io/nvidia/dynamo:latest" +hardware: + gpuSku: h200_sxm + totalGpus: 8 + numGpusPerNode: 8 +searchStrategy: thorough +overrides: + dgd: + spec: + imagePullSecrets: + - name: my-registry-secret + - name: nvcr-pull-secret diff --git a/tests/profiler/configs/9a_thorough_dsr1_sglang_overrides.yaml b/tests/profiler/configs/9a_thorough_dsr1_sglang_overrides.yaml new file mode 100644 index 000000000000..b4e0731ccb73 --- /dev/null +++ b/tests/profiler/configs/9a_thorough_dsr1_sglang_overrides.yaml @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Case 9a: Thorough sweep for DeepSeek-R1 on SGLang with DGD overrides. +# Overrides are derived by comparing the base sglang disagg template +# (examples/backends/sglang/deploy/disagg.yaml) with the DSR1 recipe +# (recipes/deepseek-r1/sglang/disagg-8gpu/deploy.yaml), excluding +# parallelization-mapping fields (--tp, --dp, --ep-size, +# --enable-dp-attention) which are swept by the profiler. +model: "deepseek-ai/DeepSeek-R1" +backend: sglang +image: "nvcr.io/nvidia/ai-dynamo/sglang-runtime:0.8.0" +hardware: + gpuSku: h200_sxm + totalGpus: 16 + numGpusPerNode: 8 +modelCache: + pvcName: model-cache + pvcMountPath: /opt/model +workload: + isl: 4000 + osl: 1000 +sla: + ttft: 2000.0 + itl: 50.0 +searchStrategy: thorough +overrides: + dgd: + spec: + services: + decode: + sharedMemory: + size: 80Gi + extraPodSpec: + mainContainer: + workingDir: /workspace + args: + - --mem-fraction-static + - "0.75" + - --prefill-round-robin-balance + - --watchdog-timeout + - "3600" + prefill: + sharedMemory: + size: 80Gi + extraPodSpec: + mainContainer: + workingDir: /workspace + args: + - --mem-fraction-static + - "0.75" + - --load-balance-method + - round_robin + - --watchdog-timeout + - "3600" diff --git a/tests/profiler/test_profile_sla_dgdr.py b/tests/profiler/test_profile_sla_dgdr.py index 2a8707cc30e0..0126e55cc0f6 100644 --- a/tests/profiler/test_profile_sla_dgdr.py +++ b/tests/profiler/test_profile_sla_dgdr.py @@ -311,28 +311,43 @@ def _save_dummy_npz(output_dir: str): ) -_THOROUGH_PATCHES = [ - patch( - "dynamo.profiler.thorough.DynamoDeploymentClient", - side_effect=lambda **kw: _mock_deployment_client(), - ), - patch("dynamo.profiler.thorough.get_prefill_ttft", return_value=50.0), - patch( - "dynamo.profiler.thorough.get_decode_itl_and_thpt_per_gpu", - return_value=(8.0, 125.0), - ), - patch("dynamo.profiler.thorough.get_num_request_range", return_value=[1, 4, 8]), - patch( - "dynamo.profiler.thorough.get_service_name_by_type", return_value="TRTLLMWorker" - ), -] - - -def _patch_kv_cache_log(modifier_module: str): +_DECODE_SVC_NAMES = { + "sglang": "decode", + "vllm": "VllmDecodeWorker", + "trtllm": "TRTLLMDecodeWorker", +} + + +def _make_thorough_patches(backend: str = "trtllm"): + """Build mock-patches for thorough mode, parameterised by backend.""" + svc_name = _DECODE_SVC_NAMES.get(backend, "TRTLLMDecodeWorker") + return [ + patch( + "dynamo.profiler.thorough.DynamoDeploymentClient", + side_effect=lambda **kw: _mock_deployment_client(), + ), + patch("dynamo.profiler.thorough.get_prefill_ttft", return_value=50.0), + patch( + "dynamo.profiler.thorough.get_decode_itl_and_thpt_per_gpu", + return_value=(8.0, 125.0), + ), + patch("dynamo.profiler.thorough.get_num_request_range", return_value=[1, 4, 8]), + patch( + "dynamo.profiler.thorough.get_service_name_by_type", + return_value=svc_name, + ), + ] + + +# Backward compat: existing tests use the trtllm-flavored list +_THOROUGH_PATCHES = _make_thorough_patches("trtllm") + + +def _patch_kv_cache_log(backend: str = "trtllm"): """Patch get_kv_cache_size_from_dynamo_log on the real config modifier.""" from dynamo.profiler.utils.config_modifiers import CONFIG_MODIFIERS - real_modifier = CONFIG_MODIFIERS["trtllm"] + real_modifier = CONFIG_MODIFIERS[backend] return patch.object( real_modifier, "get_kv_cache_size_from_dynamo_log", return_value=100000 ) @@ -352,7 +367,7 @@ def test_thorough_no_planner_with_load(self, tmp_path): dgdr = _load_dgdr(CONFIGS_DIR / "6_thorough_no_planner_with_load.yaml") ops = _make_ops(tmp_path) - with _patch_kv_cache_log("thorough"): + with _patch_kv_cache_log("trtllm"): for p in _THOROUGH_PATCHES: p.start() try: @@ -399,7 +414,7 @@ def mock_profile_decode(work_dir, *args, **kwargs): ), ] - with _patch_kv_cache_log("thorough"): + with _patch_kv_cache_log("trtllm"): all_patches = _THOROUGH_PATCHES + interp_patches for p in all_patches: p.start() @@ -429,3 +444,94 @@ def mock_profile_decode(work_dir, *args, **kwargs): ) assert prefill_npz.exists(), "Prefill interpolation data should be saved" assert decode_npz.exists(), "Decode interpolation data should be saved" + + +# --------------------------------------------------------------------------- +# Shared helper for mocked-thorough + override tests +# --------------------------------------------------------------------------- + + +def _run_mocked_thorough(dgdr, ops, backend: str): + """Run the full mocked-thorough pipeline for an arbitrary backend.""" + thorough_patches = _make_thorough_patches(backend) + kv_patch = _patch_kv_cache_log(backend) + + with kv_patch: + for p in thorough_patches: + p.start() + try: + asyncio.run(run_profile(dgdr, ops)) + finally: + for p in thorough_patches: + p.stop() + + +def _assert_overrides_applied(final_config_path: Path, dgdr): + """Assert the final DGD exists and that overrides are reflected.""" + assert final_config_path.exists(), "final_config.yaml should exist" + raw = final_config_path.read_text() + docs = list(yaml.safe_load_all(raw)) + dgd = docs[-1] if docs else {} + assert dgd and "spec" in dgd, "DGD should have a spec" + + override_spec = dgdr.overrides.dgd.get("spec", {}) + + for ovr_key in ("envs", "backendFramework"): + if ovr_key in override_spec: + assert ovr_key in dgd["spec"], f"Override field spec.{ovr_key} should exist" + + svc_overrides = override_spec.get("services", {}) + dgd_services = dgd.get("spec", {}).get("services", {}) + for svc_name, svc_ovr in svc_overrides.items(): + if svc_name in dgd_services: + dgd_svc = dgd_services[svc_name] + if "sharedMemory" in svc_ovr: + assert ( + "sharedMemory" in dgd_svc + ), f"Override sharedMemory on {svc_name} should be applied" + mc = svc_ovr.get("extraPodSpec", {}).get("mainContainer", {}) + if "args" in mc: + dgd_args = ( + dgd_svc.get("extraPodSpec", {}) + .get("mainContainer", {}) + .get("args", []) + ) + for arg in mc["args"]: + assert ( + arg in dgd_args + ), f"Override arg '{arg}' should be in {svc_name} args" + + +class TestThoroughMockedOverrides: + """Thorough + DGD overrides: verify overrides are applied end-to-end.""" + + @pytest.mark.pre_merge + @pytest.mark.gpu_0 + def test_9a_sglang_overrides(self, tmp_path): + """Case 9a: SGLang thorough sweep with DSR1 overrides.""" + dgdr = _load_dgdr(CONFIGS_DIR / "9a_thorough_dsr1_sglang_overrides.yaml") + ops = _make_ops(tmp_path) + _run_mocked_thorough(dgdr, ops, "sglang") + _assert_overrides_applied( + tmp_path / "profiling_results" / "final_config.yaml", + dgdr, + ) + + @pytest.mark.pre_merge + @pytest.mark.gpu_0 + def test_10_override_security_context(self, tmp_path): + """Case 10: imagePullSecrets injected via overrides into a new spec field.""" + dgdr = _load_dgdr(CONFIGS_DIR / "10_thorough_override_security_context.yaml") + ops = _make_ops(tmp_path) + _run_mocked_thorough(dgdr, ops, "trtllm") + + output = tmp_path / "profiling_results" / "final_config.yaml" + assert output.exists(), "final_config.yaml should exist" + config = yaml.safe_load(output.read_text()) + assert config and "spec" in config + + secrets = config["spec"].get("imagePullSecrets") + assert secrets is not None, "imagePullSecrets should be present" + secret_names = [s["name"] for s in secrets] + assert "my-registry-secret" in secret_names + assert "nvcr-pull-secret" in secret_names