Model-Optimizer/tests/_test_utils/examples/run_command.py

# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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.
"""Utility functions for running example commands reused in multiple example tests."""

import contextlib
import os
import signal
import subprocess
import threading
import time
import traceback
import warnings
from pathlib import Path

from _test_utils.torch.distributed.utils import get_free_port

MODELOPT_ROOT = Path(__file__).parents[3]

# Substrings that identify a *transient* HuggingFace Hub / dataset access failure
_HF_TRANSIENT_MARKERS = (
    "HfHubHTTPError",
    "Too Many Requests",
    "500 Server Error",
    "502 Bad Gateway",
    "503 Server Error",
    "504 Server Error",
    "Bad Gateway",
    "Service Unavailable",
    "Gateway Time-out",
    "ConnectionError",
    "ReadTimeout",
    "ConnectTimeout",
    "Max retries exceeded",
    "NewConnectionError",
    "Connection reset by peer",
    "Connection aborted",
    "Consistency check failed",  # partial / interrupted HF download
    "couldn't connect to 'https://huggingface.co'",
    "Couldn't reach",  # datasets: "Couldn't reach <repo> on the Hub"
    "Temporary failure in name resolution",  # transient DNS
    "Name or service not known",  # transient DNS
)
_HF_MAX_RETRIES = 1
_HF_RETRY_DELAY_S = 10


def extend_cmd_parts(cmd_parts: list[str], **kwargs):
    for key, value in kwargs.items():
        if value is not None:
            cmd_parts.extend([f"--{key}", str(value)])
    if kwargs.get("trust_remote_code", False):
        cmd_parts.append("--trust_remote_code")
    return cmd_parts


# Grace period for descendants to flush and close the inherited output pipe after the command exits,
# then a short bound on the post-SIGKILL drain (the fds close as the kernel tears the group down).
_ORPHAN_PIPE_TIMEOUT_S = 30
_KILLED_PIPE_TIMEOUT_S = 5


def _run_capturing(cmd_parts: list[str], cwd: Path, env: dict[str, str]) -> tuple[int, str]:
    """Run a command, streaming and capturing combined stdout/stderr to catch transient HF errors.

    Drained by a reader thread rather than ``subprocess.run``, which waits for EOF: a descendant
    outliving the command keeps the pipe open, hanging the read and discarding every log line.
    """
    chunks: list[str] = []
    process = subprocess.Popen(
        cmd_parts,
        cwd=cwd,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        start_new_session=True,
    )
    # start_new_session makes the child its own group leader, so the group id is its pid (read now:
    # os.getpgid() stops resolving it once wait() reaps the child).
    pgid = process.pid

    def _kill_process_group():
        with contextlib.suppress(ProcessLookupError, PermissionError):
            if os.name == "posix":
                os.killpg(pgid, signal.SIGKILL)
            else:  # no process groups; the command itself is the best we can reach
                process.kill()

    def _drain():
        with contextlib.suppress(ValueError):  # the stream is closed under us on the escape path
            for line in process.stdout:  # type: ignore[union-attr]
                print(line, end="")
                chunks.append(line)

    reader = threading.Thread(target=_drain, daemon=True)
    reader.start()
    try:
        if os.name == "posix":
            # Wait *without* reaping: a reaped pid can be recycled, and pgid == the child's pid, so
            # reaping before the kill below would risk signalling an unrelated group.
            os.waitid(os.P_PID, process.pid, os.WEXITED | os.WNOWAIT)
        else:  # no process groups, so nothing to protect the pid for
            process.wait()
    except BaseException:
        # Interrupted (most likely pytest-timeout) while the command runs: its own process group is
        # not swept up with pytest's, so take the descendants down rather than leak them.
        _kill_process_group()
        with contextlib.suppress(Exception):
            process.wait(timeout=_KILLED_PIPE_TIMEOUT_S)  # reap: WNOWAIT left it waitable
        raise
    reader.join(_ORPHAN_PIPE_TIMEOUT_S)

    if reader.is_alive():
        warnings.warn(
            f"{cmd_parts[0]} left descendants holding its output pipe open; killing the process "
            "group. Output may be truncated."
        )
        _kill_process_group()
        reader.join(_KILLED_PIPE_TIMEOUT_S)

    returncode = process.wait()

    with contextlib.suppress(Exception):
        process.stdout.close()  # type: ignore[union-attr]
    return returncode, "".join(chunks)


