NVIDIA · grzegorz-k-karch · Aug 7, 2026 · Aug 7, 2026
@@ -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()

@@ -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"

@@ -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

@@ -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 [])

@@ -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:

@@ -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)

@@ -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

@@ -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)

@@ -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)
@@ -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 = []