feat: add --submodel to perf and build commands

src/winml/modelkit/commands/build.py

@@ -508,6 +508,15 @@ def _validate_loader_tasks_for_model(
 @cli_utils.trust_remote_code_option(
     optional_message="Trust remote code for custom model architectures (e.g., Mu2)."
 )
+@click.option(
+    "--submodel",
+    type=str,
+    default=None,
+    help=(
+        "Build a specific sub-model from a composite model "
+        "(e.g., 'encoder', 'decoder'). Omit to build all sub-models automatically."
+    ),
+)
 @cli_utils.verbosity_options()
 @cli_utils.no_color_option()
 @click.pass_context
@@ -528,6 +537,7 @@ def build(
     max_optim_iterations: int | None,
     allow_unsupported_nodes: bool,
     trust_remote_code: bool,
+    submodel: str | None,
     verbose: int,
     quiet: bool,
 ) -> None:
@@ -884,6 +894,22 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
                         f"Composite model detection failed unexpectedly: {e}"
                     ) from e
 
+            # ── --submodel validation ──────────────────────────────────────
+            if submodel is not None:
+                if components is None:
+                    raise click.BadParameter(
+                        f"'{submodel}' was specified, but '{model}' "
+                        f"is not a composite model (no sub-models detected).",
+                        param_hint="--submodel",
+                    )
+                if submodel not in components:
+                    raise click.BadParameter(
+                        f"Unknown sub-model '{submodel}'. "
+                        f"Available: {', '.join(components.keys())}",
+                        param_hint="--submodel",
+                    )
+                components = {submodel: components[submodel]}
+
             if components:
                 if use_cache:
                     raise click.UsageError(

src/winml/modelkit/commands/perf.py

@@ -81,6 +81,7 @@ class BenchmarkConfig:
 
     model_id: str
     task: str | None = None
+    submodel: str | None = None
     device: str = "auto"
     precision: str = "auto"
     iterations: int = 100
@@ -1563,25 +1564,30 @@ def report_composite_results(
         json.dump(combined, f, indent=2)
 
 
-def generate_output_path(model_id: str, *, module_class: str | None = None) -> Path:
+def generate_output_path(
+    model_id: str, *, module_class: str | None = None, submodel: str | None = None
+) -> Path:
     r"""Generate default output path under the user's cache directory.
 
-    Returns ``~/.cache/winml/perf/<slug>[/<module_class>]/<timestamp>.json``
+    Returns ``~/.cache/winml/perf/<slug>[/<module_class>][/<submodel>]/<timestamp>.json``
     so repeated runs accumulate under a stable per-model directory without
     polluting CWD (see #551). The timestamp is generated at call time using
     local time, format ``YYYYMMDD-HHMMSS``.
 
     For ONNX inputs, the file stem is used as the slug
     (e.g., ``model.onnx`` -> ``model``). For HF model IDs, ``/`` and ``\``
     are replaced with ``_`` (e.g., ``microsoft/resnet-50`` ->
-    ``microsoft_resnet-50``).
+    ``microsoft_resnet-50``). A ``submodel`` (composite sub-component) is nested
+    under its own directory so per-sub-model reports don't collide.
     """
     p = Path(model_id)
     slug = p.stem if p.suffix.lower() == ".onnx" else model_id.replace("/", "_").replace("\\", "_")
 
     out_dir = Path.home() / ".cache" / "winml" / "perf" / slug
     if module_class:
         out_dir = out_dir / module_class
+    if submodel:
+        out_dir = out_dir / submodel
 
     timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
     return out_dir / f"{timestamp}.json"
@@ -1822,6 +1828,33 @@ def _run_genai_runtime(ctx: click.Context, *, console: Console, json_mode: bool)
     run_genai_perf(config, console=console, json_mode=json_mode)
 
 
+def _resolve_composite_components_for_perf(model: str, task: str | None) -> dict[str, str] | None:
+    """Detect a composite model's sub-components (name -> component task), else None.
+
+    Mirrors the registry-driven detection in ``winml export`` / ``winml build``
+    so ``--submodel`` resolves the same components (and the same seq2seq bridge
+    when ``--task`` is omitted). Only the "not a resolvable HF config" case
+    (``OSError``) is suppressed (fall through to "not composite"); intentional
+    loud guards (empty registry, model-task incompatibility) and any unexpected
+    failure are surfaced rather than masked.
+    """
+    from ..loader.resolution import resolve_composite_components
+
+    try:
+        return resolve_composite_components(model, task=task)
+    except click.ClickException:
+        raise
+    except ValueError as e:
+        raise click.UsageError(str(e)) from e
+    except RuntimeError:
+        raise
+    except OSError as e:
+        logger.debug("Composite detection unavailable (config not resolvable): %s", e)
+        return None
+    except Exception as e:
+        raise click.ClickException(f"Composite model detection failed unexpectedly: {e}") from e
+
+
 @click.command("perf")
 @cli_utils.model_option(required=False)
 @click.option(
@@ -1869,6 +1902,15 @@ def _run_genai_runtime(ctx: click.Context, *, console: Console, json_mode: bool)
     default=None,
     help="Explicit task (e.g., 'image-classification'). Auto-detected if not specified.",
 )
+@click.option(
+    "--submodel",
+    type=str,
+    default=None,
+    help=(
+        "Benchmark a specific sub-model of a composite model "
+        "(e.g., 'text_model', 'vision_model'). Omit to benchmark all sub-models."
+    ),
+)
 @click.option(
     "--iterations",
     type=click.IntRange(min=1),
@@ -1993,6 +2035,7 @@ def perf(
     max_new_tokens: int,
     compile_timeout: int,
     task: str | None,
+    submodel: str | None,
     iterations: int,
     warmup: int,
     device: str,
@@ -2128,6 +2171,46 @@ def perf(
     if is_onnx and model_input.local_path and not Path(model_input.local_path).exists():
         raise click.UsageError(f"ONNX file not found: {hf_model}")
 
+    # =========================================================================
+    # --submodel: narrow a composite model to one sub-component, benchmarked as
+    # a standalone single-session model. The composite is detected the same way
+    # `winml export` / `winml build` / `winml inspect` do (registry-driven, via
+    # the seq2seq bridge), so it works even when --task is omitted. The selected
+    # component is then loaded through the normal single-model path using its own
+    # component task — exactly how the composite builds that sub-model — which
+    # sidesteps the config-ambiguous pipeline task (e.g. t5 translation vs
+    # summarization) and avoids building the other sub-models just to discard.
+    # =========================================================================
+    if submodel is not None:
+        if is_onnx:
+            raise click.BadParameter(
+                "--submodel is not supported for ONNX files; a .onnx file is "
+                "already a single model.",
+                param_hint="--submodel",
+            )
+        if module_class:
+            raise click.BadParameter(
+                "--submodel cannot be combined with --module.",
+                param_hint="--submodel",
+            )
+        components = _resolve_composite_components_for_perf(hf_model, task)
+        if components is None:
+            raise click.BadParameter(
+                f"'{submodel}' was specified, but '{hf_model}' is not a "
+                f"composite model (no sub-models detected).",
+                param_hint="--submodel",
+            )
+        if submodel not in components:
+            raise click.BadParameter(
+                f"Unknown sub-model '{submodel}'. Available: {', '.join(components)}",
+                param_hint="--submodel",
+            )
+        # Load only this component, using its own task, via the single-model path.
+        task = components[submodel]
+        console.print(
+            f"[dim]Composite sub-model:[/dim] {submodel} (task={task}) [dim]from[/dim] {hf_model}"
+        )
+
     # =========================================================================
     # MODULE MODE: per-module build + benchmark
     # =========================================================================
@@ -2222,7 +2305,7 @@ def perf(
 
     # Resolve output path
     if output is None:
-        output = generate_output_path(hf_model)
+        output = generate_output_path(hf_model, submodel=submodel)
 
     # Refuse to clobber an existing report unless the user opted in.
     cli_utils.guard_output(output, overwrite)
@@ -2233,6 +2316,7 @@ def perf(
     config = BenchmarkConfig(
         model_id=hf_model,
         task=task,
+        submodel=submodel,
         device=device.lower(),
         precision=precision.lower(),
         iterations=iterations,

tests/unit/commands/test_build.py

@@ -2493,3 +2493,161 @@ def test_composite_carries_outer_quant_when_component_quant_none(
         # Outer compile carried over too, deep-copied.
         assert built_config.compile is not None
         assert built_config.compile is not outer_cfg.compile
+
+
+class TestBuildSubmodel:
+    """Test --submodel narrows a composite build to a single sub-model."""
+
+    def test_submodel_builds_only_requested_component(
+        self,
+        runner: CliRunner,
+        tmp_path: Path,
+    ) -> None:
+        """--submodel builds only the named component."""
+        from winml.modelkit.commands.build import build
+
+        components = {
+            "decoder_prefill": "feature-extraction",
+            "decoder_gen": "text2text-generation",
+        }
+        output_dir = tmp_path / "out"
+
+        fake_cfg = MagicMock()
+        fake_cfg.quant = None
+        fake_cfg.compile = None
+        fake_cfg.loader = MagicMock(task=None)
+
+        with (
+            patch(
+                "winml.modelkit.loader.resolution.resolve_composite_components",
+                return_value=components,
+            ),
+            patch(
+                "winml.modelkit.config.generate_build_config",
+                return_value=fake_cfg,
+            ),
+            patch(
+                "winml.modelkit.commands.build._run_single_build",
+            ) as mock_single_build,
+            patch(
+                "winml.modelkit.commands.build._validate_loader_tasks_for_model",
+                return_value=None,
+            ),
+        ):
+            result = runner.invoke(
+                build,
+                [
+                    "-m",
+                    "Qwen/Qwen3-0.6B",
+                    "-o",
+                    str(output_dir),
+                    "--submodel",
+                    "decoder_prefill",
+                ],
+                obj={"debug": False},
+            )
+
+        assert result.exit_code == 0, result.output
+        # Only the requested component is built.
+        assert mock_single_build.call_count == 1
+        assert mock_single_build.call_args.kwargs["cache_key"] == "decoder_prefill"
+
+    def test_submodel_rejects_unknown_name(
+        self,
+        runner: CliRunner,
+        tmp_path: Path,
+    ) -> None:
+        """--submodel with an invalid name is a clean error."""
+        from winml.modelkit.commands.build import build
+
+        components = {
+            "decoder_prefill": "feature-extraction",
+            "decoder_gen": "text2text-generation",
+        }
+
+        fake_cfg = MagicMock()
+        fake_cfg.quant = None
+        fake_cfg.compile = None
+        fake_cfg.loader = MagicMock(task=None)
+
+        with (
+            patch(
+                "winml.modelkit.loader.resolution.resolve_composite_components",
+                return_value=components,
+            ),
+            patch(
+                "winml.modelkit.config.generate_build_config",
+                return_value=fake_cfg,
+            ),
+            patch(
+                "winml.modelkit.commands.build._run_single_build",
+            ) as mock_single_build,
+            patch(
+                "winml.modelkit.commands.build._validate_loader_tasks_for_model",
+                return_value=None,
+            ),
+        ):
+            result = runner.invoke(
+                build,
+                [
+                    "-m",
+                    "Qwen/Qwen3-0.6B",
+                    "-o",
+                    str(tmp_path / "out"),
+                    "--submodel",
+                    "encoder",
+                ],
+                obj={"debug": False},
+            )
+
+        assert result.exit_code != 0
+        assert "Unknown sub-model 'encoder'" in result.output
+        assert "decoder_prefill" in result.output
+        mock_single_build.assert_not_called()
+
+    def test_submodel_rejects_non_composite(
+        self,
+        runner: CliRunner,
+        tmp_path: Path,
+    ) -> None:
+        """--submodel on a non-composite model is a clean error."""
+        from winml.modelkit.commands.build import build
+
+        fake_cfg = MagicMock()
+        fake_cfg.quant = None
+        fake_cfg.compile = None
+        fake_cfg.loader = MagicMock(task=None)
+
+        with (
+            patch(
+                "winml.modelkit.loader.resolution.resolve_composite_components",
+                return_value=None,
+            ),
+            patch(
+                "winml.modelkit.config.generate_build_config",
+                return_value=fake_cfg,
+            ),
+            patch(
+                "winml.modelkit.commands.build._run_single_build",
+            ) as mock_single_build,
+            patch(
+                "winml.modelkit.commands.build._validate_loader_tasks_for_model",
+                return_value=None,
+            ),
+        ):
+            result = runner.invoke(
+                build,
+                [
+                    "-m",
+                    "prajjwal1/bert-tiny",
+                    "-o",
+                    str(tmp_path / "out"),
+                    "--submodel",
+                    "encoder",
+                ],
+                obj={"debug": False},
+            )
+
+        assert result.exit_code != 0
+        assert "not a composite model" in result.output
+        mock_single_build.assert_not_called()