# Set by a suite's conftest to run a command without spawning a subprocess (see
# ``_test_utils.examples.megatron_example_runner``). Unset means always use a subprocess.
_in_process_runner = None


def set_in_process_runner(runner) -> None:
    """Register ``runner(cmd_parts, example_path) -> str``; it must not fall back to a subprocess."""
    global _in_process_runner
    _in_process_runner = runner


def run_example_command(
    cmd_parts: list[str],
    example_path: str,
    setup_free_port: bool = False,
    env: dict[str, str] | None = None,
    hf_max_retries: int = _HF_MAX_RETRIES,
    hf_retry_delay_s: int = _HF_RETRY_DELAY_S,
) -> str | None:
    """Run an example command, retrying transient HuggingFace access errors."""
    print(f"[{example_path}] Running command: {cmd_parts}")
    in_process = _in_process_runner if env is None else None
    if _in_process_runner is not None and env is not None:
        # The in-process runner uses the ambient environment, so a caller-supplied env would be
        # silently dropped. Fall back to a subprocess rather than run with the wrong environment.
        warnings.warn(f"[{example_path}] env= given; running this step as a subprocess")
    env = os.environ.copy() if env is None else env
    cwd = MODELOPT_ROOT / "examples" / example_path

    for attempt in range(hf_max_retries + 1):
        if setup_free_port:
            # Subprocess steps only: an in-process runner picks its own free port, since env is a
            # copy it never sees.
            env["MASTER_PORT"] = str(get_free_port())  # fresh port per attempt
        if in_process is not None:
            # Inside the loop so in-process steps get the same transient-HuggingFace retries;
            # these tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail").
            try:
                result = in_process(cmd_parts, example_path)
            except Exception as e:
                # Re-raise unless it looks transient, so a real failure keeps its traceback
                # instead of being flattened into CalledProcessError.
                # format_exc() walks __cause__/__context__: most markers are exception type
                # names that only ever appear in a formatted traceback, never in str(e) -- a Hub
                # ConnectionError typically surfaces wrapped in a DatasetGenerationError. The
                # runner also attaches what the step printed, for launcher-level failures.
                captured = getattr(e, "captured_output", "")
                text = f"{traceback.format_exc()}\n{captured}"
                if attempt == hf_max_retries or not any(
                    marker in text for marker in _HF_TRANSIENT_MARKERS
                ):
                    raise
                returncode, output = 1, text
            else:
                assert result is not None, "an in-process runner must return output, not fall back"
                return result
        else:
            returncode, output = _run_capturing(cmd_parts, cwd, env)
        if returncode == 0:
            return output
        transient = any(marker in output for marker in _HF_TRANSIENT_MARKERS)
        if not transient or attempt == hf_max_retries:
            raise subprocess.CalledProcessError(returncode, cmd_parts, output=output)
        warnings.warn(
            f"[{example_path}] transient HuggingFace access error; retrying in "
            f"{hf_retry_delay_s}s (attempt {attempt + 1}/{hf_max_retries})"
        )
        time.sleep(hf_retry_delay_s)


def run_hf_ptq_command(*, model: str, quant: str | None = None, vlm: bool = False, **kwargs):
    kwargs.update({"model": model, "quant": quant})
    kwargs.setdefault("tasks", "quant")
    kwargs.setdefault("calib", 16)

    cmd_parts = ["scripts/huggingface_example.sh", "--no-verbose"]
    if vlm:
        # VLM PTQ shares the hf_ptq entry point; --vlm runs the multimodal deploy smoke test.
        cmd_parts.append("--vlm")
    cmd_parts = extend_cmd_parts(cmd_parts, **kwargs)
    run_example_command(cmd_parts, "hf_ptq")