From 8ba76d6a8354b871acdad533ea5c7a3c8297921f Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:12:13 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20Add=20docstrings=20to=20`gkarch/?= =?UTF-8?q?add=5Fdownstream=5Feval`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Docstrings generation was requested by @grzegorz-k-karch. * https://github.com/NVIDIA/Model-Optimizer/pull/2104#issuecomment-5215641305 The following files were modified: * `modelopt/torch/puzzletron/benchmarks/aiperf.py` * `modelopt/torch/puzzletron/orchestration/adapters/pool.py` * `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` * `modelopt/torch/puzzletron/orchestration/compiler.py` * `modelopt/torch/puzzletron/orchestration/controller.py` * `modelopt/torch/puzzletron/orchestration/executors/slurm.py` * `modelopt/torch/puzzletron/orchestration/progress.py` * `modelopt/torch/puzzletron/orchestration/task_launcher.py` * `modelopt/torch/puzzletron/post_mip/builtin.py` * `modelopt/torch/puzzletron/post_mip/reporting.py` * `modelopt/torch/puzzletron/post_mip/runner.py` * `modelopt/torch/puzzletron/stages/pipeline.py` * `modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py` * `puzzletron_setup/bundle.py` * `puzzletron_setup/v2/validation.py` * `puzzletron_setup/v2/wizard.py` * `puzzletron_setup/wizard.py` --- .../torch/puzzletron/benchmarks/aiperf.py | 60 +++++- .../puzzletron/orchestration/adapters/pool.py | 14 ++ .../orchestration/adapters/post_mip.py | 13 ++ .../puzzletron/orchestration/compiler.py | 26 ++- .../puzzletron/orchestration/controller.py | 23 +++ .../orchestration/executors/slurm.py | 30 ++- .../puzzletron/orchestration/progress.py | 11 ++ .../puzzletron/orchestration/task_launcher.py | 37 +++- modelopt/torch/puzzletron/post_mip/builtin.py | 1 + .../torch/puzzletron/post_mip/reporting.py | 23 ++- modelopt/torch/puzzletron/post_mip/runner.py | 183 +++++++++++++++++- modelopt/torch/puzzletron/stages/pipeline.py | 82 +++++++- .../subblock_stats/calc_subblock_stats.py | 111 ++++++++++- puzzletron_setup/bundle.py | 28 +++ puzzletron_setup/v2/validation.py | 7 +- puzzletron_setup/v2/wizard.py | 62 +++++- puzzletron_setup/wizard.py | 68 ++++++- 17 files changed, 749 insertions(+), 30 deletions(-) diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index 5824247f8be..ebbeea66fde 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -114,7 +114,14 @@ def _canonical_topology(topology: dict[str, Any]) -> dict[str, Any]: def _topology_vllm_args(topology: dict[str, Any]) -> list[str]: - """Translate the canonical TP/PP/DP/EP/CP contract to vLLM CLI arguments.""" + """Convert canonical parallelism settings into vLLM command-line arguments. + + Parameters: + topology (dict[str, Any]): Topology configuration to normalize and convert. + + Returns: + list[str]: vLLM command-line arguments for the configured tensor, pipeline, context, data, and expert parallelism. + """ canonical = _canonical_topology(topology) args = [ @@ -164,7 +171,17 @@ def _has_vllm_option(args: Iterable[str], *options: str) -> bool: def _server_vllm_args( checkpoint_dir: Path, topology: dict[str, Any], concurrency_values: Iterable[int] ) -> list[str]: - """Build stable vLLM server args derived from topology and benchmark demand.""" + """ + Build vLLM server arguments from checkpoint configuration, topology, and concurrency demand. + + Parameters: + checkpoint_dir (Path): Directory containing the model checkpoint. + topology (dict[str, Any]): Topology and extra vLLM argument configuration. + concurrency_values (Iterable[int]): Requested concurrency levels used to set the default maximum number of sequences. + + Returns: + list[str]: Command-line arguments for the vLLM server. + """ args = _topology_vllm_args(topology) args.extend(_descriptor_vllm_args(checkpoint_dir)) @@ -178,7 +195,16 @@ def _server_vllm_args( def _exact_length_extra_inputs( extra_inputs: dict[str, Any] | None, output_tokens: int ) -> dict[str, Any]: - """Guarantee the measured OSL unless the caller chose an explicit policy.""" + """ + Configure extra input settings for exact output-length measurements. + + Parameters: + extra_inputs (dict[str, Any] | None): Optional caller-provided input settings. + + Returns: + dict[str, Any]: The input settings with ``ignore_eos`` enabled when neither + ``ignore_eos`` nor ``min_tokens`` was explicitly provided. + """ resolved = dict(extra_inputs or {}) if "ignore_eos" not in resolved and "min_tokens" not in resolved: resolved["ignore_eos"] = True @@ -409,7 +435,33 @@ def run_aiperf_sweep( benchmark_timeout: float = 600, gpu_telemetry: str | None = "pynvml", ) -> list[BenchmarkResult]: - """Run multiple concurrencies against one persistent vLLM server.""" + """ + Run multiple concurrency benchmarks against a persistent vLLM server. + + Parameters: + checkpoint_dir (str | Path): Model checkpoint directory. + artifact_dir (str | Path): Directory for benchmark outputs and logs. + concurrencies (Iterable[int]): Unique positive concurrency levels to benchmark. + input_tokens (int): Synthetic input length for each request. + output_tokens (int): Synthetic output length for each request. + gpu_ids (str): GPU visibility specification for the benchmark processes. + topology (dict[str, Any]): vLLM parallelism and environment configuration. + request_counts (dict[int, int] | None): Optional request count for each concurrency level. + solution_id (str): Identifier for the benchmark solution. + profile_id (str): Identifier for the benchmark profile. + topology_id (str | None): Optional topology identifier. + executable (str | Path): AIPerf executable to run. + endpoint_type (str): Endpoint type used by AIPerf. + extra_inputs (dict[str, Any] | None): Additional AIPerf input settings. + use_server_token_count (bool): Whether to use token counts reported by the server. + seed (int): Seed for synthetic request generation. + readiness_timeout (float): Maximum time to wait for vLLM readiness, in seconds. + benchmark_timeout (float): Maximum time allowed for each benchmark, in seconds. + gpu_telemetry (str | None): GPU telemetry backend, or None to disable telemetry. + + Returns: + list[BenchmarkResult]: Benchmark results in the original concurrency order. + """ checkpoint_dir = Path(checkpoint_dir).resolve() artifact_dir = Path(artifact_dir).resolve() diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 7abee0a5263..098c05261b6 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -224,6 +224,20 @@ def command( runner, overrides: list[str] | None = None, ) -> AttemptSpec: + """ + Build the execution command and resource allocation for a planned work item. + + Parameters: + plan (CampaignPlan): Campaign configuration and execution context. + node (StagePlanNode): Stage and resource configuration for the work item. + item (WorkItem): Work item metadata, role, and local GPU assignments. + attempt_id (str): Identifier for the execution attempt. + runner: Runner context containing the repository location. + overrides (list[str] | None): Optional configuration overrides to apply. + + Returns: + AttemptSpec: Command, environment, resource allocation, and execution metadata for the work item. + """ repo = Path(runner.contract.repository) role = item.metadata.get("role", "worker") log_dir = plan.puzzle_dir / "logs" diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 90f1f745f9a..6887a744beb 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -108,6 +108,19 @@ class PostMIPAdapter(WorkAdapter): strategy = ExecutionStrategy.SHARDED def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: + """ + Plan sharded execution for a post-MIP node and mark aggregation as required. + + Parameters: + plan (CampaignPlan): Campaign execution plan containing node configuration and candidate information. + node (StagePlanNode): Post-MIP node to plan. + + Returns: + WorkPlan: Work plan containing the node's work item and execution strategy. + + Raises: + RuntimeError: If an evaluation node has no candidate architectures to evaluate. + """ config = _node_config(plan, node.stage_id) node_type = str(config.get("type")) count = 1 if node_type in {"filter", "manual_filter"} else node.instances diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index 05944ed6288..ba915faec15 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -58,6 +58,15 @@ def _mapping(value: Any) -> dict[str, Any]: + """ + Convert a mapping to a dictionary. + + Parameters: + value (Any): The value to convert. + + Returns: + dict[str, Any]: A dictionary containing the mapping's entries, or an empty dictionary for other values. + """ return dict(value) if isinstance(value, Mapping) else {} @@ -421,7 +430,22 @@ def compile_campaign_plan( overrides: list[str] | None = None, stage_filter: str | None = None, ) -> CampaignPlan: - """Compile one campaign plan from experiment + runner + execution configs.""" + """ + Compile a campaign plan from experiment, runner, and execution configurations. + + Parameters: + experiment_config_path: Path to the experiment configuration file. + runner: Runner environment used to execute the campaign. + execution: Execution defaults and per-stage settings. + overrides: Optional experiment configuration overrides. + stage_filter: Optional stage identifier limiting the plan to one enabled stage. + + Returns: + A compiled campaign plan containing stage meshes, dependencies, resources, and GPU allocations. + + Raises: + ValueError: If the selected stage is disabled, a CPU stage requests multiple instances, or a mesh override conflicts with its topology. + """ experiment_path = Path(experiment_config_path) experiment_config = load_experiment_config(experiment_path, overrides=overrides or []) diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index 335066e7100..48deddeb2f8 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -63,6 +63,19 @@ def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor: + """ + Create an executor for the campaign plan's configured runner. + + Parameters: + plan (CampaignPlan): Campaign plan containing runner configuration. + local (bool): Whether to use a local executor instead of the configured runner. + + Returns: + Executor: Executor configured for local, Slurm, or bare-metal SSH execution. + + Raises: + ValueError: If the configured runner kind is unsupported. + """ if local: return LocalExecutor(plan.runner) if plan.runner.kind == "slurm": @@ -80,6 +93,16 @@ def _stage_dashboard_display_name( *, granularity: str | None = None, ) -> str: + """Resolve the dashboard display name for a campaign stage. + + Parameters: + config (Mapping[str, Any]): Campaign configuration containing post-MIP flow definitions. + stage_id (str): Stage identifier to format. + granularity (str | None): Optional naming granularity. + + Returns: + str: ``"Downstream Evaluation"`` for downstream evaluation post-MIP stages; otherwise, the formatted stage name. + """ if stage_id.startswith("post."): parts = stage_id.split(".", 2) if len(parts) == 3: diff --git a/modelopt/torch/puzzletron/orchestration/executors/slurm.py b/modelopt/torch/puzzletron/orchestration/executors/slurm.py index 047c2a0ce9f..b94e2c7319f 100644 --- a/modelopt/torch/puzzletron/orchestration/executors/slurm.py +++ b/modelopt/torch/puzzletron/orchestration/executors/slurm.py @@ -91,7 +91,15 @@ def render_hook_lines(commands: Sequence[str]) -> str: def _render_host_container_env(repository: str) -> str: - """Render host-side container runtime defaults for Pyxis/Enroot.""" + """ + Render shell commands that configure default Pyxis/Enroot paths and create their directories. + + Parameters: + repository (str): Repository path used to derive default Enroot cache and data paths. + + Returns: + str: Shell commands for configuring and preparing the container runtime environment. + """ cache_root = Path(repository) / ".cache" / "enroot" lines = [ @@ -125,7 +133,25 @@ def render_sbatch_script( qos: str | None, job_name: str, ) -> str: - """Render one sbatch script for an attempt.""" + """ + Render an executable Slurm batch script for an attempt, including resource + allocations, environment setup, hooks, logging, and optional container + configuration. + + Parameters: + attempt (AttemptSpec): Attempt specification containing the command and + requested task topology. + runner (RunnerEnvironment): Runner configuration used for repository, + environment, and container settings. + partition (str): Slurm partition for the job. + account (str): Slurm account for the job. + time_limit (str): Slurm time limit. + qos (str | None): Optional Slurm quality-of-service name. + job_name (str): Name assigned to the Slurm job. + + Returns: + str: The generated executable sbatch script. + """ contract = runner.contract topology = resolve_task_topology(attempt) diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index a9c943357e6..c89c276d66a 100644 --- a/modelopt/torch/puzzletron/orchestration/progress.py +++ b/modelopt/torch/puzzletron/orchestration/progress.py @@ -387,6 +387,17 @@ def _post_mip_progress( stage_id: str, config: Mapping[str, Any] | None, ) -> str | None: + """ + Summarize post-MIP candidate processing progress for a configured node. + + Parameters: + puzzle_dir (Path): Root directory containing post-MIP artifacts. + stage_id (str): Identifier of the post-MIP stage and node. + config (Mapping[str, Any] | None): Configuration containing post-MIP flow and node definitions. + + Returns: + str | None: Progress summary with completed, failed, and timed-out candidate counts, or None when progress data is unavailable or the node is not applicable. + """ parts = stage_id.split(".", 2) if len(parts) != 3: return None diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 42606b6180b..6428ada478e 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -134,7 +134,18 @@ def build_task_command( binding: TaskBinding, gpus_per_task: int, ) -> tuple[str, ...]: - """Wrap an application payload in torchrun when the topology requests it.""" + """ + Build the command used to launch the application for the selected launcher and task topology. + + Parameters: + payload (Sequence[str]): Application command and its arguments. + launcher (TaskLauncher): Launcher mode that determines whether to wrap the payload. + binding (TaskBinding): Resolved task placement and rendezvous information. + gpus_per_task (int): Number of processes to launch per node when using distributed execution. + + Returns: + tuple[str, ...]: The original payload for direct execution, or a torchrun command configured for the task topology. + """ command = tuple(str(part) for part in payload) if launcher is TaskLauncher.DIRECT: @@ -174,6 +185,20 @@ def _direct_distributed_env(binding: TaskBinding) -> dict[str, str]: def _required_index(env: Mapping[str, str], primary: str, fallback: str) -> int: + """ + Read a task index from the primary environment variable or its fallback. + + Parameters: + env (Mapping[str, str]): Environment variables containing the task index. + primary (str): Preferred environment variable name. + fallback (str): Alternate environment variable name. + + Returns: + int: The task index parsed from the selected environment variable. + + Raises: + RuntimeError: If neither environment variable is set. + """ value = env.get(primary, env.get(fallback)) if value is None: raise RuntimeError(f"missing task identity: set {primary} or {fallback}") @@ -196,7 +221,15 @@ def _parser() -> argparse.ArgumentParser: def main(argv: Sequence[str] | None = None) -> int: - """Resolve this task's binding and replace the launcher with its payload.""" + """ + Resolve the task's distributed binding, prepare its execution environment, and replace the current process with the payload command. + + Parameters: + argv (Sequence[str] | None): Optional command-line arguments to parse instead of the process arguments. + + Returns: + int: Zero after replacing the current process with the payload command. + """ args = _parser().parse_args(argv) payload = tuple(args.payload[1:] if args.payload[:1] == ["--"] else args.payload) diff --git a/modelopt/torch/puzzletron/post_mip/builtin.py b/modelopt/torch/puzzletron/post_mip/builtin.py index 050352af445..52ab60509d5 100644 --- a/modelopt/torch/puzzletron/post_mip/builtin.py +++ b/modelopt/torch/puzzletron/post_mip/builtin.py @@ -125,4 +125,5 @@ class DownstreamEvaluationNode(PostMIPNode): @classmethod def render_report(cls, node, payload): + """Render the downstream evaluation report for the payload's section.""" return render_downstream_evaluation_report(str(payload["section_id"]), payload) diff --git a/modelopt/torch/puzzletron/post_mip/reporting.py b/modelopt/torch/puzzletron/post_mip/reporting.py index bc287990805..c21d61ce4a9 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -319,7 +319,17 @@ def render_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str def render_aiperf_report(section_id: str, payload: Mapping[str, Any]) -> str: - """Render AIPerf throughput/latency observations and timeout evidence.""" + """ + Render AIPerf candidate status, performance metrics, selection markers, and errors. + + Parameters: + section_id (str): Identifier used to scope the throughput chart element. + payload (Mapping[str, Any]): AIPerf observations and status data. + + Returns: + str: HTML fragment containing the status summary, throughput chart placeholder, + and candidate metrics table. + """ observations = list(payload.get("observations") or ()) rows = [] @@ -368,7 +378,16 @@ def render_downstream_evaluation_report(section_id: str, payload: Mapping[str, A def render_global_kd_report(section_id: str, payload: Mapping[str, Any]) -> str: - """Render several candidate KD histories on shared, lineage-colored plots.""" + """ + Render the Short KD comparison with candidate statuses, loss plots, and run summaries. + + Parameters: + section_id (str): Identifier used to generate unique plot element IDs. + payload (Mapping[str, Any]): Short KD runs and status data to display. + + Returns: + str: HTML fragment containing the comparison summary, plot placeholders, and run table. + """ runs = list(payload.get("runs") or ()) rows = [] diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 1bbe463d0b7..3d23e0d05b0 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -470,6 +470,17 @@ def _evaluate( def _aiperf( config: dict[str, Any], node: CompiledPostMIPNode, source, execution_identity: str ) -> dict[str, Any]: + """Run an AI performance sweep for a checkpoint across configured concurrency levels. + + Parameters: + config (dict[str, Any]): Workflow configuration used to determine the execution directory. + node (CompiledPostMIPNode): Compiled post-MIP node containing benchmark settings. + source: Candidate source containing the checkpoint and architecture identifier. + execution_identity (str): Identifier for the current node execution. + + Returns: + dict[str, Any]: Benchmark metrics and paths to the raw result artifacts. + """ from ..benchmarks import run_aiperf_sweep settings = dict(node.config.get("config") or {}) @@ -541,10 +552,27 @@ def _aiperf( def _as_cli_bool(value: bool) -> str: + """Convert a Boolean value to the CLI-compatible ``"True"`` or ``"False"`` string. + + Parameters: + value (bool): The Boolean value to convert. + + Returns: + str: ``"True"`` for true values and ``"False"`` for false values. + """ return "True" if value else "False" def _as_lmms_eval_arg(value: Any) -> str: + """ + Convert a value to the command-line argument format expected by lmms-eval. + + Parameters: + value (Any): The value to convert. + + Returns: + str: The formatted command-line argument value. + """ if isinstance(value, bool): return _as_cli_bool(value) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -555,6 +583,20 @@ def _as_lmms_eval_arg(value: Any) -> str: def _join_cli_values(value: Any, *, path: str) -> str: + """ + Convert a string or sequence of values into a comma-separated CLI value. + + Parameters: + value (Any): String or sequence of values to normalize. + path (str): Configuration path used in validation errors. + + Returns: + str: The normalized comma-separated value. + + Raises: + TypeError: If value is neither a string nor a sequence. + ValueError: If value is empty or contains an empty item. + """ if isinstance(value, str): text = value.strip() if not text: @@ -569,6 +611,18 @@ def _join_cli_values(value: Any, *, path: str) -> str: def _model_arg_string(values: Mapping[str, Any]) -> str: + """ + Convert model arguments to lmms-eval's comma-separated argument format. + + Parameters: + values (Mapping[str, Any]): Model argument names and values. + + Returns: + str: A comma-separated string of rendered key-value arguments. + + Raises: + ValueError: If an argument key or value is invalid, or if no arguments are provided. + """ parts = [] for key, value in values.items(): if value is None: @@ -589,6 +643,16 @@ def _model_arg_string(values: Mapping[str, Any]) -> str: def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> str: + """ + Merge checkpoint, topology, and supported model settings into lmms-eval model arguments. + + Parameters: + settings (Mapping[str, Any]): Downstream evaluation settings containing optional model arguments and configuration overrides. + checkpoint (str): Path to the checkpoint used for evaluation. + + Returns: + str: Comma-separated lmms-eval model arguments. + """ raw = settings.get("model_args") checkpoint_arg = str(settings.get("checkpoint_arg", "model")) topology = dict(settings.get("topology") or {}) @@ -625,6 +689,18 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> def _command_prefix(settings: Mapping[str, Any]) -> list[str]: + """ + Resolve the command prefix used to invoke lmms-eval. + + Parameters: + settings (Mapping[str, Any]): Downstream evaluation settings containing an optional command prefix. + + Returns: + list[str]: The configured command prefix, or the current Python interpreter followed by the lmms-eval module. + + Raises: + ValueError: If the configured command prefix is empty or contains an empty value. + """ raw = settings.get("command_prefix") if raw is None: return [sys.executable, "-m", "lmms_eval"] @@ -643,7 +719,17 @@ def _lmms_eval_command( checkpoint: str, output_path: Path, ) -> tuple[list[str], dict[str, str], float | None]: - """Build a deterministic lmms-eval CLI invocation for one realized checkpoint.""" + """ + Builds an lmms-eval command, environment, and optional timeout for a realized checkpoint. + + Parameters: + settings (Mapping[str, Any]): Downstream evaluation settings. + checkpoint (str): Path to the realized checkpoint. + output_path (Path): Directory for lmms-eval output. + + Returns: + tuple[list[str], dict[str, str], float | None]: The command arguments, environment variables, and timeout in seconds. + """ tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") argv = [ @@ -697,6 +783,15 @@ def _lmms_eval_command( def _metric_key(value: Any) -> str: + """ + Normalize a metric name component for use in metric keys. + + Parameters: + value (Any): The value to convert into a normalized metric name component. + + Returns: + str: The stripped string representation with spaces, commas, and slashes replaced by underscores. + """ return ( str(value) .strip() @@ -708,6 +803,15 @@ def _metric_key(value: Any) -> str: def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: + """ + Flatten finite numeric task metrics from an lmms-eval result payload. + + Parameters: + payload (Mapping[str, Any]): Result payload containing task metrics under the ``results`` key. + + Returns: + dict[str, float]: Metric names mapped to finite numeric values, or an empty dictionary when no valid results are present. + """ results = payload.get("results") if not isinstance(results, Mapping): return {} @@ -726,6 +830,18 @@ def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: + """ + Finds the newest valid lmms-eval result payload under an output directory. + + Parameters: + output_path (Path): Directory containing lmms-eval output files. + + Returns: + tuple[dict[str, Any], Path]: The result payload and path of the newest JSON file containing a `results` mapping. + + Raises: + FileNotFoundError: If no valid result JSON file is found. + """ candidates = [] for path in sorted(output_path.rglob("*.json")): try: @@ -743,6 +859,16 @@ def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: def _write_lmms_eval_streams( output_path: Path, result: subprocess.CompletedProcess[str] ) -> dict[str, str]: + """ + Persist non-empty lmms-eval subprocess output streams and return their artifact paths. + + Parameters: + output_path (Path): Directory where stream files are written. + result (subprocess.CompletedProcess[str]): Completed subprocess result containing captured output. + + Returns: + dict[str, str]: Mapping of stream path keys to the paths of written output files. + """ stream_paths = {} for stream_name, text in (("stdout", result.stdout), ("stderr", result.stderr)): if not text: @@ -754,6 +880,16 @@ def _write_lmms_eval_streams( def _lmms_eval_output_tail(result: subprocess.CompletedProcess[str], *, max_lines: int = 20) -> str: + """ + Format the most recent subprocess output lines from stderr and stdout. + + Parameters: + result (subprocess.CompletedProcess[str]): Completed process containing captured output. + max_lines (int): Maximum number of lines to include from each stream. + + Returns: + str: Formatted stderr and stdout output tails. + """ sections = [] for stream_name, text in (("stderr", result.stderr), ("stdout", result.stdout)): lines = (text or "").strip().splitlines() @@ -769,6 +905,23 @@ def _downstream_evaluation( source, execution_identity: str, ) -> dict[str, Any]: + """ + Run downstream lmms-eval benchmarking for a materialized checkpoint. + + Parameters: + config (dict[str, Any]): Campaign configuration used to determine execution paths. + node (CompiledPostMIPNode): Post-MIP node containing lmms-eval settings. + source: Checkpoint artifact to evaluate. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: Paths to the evaluation summary, raw result, command record, and captured streams, together with numeric metrics. + + Raises: + ValueError: If the source is not a checkpoint artifact. + RuntimeError: If lmms-eval fails or produces no numeric task metrics. + FileNotFoundError: If no valid lmms-eval result file is produced. + """ if source.artifact_kind is not ArtifactKind.CHECKPOINT: raise ValueError("downstream_evaluation requires materialized checkpoint artifacts") settings = dict(node.config.get("config") or {}) @@ -891,6 +1044,22 @@ def _run_candidate( input_revision_id: str, execution_identity: str, ) -> dict[str, Any]: + """ + Execute a candidate according to the node type and return its execution result. + + Parameters: + config (dict[str, Any]): Runtime configuration for the candidate execution. + node (CompiledPostMIPNode): Compiled node defining the execution type and model source. + ledger (CandidateLedger): Ledger containing the input candidate revision. + input_revision_id (str): Identifier of the candidate revision to execute. + execution_identity (str): Identifier for the current node execution. + + Returns: + dict[str, Any]: A successful result containing the input and source revision identifiers, architecture identifier, and executor-specific metadata. + + Raises: + ValueError: If the node type is not a supported candidate executor. + """ source = ledger.source_revision(input_revision_id, node.model_source) if node.node_type == "materialize": result = _materialize(config, node, ledger, input_revision_id, source, execution_identity) @@ -934,6 +1103,18 @@ def _distributed_shard(config: dict[str, Any], node: CompiledPostMIPNode) -> Ite def run_post_mip_node_shard( config: dict[str, Any], stage_id: str, *, shard_index: int = 0, shard_count: int = 1 ) -> Path: + """ + Execute the assigned candidate revisions for a post-MIP node shard and persist the results. + + Parameters: + config (dict[str, Any]): Post-MIP configuration. + stage_id (str): Identifier of the compiled node to execute. + shard_index (int): Zero-based index of this shard. + shard_count (int): Total number of shards distributing the candidate revisions. + + Returns: + Path: Path to the shard result artifact. + """ node = _compiled_node(config, stage_id) ledger = _ledger(config) ledger.ingest_mip(_puzzle_dir(config)) diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index 45876805f93..b723526cc1b 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -207,7 +207,13 @@ def _vllm_stats_is_explicit(config: dict[str, Any]) -> bool: def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> None: - """Append one analytical memory profile for every configured MIP workload.""" + """ + Append an analytical memory profile for each configured MIP workload. + + Parameters: + config (dict[str, Any]): Pipeline configuration containing optional MIP workloads. + hydra_cfg (Any): Base subblock-statistics configuration to customize for each workload. + """ from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats workloads = dict((config.get("mip") or {}).get("workloads") or {}) @@ -243,6 +249,15 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> def _scenario_hidden_width(puzzle_dir: Path) -> int | None: + """ + Read the hidden width from a scenario manifest. + + Parameters: + puzzle_dir (Path): Directory containing the scenario manifest. + + Returns: + int | None: The manifest's hidden width, or `None` when the manifest or value is absent. + """ manifest_path = puzzle_dir / "scenario_manifest.json" if not manifest_path.is_file(): return None @@ -258,6 +273,18 @@ def _has_runtime_measurement( measurement: Any, allow_missing_workload_id: bool = False, ) -> bool: + """ + Determine whether a statistics file contains a compatible runtime measurement. + + Parameters: + path (Path): Statistics file to inspect. + hidden_width (int): Model hidden width expected by the measurement. + measurement (Any): Runtime measurement configuration to match. + allow_missing_workload_id (bool): Whether entries without a workload identifier may match. + + Returns: + bool: `True` if a compatible runtime measurement is present, `False` otherwise. + """ try: payload = json.loads(path.read_text()) except (OSError, ValueError): @@ -303,6 +330,18 @@ def _runtime_measurement_candidate_paths( stats_path: Path, measurement: Any, ) -> list[tuple[Path, bool]]: + """ + Builds candidate paths for locating reusable runtime measurement statistics. + + Parameters: + config (dict[str, Any]): Configuration containing the statistics filename. + puzzle_dir (Path): Directory associated with the current scenario. + stats_path (Path): Primary statistics file path. + measurement (Any): Measurement configuration that may specify a relative statistics path. + + Returns: + list[tuple[Path, bool]]: Candidate statistics paths paired with a flag indicating whether each path came from a configured relative path. + """ stats_name = str( (config.get("vllm_stats") or {}).get("subblock_stats_filename", stats_path.name) ) @@ -329,6 +368,22 @@ def _runtime_reuse_source_path( hidden_width: int, measurement: Any, ) -> Path: + """ + Finds a reusable vLLM measurement file matching the requested hidden width and workload. + + Parameters: + config (dict[str, Any]): Runtime configuration used to resolve candidate measurement paths. + puzzle_dir (Path): Experiment directory containing scenario-specific measurement files. + stats_path (Path): Configured statistics file path. + hidden_width (int): Hidden width required for the reusable measurement. + measurement (Any): Workload measurement whose identity must match. + + Returns: + Path: The first candidate measurement file containing a compatible runtime measurement. + + Raises: + RuntimeError: If no candidate contains a matching reusable measurement. + """ candidates = _runtime_measurement_candidate_paths( config=config, puzzle_dir=puzzle_dir, @@ -355,7 +410,7 @@ def _refresh_scenario_runtime_workload_stats( hydra_cfg: Any, stats_path: Path, ) -> None: - """Refresh width-scenario runtime rows with the local parameter inventory identity.""" + """Refresh scenario-specific runtime statistics using measurements for the local hidden width.""" from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats puzzle_dir = _puzzle_dir(config, hydra_cfg) @@ -395,7 +450,13 @@ def _refresh_scenario_runtime_workload_stats( def _write_runtime_subblock_library(path: Path, block_configs: tuple[Any, ...]) -> None: - """Write the legacy subblock-library input without assembling a replacement library.""" + """ + Write runtime subblock configurations to a JSON library file. + + Parameters: + path (Path): Destination path for the library file. + block_configs (tuple[Any, ...]): Block configurations to serialize. + """ rows = [] for block_config in block_configs: row = { @@ -1123,6 +1184,21 @@ def bypass_overfit_stage(config: dict[str, Any], manifest: StageManifest): def build_library_stage(config: dict[str, Any], manifest: StageManifest): + """ + Build the replacement and candidate libraries and record their associated statistics. + + The stage validates and shares the resolved scoring parent, optionally refreshes runtime + statistics, calculates static workload statistics, and publishes the resulting artifact + paths and execution metadata. + + Parameters: + config (dict[str, Any]): Pipeline configuration. + manifest (StageManifest): Manifest used to record stage completion and outputs. + + Returns: + StageManifest: Updated manifest containing the generated library paths, statistics + metadata, and scoring-parent information. + """ hydra_cfg = load_runtime_hydra_config(config) puzzle_dir = _puzzle_dir(config, hydra_cfg) candidate_library_path = puzzle_dir / "candidate_library.json" diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 38580f156b7..9677b5c55c1 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -276,7 +276,16 @@ def _runtime_reuse_key_from_args( *, fallback_workload_id: str | None = None, ) -> tuple | None: - """Return the exact runtime-reuse identity represented by persisted args.""" + """ + Builds the identity used to match reusable runtime measurements. + + Parameters: + args (Mapping): Persisted calculation arguments containing runtime and workload settings. + fallback_workload_id (str | None): Workload identifier to use when `args` does not provide one. + + Returns: + tuple | None: Runtime-reuse identity, or `None` when runtime statistics are unavailable or the dtype is not bfloat16. + """ if not args.get("runtime_stats") or args.get("weights_dtype") != str(torch.bfloat16): return None @@ -307,7 +316,15 @@ def _runtime_reuse_key( generation_seq_len: int, runtime_stats_config: Mapping, ) -> tuple: - """Return the exact runtime-reuse identity requested by this calculation.""" + """ + Builds an exact identity for matching reusable runtime measurements. + + Parameters: + runtime_stats_config (Mapping): Runtime settings that determine measurement compatibility. + + Returns: + tuple: Identity containing model dimensions, runtime settings, vLLM arguments, and workload identity. + """ return ( int(width), @@ -332,7 +349,21 @@ def _reuse_runtime_stats( source_path: str, fallback_workload_id: str | None = None, ) -> dict: - """Overlay immutable measured latency onto refreshed static statistics.""" + """ + Reuse measured runtime statistics from a compatible source entry in refreshed statistics. + + Parameters: + target (dict): Statistics entry to update with reusable runtime data. + source (dict): Statistics entry containing the measured runtime data. + source_path (str): Path identifying the source statistics entry. + fallback_workload_id (str | None): Workload identifier to use when the source does not provide one. + + Returns: + dict: The updated target statistics entry. + + Raises: + KeyError: If a target subblock has no matching runtime statistics in the source entry. + """ # Synthetic vLLM timings are layer-independent: collection benchmarks the # set of unique subblock configs, while a post-scoring replacement library @@ -826,6 +857,33 @@ def calculate_subblock_stats( runtime_selection_identity: str | None = None, parameter_inventory: Mapping | None = None, ) -> dict: + """ + Compute parameter, memory, additive-metric, and optional runtime statistics for subblock configurations. + + Parameters: + calc_subblock_stats_config (DictConfig): Runtime measurement and calculation settings. + teacher_dir (Path): Directory containing the teacher model or checkpoint. + model_config (PretrainedConfig): Model configuration used for metric calculations. + descriptor (Type[ModelDescriptor]): Model descriptor defining architecture-specific behavior. + master_puzzle_dir (Path): Puzzle directory used for runtime measurement caches. + subblock_configs (list[immutabledict[str, SubblockConfig]]): Subblock configurations and their parent layer indices. + batch_size (int): Number of sequences in the workload. + prefill_seq_len (int): Input sequence length used for prefill calculations. + generation_seq_len (int): Number of generated tokens used for decode calculations. + n_embd (int): Model hidden size. + n_head (int): Number of attention heads. + vocab_size (int): Model vocabulary size. + runtime_stats_enabled (bool): Whether to measure runtime statistics. + use_cuda_graph (bool): Whether to use CUDA graphs during runtime measurement. + weights_dtype (torch.dtype): Data type used for model weights. + activations_dtype (torch.dtype): Data type used for activations. + kv_cache_dtype (torch.dtype): Data type used for the key-value cache. + runtime_selection_identity (str | None): Identity of the runtime subblock selection. + parameter_inventory (Mapping | None): Precomputed parameter inventory to use for parameter counts. + + Returns: + dict: Statistics for the requested workload, including calculation arguments, non-block statistics, and per-subblock metrics. + """ runtime_granularity = "subblock" runtime_stats_config = ( calc_subblock_stats_config.get("runtime_stats", {}) if runtime_stats_enabled else {} @@ -1246,12 +1304,26 @@ def _subblock_stats_already_complete( prefill_seq_len: int = 2048, generation_seq_len: int = 2048, ) -> bool: - """Whether ``existing_stats`` already covers every configuration this run would compute. - - When runtime benchmarking is enabled, the bf16 entries (the only ones for - which runtime is ever measured) must additionally already carry runtime - measurements **at the requested granularity** — switching subblock<->block must trigger a - recompute rather than silently reusing the other granularity's numbers. + """ + Determine whether existing statistics cover all requested configurations. + + Parameters: + existing_stats (list): Previously calculated statistics entries. + subblock_configs (list): Subblock configurations required for coverage. + batch_sizes (Iterable[int]): Batch sizes to verify. + data_types (list): Weight, activation, and KV-cache dtype combinations. + model_hidden_sizes (Iterable[int]): Model widths to verify. + runtime_stats_enabled (bool): Whether runtime measurements are required. + runtime_granularity (str): Required runtime measurement granularity. + runtime_max_num_seqs (int | None): Required maximum number of runtime sequences. + runtime_workload_id (str | None): Required runtime workload identity. + runtime_selection_identity (str | None): Required runtime subblock-selection identity. + parameter_inventory_identities (Mapping[int, str] | None): Inventory identity required for each model width. + prefill_seq_len (int): Required prefill sequence length. + generation_seq_len (int): Required generation sequence length. + + Returns: + bool: True if every requested configuration and required measurement is present, False otherwise. """ by_signature = {_arg_signature(entry["args"]): entry for entry in existing_stats} required_subblock_keys = { @@ -1354,6 +1426,27 @@ def calculate_subblock_stats_for_puzzle_dir( # from attach_helper import debugging_setup # debugging_setup() # You can optionally pass a name to identify the job (e.g. `debugging_setup(name="my_script")`) # ==== END === Setup for attach-helper ==== + """ + Compute and persist subblock statistics for all requested batch sizes, data types, and model widths. + + Parameters: + calc_subblock_stats_config (DictConfig): Configuration for statistics calculation and optional runtime measurement. + master_puzzle_dir (Path | str): Puzzle directory containing subblock configurations and output files. + teacher_dir (Path | str): Teacher checkpoint directory used for model metadata and parameter inventories. + descriptor (Type[ModelDescriptor]): Model descriptor defining architecture-specific behavior. + model_hidden_sizes (ListConfig): Hidden sizes to evaluate; the teacher hidden size is always included. + ffn_hidden_sizes (ListConfig): Additional FFN sizes to include in the subblock configurations. + batch_sizes (Iterable[int]): Batch sizes to evaluate. + prefill_seq_len (int): Number of prompt tokens used for runtime measurements. + generation_seq_len (int): Number of generated tokens used for runtime measurements. + runtime_stats_enabled (bool): Whether to compute or reuse runtime statistics. + merge_with_existing_stats (bool): Whether to update an existing incomplete statistics file. + subblock_stats_filename (str): Name of the JSON file used to persist statistics. + + Raises: + FileNotFoundError: If a configured runtime manifest or reusable runtime statistics file cannot be found. + ValueError: If runtime settings or reusable runtime statistics do not cover the requested configurations. + """ if isinstance(batch_sizes, str): batch_sizes = [ int(batch_size) for batch_size in batch_sizes.strip("[]").replace(" ", "").split(",") diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 3ee3e8f0936..13bad12f5ef 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -329,6 +329,19 @@ def _post_mip_flows( global_kd_mesh: Mapping[str, Any], default_serving_topology: Mapping[str, Any], ) -> dict[str, Any]: + """ + Prepare post-MIP flow configurations with mesh settings, serving defaults, and smoke-run limits. + + Parameters: + state (Mapping[str, Any]): Campaign state containing post-MIP flow definitions. + smoke (bool): Whether to apply reduced settings for a smoke run. + common_mesh (Mapping[str, Any]): Mesh used by evaluation and materialization nodes. + global_kd_mesh (Mapping[str, Any]): Mesh used by global knowledge-distillation nodes. + default_serving_topology (Mapping[str, Any]): Default topology for serving-based nodes. + + Returns: + dict[str, Any]: The normalized post-MIP flow configurations. + """ flows = deepcopy(_mapping(_answers(state, "post_mip").get("flows"))) for flow in flows.values(): for node in _mapping(flow.get("nodes")).values(): @@ -795,6 +808,21 @@ def _dynamic_stage_entries( *, pool_source_evaluations: bool, ) -> dict[str, Any]: + """ + Builds scheduler entries for dynamic post-MIP stages. + + Parameters: + experiment (Mapping[str, Any]): Experiment configuration containing post-MIP flows. + workers (Mapping[str, Any]): Worker limits for pooled and sharded stages. + gpus_per_node (int): Number of GPUs assigned to each node. + common (Mapping[str, Any]): Parallel configuration for evaluation stages. + single_gpu (Mapping[str, Any]): Parallel configuration for materialization stages. + cpu_partition (str | None): CPU partition to assign to CPU stages. + pool_source_evaluations (bool): Whether source evaluations should use pooled workers. + + Returns: + dict[str, Any]: Scheduler entries keyed by post-MIP flow and node identifiers. + """ entries = {} candidate_limits = _post_mip_candidate_limits(experiment) for flow_id, flow in _mapping(_mapping(experiment.get("post_mip")).get("flows")).items(): diff --git a/puzzletron_setup/v2/validation.py b/puzzletron_setup/v2/validation.py index d9c45c143e6..4000bb82189 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -186,7 +186,12 @@ def _dataset_subset_issues(state: WizardState) -> list[ValidationIssue]: def validate_state(state: WizardState) -> tuple[ValidationIssue, ...]: - """Return actionable authoring issues before canonical compilation.""" + """ + Validate wizard state and identify issues that prevent canonical compilation. + + Returns: + tuple[ValidationIssue, ...]: Validation issues sorted by configuration section and path. + """ issues: list[ValidationIssue] = [] required = ( "model.source", diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index af87270c860..766aefd04ac 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3597,6 +3597,17 @@ def _post_mip_strategy(node: NodeDraft) -> str: def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context: dict) -> bool: + """ + Configure post-MIP execution flows for each MIP run, using recommended or custom nodes. + + Parameters: + session (WizardSession): Wizard session used to read state and collect configuration. + resolver (DefaultsResolver): Resolver for stage resource defaults. + context (dict): Model and pruning context required to configure serving and evaluation nodes. + + Returns: + bool: `True` when post-MIP flows are configured, `False` when the section is exited through back navigation. + """ mip = _mapping_copy(session.state.collection("mip_config")) runs = _mapping_copy(mip.get("runs")) sequence = int(session.state.get_field("data.sequence_length", 4096)) @@ -4023,7 +4034,22 @@ def _serving_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """Ask the complete AIPerf workload and serving-only parallel setting.""" + """ + Collect AIPerf serving workload settings and the vLLM serving topology. + + Parameters: + session (WizardSession): Wizard session used to collect and validate responses. + prefix (str): State key prefix for the serving settings. + defaults (Mapping[str, Any]): Default workload and topology values. + inventory (Any): Model inventory used to validate the topology. + pruning (Mapping[str, Any]): Pruning configuration relevant to topology validation. + stage_id (str): Pruning stage associated with the serving configuration. + + Returns: + Any: A mapping containing input and output sequence lengths, concurrency values, + request count, model selection mode, and topology, or the `BACK` sentinel when + the user navigates to the previous prompt. + """ values = {} for name, label, default in ( ("input_tokens", "Serving input sequence length (ISL):", defaults["input_tokens"]), @@ -4100,9 +4126,30 @@ def _downstream_evaluation_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """Ask lmms-eval task settings and the vLLM topology used to run them.""" + """ + Collect lmms-eval tasks, execution settings, model arguments, and vLLM topology. + + Parameters: + session (WizardSession): Wizard session used to prompt for settings. + prefix (str): State-key prefix for the prompted values. + defaults (Mapping[str, Any]): Existing values used as prompt defaults. + inventory (Any): Model inventory used to validate the vLLM topology. + pruning (Mapping[str, Any]): Pruning configuration relevant to topology validation. + stage_id (str): Identifier of the stage using the evaluation settings. + + Returns: + Any: A mapping containing lmms-eval tasks, sample and batch limits, timeout, model arguments, logging settings, and vLLM topology, or `BACK` if prompting is cancelled. + """ def validate_tasks(value: str) -> bool | str: + """Validate a comma-separated list of lmms-eval tasks. + + Parameters: + value (str): Comma-separated task names. + + Returns: + bool | str: `True` if at least one task is provided, otherwise an error message. + """ tasks = [item.strip() for item in value.split(",") if item.strip()] return True if tasks else "Enter at least one lmms-eval task." @@ -4176,7 +4223,16 @@ def _configure_dynamic_resources( *, ask: bool, ) -> Any: - """Attach an independent resource/batch card to every node in one flow.""" + """ + Configure independent resource assignments for all nodes in a post-MIP flow. + + Parameters: + flow_id (str): Identifier of the flow whose nodes are configured. + ask (bool): Whether to prompt for resource and batch customizations. + + Returns: + True when configuration completes, or `BACK` when navigation is requested. + """ registry = ResourceProfileRegistry.from_dict( session.state.collection("parallel_profiles") or {} ) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 6938eeeee22..dcd7b39b898 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -635,7 +635,18 @@ def _ask_aiperf_config( runtime: Mapping[str, Any], defaults: Mapping[str, Any] | None = None, ) -> dict[str, Any]: - """Ask for one AIPerf node's independent Serving topology and workload.""" + """ + Configure an AIPerf serving node's parallel topology and workload settings. + + Parameters: + detailed (bool): Whether to prompt for workload and timeout values. + moe (bool): Whether to configure expert parallelism for a mixture-of-experts model. + runtime (Mapping[str, Any]): Runtime defaults for input length, output length, and concurrency. + defaults (Mapping[str, Any] | None): Previously saved configuration values. + + Returns: + dict[str, Any]: The configured AIPerf topology, workload, and timeout settings. + """ defaults = dict(defaults or {}) topology_defaults = dict(defaults.get("topology") or {}) checkpoint = prompts.checkpoint() @@ -735,7 +746,17 @@ def _ask_downstream_evaluation_config( moe: bool, defaults: Mapping[str, Any] | None = None, ) -> dict[str, Any]: - """Ask for lmms-eval task and vLLM settings.""" + """ + Collect lmms-eval tasks, sampling settings, vLLM topology, and evaluation timeout. + + Parameters: + detailed (bool): Whether to prompt for the per-candidate timeout. + moe (bool): Whether to allow configuring expert parallelism. + defaults (Mapping[str, Any] | None): Previously saved settings used as prompt defaults. + + Returns: + dict[str, Any]: The configured downstream evaluation settings. + """ defaults = defaults or {} tasks = prompts.text( @@ -846,6 +867,21 @@ def _default_flow( objective: Mapping[str, Any] | None = None, include_initial_filter: bool = True, ) -> dict[str, Any]: + """ + Build the standard post-MIP evaluation and selection flow. + + Parameters: + run_id (str): Identifier of the MIP run. + run (Mapping[str, Any]): MIP run configuration. + runtime (Mapping[str, Any]): Runtime settings for serving evaluation. + data (Mapping[str, Any]): Dataset settings, including sequence length. + prefix (str): Prefix applied to generated node identifiers. + objective (Mapping[str, Any] | None): Objective used to configure ranking; the run's first objective is used when omitted. + include_initial_filter (bool): Whether to include the initial MIP-score filter. + + Returns: + dict[str, Any]: Flow configuration containing the source metadata and ordered post-MIP nodes. + """ def node_id(name: str) -> str: return f"{prefix}{name}" @@ -1005,6 +1041,20 @@ def _custom_flow( detailed: bool, moe: bool, ) -> dict[str, Any]: + """ + Build a custom post-MIP evaluation flow through interactive configuration. + + Parameters: + run_id (str): Identifier of the MIP run supplying candidate models. + runtime (Mapping[str, Any]): Runtime settings used by serving evaluations. + data (Mapping[str, Any]): Dataset settings used by evaluation nodes. + used_ids (set[str]): Node IDs already in use; newly configured IDs are added. + detailed (bool): Whether to collect detailed evaluation settings. + moe (bool): Whether to enable mixture-of-experts configuration options. + + Returns: + dict[str, Any]: A flow definition containing the MIP source and configured nodes. + """ nodes: OrderedDict[str, Any] = OrderedDict() available_metrics = ["mip.score"] transformer_nodes = [] @@ -1221,6 +1271,20 @@ def _resource_rows( gpus_per_node: int, workers: Mapping[str, int], ) -> list[dict[str, Any]]: + """ + Calculate resource requirements for each campaign execution stage. + + Parameters: + state (AnswerState): Campaign configuration containing post-MIP flows and execution details. + common (Mapping[str, int]): Parallel mesh dimensions shared by common stages. + bypass (Mapping[str, int]): Parallel mesh dimensions for bypass processing. + global_kd (Mapping[str, int]): Parallel mesh dimensions for global knowledge distillation. + gpus_per_node (int): Number of GPUs available on each node. + workers (Mapping[str, int]): Worker limits for pool and sharded stages. + + Returns: + list[dict[str, Any]]: Resource rows containing each stage's name, instance count, GPUs per instance, and required node count. + """ from .bundle import _post_mip_candidate_limits, _serving_parallel rows = []