diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 858243686d0..910282b0ac2 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -100,6 +100,15 @@ MODELOPT_STATE_PATH= python vllm_serve_fakequant.py QUANT_CFG= QUANT_FILE_PATH= python vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 ``` +For scale-free E5M2 fake quantization, use the built-in `E5M2_DEFAULT_CFG`: + +```bash +QUANT_CFG=E5M2_DEFAULT_CFG python vllm_serve_fakequant.py Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 +``` + +This simulates E5M2 by casting weights and activations directly to `torch.float8_e5m2` +and back to their original dtype, without calibration or per-tensor scaling. + ## Serve a model with sparse attention in vLLM Apply ModelOpt sparse attention at serve time. Right after model load, the launcher replaces each native attention implementation with its matching ModelOpt adapter: `ModelOptSparseAttentionImpl` for FlashAttention or `ModelOptSparseFlashInferImpl` for FlashInfer. Both adapters use the same Triton kernel with paged KV cache support. diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 6257623d108..0f10b6c5225 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -428,6 +428,11 @@ def validate_num_bits(self): raise ValueError( "Supported FPx quantization formats: FP8 (E4M3, E5M2), FP6(E3M2, E2M3), FP4(E2M1)." ) + elif num_bits == (5, 2) and block_sizes is None: + if self.type != "dynamic" or self.axis is not None: + raise ValueError( + "Unscaled E5M2 quantization requires type='dynamic' and axis=None." + ) elif num_bits not in [(4, 3), (2, 1)] and ( block_sizes is None or block_sizes.get("type", None) != "dynamic" ): @@ -1659,6 +1664,7 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]: "configs/ptq/presets/model/int8_weight_only" ) FP8_DEFAULT_CFG: dict[str, Any] = _load_quantize_config_dict("configs/ptq/presets/model/fp8") +E5M2_DEFAULT_CFG: dict[str, Any] = _load_quantize_config_dict("configs/ptq/presets/model/e5m2") MAMBA_MOE_FP8_AGGRESSIVE_CFG: dict[str, Any] = _load_quantize_config_dict( "configs/ptq/presets/model/mamba_moe_fp8_aggressive" ) @@ -1755,6 +1761,7 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]: # DO NOT ADD NEW CONFIGS HERE. If you want to add a new general recipe, add it to # modelopt_recipes/general/ptq/ as a yaml file choices: set[str] = { + "E5M2_DEFAULT_CFG", "FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG", "FP8_AFFINE_KV_CFG", "FP8_DEFAULT_CFG", diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 0e313dd97e1..b800cb7362e 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -56,6 +56,7 @@ ) from ...tensor_quant import ( dynamic_block_quant, + fake_fp8, fake_tensor_quant, fp4_cast_ste, int_cast_ste, @@ -895,7 +896,7 @@ def _fake_quantize(self, inputs): return entrypoint(inputs, self) amax = None - if not self.is_mx_format: + if not self.is_mx_format and self._num_bits != (5, 2): amax = self._get_amax(inputs) if self.block_sizes is not None and self.block_sizes.get("type", "static") == "dynamic": @@ -921,7 +922,7 @@ def _fake_quantize(self, inputs): # Float-point quantization, e.g., FP8 E, M = self._num_bits # noqa: N806 - outputs = scaled_e4m3( + outputs = fake_fp8( inputs, amax, self._get_bias(inputs), @@ -949,6 +950,9 @@ def _fake_quantize(self, inputs): def _check_onnx_readiness(self, inputs): """Check if quantizer is ready for ONNX export.""" + if self._num_bits == (5, 2): + raise NotImplementedError("ONNX export does not support E5M2 fake quantization.") + if not self.block_sizes or self.block_sizes.get("scale_bits", None) != (8, 0): assert hasattr(self, "_amax"), ( "Quantizer has not been calibrated. ONNX export requires the quantizer to be" diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..3c4b79a0d1a 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -64,6 +64,11 @@ def fp8_eager(x, amax): return _fp8_eager(x, amax) +def e5m2_eager(x): + """Eager mode implementation of unscaled E5M2 quantization.""" + return x.to(torch.float8_e5m2).to(x.dtype) + + def scaled_e4m3_impl( inputs: torch.Tensor, amax: torch.Tensor | None = None, @@ -114,7 +119,7 @@ def fake_quant_impl( def _quantize_impl( inputs: torch.Tensor, - amax: torch.Tensor, + amax: torch.Tensor | None, num_bits: int = 8, exponent_bits: int = 0, unsigned: bool = False, @@ -122,6 +127,10 @@ def _quantize_impl( ): if num_bits == 8 and exponent_bits == 4: return scaled_e4m3_impl(inputs=inputs, amax=amax) + elif num_bits == 8 and exponent_bits == 5: + if amax is not None: + raise ValueError("E5M2 fake quantization does not support scaling.") + return e5m2_eager(inputs) elif isinstance(num_bits, int): return fake_quant_impl( inputs=inputs, @@ -138,7 +147,7 @@ def _quantize_impl( def _quantize_impl_abstract( input: torch.Tensor, - amax: torch.Tensor, + amax: torch.Tensor | None, num_bits: int = 8, exponent_bits: int = 0, unsigned: bool = False, @@ -222,7 +231,7 @@ def _dynamic_block_quantize_impl_abstract( try: torch.library.define( "tensorrt::quantize_op", - "(Tensor input, Tensor amax, int num_bits, int exponent_bits, " + "(Tensor input, Tensor? amax, int num_bits, int exponent_bits, " "bool unsigned, bool narrow_range) -> Tensor", ) torch.library.define( @@ -399,8 +408,8 @@ def backward(ctx, grad_outputs): return _fake_quant_backward_function(ctx, grad_outputs, num_args=10) -class ScaledE4M3Function(Function): - """E4M3fy input with scale.""" +class FakeFP8Function(Function): + """Fake quantize input to a supported FP8 format.""" @staticmethod @symbolic_helper.parse_args("v", "t", "t", "i", "i", "s", "b") @@ -415,6 +424,8 @@ def symbolic( pass_through_bwd=False, ): """ONNX symbolic function.""" + if E != 4 or M != 3: + raise NotImplementedError("ONNX export only supports E4M3 FP8 quantization.") from .export_onnx import export_fp8 return export_fp8(g, inputs, amax, trt_high_precision_dtype) @@ -432,8 +443,8 @@ def forward( pass_through_bwd=False, ): """Forward method.""" - if E != 4 or M != 3: - raise NotImplementedError("Only support E=4 & M=3 for now.") + if (E, M) not in ((4, 3), (5, 2)): + raise NotImplementedError("Only E4M3 and E5M2 FP8 formats are supported.") if bias is not None: inputs = inputs - bias @@ -444,7 +455,7 @@ def forward( inputs, amax, num_bits=8, - exponent_bits=4, + exponent_bits=E, unsigned=False, narrow_range=False, ) @@ -700,7 +711,8 @@ def backward(ctx, grad_outputs): fake_tensor_quant = FakeTensorQuantFunction.apply -scaled_e4m3 = ScaledE4M3Function.apply +fake_fp8 = FakeFP8Function.apply +scaled_e4m3 = fake_fp8 dynamic_block_quant = DynamicBlockQuantizationFunction.apply static_blockwise_fp4_fake_quant = StaticBlockwiseFP4FakeQuantFunction.apply fp4_cast_ste = FP4CastSTEFunction.apply diff --git a/modelopt_recipes/configs/numerics/e5m2.yaml b/modelopt_recipes/configs/numerics/e5m2.yaml new file mode 100644 index 00000000000..7e67e8a9295 --- /dev/null +++ b/modelopt_recipes/configs/numerics/e5m2.yaml @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Per-tensor unscaled FP8 E5M2 quantizer attributes. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: e5m2 +axis: +type: dynamic diff --git a/modelopt_recipes/configs/ptq/presets/model/e5m2.yaml b/modelopt_recipes/configs/ptq/presets/model/e5m2.yaml new file mode 100644 index 00000000000..7d7144330d5 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/e5m2.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for W8A8 FP8 E5M2 without scaling. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + w8a8_e5m2_e5m2: configs/ptq/units/w8a8_e5m2_e5m2 + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + +algorithm: max +quant_cfg: + - $import: base_disable_all + - $import: w8a8_e5m2_e5m2 + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/units/w8a8_e5m2_e5m2.yaml b/modelopt_recipes/configs/ptq/units/w8a8_e5m2_e5m2.yaml new file mode 100644 index 00000000000..2788865389b --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/w8a8_e5m2_e5m2.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizerCfgList snippet that enables unscaled FP8 E5M2 on weight and input quantizers. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig +imports: + e5m2: configs/numerics/e5m2 +--- + - quantizer_name: '*weight_quantizer' + cfg: + $import: e5m2 + - quantizer_name: '*input_quantizer' + cfg: + $import: e5m2 diff --git a/tests/gpu/torch/quantization/test_quantize_cuda.py b/tests/gpu/torch/quantization/test_quantize_cuda.py index 8fb4c096030..c1a2df260ed 100644 --- a/tests/gpu/torch/quantization/test_quantize_cuda.py +++ b/tests/gpu/torch/quantization/test_quantize_cuda.py @@ -84,6 +84,7 @@ "config", [ mtq.INT8_DEFAULT_CFG, + mtq.E5M2_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG, mtq.W4A8_AWQ_BETA_CFG, mtq.INT8_SMOOTHQUANT_CFG, diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cf802abaf40..877ca7e71e3 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -179,6 +179,17 @@ def test_zero_amax_per_channel_is_finite(self): assert torch.isfinite(xq).all() +class TestUnscaledE5M2: + @pytest.mark.parametrize("device", ["cuda", "cpu"]) + def test_matches_native_cast(self, device): + x = torch.randn(4, 4, device=device, dtype=torch.float32) + expected = x.to(torch.float8_e5m2).to(x.dtype) + + actual = tensor_quant.fake_fp8(x, None, None, 5, 2) + + assert torch.equal(actual, expected) + + class Testfp4: @pytest.mark.skipif(not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute >= 8.9") def test_native_block_scale_underflows_to_zero(self): diff --git a/tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py b/tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py index af884c878b2..23bc216a2f9 100644 --- a/tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py @@ -41,19 +41,30 @@ class TestBlockQuantCuda(BlockQuantTester): class TestTensorQuantizerE4M3: @pytest.mark.parametrize( - ("E", "M", "axis"), [(5, 2, None), (4, 3, None), (4, 3, 1), (7, 3, None)] + ("E", "M", "axis", "quant_type"), + [ + (5, 2, None, "dynamic"), + (5, 2, None, "static"), + (5, 2, 1, "dynamic"), + (4, 3, None, "static"), + (4, 3, 1, "static"), + (7, 3, None, "static"), + ], ) - def test_e4m3(self, E, M, axis): # noqa: N803 - is_error_expected = E != 4 or M != 3 + def test_fp8(self, E, M, axis, quant_type): # noqa: N803 + is_e4m3 = (E, M) == (4, 3) + is_e5m2 = (E, M) == (5, 2) and axis is None and quant_type == "dynamic" + is_error_expected = not (is_e4m3 or is_e5m2) with pytest.raises(ValidationError) if is_error_expected else contextlib.nullcontext(): - e4m3_desc = QuantizerAttributeConfig(num_bits=(E, M), axis=axis) - e4m3_quantizer = tensor_quantizer.TensorQuantizer(e4m3_desc).cuda() + fp8_desc = QuantizerAttributeConfig(num_bits=(E, M), axis=axis, type=quant_type) + fp8_quantizer = tensor_quantizer.TensorQuantizer(fp8_desc).cuda() x = torch.rand(3, 6, 7, 7).cuda() - e4m3_x = e4m3_quantizer(x) - ref = tensor_quant.scaled_e4m3(x, e4m3_quantizer._get_amax(x), None, E, M) - assert torch.allclose(e4m3_x, ref) + fp8_x = fp8_quantizer(x) + amax = None if is_e5m2 else fp8_quantizer._get_amax(x) + ref = tensor_quant.fake_fp8(x, amax, None, E, M) + assert torch.allclose(fp8_x, ref) def test_non_current_gpu(self, need_2_gpus): x = torch.randn(3, 4) diff --git a/tests/unit/torch/quantization/test_config_validation.py b/tests/unit/torch/quantization/test_config_validation.py index 4b969d3259c..6a5058eb000 100644 --- a/tests/unit/torch/quantization/test_config_validation.py +++ b/tests/unit/torch/quantization/test_config_validation.py @@ -23,6 +23,7 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization.algorithms import _match_quantizer_cfg from modelopt.torch.quantization.config import ( + E5M2_DEFAULT_CFG, FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, FP8_DEFAULT_CFG, FP8_PER_CHANNEL_PER_TOKEN_CFG, @@ -45,6 +46,7 @@ def test_need_calibration(): + assert not need_calibration(E5M2_DEFAULT_CFG) assert need_calibration(FP8_DEFAULT_CFG) assert not need_calibration(FP8_PER_CHANNEL_PER_TOKEN_CFG) assert not need_calibration(FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG) diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index ba352ec2162..2b8b9ff85be 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -23,6 +23,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module +import modelopt.torch.quantization.tensor_quant as tensor_quant from modelopt.torch.quantization import QuantModuleRegistry from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig from modelopt.torch.quantization.nn import ( @@ -37,6 +38,15 @@ class TestFakeTensorQuantCPU(FakeTensorQuantTester): device = "cpu" +def test_unscaled_e5m2_matches_native_cast(): + inputs = torch.tensor([-60000.0, -1.3, 0.0, 1.3, 60000.0], dtype=torch.float32) + expected = inputs.to(torch.float8_e5m2).to(inputs.dtype) + + actual = tensor_quant.fake_fp8(inputs, None, None, 5, 2) + + assert torch.equal(actual, expected) + + class TestQuantizerAttributeConfig: def test_scaled_mode(self): num_bits = np.random.randint(1, 16) @@ -110,6 +120,25 @@ def test_num_bits(self): ): QuantizerAttributeConfig(enable=True, num_bits=(-1, 2)) + def test_unscaled_e5m2_config(self): + config = QuantizerAttributeConfig(num_bits=(5, 2), axis=None, type="dynamic") + assert config.num_bits == (5, 2) + + with pytest.raises(ValueError, match="requires type='dynamic' and axis=None"): + QuantizerAttributeConfig(num_bits=(5, 2), axis=None) + + with pytest.raises(ValueError, match="requires type='dynamic' and axis=None"): + QuantizerAttributeConfig(num_bits=(5, 2), axis=0, type="dynamic") + + +def test_unscaled_e5m2_tensor_quantizer(): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(5, 2), axis=None, type="dynamic") + ) + inputs = torch.tensor([-60000.0, -1.3, 0.0, 1.3, 60000.0], dtype=torch.float32) + + assert torch.equal(quantizer(inputs), inputs.to(torch.float8_e5m2).to(inputs.dtype)) + def _run_rotated_backend(monkeypatch, rotate): calls = [] diff --git a/tools/launcher/common/vllm/fakequant_eval.sh b/tools/launcher/common/vllm/fakequant_eval.sh new file mode 100755 index 00000000000..bbfee45cdce --- /dev/null +++ b/tools/launcher/common/vllm/fakequant_eval.sh @@ -0,0 +1,82 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euo pipefail + +MODEL=${1:?Usage: fakequant_eval.sh MODEL [QUANT_CFG] [TASKS] [LIMIT]} +QUANT_CFG=${2:-E5M2_DEFAULT_CFG} +TASKS=${3:-gsm8k} +LIMIT=${4:-20} +PORT=${PORT:-8000} + +REPO=./modules/Model-Optimizer +SERVER_LOG="$PWD/vllm_fakequant_server.log" + +if command -v python3 >/dev/null 2>&1; then + PYTHON_BIN=python3 +elif command -v python >/dev/null 2>&1; then + PYTHON_BIN=python +else + echo "Neither python3 nor python is available in the workload container." >&2 + exit 127 +fi + +test -f "$MODEL/config.json" +test -f "$REPO/examples/vllm_serve/vllm_serve_fakequant.py" + +"$PYTHON_BIN" -m pip install --editable "$REPO" +"$PYTHON_BIN" -m pip install "lm_eval[api]>=0.4.8" + +export QUANT_CFG +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export OPENAI_API_KEY=${OPENAI_API_KEY:-token-abc123} + +"$PYTHON_BIN" "$REPO/examples/vllm_serve/vllm_serve_fakequant.py" \ + "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --enforce-eager \ + --gpu-memory-utilization 0.85 \ + --max-model-len 4096 \ + >"$SERVER_LOG" 2>&1 & +server_pid=$! + +cleanup() { + kill "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true +} +trap cleanup EXIT + +for _ in $(seq 1 120); do + if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null; then + break + fi + if ! kill -0 "$server_pid" 2>/dev/null; then + tail -200 "$SERVER_LOG" + exit 1 + fi + sleep 5 +done +curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null + +lm_eval \ + --model local-completions \ + --tasks "$TASKS" \ + --limit "$LIMIT" \ + --batch_size 1 \ + --output_path "$PWD/lm_eval_fakequant" \ + --model_args \ + "model=$MODEL,base_url=http://127.0.0.1:$PORT/v1/completions,num_concurrent=1,max_retries=3,tokenized_requests=False,tokenizer_backend=None" diff --git a/tools/launcher/common/vllm/fakequant_nemo_gym_eval.sh b/tools/launcher/common/vllm/fakequant_nemo_gym_eval.sh new file mode 100755 index 00000000000..1b626a279f3 --- /dev/null +++ b/tools/launcher/common/vllm/fakequant_nemo_gym_eval.sh @@ -0,0 +1,134 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euo pipefail + +MODEL=${1:?Usage: fakequant_nemo_gym_eval.sh MODEL [QUANT_CFG] [BENCHMARK] [LIMIT]} +QUANT_CFG=${2:-E5M2_DEFAULT_CFG} +BENCHMARK=${3:-gpqa} +LIMIT=${4:-5} +PORT=${PORT:-8000} +NEMO_GYM_REV=${NEMO_GYM_REV:-a85670eb167ba9b48cc53a36a070eed815e6c40d} + +REPO=./modules/Model-Optimizer +RESULTS_DIR="$PWD/nemo_gym_${BENCHMARK}" +SERVER_LOG="$PWD/vllm_fakequant_server.log" +GYM_LOG="$PWD/nemo_gym_servers.log" + +if command -v python3 >/dev/null 2>&1; then + PYTHON_BIN=python3 +elif command -v python >/dev/null 2>&1; then + PYTHON_BIN=python +else + echo "Neither python3 nor python is available in the workload container." >&2 + exit 127 +fi + +test -f "$MODEL/config.json" +test -f "$REPO/examples/vllm_serve/vllm_serve_fakequant.py" + +"$PYTHON_BIN" -c \ + 'import sys; assert sys.version_info >= (3, 12), "NeMo Gym requires Python 3.12 or newer"' +"$PYTHON_BIN" -m pip install --editable "$REPO" +"$PYTHON_BIN" -m pip install "uv>=0.9.30" + +GYM_DIR=$(mktemp -d) +git -C "$GYM_DIR" init +git -C "$GYM_DIR" remote add origin https://github.com/NVIDIA-NeMo/Gym.git +git -C "$GYM_DIR" fetch --depth 1 origin "$NEMO_GYM_REV" +git -C "$GYM_DIR" checkout --detach FETCH_HEAD +uv sync --directory "$GYM_DIR" --frozen --no-dev + +mkdir -p "$RESULTS_DIR" +( + cd "$GYM_DIR" + uv run --frozen gym eval prepare \ + --benchmark "$BENCHMARK" \ + '+hf_token=${oc.env:HF_TOKEN,null}' +) + +export QUANT_CFG +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export OPENAI_API_KEY=${OPENAI_API_KEY:-token-abc123} +export PYTHONUNBUFFERED=1 + +"$PYTHON_BIN" "$REPO/examples/vllm_serve/vllm_serve_fakequant.py" \ + "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --enforce-eager \ + --gpu-memory-utilization 0.85 \ + --max-model-len 4096 \ + >"$SERVER_LOG" 2>&1 & +server_pid=$! +gym_pid= + +cleanup() { + if [[ -n "$gym_pid" ]]; then + kill "$gym_pid" 2>/dev/null || true + wait "$gym_pid" 2>/dev/null || true + fi + kill "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true +} +trap cleanup EXIT + +for _ in $(seq 1 120); do + if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null; then + break + fi + if ! kill -0 "$server_pid" 2>/dev/null; then + tail -200 "$SERVER_LOG" + exit 1 + fi + sleep 5 +done +curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null + +( + cd "$GYM_DIR" + uv run --frozen gym env start \ + --benchmark "$BENCHMARK" \ + --model-type vllm_model \ + --model "$MODEL" \ + --model-url "http://127.0.0.1:$PORT/v1" \ + --model-api-key "$OPENAI_API_KEY" +) >"$GYM_LOG" 2>&1 & +gym_pid=$! + +for _ in $(seq 1 120); do + if grep -qE 'All [0-9]+ / [0-9]+ servers ready!' "$GYM_LOG"; then + break + fi + if ! kill -0 "$gym_pid" 2>/dev/null; then + tail -200 "$GYM_LOG" + exit 1 + fi + sleep 5 +done +grep -qE 'All [0-9]+ / [0-9]+ servers ready!' "$GYM_LOG" + +( + cd "$GYM_DIR" + uv run --frozen gym eval run --no-serve \ + --agent "${BENCHMARK}_mcqa_simple_agent" \ + --input "benchmarks/$BENCHMARK/data/${BENCHMARK}_diamond_benchmark.jsonl" \ + --output "$RESULTS_DIR/rollouts.jsonl" \ + --prompt-config benchmarks/prompts/eval/aai/mcq-4choices.yaml \ + --limit "$LIMIT" \ + --num-repeats 1 \ + --concurrency 1 +) diff --git a/tools/launcher/examples/Qwen/Qwen3-4B/vllm_fakequant_eval.yaml b/tools/launcher/examples/Qwen/Qwen3-4B/vllm_fakequant_eval.yaml new file mode 100644 index 00000000000..e0c7395b6c5 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-4B/vllm_fakequant_eval.yaml @@ -0,0 +1,22 @@ +job_name: qwen3_4b_e5m2_vllm_eval + +pipeline: + skip: false + allow_to_fail: false + note: "Qwen3-4B scale-free E5M2 vLLM fake-quant smoke evaluation" + + task_0: + script: common/vllm/fakequant_eval.sh + args: + - /hf-local/Qwen/Qwen3-4B + - E5M2_DEFAULT_CFG + - gsm8k + - "20" + environment: + - VLLM_STARTUP_TIMEOUT: "600" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: "01:00:00" diff --git a/tools/launcher/examples/Qwen/Qwen3-4B/vllm_fakequant_nemo_gym_gpqa.yaml b/tools/launcher/examples/Qwen/Qwen3-4B/vllm_fakequant_nemo_gym_gpqa.yaml new file mode 100644 index 00000000000..c011a481ebe --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-4B/vllm_fakequant_nemo_gym_gpqa.yaml @@ -0,0 +1,22 @@ +job_name: qwen3_4b_e5m2_vllm_nemo_gym_gpqa + +pipeline: + skip: false + allow_to_fail: false + note: "Qwen3-4B scale-free E5M2 vLLM fake-quant NeMo Gym GPQA smoke evaluation" + + task_0: + script: common/vllm/fakequant_nemo_gym_eval.sh + args: + - /hf-local/Qwen/Qwen3-4B + - E5M2_DEFAULT_CFG + - gpqa + - "5" + environment: + - VLLM_STARTUP_TIMEOUT: "600" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: "01:00:00" diff --git a/tools/launcher/launch.py b/tools/launcher/launch.py index 0f3313aaf4a..3b9d527ab86 100644 --- a/tools/launcher/launch.py +++ b/tools/launcher/launch.py @@ -104,6 +104,8 @@ def _add_package_glob(pattern: str) -> None: _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Megatron-LM/megatron")) _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Megatron-LM/examples")) _add_package_glob(os.path.join(LAUNCHER_DIR, "modules/Megatron-LM/*.py")) + _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/pyproject.toml")) + _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/LICENSE_HEADER")) _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/modelopt")) _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/modelopt_recipes")) _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/examples")) diff --git a/tools/mcp/modelopt_mcp/bridge.py b/tools/mcp/modelopt_mcp/bridge.py index 5ac251e83b1..dd3262b8362 100644 --- a/tools/mcp/modelopt_mcp/bridge.py +++ b/tools/mcp/modelopt_mcp/bridge.py @@ -37,6 +37,7 @@ import json import os import re +import shlex import shutil import subprocess # nosec B404 - fixed-argv CLI probes are required; shell=True is not used. import tempfile @@ -69,6 +70,14 @@ _GIT_SHA_RE = re.compile(r"^[0-9a-fA-F]{7,40}$") _SAFE_PATH_TOKEN_RE = re.compile(r"[^A-Za-z0-9_.-]+") +_LAUNCHER_CREDENTIAL_ENV_KEYS: frozenset[str] = frozenset( + { + "HF_TOKEN", + "SPECDEC_BENCH_S3_ENDPOINT", + "SPECDEC_BENCH_S3_KEY_ID", + "SPECDEC_BENCH_S3_SECRET", + } +) _SLURM_TERMINAL_STATES: frozenset[str] = frozenset( { "COMPLETED", @@ -966,6 +975,8 @@ def _launcher_overrides( account: str | None, partition: str | None, container: str | None, + container_mounts: list[str] | None, + srun_args: list[str] | None, gpus_per_node: int | None, ntasks_per_node: int | None, ) -> dict[str, str]: @@ -979,6 +990,16 @@ def _launcher_overrides( overrides.setdefault(f"{slurm_prefix}partition", partition) if container: overrides.setdefault(f"{slurm_prefix}container", container) + if container_mounts is not None: + overrides.setdefault( + f"{slurm_prefix}container_mounts", + json.dumps(container_mounts, separators=(",", ":")), + ) + if srun_args is not None: + overrides.setdefault( + f"{slurm_prefix}srun_args", + json.dumps(srun_args, separators=(",", ":")), + ) if gpus_per_node is not None: overrides.setdefault(f"{slurm_prefix}gpus_per_node", str(gpus_per_node)) if ntasks_per_node is not None: @@ -998,7 +1019,8 @@ def _apply_launcher_env( reconnect_command: str | None, ) -> None: """Apply launcher env shared by live submit and dry-run.""" - env.setdefault("NEMORUN_HOME", os.getcwd()) + _hydrate_launcher_credentials_from_cache(env) + env.setdefault("NEMORUN_HOME", str(_nemorun_home())) if checkout is not None: env["MODELOPT_MCP_SOURCE_ROOT"] = str(checkout.root) env["MODELOPT_MCP_SOURCE_REF"] = checkout.ref @@ -1015,6 +1037,40 @@ def _apply_launcher_env( env["MODELOPT_LAUNCHER_SSH_RECONNECT_COMMAND"] = reconnect_command +def _hydrate_launcher_credentials_from_cache(env: dict[str, str]) -> None: + """Fill missing launcher credentials from Pensieve's ephemeral env cache. + + Hosted runtimes may inject credentials into the agent shell after the + long-lived MCP server starts. The Pensieve credential cache is hydrated + from the user's carrier for this case. Only credentials already supported + by the launcher are read, process-environment values remain authoritative, + and values are passed only through the child process environment. + """ + configured_path = env.get("PENSIEVE_ENV_FILE") or os.environ.get("PENSIEVE_ENV_FILE") + cache_path = ( + Path(configured_path).expanduser() + if configured_path + else Path.home() / ".config" / "pensieve" / "env" + ) + try: + lines = cache_path.read_text().splitlines() + except OSError: + return + + for line in lines: + try: + fields = shlex.split(line, comments=True, posix=True) + except ValueError: + continue + if fields[:1] == ["export"]: + fields = fields[1:] + if len(fields) != 1 or "=" not in fields[0]: + continue + key, value = fields[0].split("=", 1) + if key in _LAUNCHER_CREDENTIAL_ENV_KEYS and not env.get(key): + env[key] = value + + def submit_job_impl( *, yaml_path: str, @@ -1029,6 +1085,8 @@ def submit_job_impl( account: str | None = None, partition: str | None = None, container: str | None = None, + container_mounts: list[str] | None = None, + srun_args: list[str] | None = None, gpus_per_node: int | None = None, ntasks_per_node: int | None = None, control_socket: str | None = None, @@ -1088,6 +1146,8 @@ def submit_job_impl( account=account, partition=partition, container=container, + container_mounts=container_mounts, + srun_args=srun_args, gpus_per_node=gpus_per_node, ntasks_per_node=ntasks_per_node, control_socket=control_socket, @@ -1202,6 +1262,8 @@ def submit_job_impl( account=account, partition=partition, container=container, + container_mounts=container_mounts, + srun_args=srun_args, gpus_per_node=gpus_per_node, ntasks_per_node=ntasks_per_node, ).items(): @@ -1399,6 +1461,8 @@ def _submit_job_dry_run( account: str | None, partition: str | None, container: str | None, + container_mounts: list[str] | None, + srun_args: list[str] | None, gpus_per_node: int | None, ntasks_per_node: int | None, control_socket: str | None, @@ -1473,6 +1537,8 @@ def _submit_job_dry_run( account=account, partition=partition, container=container, + container_mounts=container_mounts, + srun_args=srun_args, gpus_per_node=gpus_per_node, ntasks_per_node=ntasks_per_node, ).items(): @@ -1578,12 +1644,13 @@ def _resolve_experiment_dir(experiment_id: str) -> Path | None: """Map an experiment_id to its on-disk directory. nemo_run lays experiments out under ``$NEMORUN_HOME/experiments//`` - by default; ``NEMORUN_HOME`` falls back to cwd. We check several + by default. We check several candidate roots in order: - 1. ``$NEMORUN_HOME/experiments/`` — what submit_job_impl pins via env. + 1. NeMo Run's configured/default home. 2. cwd's ``experiments/`` + ``local_experiments/`` — for operators - running the MCP server from their own checkout. + running the MCP server from their own checkout and for experiments + created by older modelopt-mcp versions. 3. The launcher's own ``experiments/`` directory — belt-and-braces for the case where the operator didn't set NEMORUN_HOME at all AND the MCP server's cwd differs from where launch.py ran. @@ -1600,10 +1667,7 @@ def _resolve_experiment_dir(experiment_id: str) -> Path | None: def _experiment_search_roots() -> list[Path]: """Return experiment roots searched by status/log tools.""" - roots = [] - nemorun_home = os.environ.get("NEMORUN_HOME") - if nemorun_home: - roots.append(Path(nemorun_home) / "experiments") + roots = [_nemorun_home() / "experiments"] roots.append(Path.cwd() / "experiments") roots.append(Path.cwd() / "local_experiments") launcher_dir = _find_launcher_package_dir() @@ -1612,6 +1676,14 @@ def _experiment_search_roots() -> list[Path]: return roots +def _nemorun_home() -> Path: + """Return the explicit NeMo Run home or NeMo Run's native default.""" + configured = os.environ.get("NEMORUN_HOME") + if configured: + return Path(configured) + return Path.home() / ".nemo_run" + + def _experiment_not_found_diagnostic() -> str: """Describe all experiment roots used by _resolve_experiment_dir.""" roots = ", ".join(str(root) for root in _experiment_search_roots()) @@ -1923,11 +1995,15 @@ def wait_for_experiment_impl( experiment_id: str, timeout_sec: int, poll_interval_sec: int, + include_log: bool = False, + log_job_idx: int = 0, ) -> dict: """Block until ``experiment_id`` reaches a terminal status or the timeout elapses. Returns the same dict shape as ``job_status_impl`` plus a - ``waited_seconds`` field. On timeout, returns + ``waited_seconds`` field. When ``include_log`` is true, terminal + results also include ``log_tail`` or a structured ``log_error``. + On timeout, returns ``{ok: False, reason: "wait_timeout", last_status: }`` instead of raising — same structured-failure convention as the other tools. @@ -1947,7 +2023,18 @@ def wait_for_experiment_impl( # dir doesn't exist. return {**status, "waited_seconds": time.monotonic() - started} if status["status"] in ("done", "failed"): - return {**status, "waited_seconds": time.monotonic() - started} + result = {**status, "waited_seconds": time.monotonic() - started} + if include_log: + log_result = read_cluster_artifact_impl( + experiment_id=experiment_id, + path=None, + job_idx=log_job_idx, + ) + if log_result.get("ok"): + result["log_tail"] = log_result["content"] + else: + result["log_error"] = log_result + return result if time.monotonic() - started > timeout_sec: return { "ok": False, @@ -2109,64 +2196,19 @@ def read_cluster_artifact_impl( The tool returns ``{ok, content, ...}`` with the file content as a text string (8 KB max — same as the launcher's log_excerpt cap). """ - if not path: - # Mode 1: fetch log via `nemo experiment logs`. Subprocess - # because the CLI handles tunnel auth and remote-path - # resolution. - argv = [ - "uv", - "run", - "nemo", - "experiment", - "logs", - experiment_id, - str(job_idx), - ] - try: - proc = subprocess.run( # nosec B603 B607 - fixed nemo CLI argv; no shell. - argv, - capture_output=True, - text=True, - timeout=60, - check=False, - ) - except subprocess.TimeoutExpired: - return { - "ok": False, - "experiment_id": experiment_id, - "job_idx": job_idx, - "reason": "logs_fetch_timeout", - "diagnostic": ( - f"`nemo experiment logs {experiment_id} {job_idx}` " - f"did not return within 60s — tunnel may be slow " - f"or unreachable." - ), - } - if proc.returncode != 0: - return { - "ok": False, - "experiment_id": experiment_id, - "job_idx": job_idx, - "reason": "logs_fetch_failed", - "exit_code": proc.returncode, - "diagnostic": ( - f"`nemo experiment logs` exited with code " - f"{proc.returncode}. stderr: {proc.stderr.strip()[-400:]}" - ), - } - content = (proc.stdout or "")[-8192:] + invalid = _validate_experiment_id(experiment_id) + if invalid: + return invalid + + exp_dir = _resolve_experiment_dir(experiment_id) + if exp_dir is None: return { - "ok": True, + "ok": False, "experiment_id": experiment_id, - "job_idx": job_idx, - "mode": "logs", - "content": content, - "bytes": len(content), + "reason": "experiment_dir_not_found", + "diagnostic": _experiment_not_found_diagnostic(), } - # Mode 2: arbitrary path via the experiment's tunnel. nemo_run's - # Experiment loads the executor + tunnel from disk; we rsync the - # named relative path into a local tmp dir, then read it back. try: from nemo_run.run.experiment import Experiment except ImportError: @@ -2183,8 +2225,11 @@ def read_cluster_artifact_impl( ), } try: - exp = Experiment.from_id(experiment_id) - except (FileNotFoundError, ValueError) as e: + exp = Experiment._from_config(str(exp_dir)) + job = exp.jobs[job_idx] + executor = job.executor + tunnel = getattr(executor, "tunnel", None) + except (AssertionError, FileNotFoundError, ValueError) as e: return { "ok": False, "experiment_id": experiment_id, @@ -2196,13 +2241,6 @@ def read_cluster_artifact_impl( ), } - # The executor exposes a tunnel; tunnel.run() runs a remote shell - # command. Use `cat` for small files; rsync for large ones is a - # Phase-2.1 follow-up. - try: - task = exp.tasks[job_idx] - executor = task.executor - tunnel = getattr(executor, "tunnel", None) except (IndexError, AttributeError) as e: return { "ok": False, @@ -2223,10 +2261,7 @@ def read_cluster_artifact_impl( "diagnostic": "Experiment executor has no tunnel attribute.", } - # Resolve the remote path. The experiment dir on the cluster is - # exposed via tunnel.job_dir or executor.job_dir; for the launcher's - # SlurmExecutor, this is set at submit time. - remote_dir = getattr(executor, "job_dir", None) or getattr(tunnel, "job_dir", None) + remote_dir = getattr(tunnel, "job_dir", None) if not remote_dir: return { "ok": False, @@ -2238,7 +2273,44 @@ def read_cluster_artifact_impl( "of a relative one." ), } - remote_path = path if path.startswith("/") else f"{remote_dir}/{experiment_id}/{path}" + if not path: + task_dir = f"{remote_dir}/{job.id}" + quoted_task_dir = shlex.quote(task_dir) + command = ( + f"log_path=$(find {quoted_task_dir} -maxdepth 1 -type f " + "\\( -name 'log-*.out' -o -name 'sbatch_*.out' \\) " + "-print | sort | head -n 1); " + 'test -n "$log_path" && tail -c 8192 "$log_path"' + ) + try: + result = tunnel.run(command, warn=True) + except Exception as e: + return { + "ok": False, + "experiment_id": experiment_id, + "job_idx": job_idx, + "reason": "logs_fetch_failed", + "diagnostic": f"Remote log read failed: {type(e).__name__}: {e}", + } + stdout = getattr(result, "stdout", "") or "" + if not stdout: + return { + "ok": False, + "experiment_id": experiment_id, + "job_idx": job_idx, + "reason": "task_log_not_found", + "diagnostic": f"No task log was found under {task_dir}.", + } + return { + "ok": True, + "experiment_id": experiment_id, + "job_idx": job_idx, + "mode": "logs", + "content": stdout[-8192:], + "bytes": len(stdout), + } + + remote_path = path if path.startswith("/") else f"{remote_dir}/{path}" try: result = tunnel.run(f"cat {remote_path}", warn=True) diff --git a/tools/mcp/modelopt_mcp/server.py b/tools/mcp/modelopt_mcp/server.py index 2506b18a33c..ef4658c1ff5 100644 --- a/tools/mcp/modelopt_mcp/server.py +++ b/tools/mcp/modelopt_mcp/server.py @@ -210,6 +210,20 @@ def submit_job( description=("Container image override for pipeline.task_0.slurm_config.container.") ), ] = None, + container_mounts: Annotated[ + list[str] | None, + Field( + description=( + "Slurm container mount overrides, usually from resolved cluster config." + ) + ), + ] = None, + srun_args: Annotated[ + list[str] | None, + Field( + description=("Slurm srun argument overrides, usually from resolved cluster config.") + ), + ] = None, gpus_per_node: Annotated[ int | None, Field( @@ -329,6 +343,8 @@ def submit_job( account=account, partition=partition, container=container, + container_mounts=container_mounts, + srun_args=srun_args, gpus_per_node=gpus_per_node, ntasks_per_node=ntasks_per_node, control_socket=control_socket, @@ -399,7 +415,8 @@ def job_logs( description=( "Block until an experiment reaches a terminal status " "('done' or 'failed') or the timeout elapses. Returns the " - "same shape as job_status plus a `waited_seconds` field. " + "same shape as job_status plus a `waited_seconds` field " + "and, by default, the final remote task `log_tail`. " "Use this instead of writing your own polling while-loop " "around job_status." ), @@ -432,11 +449,29 @@ def wait_for_experiment( ), ), ] = 30, + include_log: Annotated[ + bool, + Field( + description=( + "Include the final remote log tail when the experiment reaches " + "a terminal state." + ) + ), + ] = True, + log_job_idx: Annotated[ + int, + Field( + ge=0, + description=("Zero-based job index whose final log should be returned."), + ), + ] = 0, ) -> dict: return bridge.wait_for_experiment_impl( experiment_id, timeout_sec, poll_interval_sec, + include_log, + log_job_idx, ) @mcp.tool( diff --git a/tools/mcp/tests/test_bridge.py b/tools/mcp/tests/test_bridge.py index 69d57822119..5351dcc7563 100644 --- a/tools/mcp/tests/test_bridge.py +++ b/tools/mcp/tests/test_bridge.py @@ -19,6 +19,7 @@ import json import subprocess +from types import SimpleNamespace import pytest @@ -745,6 +746,8 @@ def fake_run(argv, **kwargs): account="project_account", partition="batch", container="ubuntu:24.04", + container_mounts=["/host/hf:/hf-local", "/host/hf:/models"], + srun_args=["--no-container-mount-home --mpi=pmix"], ntasks_per_node=1, control_socket="~/.ssh/mfa.sock", reconnect_command="ssh mfa-cluster", @@ -757,6 +760,14 @@ def fake_run(argv, **kwargs): assert "pipeline.task_0.slurm_config.account=project_account" in captured["argv"] assert "pipeline.task_0.slurm_config.partition=batch" in captured["argv"] assert "pipeline.task_0.slurm_config.container=ubuntu:24.04" in captured["argv"] + assert ( + 'pipeline.task_0.slurm_config.container_mounts=["/host/hf:/hf-local","/host/hf:/models"]' + in captured["argv"] + ) + assert ( + 'pipeline.task_0.slurm_config.srun_args=["--no-container-mount-home --mpi=pmix"]' + in captured["argv"] + ) assert "pipeline.task_0.slurm_config.ntasks_per_node=1" in captured["argv"] assert captured["env"]["SLURM_ACCOUNT"] == "project_account" assert captured["env"]["SLURM_PARTITION"] == "batch" @@ -916,6 +927,8 @@ def fake_run(argv, **kwargs): account="project_account", partition="batch", container="ubuntu:24.04", + container_mounts=["/host/hf:/hf-local", "/host/hf:/models"], + srun_args=["--no-container-mount-home --mpi=pmix"], ntasks_per_node=1, control_socket="~/.ssh/mfa.sock", ssh_alias="mfa-cluster", @@ -927,6 +940,14 @@ def fake_run(argv, **kwargs): assert "pipeline.task_0.slurm_config.account=project_account" in captured["argv"] assert "pipeline.task_0.slurm_config.partition=batch" in captured["argv"] assert "pipeline.task_0.slurm_config.container=ubuntu:24.04" in captured["argv"] + assert ( + 'pipeline.task_0.slurm_config.container_mounts=["/host/hf:/hf-local","/host/hf:/models"]' + in captured["argv"] + ) + assert ( + 'pipeline.task_0.slurm_config.srun_args=["--no-container-mount-home --mpi=pmix"]' + in captured["argv"] + ) assert "pipeline.task_0.slurm_config.ntasks_per_node=1" in captured["argv"] assert captured["env"]["SLURM_HOST"] == "mfa-login.example.com" assert captured["env"]["SLURM_ACCOUNT"] == "project_account" @@ -1294,6 +1315,76 @@ def test_job_status_launcher_experiments_fallback(tmp_path, monkeypatch): assert result["status"] == "running" +def test_launcher_and_status_share_native_nemorun_home(tmp_path, monkeypatch): + """Without an override, submission and status use NeMo Run's native home.""" + native_home = tmp_path / ".nemo_run" + monkeypatch.delenv("NEMORUN_HOME", raising=False) + monkeypatch.setattr(bridge, "_nemorun_home", lambda: native_home) + env = {} + + bridge._apply_launcher_env( + env, + checkout=None, + executor="slurm", + cluster_host="cluster.example.com", + account=None, + partition=None, + control_socket=None, + reconnect_command=None, + ) + + assert env["NEMORUN_HOME"] == str(native_home) + assert bridge._experiment_search_roots()[0] == native_home / "experiments" + + +def test_launcher_env_hydrates_credentials_from_pensieve_cache(tmp_path, monkeypatch): + """Carrier-backed credentials reach launcher subprocesses without entering argv.""" + cache = tmp_path / "pensieve.env" + cache.write_text( + "export HF_TOKEN='carrier hf token'\n" + "SPECDEC_BENCH_S3_SECRET=carrier-s3-secret\n" + "UNRELATED_SECRET=do-not-forward\n" + ) + monkeypatch.setenv("PENSIEVE_ENV_FILE", str(cache)) + env = {} + + bridge._apply_launcher_env( + env, + checkout=None, + executor="slurm", + cluster_host="cluster.example.com", + account=None, + partition=None, + control_socket=None, + reconnect_command=None, + ) + + assert env["HF_TOKEN"] == "carrier hf token" + assert env["SPECDEC_BENCH_S3_SECRET"] == "carrier-s3-secret" + assert "UNRELATED_SECRET" not in env + + +def test_launcher_env_prefers_process_credentials(tmp_path, monkeypatch): + """An inherited non-empty credential takes precedence over the cache.""" + cache = tmp_path / "pensieve.env" + cache.write_text("HF_TOKEN=cache-token\n") + monkeypatch.setenv("PENSIEVE_ENV_FILE", str(cache)) + env = {"HF_TOKEN": "process-token"} + + bridge._apply_launcher_env( + env, + checkout=None, + executor="slurm", + cluster_host="cluster.example.com", + account=None, + partition=None, + control_socket=None, + reconnect_command=None, + ) + + assert env["HF_TOKEN"] == "process-token" + + def test_job_status_rejects_unsafe_experiment_id(): """Experiment ids are path tokens, not filesystem paths or globs.""" result = bridge.job_status_impl("../exp_1781000008") @@ -1392,6 +1483,29 @@ def test_wait_for_experiment_returns_terminal_immediately(tmp_path, monkeypatch) assert result["waited_seconds"] < 1 # didn't actually wait +def test_wait_for_experiment_includes_terminal_log(tmp_path, monkeypatch): + """The MCP wait surface can attach the final remote log tail.""" + exp = tmp_path / "experiments" / "cicd" / "cicd_done" + exp.mkdir(parents=True) + (exp / "_DONE").touch() + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + monkeypatch.setattr( + bridge, + "read_cluster_artifact_impl", + lambda **_: {"ok": True, "content": "final output\n"}, + ) + + result = bridge.wait_for_experiment_impl( + "cicd_done", + timeout_sec=10, + poll_interval_sec=1, + include_log=True, + ) + + assert result["status"] == "done" + assert result["log_tail"] == "final output\n" + + def test_wait_for_experiment_polls_until_done(tmp_path, monkeypatch): """Spin through running → done.""" exp = tmp_path / "experiments" / "exp_in_flight" @@ -1578,24 +1692,27 @@ def test_provision_ssh_explicit_identity_overrides_default(tmp_path, monkeypatch # --------------------------------------------------------------------------- -# read_cluster_artifact — logs mode (subprocess mocked) +# read_cluster_artifact — logs mode # --------------------------------------------------------------------------- -def test_read_cluster_artifact_logs_mode_ok(monkeypatch): - """path=None → wraps `nemo experiment logs `.""" - captured = {} - - def fake_run(argv, **kwargs): - captured["argv"] = argv - return subprocess.CompletedProcess( - args=argv, - returncode=0, - stdout="line 1\nline 2\nline 3\n", - stderr="", - ) +def test_read_cluster_artifact_logs_mode_ok(tmp_path, monkeypatch): + """path=None reads the task log through the stored remote tunnel.""" + exp = tmp_path / "experiments" / "cicd" / "cicd_42" + exp.mkdir(parents=True) + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + from nemo_run.run.experiment import Experiment - monkeypatch.setattr(subprocess, "run", fake_run) + tunnel = SimpleNamespace( + job_dir="/remote/cicd/cicd_42", + run=lambda command, warn: SimpleNamespace(stdout="line 1\nline 2\nline 3\n"), + ) + job = SimpleNamespace(id="task_0", executor=SimpleNamespace(tunnel=tunnel)) + monkeypatch.setattr( + Experiment, + "_from_config", + lambda _: SimpleNamespace(jobs=[job]), + ) result = bridge.read_cluster_artifact_impl( experiment_id="cicd_42", path=None, @@ -1604,49 +1721,32 @@ def fake_run(argv, **kwargs): assert result["ok"] is True assert result["mode"] == "logs" assert "line 1" in result["content"] - # Verify the wrapped command - assert "nemo" in captured["argv"] - assert "experiment" in captured["argv"] - assert "logs" in captured["argv"] - assert "cicd_42" in captured["argv"] -def test_read_cluster_artifact_logs_mode_subprocess_failed(monkeypatch): - """Nemo cli non-zero → structured logs_fetch_failed.""" +def test_read_cluster_artifact_logs_mode_subprocess_failed(tmp_path, monkeypatch): + """Tunnel failures produce structured logs_fetch_failed.""" + exp = tmp_path / "experiments" / "cicd" / "cicd_42" + exp.mkdir(parents=True) + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + from nemo_run.run.experiment import Experiment - def fake_run(argv, **kwargs): - return subprocess.CompletedProcess( - args=argv, - returncode=1, - stdout="", - stderr="experiment cicd_42 not found\n", - ) + def fail_run(command, warn): + raise RuntimeError("experiment log unavailable") - monkeypatch.setattr(subprocess, "run", fake_run) - result = bridge.read_cluster_artifact_impl( - experiment_id="cicd_42", - path=None, - job_idx=0, + tunnel = SimpleNamespace(job_dir="/remote/cicd/cicd_42", run=fail_run) + job = SimpleNamespace(id="task_0", executor=SimpleNamespace(tunnel=tunnel)) + monkeypatch.setattr( + Experiment, + "_from_config", + lambda _: SimpleNamespace(jobs=[job]), ) - assert result["ok"] is False - assert result["reason"] == "logs_fetch_failed" - assert result["exit_code"] == 1 - - -def test_read_cluster_artifact_logs_mode_timeout(monkeypatch): - """Hanging tunnel → structured logs_fetch_timeout, no exception.""" - - def fake_run(argv, **kwargs): - raise subprocess.TimeoutExpired(cmd=argv, timeout=60) - - monkeypatch.setattr(subprocess, "run", fake_run) result = bridge.read_cluster_artifact_impl( experiment_id="cicd_42", path=None, job_idx=0, ) assert result["ok"] is False - assert result["reason"] == "logs_fetch_timeout" + assert result["reason"] == "logs_fetch_failed" # --------------------------------------------------------------------------- diff --git a/tools/mcp/tests/test_server.py b/tools/mcp/tests/test_server.py new file mode 100644 index 00000000000..0e6723bd993 --- /dev/null +++ b/tools/mcp/tests/test_server.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the modelopt-mcp tool schema.""" + +from modelopt_mcp.server import _build_server + + +def test_submit_job_exposes_resolved_slurm_list_fields(): + """Cluster inventory lists must be accepted by the typed submit tool.""" + tool = _build_server()._tool_manager.get_tool("submit_job") + properties = tool.parameters["properties"] + + for field in ("container_mounts", "srun_args"): + assert properties[field]["anyOf"][0] == { + "items": {"type": "string"}, + "type": "array", + } + + +def test_wait_for_experiment_includes_logs_by_default(): + """Terminal waits should return the final task log without a second tool call.""" + tool = _build_server()._tool_manager.get_tool("wait_for_experiment") + properties = tool.parameters["properties"] + + assert properties["include_log"]["default"] is True + assert properties["log_job_idx"]["default"] == 0