From d834842b409a73f51a2a171f738ce3c87ca71c4b Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:02:20 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20Add=20docstrings=20to=20`gkarch/?= =?UTF-8?q?add=5Flmms=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/2034#issuecomment-5124533292 The following files were modified: * `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` * `modelopt/torch/puzzletron/orchestration/compiler.py` * `modelopt/torch/puzzletron/orchestration/progress.py` * `modelopt/torch/puzzletron/post_mip/builtin.py` * `modelopt/torch/puzzletron/post_mip/reporting.py` * `modelopt/torch/puzzletron/post_mip/runner.py` * `puzzletron_setup/bundle.py` * `puzzletron_setup/v2/validation.py` * `puzzletron_setup/v2/wizard.py` * `puzzletron_setup/wizard.py` --- .../orchestration/adapters/post_mip.py | 13 ++ .../puzzletron/orchestration/compiler.py | 9 + .../puzzletron/orchestration/progress.py | 11 ++ modelopt/torch/puzzletron/post_mip/builtin.py | 9 + .../torch/puzzletron/post_mip/reporting.py | 22 ++- modelopt/torch/puzzletron/post_mip/runner.py | 175 +++++++++++++++++- puzzletron_setup/bundle.py | 28 +++ puzzletron_setup/v2/validation.py | 10 +- puzzletron_setup/v2/wizard.py | 63 ++++++- puzzletron_setup/wizard.py | 70 ++++++- 10 files changed, 400 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 90f1f745f9a..4f993841183 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: + """ + Create the work plan for a post-MIP node. + + Parameters: + plan (CampaignPlan): Campaign execution plan containing node configuration and artifacts. + node (StagePlanNode): Post-MIP node whose workers should be planned. + + Returns: + WorkPlan: A work plan with the required worker allocation and aggregation enabled. + + Raises: + RuntimeError: If an evaluation node has no candidate architectures available. + """ 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..236dcdc6d6e 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): Value to normalize. + + Returns: + dict[str, Any]: A dictionary copy of the value, or an empty dictionary if the value is not a mapping. + """ return dict(value) if isinstance(value, Mapping) else {} diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index a9c943357e6..05484b49b70 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): Post-MIP stage identifier containing the flow and node IDs. + config (Mapping[str, Any] | None): Configuration containing post-MIP flow definitions. + + Returns: + str | None: A candidate completion summary with failure or timeout counts, or `None` when progress data is unavailable. + """ parts = stage_id.split(".", 2) if len(parts) != 3: return None diff --git a/modelopt/torch/puzzletron/post_mip/builtin.py b/modelopt/torch/puzzletron/post_mip/builtin.py index 050352af445..30c70442598 100644 --- a/modelopt/torch/puzzletron/post_mip/builtin.py +++ b/modelopt/torch/puzzletron/post_mip/builtin.py @@ -125,4 +125,13 @@ class DownstreamEvaluationNode(PostMIPNode): @classmethod def render_report(cls, node, payload): + """Render the downstream evaluation report for a node. + + Parameters: + node: The node associated with the report. + payload: Report data containing the section identifier. + + Returns: + str: The rendered downstream evaluation report. + """ 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..20320c22775 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -319,7 +319,16 @@ 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 an HTML report of AIPerf candidate throughput, latency, and selection results. + + Parameters: + section_id (str): Identifier used to generate the report's chart element ID. + payload (Mapping[str, Any]): Report data containing observations and status information. + + Returns: + str: HTML containing a status summary, throughput chart, and candidate results table. + """ observations = list(payload.get("observations") or ()) rows = [] @@ -368,7 +377,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 a shared comparison of candidate short KD runs and their loss histories. + + Parameters: + section_id (str): Identifier used to construct the chart element IDs. + payload (Mapping[str, Any]): Report data containing candidate KD runs and status information. + + Returns: + str: HTML containing the comparison heading, status summary, loss chart containers, 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 3a1ee61afab..9479c84b687 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -470,6 +470,18 @@ def _evaluate( def _aiperf( config: dict[str, Any], node: CompiledPostMIPNode, source, execution_identity: str ) -> dict[str, Any]: + """ + Run an AIPerf benchmark for a candidate checkpoint across configured concurrency levels. + + Parameters: + config (dict[str, Any]): Runtime configuration used to determine the artifact output directory. + node (CompiledPostMIPNode): Node configuration containing benchmark settings. + source: Candidate source containing the checkpoint to benchmark. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: A mapping containing aggregated benchmark metrics and raw artifact paths. + """ from ..benchmarks import run_aiperf_sweep settings = dict(node.config.get("config") or {}) @@ -541,10 +553,27 @@ def _aiperf( def _as_cli_bool(value: bool) -> str: + """Convert a boolean value to its CLI string representation. + + Parameters: + value (bool): The boolean value to convert. + + Returns: + str: "True" for a true value or "False" for a false value. + """ return "True" if value else "False" def _as_lmms_eval_arg(value: Any) -> str: + """ + Convert a supported value to its deterministic LMMS-Eval argument representation. + + Parameters: + value (Any): The value to convert. + + Returns: + str: A string representation suitable for an LMMS-Eval argument. + """ if isinstance(value, bool): return _as_cli_bool(value) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -555,6 +584,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 string. + + Parameters: + value (Any): String or sequence of values to join. + path (str): Configuration path used in validation error messages. + + Returns: + str: The stripped input string or comma-separated sequence values. + + Raises: + TypeError: If value is neither a string nor a sequence. + ValueError: If value or any sequence item is empty. + """ if isinstance(value, str): text = value.strip() if not text: @@ -569,6 +612,18 @@ def _join_cli_values(value: Any, *, path: str) -> str: def _model_arg_string(values: Mapping[str, Any]) -> str: + """ + Build a comma-separated lmms-eval model arguments string. + + Parameters: + values (Mapping[str, Any]): Model argument names and values to serialize. + + Returns: + str: The serialized model 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 +644,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 configured values into an LMMS-Eval model-arguments string. + + Parameters: + settings (Mapping[str, Any]): LMMS-Eval settings containing optional model arguments and topology. + checkpoint (str): Path to the checkpoint supplied to the evaluator. + + 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 +690,18 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> def _command_prefix(settings: Mapping[str, Any]) -> list[str]: + """ + Build 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 arguments, or the default Python module invocation. + + Raises: + ValueError: If the configured command prefix is empty. + """ raw = settings.get("command_prefix") if raw is None: return [sys.executable, "-m", "lmms_eval"] @@ -643,7 +720,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.""" + """ + Build the LMMS-Eval command, environment, and optional timeout for a checkpoint evaluation. + + Parameters: + settings (Mapping[str, Any]): Evaluation settings used to construct the invocation. + checkpoint (str): Path to the checkpoint to evaluate. + output_path (Path): Directory where LMMS-Eval writes its outputs. + + 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 +784,15 @@ def _lmms_eval_command( def _metric_key(value: Any) -> str: + """ + Normalize a metric name for use as a stable key. + + Parameters: + value (Any): Value to convert into a normalized metric key. + + Returns: + str: The stripped string representation with spaces, commas, slashes, and backslashes replaced by underscores. + """ return ( str(value) .strip() @@ -708,6 +804,15 @@ def _metric_key(value: Any) -> str: def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: + """ + Flatten finite numeric LMMS-Eval task metrics into normalized metric names. + + Parameters: + payload (Mapping[str, Any]): LMMS-Eval output containing task results. + + Returns: + dict[str, float]: Numeric metrics keyed by normalized task and metric names. + """ results = payload.get("results") if not isinstance(results, Mapping): return {} @@ -726,6 +831,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 results payload under a directory. + + Parameters: + output_path (Path): Directory containing LMMS-Eval output files. + + Returns: + tuple[dict[str, Any], Path]: The results payload and the path of its JSON file. + + Raises: + FileNotFoundError: If no JSON file containing a mapping-valued ``results`` field is found. + """ candidates = [] for path in sorted(output_path.rglob("*.json")): try: @@ -746,6 +863,22 @@ def _downstream_evaluation( source, execution_identity: str, ) -> dict[str, Any]: + """ + Run LMMS-Eval on a materialized checkpoint and collect its numeric task metrics. + + Parameters: + config (dict[str, Any]): Runtime configuration used to determine output locations. + node (CompiledPostMIPNode): Compiled downstream evaluation node configuration. + source: Candidate source containing the checkpoint to evaluate. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: Evaluation metrics and paths to the summary, raw result, and recorded command. + + Raises: + ValueError: If the source does not contain a checkpoint artifact. + RuntimeError: If LMMS-Eval fails or produces no numeric task metrics. + """ if source.artifact_kind is not ArtifactKind.CHECKPOINT: raise ValueError("downstream_evaluation requires materialized checkpoint artifacts") settings = dict(node.config.get("config") or {}) @@ -815,7 +948,17 @@ def _post_mip_kd_settings( config: Mapping[str, Any], node_settings: Mapping[str, Any], ) -> dict[str, Any]: - """Resolve KD settings while preserving the transformer's checkpoint contract.""" + """ + Resolve global distillation settings with node-specific overrides. + + Parameters: + config (Mapping[str, Any]): Configuration containing optional global distillation settings. + node_settings (Mapping[str, Any]): Node-specific distillation settings that override global values. + + Returns: + dict[str, Any]: The merged distillation settings, with disabled string values for + ``save_consolidated`` normalized to ``True``. + """ settings = copy.deepcopy(dict(config.get("global_distillation") or {})) settings.update(copy.deepcopy(dict(node_settings))) @@ -863,6 +1006,22 @@ def _run_candidate( input_revision_id: str, execution_identity: str, ) -> dict[str, Any]: + """ + Execute a candidate through the configured post-MIP node and assemble its result. + + Parameters: + config (dict[str, Any]): Pipeline configuration. + node (CompiledPostMIPNode): Compiled node defining the candidate operation. + ledger (CandidateLedger): Ledger containing candidate revisions. + input_revision_id (str): Revision identifier to process. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: Successful result containing input, source, architecture, and operation-specific fields. + + 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) @@ -906,6 +1065,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 a shard of a compiled post-MIP node and persist its per-candidate results. + + Parameters: + config (dict[str, Any]): Runtime configuration for the post-MIP execution. + 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 candidates. + + Returns: + Path: Path to the shard results JSON file. + """ node = _compiled_node(config, stage_id) ledger = _ledger(config) ledger.ingest_mip(_puzzle_dir(config)) diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 3ee3e8f0936..c11bbba398e 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 for the selected execution budget. + + Parameters: + state (Mapping[str, Any]): Setup state containing post-MIP flow definitions. + smoke (bool): Whether to apply smoke-budget limits. + 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]: Deep-copied post-MIP flows with normalized node configurations and budget-specific limits. + """ 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]: + """ + Build execution stage entries for post-MIP flow nodes. + + 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 allocated per node. + common (Mapping[str, Any]): Parallelization settings for evaluation stages. + single_gpu (Mapping[str, Any]): Parallelization settings for materialization stages. + cpu_partition (str | None): Optional CPU partition for selector and materialization stages. + pool_source_evaluations (bool): Whether source evaluations use a persistent worker pool. + + Returns: + dict[str, Any]: Mapping of post-MIP stage identifiers to their execution configurations. + """ 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..a0dc43e6680 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -186,7 +186,15 @@ def _dataset_subset_issues(state: WizardState) -> list[ValidationIssue]: def validate_state(state: WizardState) -> tuple[ValidationIssue, ...]: - """Return actionable authoring issues before canonical compilation.""" + """ + Validate persisted wizard state and return actionable cross-field authoring issues. + + Parameters: + state (WizardState): Persisted wizard state to validate. + + Returns: + tuple[ValidationIssue, ...]: Issues sorted by validation category and path. + """ issues: list[ValidationIssue] = [] required = ( "model.source", diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index af87270c860..b9fcabc7066 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 flows for each configured MIP run. + + Parameters: + session (WizardSession): Interactive wizard session used to read and store configuration. + resolver (DefaultsResolver): Resolver for stage resource defaults. + context (dict): Wizard context containing the inspected model and its inventory. + + Returns: + bool: `True` when post-MIP flows are configured, `False` if the user backs out. + """ 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,19 @@ def _serving_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """Ask the complete AIPerf workload and serving-only parallel setting.""" + """ + Collect serving workload parameters and its vLLM parallel topology. + + Parameters: + prefix (str): State path prefix used for the interactive prompts. + defaults (Mapping[str, Any]): Default serving settings. + inventory (Any): Model inventory used to validate the topology. + pruning (Mapping[str, Any]): Pruning configuration used to validate the topology. + stage_id (str): Stage associated with the serving topology. + + Returns: + Any: The configured serving settings, or `BACK` if the user navigates backward. + """ values = {} for name, label, default in ( ("input_tokens", "Serving input sequence length (ISL):", defaults["input_tokens"]), @@ -4100,9 +4123,31 @@ 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.""" + """ + Prompt for lmms-eval tasks and the vLLM configuration used to run them. + + Parameters: + session (WizardSession): Interactive wizard session. + prefix (str): State path prefix for the prompted settings. + defaults (Mapping[str, Any]): Default downstream evaluation settings. + inventory (Any): Model inventory used to validate the vLLM topology. + pruning (Mapping[str, Any]): Pruning configuration used to validate the vLLM topology. + stage_id (str): Stage identifier associated with the vLLM topology. + + Returns: + dict: Configured lmms-eval settings, including tasks, sampling limits, timeout, vLLM topology, and model arguments; `BACK` if the user backs out. + """ def validate_tasks(value: str) -> bool | str: + """ + Validate a comma-separated list of lmms-eval tasks. + + Parameters: + value (str): The task list to validate. + + 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 +4221,19 @@ def _configure_dynamic_resources( *, ask: bool, ) -> Any: - """Attach an independent resource/batch card to every node in one flow.""" + """ + Configure independent resources for each node in a post-MIP flow. + + Parameters: + session (WizardSession): Interactive session used to read state and collect resource settings. + editor (PostMIPFlowEditor): Flow editor containing the nodes to configure. + flow_id (str): Identifier of the post-MIP flow. + model (Any): Model information used to validate parallel profiles. + ask (bool): Whether to prompt for per-node resource and batch customization. + + Returns: + True when configuration completes, or BACK if the user exits during a prompt. + """ registry = ResourceProfileRegistry.from_dict( session.state.collection("parallel_profiles") or {} ) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 6938eeeee22..79e05c60372 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -635,7 +635,19 @@ 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 topology and workload. + + Parameters: + prompts (PromptSession): Interactive prompt session. + detailed (bool): Whether to collect detailed workload settings. + 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): Existing topology and workload values to use as defaults. + + Returns: + dict[str, Any]: AIPerf workload configuration containing serving topology, token counts, concurrency, and timeout settings. + """ defaults = dict(defaults or {}) topology_defaults = dict(defaults.get("topology") or {}) checkpoint = prompts.checkpoint() @@ -735,7 +747,18 @@ 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: + prompts (PromptSession): Interactive prompt session. + detailed (bool): Whether to prompt for the per-candidate timeout. + moe (bool): Whether to offer expert parallelism for the vLLM topology. + defaults (Mapping[str, Any] | None): Previously saved configuration values to use as prompt defaults. + + Returns: + dict[str, Any]: Downstream evaluation configuration containing tasks, sampling settings, vLLM topology, model arguments, and timeout. + """ defaults = defaults or {} tasks = prompts.text( @@ -846,6 +869,21 @@ def _default_flow( objective: Mapping[str, Any] | None = None, include_initial_filter: bool = True, ) -> dict[str, Any]: + """ + Build an opinionated post-MIP evaluation and selection flow. + + Parameters: + run_id (str): Identifier of the MIP run used as the flow source. + run (Mapping[str, Any]): MIP run configuration, including objectives and homogeneous candidate settings. + runtime (Mapping[str, Any]): Runtime workload settings for the serving evaluation. + data (Mapping[str, Any]): Dataset settings used to configure evaluation block size. + prefix (str): Prefix applied to generated node identifiers. + objective (Mapping[str, Any] | None): Objective to use for the flow; defaults to the first objective in the run. + include_initial_filter (bool): Whether to include the initial MIP-score filtering node. + + Returns: + dict[str, Any]: Flow definition containing its source specification and ordered orchestration nodes. + """ def node_id(name: str) -> str: return f"{prefix}{name}" @@ -1005,6 +1043,20 @@ def _custom_flow( detailed: bool, moe: bool, ) -> dict[str, Any]: + """ + Interactively build a post-MIP orchestration flow for a run. + + Parameters: + run_id (str): Identifier of the MIP run associated with the flow. + runtime (Mapping[str, Any]): Runtime settings used to configure serving evaluations. + data (Mapping[str, Any]): Data settings used to configure model evaluations. + used_ids (set[str]): Node identifiers already used by other flows. + detailed (bool): Whether to collect detailed configuration options. + moe (bool): Whether the selected model is a mixture-of-experts model. + + Returns: + dict[str, Any]: A flow definition containing the run source and configured post-MIP nodes. + """ nodes: OrderedDict[str, Any] = OrderedDict() available_metrics = ["mip.score"] transformer_nodes = [] @@ -1221,6 +1273,20 @@ def _resource_rows( gpus_per_node: int, workers: Mapping[str, int], ) -> list[dict[str, Any]]: + """ + Builds a resource plan for each campaign execution stage. + + Parameters: + state (AnswerState): Campaign state containing post-MIP flows and execution mode. + common (Mapping[str, int]): Common parallel mesh dimensions. + bypass (Mapping[str, int]): Bypass parallel mesh dimensions. + global_kd (Mapping[str, int]): Global knowledge-distillation mesh dimensions. + gpus_per_node (int): Number of GPUs available on each node. + workers (Mapping[str, int]): Worker limits for pool and sharded workloads. + + Returns: + list[dict[str, Any]]: Resource rows containing each stage name, instance count, GPUs per instance, and required node count. + """ from .bundle import _post_mip_candidate_limits, _serving_parallel rows = []