From f1bd41a74b75e6b5fbc607fd3ada4f874bd50e2f Mon Sep 17 00:00:00 2001 From: JOYSTON <116254639+Joystonm@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:14:41 +0530 Subject: [PATCH] commit --- .github/workflows/tests.yaml | 5 +- docs/user_manual/mcp_server.rst | 138 +++++++ pyproject.toml | 4 + src/pruna/mcp_server/__init__.py | 60 ++++ src/pruna/mcp_server/introspection.py | 347 ++++++++++++++++++ src/pruna/mcp_server/server.py | 480 +++++++++++++++++++++++++ src/pruna/mcp_server/validation.py | 369 +++++++++++++++++++ tests/conftest.py | 1 + tests/mcp_server/__init__.py | 13 + tests/mcp_server/test_introspection.py | 127 +++++++ tests/mcp_server/test_server.py | 292 +++++++++++++++ tests/mcp_server/test_validation.py | 307 ++++++++++++++++ 12 files changed, 2142 insertions(+), 1 deletion(-) create mode 100644 docs/user_manual/mcp_server.rst create mode 100644 src/pruna/mcp_server/__init__.py create mode 100644 src/pruna/mcp_server/introspection.py create mode 100644 src/pruna/mcp_server/server.py create mode 100644 src/pruna/mcp_server/validation.py create mode 100644 tests/mcp_server/__init__.py create mode 100644 tests/mcp_server/test_introspection.py create mode 100644 tests/mcp_server/test_server.py create mode 100644 tests/mcp_server/test_validation.py diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 72684514..4d048cc0 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -66,7 +66,7 @@ jobs: strategy: matrix: python-version: ["3.11"] - name: ["base", 'lmharness', 'rapidata', 'llamacpp'] + name: ["base", 'lmharness', 'rapidata', 'llamacpp', 'mcp'] include: - name: base extras: "dev" @@ -80,6 +80,9 @@ jobs: - name: llamacpp extras: "dev,llamacpp" mark_filter: "requires_llamacpp" + - name: mcp + extras: "dev,mcp" + mark_filter: "requires_mcp" env: HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/docs/user_manual/mcp_server.rst b/docs/user_manual/mcp_server.rst new file mode 100644 index 00000000..e49e2cce --- /dev/null +++ b/docs/user_manual/mcp_server.rst @@ -0,0 +1,138 @@ +:title: MCP Server - Pruna AI Framework Documentation +:description: Serve the Pruna algorithm catalogue to an LLM over the Model Context Protocol, with model-free configuration validation. + +MCP Server +========== + +|pruna| ships a `Model Context Protocol `_ server that lets an LLM +discover optimization algorithms, inspect their hyperparameters and validate configurations without +writing code by hand. + +The server has two tiers of tools. + +**Discovery and validation** tools are read-only. They need no GPU, load no model and answer +immediately, so they are always registered. They cover the whole algorithm catalogue, including the +hyperparameters each algorithm accepts and the algorithms it may be combined with. + +**Execution** tools load a model and optimize it. They are opt-in and are not registered unless the +server is started with ``--allow-execute``. + +Installation +------------ + +The server lives in the ``pruna[mcp]`` extra: + +.. code-block:: bash + :class: noextract + + pip install "pruna[mcp]" + +This registers the ``pruna-mcp`` console script. + +Run the Server +-------------- + +Start it over stdio, which is what Claude Desktop and Claude Code use: + +.. code-block:: bash + :class: noextract + + pruna-mcp + +To additionally expose the tools that load and optimize models: + +.. code-block:: bash + :class: noextract + + pruna-mcp --allow-execute + +Because execution tools download weights and need a device, they are disabled by default. Leaving +them off keeps the server cheap to run and safe to point at from any client. + +Connect a Client +---------------- + +.. code-block:: json + :class: noextract + + { + "mcpServers": { + "pruna": { + "command": "pruna-mcp" + } + } + } + +Add ``"args": ["--allow-execute"]`` to the same entry to enable the execution tools. + +Available Tools +--------------- + +``list_tags`` + List the algorithm categories, such as ``quantizer`` or ``cacher``, with a description of each. + +``list_algorithms`` + List every optimization algorithm, optionally filtered by category. Reports the devices each + algorithm runs on and whether it needs a tokenizer, processor or dataset. + +``describe_algorithm`` + Describe one algorithm in full: every hyperparameter with its allowed values and default, the + algorithms it must run before or after, and its references. + +``validate_config`` + Check a proposed configuration **without loading a model**. Reports missing pip extras, missing + tokenizers or datasets, and algorithm pairs that cannot be combined. + +``list_benchmarks``, ``list_metrics`` and ``list_datasets`` + List the available evaluation benchmarks, metrics and datasets. ``list_datasets`` returns the + exact names accepted by the ``data`` argument of ``evaluate_model``. + +``smash_model`` and ``evaluate_model`` + Load a model and optimize or evaluate it. Registered only with ``--allow-execute``. + +A Typical Session +----------------- + +The intended workflow is discovery, then validation, then execution: + +.. code-block:: text + :class: noextract + + list_tags() + -> {"tags": [{"name": "quantizer", ...}], "count": 11} + + list_algorithms(tag="quantizer") + -> {"algorithms": [{"name": "hqq", ...}], "tag": "quantizer", "count": 12} + + describe_algorithm("hqq") + -> {"hyperparameters": [{"qualified_name": "hqq_weight_bits", + "default": 8, "allowed": [2, 4, 8]}, ...], ...} + + validate_config(algorithms=["hqq", "torch_compile"], + hyperparameters={"hqq_weight_bits": 4}) + -> {"valid": true, "errors": [], "execution_order": ["hqq", "torch_compile"], ...} + + smash_model(model_path="...", algorithms=["hqq", "torch_compile"], save_to="./optimized") + +Every tool answers with a single JSON object, so a client parses one payload rather than a list of +fragments. Hyperparameters are always keyed by their fully qualified name, prefixed with the +algorithm name, for example ``hqq_weight_bits``. Any hyperparameter that is omitted keeps its +default. + +What Validation Covers +---------------------- + +``validate_config`` reuses the same routines ``smash`` runs, so it catches configuration mistakes +before any weights are downloaded. It reports: + +* algorithms whose pip extra is not installed, +* algorithms that require a tokenizer, processor or dataset that has not been attached, +* pairs of algorithms that declare no compatibility with each other. + +A valid report also carries the resolved value of every hyperparameter and the ``execution_order`` +the algorithms will be applied in. The report keeps the same keys whether or not the configuration +was valid, so a rejected hyperparameter value is reported in the same shape as a missing dependency. + +Two checks cannot run without a model and are reported as such in the ``notes`` field: whether the +model itself is compatible, and whether two algorithms target overlapping modules. Both are performed +by ``smash`` once the model is loaded. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index c299a0dc..e3e5446b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -156,7 +156,11 @@ dependencies = [ "termcolor==2.3.0", ] +[project.scripts] +pruna-mcp = "pruna.mcp_server.server:main" + [project.optional-dependencies] +mcp = ["mcp>=2.0,<3.0"] transformers-legacy = [ "transformers<5.0.0", ] diff --git a/src/pruna/mcp_server/__init__.py b/src/pruna/mcp_server/__init__.py new file mode 100644 index 00000000..9541438b --- /dev/null +++ b/src/pruna/mcp_server/__init__.py @@ -0,0 +1,60 @@ +# Copyright 2025 - Pruna AI GmbH. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model Context Protocol server exposing pruna to an LLM. + +The server itself lives in :mod:`pruna.mcp_server.server` and is imported lazily, so that +importing this package does not require the ``mcp`` dependency to be installed. The +metadata helpers in :mod:`pruna.mcp_server.introspection` and +:mod:`pruna.mcp_server.validation` are always importable. +""" + +from typing import Any + +# Both names are provided lazily by __getattr__, which ruff cannot follow. +__all__ = ["build_server", "main"] # noqa: F822 + +_LAZY: dict[str, tuple[str, str]] = { + "build_server": ("pruna.mcp_server.server", "build_server"), + "main": ("pruna.mcp_server.server", "main"), +} + + +def __getattr__(name: str) -> Any: + """ + Resolve a server entry point lazily. + + Parameters + ---------- + name : str + The attribute being accessed. + + Returns + ------- + Any + The requested server attribute. + + Raises + ------ + AttributeError + If the attribute is not a lazy entry point. + """ + if name not in _LAZY: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + from importlib import import_module + + module_name, attribute = _LAZY[name] + return getattr(import_module(module_name), attribute) diff --git a/src/pruna/mcp_server/introspection.py b/src/pruna/mcp_server/introspection.py new file mode 100644 index 00000000..88e25eb6 --- /dev/null +++ b/src/pruna/mcp_server/introspection.py @@ -0,0 +1,347 @@ +# Copyright 2025 - Pruna AI GmbH. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Introspection helpers that expose pruna's algorithm catalogue as plain, JSON-serializable data. + +Every function in this module imports pruna lazily so that merely importing +``pruna.mcp_server`` does not pull in torch and the rest of the heavy runtime. +""" + +from __future__ import annotations + +from typing import Any + +# Attributes of a ConfigSpace hyperparameter that carry its allowed value range. +_RANGE_ATTRIBUTES = ("choices", "sequence", "lower", "upper", "q", "log") + + +def _registry() -> Any: + """ + Import and return the algorithm registry. + + Returns + ------- + Any + The ``AlgorithmRegistry`` class from pruna. + """ + from pruna import AlgorithmRegistry + + return AlgorithmRegistry + + +def _to_jsonable(value: Any) -> Any: + """ + Convert an arbitrary ConfigSpace value into something JSON-serializable. + + Parameters + ---------- + value : Any + The value to convert. + + Returns + ------- + Any + A JSON-serializable representation of the value. + """ + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, dict): + return {str(key): _to_jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [_to_jsonable(item) for item in value] + return str(value) + + +def _serialize_hyperparameter(hyperparameter: Any, algorithm_name: str) -> dict[str, Any]: + """ + Serialize a single ConfigSpace hyperparameter of an algorithm. + + Hyperparameter names are reported both unprefixed (as the algorithm class declares them) + and fully qualified with the algorithm prefix, because the prefix is what appears in a + ``SmashConfig``. + + Parameters + ---------- + hyperparameter : Any + A ConfigSpace hyperparameter object. + algorithm_name : str + Name of the algorithm the hyperparameter belongs to. + + Returns + ------- + dict[str, Any] + JSON-serializable description of the hyperparameter. + """ + raw_name = str(hyperparameter.name) + prefix = f"{algorithm_name}_" + short_name = raw_name[len(prefix) :] if raw_name.startswith(prefix) else raw_name + meta = getattr(hyperparameter, "meta", None) or {} + + described: dict[str, Any] = { + "name": short_name, + "qualified_name": f"{algorithm_name}_{short_name}", + "type": type(hyperparameter).__name__, + "default": _to_jsonable(getattr(hyperparameter, "default_value", None)), + } + if meta.get("desc"): + described["description"] = str(meta["desc"]) + + for attribute in _RANGE_ATTRIBUTES: + value = getattr(hyperparameter, attribute, None) + if value is not None: + described[attribute] = _to_jsonable(value) + + described["allowed"] = _allowed_values(hyperparameter, described["default"]) + return described + + +def _allowed_values(hyperparameter: Any, default: Any) -> list[Any]: + """ + Enumerate the values a hyperparameter will accept. + + The ConfigSpace hyperparameters spell their domain differently: categoricals expose + ``choices``, ordinals expose ``sequence`` and a constant exposes nothing at all. A caller + should not have to know which spelling applies, so the enumerable domains are collected + under one key. A continuous domain cannot be enumerated and is reported as an empty list, + with ``lower`` and ``upper`` left in place. + + Parameters + ---------- + hyperparameter : Any + A ConfigSpace hyperparameter object. + default : Any + The already-serialized default value, reused for a constant. + + Returns + ------- + list[Any] + Every accepted value, or an empty list when the domain is continuous. + """ + if type(hyperparameter).__name__ == "Constant": + return [default] + + for attribute in ("choices", "sequence", "q"): + value = getattr(hyperparameter, attribute, None) + if value is not None: + return [_to_jsonable(item) for item in value] + + return [] + + +def _requirements(algorithm: Any) -> dict[str, Any]: + """ + Summarize what an algorithm needs before it can run. + + Parameters + ---------- + algorithm : Any + A registered algorithm instance. + + Returns + ------- + dict[str, Any] + JSON-serializable requirements of the algorithm. + """ + return { + "tokenizer": bool(getattr(algorithm, "tokenizer_required", False)), + "processor": bool(getattr(algorithm, "processor_required", False)), + "dataset": bool(getattr(algorithm, "dataset_required", False)), + "pip_install": getattr(algorithm, "required_install", None), + } + + +def _summary(algorithm: Any) -> dict[str, Any]: + """ + Build the compact algorithm summary used by ``list_algorithms``. + + Parameters + ---------- + algorithm : Any + A registered algorithm instance. + + Returns + ------- + dict[str, Any] + JSON-serializable summary of the algorithm. + """ + save_fn = getattr(algorithm, "save_fn", None) + return { + "name": algorithm.algorithm_name, + "tags": [str(tag) for tag in getattr(algorithm, "group_tags", [])], + "runs_on": list(getattr(algorithm, "runs_on", [])), + "requires": _requirements(algorithm), + "save_fn": getattr(save_fn, "name", None), + "references": _to_jsonable(getattr(algorithm, "references", None)), + } + + +def list_tags() -> list[dict[str, str]]: + """ + List every algorithm category with its human-readable description. + + Returns + ------- + list[dict[str, str]] + One entry per algorithm tag, each with a ``name`` and a ``description``. + """ + from pruna.algorithms.base.tags import AlgorithmTag + + return [tag.get_documentation() for tag in AlgorithmTag] + + +def list_algorithms(tag: str | None = None) -> list[dict[str, Any]]: + """ + List all registered optimization algorithms, optionally filtered by category. + + Parameters + ---------- + tag : str | None + Restrict the result to algorithms in this category, for example ``"quantizer"`` or + ``"cacher"``. Use ``list_tags`` for the full set of valid categories. If None, every + algorithm is returned. + + Returns + ------- + list[dict[str, Any]] + One summary entry per algorithm. + + Raises + ------ + ValueError + If ``tag`` is not a known algorithm category. + """ + from pruna.algorithms.base.tags import AlgorithmTag + + registry = _registry() + if tag is None: + return [_summary(algorithm) for algorithm in registry.get_all_algorithms()] + + try: + algorithm_tag = AlgorithmTag[tag.upper()] + except KeyError: + valid = ", ".join(str(member) for member in AlgorithmTag) + raise ValueError(f"Unknown algorithm tag '{tag}'. Valid tags are: {valid}.") from None + + return [_summary(registry[name]) for name in registry.get_algorithms_by_tag(algorithm_tag)] + + +def describe_algorithm(name: str) -> dict[str, Any]: + """ + Describe a single algorithm in full detail. + + This is the tool to call before composing a configuration: it returns every + hyperparameter with its allowed values and default, the order constraints, and the + requirements the algorithm imposes. + + Parameters + ---------- + name : str + Name of the algorithm, as returned by ``list_algorithms``. + + Returns + ------- + dict[str, Any] + Full description of the algorithm. + + Raises + ------ + KeyError + If no algorithm with that name is registered. + """ + registry = _registry() + if name not in {algorithm.algorithm_name for algorithm in registry.get_all_algorithms()}: + raise KeyError(f"Unknown algorithm '{name}'. Use list_algorithms to see valid names.") + + algorithm = registry[name] + described = _summary(algorithm) + described.update( + { + "compatible_before": _to_jsonable(algorithm.get_algorithms_to_run_before()), + "compatible_after": _to_jsonable(algorithm.get_algorithms_to_run_after()), + "compatible_with": _to_jsonable(algorithm.get_compatible_algorithms()), + "disjointly_compatible_with": _to_jsonable(algorithm.get_disjointly_compatible_algorithms()), + "hyperparameters": [ + _serialize_hyperparameter(hyperparameter, name) for hyperparameter in algorithm.hyperparameters + ], + } + ) + return described + + +def list_metrics() -> list[str]: + """ + List every evaluation metric name that can be requested from an evaluation task. + + Returns + ------- + list[str] + Sorted metric names available through ``EvaluationAgent``. + """ + from pruna.evaluation.metrics import MetricRegistry + + # MetricRegistry exposes registration and lookup but no public enumeration helper, + # so the underlying registry mapping is read directly. + registered = getattr(MetricRegistry, "_registry", {}) + return sorted(str(key) for key in registered) + + +def list_datasets() -> list[str]: + """ + List the named datasets that an evaluation task can be run on. + + These are the exact names ``PrunaDataModule.from_string`` accepts, so a caller never has to + guess the spelling of a dataset. + + Returns + ------- + list[str] + The registered dataset names. + """ + from pruna.data import base_datasets + + return sorted(base_datasets.keys()) + + +def list_benchmarks(task_type: str | None = None) -> list[dict[str, Any]]: + """ + List the built-in evaluation benchmarks, optionally filtered by task type. + + Parameters + ---------- + task_type : str | None + Restrict the result to benchmarks for this task type, for example + ``"text_to_image"``. If None, every benchmark is returned. + + Returns + ------- + list[dict[str, Any]] + One entry per benchmark, including the metrics it evaluates. + """ + from pruna.evaluation.benchmarks import BenchmarkRegistry + + described = [] + for key in BenchmarkRegistry.list(task_type=task_type): + benchmark = BenchmarkRegistry.get(key) + described.append( + { + "name": benchmark.name, + "lookup_key": benchmark.lookup_key, + "task_type": benchmark.task_type, + "description": benchmark.description, + "metrics": list(benchmark.metrics), + "reference": benchmark.reference, + } + ) + return described diff --git a/src/pruna/mcp_server/server.py b/src/pruna/mcp_server/server.py new file mode 100644 index 00000000..fbd26609 --- /dev/null +++ b/src/pruna/mcp_server/server.py @@ -0,0 +1,480 @@ +# Copyright 2025 - Pruna AI GmbH. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model Context Protocol server exposing pruna's optimization catalogue to an LLM. + +The server has two tiers of tools. Discovery and validation tools are cheap, read-only and +need neither a GPU nor a loaded model, so they are always available. Tools that actually +optimize a model are opt-in and must be enabled explicitly with ``--allow-execute``. + +Run it over stdio with:: + + pruna-mcp +""" + +from __future__ import annotations + +import argparse +import sys +from typing import Any + +try: + from mcp.server import MCPServer + from mcp.server.mcpserver.exceptions import ToolError + from mcp.types import ToolAnnotations + + _MCP_IMPORT_ERROR: ImportError | None = None +except ImportError as error: # pragma: no cover - depends on the installed extras + # Deferred rather than raised, so that importing this module stays possible without the + # extra. The console script imports it before main() runs, so raising here would bury the + # explanation under a traceback that a client shows verbatim to the user. + _MCP_IMPORT_ERROR = error + +from . import introspection, validation + +INSTRUCTIONS = """ +pruna is a model-optimization framework. Its entry point is `smash(model, SmashConfig)`, +where a `SmashConfig` activates one or more optimization algorithms. + +Recommended workflow: + +1. Call `list_tags` and `list_algorithms` to find the algorithms relevant to a goal such as + quantization or caching. +2. Call `describe_algorithm` on each candidate. It returns every hyperparameter with its + allowed values and default, plus which algorithms it may run before or after. +3. Call `validate_config` with the intended combination. This needs no GPU and reports + missing pip extras, missing tokenizers and incompatible algorithm pairs immediately. +4. Only then call `smash_model`, which loads a model and optimizes it. This tool exists only + when the server was started with --allow-execute. + +Hyperparameters are always keyed by their fully qualified name, prefixed with the algorithm +name, for example `hqq_weight_bits`. Omit any hyperparameter to keep its default. +""".strip() + + +def _require_mcp() -> None: + """ + Fail with actionable guidance when the mcp extra is missing. + + Raises + ------ + ImportError + If the server was imported without pruna[mcp] installed. + """ + if _MCP_IMPORT_ERROR is not None: + raise ImportError( + "The pruna MCP server requires the 'mcp' package, which is not installed. " + f"Install it with: pip install 'pruna[mcp]' (original error: {_MCP_IMPORT_ERROR})" + ) from _MCP_IMPORT_ERROR + + +def _require_execution(allow_execute: bool) -> None: + """ + Guard the execution tier behind an explicit opt-in flag. + + Parameters + ---------- + allow_execute : bool + Whether the server was started with execution enabled. + + Raises + ------ + ToolError + If execution tools were called on a read-only server. + """ + if not allow_execute: + raise ToolError( + "This server runs in read-only mode. Restart it with --allow-execute to enable " + "smash_model and evaluate_model. Use validate_config to check a configuration first." + ) + + +def build_server(allow_execute: bool = False) -> MCPServer: + """ + Construct the MCP server and register every tool. + + Parameters + ---------- + allow_execute : bool + Whether to register the tools that load and optimize models. These are omitted + entirely when False, so a client can discover that execution is unavailable. + + Returns + ------- + MCPServer + The configured server, ready to be run over stdio. + """ + _require_mcp() + server = MCPServer(name="pruna", instructions=INSTRUCTIONS) + read_only = ToolAnnotations(read_only_hint=True, destructive_hint=False, idempotent_hint=True) + executing = ToolAnnotations(read_only_hint=False, destructive_hint=True, idempotent_hint=False) + + @server.tool( + annotations=read_only, + description="List the algorithm categories, such as quantizer or cacher, with descriptions.", + ) + def list_tags() -> dict[str, Any]: + """ + List every algorithm category with a description of what it does. + + Returns + ------- + dict[str, Any] + The categories under the key "tags", each with a name and a description. + """ + tags = introspection.list_tags() + return {"tags": tags, "count": len(tags)} + + @server.tool( + annotations=read_only, + description="List all optimization algorithms. Filter by category, e.g. tag='quantizer'.", + ) + def list_algorithms(tag: str | None = None) -> dict[str, Any]: + """ + List the available optimization algorithms and what each one needs to run. + + Parameters + ---------- + tag : str | None + Optional category filter, for example "quantizer" or "cacher". Call list_tags for + the full set. If None, every algorithm is returned. + + Returns + ------- + dict[str, Any] + The algorithm summaries under the key "algorithms", the applied filter under "tag" + and their number under "count". + """ + try: + algorithms = introspection.list_algorithms(tag) + except ValueError as error: + raise ToolError(str(error)) from None + return {"algorithms": algorithms, "tag": tag, "count": len(algorithms)} + + @server.tool( + annotations=read_only, + description=( + "Describe one algorithm in full: every hyperparameter with allowed values and " + "defaults, order constraints, and requirements. Call this before composing a config." + ), + ) + def describe_algorithm(name: str) -> dict[str, Any]: + """ + Describe a single algorithm, including all of its hyperparameters. + + Parameters + ---------- + name : str + The algorithm name, as returned by list_algorithms. + + Returns + ------- + dict[str, Any] + The full algorithm description, including per-hyperparameter names, types, + defaults and allowed ranges. + """ + try: + return introspection.describe_algorithm(name) + except KeyError: + available = len(introspection.list_algorithms()) + raise ToolError( + f"Unknown algorithm '{name}'. Call list_algorithms to see the {available} available names." + ) from None + + @server.tool( + annotations=read_only, + description="List evaluation benchmarks and the metrics each one computes.", + ) + def list_benchmarks(task_type: str | None = None) -> dict[str, Any]: + """ + List the built-in evaluation benchmarks. + + Parameters + ---------- + task_type : str | None + Optional filter such as "text_to_image". If None, every benchmark is returned. + + Returns + ------- + dict[str, Any] + The benchmarks under the key "benchmarks", each including the metrics it evaluates. + """ + benchmarks = introspection.list_benchmarks(task_type) + return {"benchmarks": benchmarks, "count": len(benchmarks)} + + @server.tool( + annotations=read_only, + description=( + "List the named datasets an evaluation can run on, e.g. 'LAION256' or 'WikiText'. " + "Pass one of these names as the 'data' argument of evaluate_model." + ), + ) + def list_datasets() -> dict[str, Any]: + """ + List the named datasets available to an evaluation task. + + Returns + ------- + dict[str, Any] + The dataset names under the key "datasets", accepted verbatim by evaluate_model. + """ + datasets = introspection.list_datasets() + return {"datasets": datasets, "count": len(datasets)} + + @server.tool( + annotations=read_only, + description="List every evaluation metric name available to an evaluation task.", + ) + def list_metrics() -> dict[str, Any]: + """ + List the available evaluation metrics. + + Returns + ------- + dict[str, Any] + The sorted metric names usable with evaluate_model, under the key "metrics". + """ + metrics = introspection.list_metrics() + return {"metrics": metrics, "count": len(metrics)} + + @server.tool( + annotations=read_only, + description=( + "Validate a proposed configuration WITHOUT loading a model. Reports missing pip " + "extras, missing tokenizers/processors/datasets and incompatible algorithm pairs. " + "Cheap and safe to call repeatedly while iterating." + ), + ) + def validate_config( + algorithms: list[str], + hyperparameters: dict[str, Any] | None = None, + device: str | None = None, + batch_size: int = 1, + ) -> dict[str, Any]: + """ + Check whether a proposed configuration will work, without loading a model. + + Parameters + ---------- + algorithms : list[str] + Algorithm names to activate, as returned by list_algorithms. + hyperparameters : dict[str, Any] | None + Hyperparameter values keyed by fully qualified name, e.g. {"hqq_weight_bits": 4}. + Omitted hyperparameters keep their defaults. + device : str | None + Optional target device such as "cpu" or "cuda". + batch_size : int + Inference batch size to record on the configuration. + + Returns + ------- + dict[str, Any] + A report with a valid flag, any errors with hints, the resolved configuration, + the requirements the algorithms impose, and the checks that still need a model. + """ + return validation.validate_config(algorithms, hyperparameters, device, batch_size) + + if allow_execute: + + @server.tool( + annotations=executing, + description=( + "Load a base model and optimize it. This downloads weights and needs a device, " + "so it is slow. Call validate_config first." + ), + ) + def smash_model( + model_path: str, + algorithms: list[str], + hyperparameters: dict[str, Any] | None = None, + loader: str = "auto", + device: str | None = None, + experimental: bool = False, + save_to: str | None = None, + batch_size: int = 1, + ) -> dict[str, Any]: + """ + Load a base model, optimize it with the given algorithms, and optionally save it. + + The optimized model only exists in memory, so pass ``save_to`` to write it to disk + for later use. + + Parameters + ---------- + model_path : str + Hugging Face model id or local directory of the base model. + algorithms : list[str] + Algorithm names to apply, as returned by list_algorithms. + hyperparameters : dict[str, Any] | None + Hyperparameter values keyed by fully qualified name. Omitted hyperparameters + keep their defaults. + loader : str + How to load the model: "transformers", "diffusers", or "auto" to try + transformers first and fall back to diffusers. + device : str | None + Optional target device such as "cpu" or "cuda". + experimental : bool + Skip the strict compatibility checks that ``smash`` normally enforces. + save_to : str | None + Directory to save the optimized model to. The result is discarded if None. + batch_size : int + Inference batch size to record on the configuration. + + Returns + ------- + dict[str, Any] + The algorithms applied, the resolved configuration, and the save location if + the model was written to disk. + """ + _require_execution(allow_execute) + from pruna.engine.load import load_diffusers_model, load_transformers_model + from pruna.engine.save import save_pruna_model + from pruna.smash import smash + + loaders = { + "transformers": load_transformers_model, + "diffusers": load_diffusers_model, + } + if loader == "auto": + try: + model = load_transformers_model(model_path) + except Exception: # A diffusers checkpoint is not loadable by transformers + model = load_diffusers_model(model_path) + elif loader in loaders: + model = loaders[loader](model_path) + else: + raise ValueError(f"Unknown loader '{loader}'. Use 'auto', 'transformers' or 'diffusers'.") + + config = validation.build_config(algorithms, hyperparameters, device, batch_size) + optimized = smash(model, config, experimental=experimental) + + saved_to = None + if save_to is not None: + save_pruna_model(optimized, save_to, config) + saved_to = save_to + + return { + "applied_algorithms": config.get_active_algorithms(), + "resolved_hyperparameters": { + name: { + parameter.name: validation.resolve_hyperparameter(config, name, parameter) + for parameter in _algorithm_registry()[name].hyperparameters + } + for name in config.get_active_algorithms() + }, + "saved_to": saved_to, + } + + @server.tool( + annotations=executing, + description="Evaluate a model with the pruna evaluation agent. Needs a device and is slow.", + ) + def evaluate_model( + model_path: str, + request: str | list[str], + data: str, + device: str | None = None, + ) -> dict[str, Any]: + """ + Load a model and evaluate it, returning the computed metric results. + + Parameters + ---------- + model_path : str + Hugging Face model id or local directory of the model to evaluate. + request : str | list[str] + A metric name, a list of metric names, or one of the special requests + "image_generation_quality" and "text_generation_quality". + data : str + Name of the dataset to evaluate on, for example "LAION256", "ImageNet" or + "WikiText". + device : str | None + Optional device to evaluate on. + + Returns + ------- + dict[str, Any] + The metric results reported by the evaluation agent. + """ + _require_execution(allow_execute) + from pruna.data.pruna_datamodule import PrunaDataModule + from pruna.engine.load import load_diffusers_model, load_transformers_model + from pruna.evaluation.evaluation_agent import EvaluationAgent + + try: + model = load_transformers_model(model_path) + except Exception: # A diffusers checkpoint is not loadable by transformers + model = load_diffusers_model(model_path) + + # from_string is pruna's own entry point for a named dataset; the constructor takes + # dataset objects and has no data_name argument. An unknown name would raise a plain + # KeyError whose text the SDK discards, so it is turned into guidance here. + try: + datamodule = PrunaDataModule.from_string(data) + except Exception as error: + raise ToolError( + f"Could not load dataset '{data}': {error}. " + f"Call list_datasets for the {len(introspection.list_datasets())} available names." + ) from None + + agent = EvaluationAgent(request=request, datamodule=datamodule, device=device) + results = agent.evaluate(model, model_name=model_path) + + return {"metrics": [str(result) for result in results]} + + return server + + +def _algorithm_registry() -> Any: + """ + Import and return the algorithm registry. + + Returns + ------- + Any + The ``AlgorithmRegistry`` class from pruna. + """ + from pruna import AlgorithmRegistry + + return AlgorithmRegistry + + +def main() -> None: + """ + Command line entry point for the pruna MCP server. + + Parses arguments and runs the server over stdio. + """ + parser = argparse.ArgumentParser(prog="pruna-mcp", description="Model Context Protocol server for pruna.") + parser.add_argument( + "--allow-execute", + action="store_true", + help="Register smash_model and evaluate_model. Off by default because they load models and need a device.", + ) + arguments = parser.parse_args() + + try: + server = build_server(allow_execute=arguments.allow_execute) + except ImportError as error: + # A client launches this as a subprocess and shows whatever lands on stderr, so a clean + # one-line explanation is far more useful than an import traceback. + sys.stderr.write(f"pruna-mcp: {error}\n") + raise SystemExit(1) from None + + server.run(transport="stdio") + + +if __name__ == "__main__": + main() diff --git a/src/pruna/mcp_server/validation.py b/src/pruna/mcp_server/validation.py new file mode 100644 index 00000000..cc4ea67e --- /dev/null +++ b/src/pruna/mcp_server/validation.py @@ -0,0 +1,369 @@ +# Copyright 2025 - Pruna AI GmbH. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Configuration validation that runs without loading a model. + +``pruna.smash`` performs its checks as part of the optimization run, which means a bad +configuration is only discovered after a model has been loaded onto a device. The helpers +here reuse the same routines from :mod:`pruna.config.pre_smash_routines` to surface those +problems earlier, plus a static compatibility check that ``smash`` can only perform once a +model is in hand. +""" + +from __future__ import annotations + +from typing import Any + +from .introspection import _to_jsonable + + +def build_config( + algorithms: list[str], + hyperparameters: dict[str, Any] | None = None, + device: str | None = None, + batch_size: int = 1, +) -> Any: + """ + Build a ``SmashConfig`` from the flat arguments the MCP tools accept. + + Parameters + ---------- + algorithms : list[str] + Names of the algorithms to activate. + hyperparameters : dict[str, Any] | None + Hyperparameter values keyed by their fully qualified name, for example + ``{"hqq_weight_bits": 4}``. If None, every hyperparameter keeps its default. + device : str | None + Target device, for example ``"cpu"`` or ``"cuda"``. If None, the best available + device is chosen by pruna. + batch_size : int + Inference batch size recorded on the configuration. + + Returns + ------- + Any + The constructed ``SmashConfig``. + """ + from pruna import SmashConfig + + payload: dict[str, Any] = dict.fromkeys(algorithms, True) + payload.update(hyperparameters or {}) + config = SmashConfig(payload) + if device is not None: + config.device = device + config.batch_size = batch_size + return config + + +def collect_requirements(algorithms: list[str]) -> dict[str, Any]: + """ + Report what a set of algorithms will need before it can run. + + Parameters + ---------- + algorithms : list[str] + Names of the activated algorithms. + + Returns + ------- + dict[str, Any] + Lists of algorithms requiring a tokenizer, a processor or a dataset, together with the + union of the devices they support. + """ + from pruna import AlgorithmRegistry + + registry = AlgorithmRegistry + tokenizer: list[str] = [] + processor: list[str] = [] + dataset: list[str] = [] + devices: list[str] = [] + pip_install: list[str] = [] + + for name in algorithms: + algorithm = registry[name] + if getattr(algorithm, "tokenizer_required", False): + tokenizer.append(name) + if getattr(algorithm, "processor_required", False): + processor.append(name) + if getattr(algorithm, "dataset_required", False): + dataset.append(name) + for device in getattr(algorithm, "runs_on", []): + if device not in devices: + devices.append(device) + required_install = getattr(algorithm, "required_install", None) + if required_install and required_install not in pip_install: + pip_install.append(required_install) + + return { + "tokenizer": tokenizer, + "processor": processor, + "dataset": dataset, + "supported_devices": devices, + "pip_install": pip_install, + } + + +def find_incompatible_pairs(algorithms: list[str]) -> list[tuple[str, str]]: + """ + Statically detect algorithm pairs that cannot be combined. + + This mirrors the pairwise branch of ``check_algorithm_cross_compatibility`` without + needing a model. It therefore cannot detect overlapping ``target_modules`` between + algorithms that are only disjointly compatible; that check still needs a loaded model. + + Parameters + ---------- + algorithms : list[str] + Names of the activated algorithms. + + Returns + ------- + list[tuple[str, str]] + The incompatible pairs, each pair ordered deterministically. + """ + from pruna import AlgorithmRegistry + + incompatible: list[tuple[str, str]] = [] + for index, first in enumerate(algorithms): + for second in algorithms[index + 1 :]: + other = AlgorithmRegistry[second] + if first in other.get_compatible_algorithms(): + continue + if first in other.get_disjointly_compatible_algorithms(): + continue + incompatible.append((first, second)) + return incompatible + + +def determine_execution_order(config: Any) -> list[str] | None: + """ + Report the order the active algorithms will actually be applied in. + + This delegates to pruna's own ordering routine so the MCP server does not reimplement the + topological sort. A ``None`` result means the order could not be determined without a model, + which happens when an algorithm is only disjointly compatible with another and the decision + depends on their target modules. + + Parameters + ---------- + config : Any + The ``SmashConfig`` to order. + + Returns + ------- + list[str] | None + The algorithms in the order they will run, or None if it cannot be determined + without a model. + """ + from pruna.config.pre_smash_routines import determine_algorithm_order + + try: + return determine_algorithm_order(None, config) + except Exception: + return None + + +def resolve_hyperparameter(config: Any, algorithm_name: str, hyperparameter: Any) -> Any: + """ + Read the effective value of one hyperparameter from a configuration. + + ConfigSpace registers hyperparameters under their prefixed name, but the prefixing + happens as a side effect of algorithm discovery. Both the prefixed and the bare name are + therefore tried, so an explicitly configured value is never silently replaced by the + default. + + Parameters + ---------- + config : Any + The ``SmashConfig`` to read from. + algorithm_name : str + Name of the algorithm the hyperparameter belongs to. + hyperparameter : Any + The ConfigSpace hyperparameter object. + + Returns + ------- + Any + The configured value, or the hyperparameter's default when it was never set. + """ + names = [str(hyperparameter.name)] + prefixed = f"{algorithm_name}_{names[0]}" + if prefixed != names[0]: + names.append(prefixed) + + for candidate in names: + value = config.get(candidate) + if value is not None: + return value + return getattr(hyperparameter, "default_value", None) + + +def _requirements_or_empty(algorithms: list[str]) -> dict[str, Any]: + """ + Collect requirements, tolerating names that the registry does not know. + + Parameters + ---------- + algorithms : list[str] + Names of the algorithms the caller asked for, which may include unknown ones. + + Returns + ------- + dict[str, Any] + The requirements of every recognised algorithm, or an empty block if none is known. + """ + from pruna import AlgorithmRegistry + + known: list[str] = [] + for name in algorithms: + try: + AlgorithmRegistry[name] + except KeyError: + continue + known.append(name) + return collect_requirements(known) if known else {} + + +def validate_config( + algorithms: list[str], + hyperparameters: dict[str, Any] | None = None, + device: str | None = None, + batch_size: int = 1, +) -> dict[str, Any]: + """ + Validate a proposed optimization configuration without loading a model. + + Runs the same argument and dependency checks that ``pruna.smash`` runs, plus a static + pairwise compatibility check, so an invalid configuration is caught before any device + is touched. + + Parameters + ---------- + algorithms : list[str] + Names of the algorithms to activate, as returned by ``list_algorithms``. + hyperparameters : dict[str, Any] | None + Hyperparameter values keyed by their fully qualified name, for example + ``{"hqq_weight_bits": 4}``. If None, every hyperparameter keeps its default. + device : str | None + Target device, for example ``"cpu"`` or ``"cuda"``. If None, the best available + device is chosen by pruna. + batch_size : int + Inference batch size to record on the configuration. + + Returns + ------- + dict[str, Any] + A report with the ``valid`` flag, a list of ``errors``, the resolved configuration, + the requirements the algorithms impose, the ``execution_order`` they will be applied + in, and a list of ``notes`` describing what could not be checked without a model. + """ + from pruna.config.pre_smash_routines import ( + check_algorithm_packages_availability, + check_argument_compatibility, + ) + + errors: list[dict[str, str]] = [] + notes: list[str] = [] + + try: + config = build_config(algorithms, hyperparameters, device, batch_size) + except Exception as error: + # SmashConfig raises on an unknown algorithm or a rejected hyperparameter value. The + # report has to keep the same shape as the success path, so a client reading a fixed + # set of keys does not break on exactly the configurations it most needs to see. + return { + "valid": False, + "errors": [ + { + "check": "build_config", + "message": str(error), + "hint": ( + "Check the algorithm names against list_algorithms and the hyperparameter " + "values against the allowed values describe_algorithm reports." + ), + } + ], + "active_algorithms": [], + "resolved_hyperparameters": {}, + "requirements": _requirements_or_empty(algorithms), + "execution_order": None, + "notes": ["No configuration could be built, so nothing was checked against the registry."], + } + + active = config.get_active_algorithms() + + try: + check_argument_compatibility(config) + except Exception as error: + errors.append( + { + "check": "check_argument_compatibility", + "message": str(error), + "hint": "Attach the missing object to the SmashConfig before smashing.", + } + ) + + try: + check_algorithm_packages_availability(config) + except ImportError as error: + errors.append( + { + "check": "check_algorithm_packages_availability", + "message": str(error), + "hint": "Install the missing pip extra named in the message.", + } + ) + except Exception as error: + errors.append( + { + "check": "check_algorithm_packages_availability", + "message": str(error), + "hint": "", + } + ) + + for first, second in find_incompatible_pairs(active): + errors.append( + { + "check": "cross_compatibility", + "message": f"Algorithm '{first}' is incompatible with '{second}'.", + "hint": ( + "Call describe_algorithm on both and keep only pairs where one appears in the " + "other's 'compatible_with' list." + ), + } + ) + + notes.append("Model compatibility and target-module overlap are not checked here; they require a loaded model.") + + resolved: dict[str, Any] = {} + from pruna import AlgorithmRegistry + + for name in active: + algorithm = AlgorithmRegistry[name] + resolved[name] = { + hyperparameter.name: _to_jsonable(resolve_hyperparameter(config, name, hyperparameter)) + for hyperparameter in algorithm.hyperparameters + } + + return { + "valid": not errors, + "errors": errors, + "active_algorithms": active, + "resolved_hyperparameters": resolved, + "requirements": collect_requirements(active), + "execution_order": determine_execution_order(config), + "notes": notes, + } diff --git a/tests/conftest.py b/tests/conftest.py index dec0a636..e0f6cd3e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,6 +24,7 @@ "requires_kvpress": "mark test that needs pruna[kvpress]", "requires_llamacpp": "mark test that needs pruna[llamacpp]", "requires_transformers_legacy": "mark test that needs pruna[transformers-legacy]", + "requires_mcp": "mark test that needs pruna[mcp]", } diff --git a/tests/mcp_server/__init__.py b/tests/mcp_server/__init__.py new file mode 100644 index 00000000..38e0d7e5 --- /dev/null +++ b/tests/mcp_server/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 - Pruna AI GmbH. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/mcp_server/test_introspection.py b/tests/mcp_server/test_introspection.py new file mode 100644 index 00000000..44c61a76 --- /dev/null +++ b/tests/mcp_server/test_introspection.py @@ -0,0 +1,127 @@ +"""Tests for the MCP introspection helpers that expose the pruna algorithm catalogue.""" + +import pytest + +from pruna.mcp_server import introspection + + +@pytest.mark.cpu +class TestListTags: + """Tests for listing algorithm categories.""" + + def test_returns_every_tag_with_a_description(self) -> None: + """Every algorithm category should be reported with a human-readable description.""" + tags = introspection.list_tags() + assert len(tags) > 0 + for tag in tags: + assert set(tag) == {"name", "description"} + assert tag["name"] and tag["description"] + + +@pytest.mark.cpu +class TestListAlgorithms: + """Tests for enumerating the algorithm catalogue.""" + + def test_lists_every_registered_algorithm(self) -> None: + """The unfiltered listing should cover the whole registry.""" + from pruna import AlgorithmRegistry + + algorithms = introspection.list_algorithms() + assert len(algorithms) == len(AlgorithmRegistry.get_all_algorithms()) + + def test_summary_reports_name_devices_and_requirements(self) -> None: + """Each summary should carry the fields a caller needs to plan a configuration.""" + for algorithm in introspection.list_algorithms(): + assert algorithm["name"] + assert isinstance(algorithm["runs_on"], list) + assert set(algorithm["requires"]) == {"tokenizer", "processor", "dataset", "pip_install"} + + def test_tag_filter_returns_a_subset(self) -> None: + """Filtering by a category should return fewer algorithms than the full listing.""" + everything = introspection.list_algorithms() + quantizers = introspection.list_algorithms(tag="quantizer") + assert len(quantizers) < len(everything) + assert all("quantizer" in algorithm["tags"] for algorithm in quantizers) + + def test_tag_filter_is_case_insensitive(self) -> None: + """Category names should be accepted in either case.""" + assert introspection.list_algorithms(tag="QUANTIZER") == introspection.list_algorithms(tag="quantizer") + + def test_unknown_tag_raises_with_the_valid_options(self) -> None: + """An unknown category should fail loudly and list what is accepted.""" + with pytest.raises(ValueError, match="Unknown algorithm tag"): + introspection.list_algorithms(tag="not_a_real_tag") + + +@pytest.mark.cpu +class TestDescribeAlgorithm: + """Tests for the single-algorithm detail view.""" + + def test_describes_hyperparameters_with_qualified_names(self) -> None: + """Hyperparameters should be reported both bare and algorithm-prefixed.""" + described = introspection.describe_algorithm("hqq") + assert described["name"] == "hqq" + assert described["hyperparameters"] + + for hyperparameter in described["hyperparameters"]: + assert hyperparameter["qualified_name"] == f"hqq_{hyperparameter['name']}" + assert hyperparameter["type"] + assert "default" in hyperparameter + + def test_reports_compatibility_and_requirements(self) -> None: + """Order constraints and requirements should be present for configuration planning.""" + described = introspection.describe_algorithm("hqq") + for key in ("compatible_before", "compatible_after", "compatible_with", "requires", "runs_on"): + assert key in described + + def test_unknown_algorithm_raises(self) -> None: + """An unknown algorithm name should raise rather than return an empty description.""" + with pytest.raises(KeyError): + introspection.describe_algorithm("definitely_not_an_algorithm") + + +@pytest.mark.cpu +class TestEvaluationCatalogue: + """Tests for the benchmark and metric listings.""" + + def test_list_metrics_returns_sorted_names(self) -> None: + """Metrics should be reported as a sorted list of names.""" + metrics = introspection.list_metrics() + assert metrics == sorted(metrics) + assert all(isinstance(name, str) for name in metrics) + + def test_list_benchmarks_reports_metrics_per_benchmark(self) -> None: + """Each benchmark should expose the metrics it evaluates.""" + benchmarks = introspection.list_benchmarks() + assert len(benchmarks) > 0 + for benchmark in benchmarks: + assert benchmark["name"] + assert isinstance(benchmark["metrics"], list) + + def test_list_benchmarks_filters_by_task_type(self) -> None: + """Filtering by task type should return a subset carrying that task type.""" + filtered = introspection.list_benchmarks(task_type="text_to_image") + assert len(filtered) <= len(introspection.list_benchmarks()) + for benchmark in filtered: + assert benchmark["task_type"] == "text_to_image" + +@pytest.mark.cpu +class TestDatasets: + """Tests for the named dataset catalogue.""" + + def test_returns_sorted_names(self) -> None: + """The names should come back in a stable order so clients can rely on them.""" + names = introspection.list_datasets() + assert names == sorted(names) + assert "LAION256" in names + + def test_every_name_is_accepted_by_the_datamodule_factory(self) -> None: + """ + Each advertised name must be one PrunaDataModule.from_string can actually resolve. + + The list is read straight from pruna's own dataset registry, so this guards against the + registry and the factory drifting apart, which would hand evaluate_model a bad name. + """ + from pruna.data import base_datasets + + assert introspection.list_datasets() == sorted(base_datasets.keys()) diff --git a/tests/mcp_server/test_server.py b/tests/mcp_server/test_server.py new file mode 100644 index 00000000..fa2305d8 --- /dev/null +++ b/tests/mcp_server/test_server.py @@ -0,0 +1,292 @@ +"""Tests for MCP tool registration and the read-only/execution split.""" + +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + +mcp = pytest.importorskip("mcp") + +from pruna.mcp_server.server import build_server # noqa: E402 + +READ_ONLY_TOOLS = { + "list_tags", + "list_algorithms", + "describe_algorithm", + "list_benchmarks", + "list_metrics", + "list_datasets", + "validate_config", +} +EXECUTION_TOOLS = {"smash_model", "evaluate_model"} + + +def _tool_names(allow_execute: bool = False) -> set[str]: + """Return the names of every tool registered on a server.""" + import anyio + + server = build_server(allow_execute=allow_execute) + return {tool.name for tool in anyio.run(server.list_tools)} + + +@pytest.mark.cpu +@pytest.mark.requires_mcp +class TestToolRegistration: + """Tests for which tools the server exposes.""" + + def test_every_documented_tool_is_registered(self) -> None: + """The documented tool list and the code must not drift apart.""" + import re + from pathlib import Path + + docs = Path(__file__).parents[2] / "docs" / "user_manual" / "mcp_server.rst" + documented = set(re.findall(r"^``([a-z_]+)``$", docs.read_text(encoding="utf-8"), re.M)) + assert documented <= _tool_names(allow_execute=True), ( + f"documented but not registered: {sorted(documented - _tool_names(allow_execute=True))}" + ) + + def test_read_only_mode_registers_the_discovery_tier(self) -> None: + """The safe tools should always be available.""" + assert READ_ONLY_TOOLS <= _tool_names(allow_execute=False) + + def test_execution_tools_are_absent_by_default(self) -> None: + """Tools that load models must not be exposed without an explicit opt-in.""" + assert not (EXECUTION_TOOLS & _tool_names(allow_execute=False)) + + def test_execution_tools_appear_when_enabled(self) -> None: + """Opting in should add the execution tier alongside the discovery tier.""" + names = _tool_names(allow_execute=True) + assert EXECUTION_TOOLS <= names + assert READ_ONLY_TOOLS <= names + + def test_tools_carry_input_schemas(self) -> None: + """Every tool must expose a schema so a client can call it.""" + import anyio + + server = build_server(allow_execute=False) + for tool in anyio.run(server.list_tools): + assert tool.input_schema["type"] == "object" + + def test_discovery_tools_are_annotated_read_only(self) -> None: + """Clients rely on the read-only hint to avoid prompting on safe calls.""" + import anyio + + server = build_server(allow_execute=False) + for tool in anyio.run(server.list_tools): + assert tool.annotations is not None + assert tool.annotations.read_only_hint is True + + +@pytest.mark.cpu +@pytest.mark.requires_mcp +class TestToolExecution: + """Tests for tools actually returning data.""" + + def test_list_tags_returns_the_catalogue(self) -> None: + """Calling the server should yield the algorithm categories.""" + import anyio + + result = anyio.run(lambda: build_server().call_tool("list_tags", {})) + assert result.is_error is False + assert result.structured_content["count"] > 0 + assert all({"name", "description"} <= set(tag) for tag in result.structured_content["tags"]) + + def test_every_tool_returns_one_coherent_json_object(self) -> None: + """ + Each tool must answer with a single content block holding a JSON object. + + A bare list return is serialized by the SDK as one content block per element, so a + catalogue of 48 algorithms would arrive as 48 fragments for the model to reassemble. + Returning a named object keeps every tool's shape uniform. + """ + import anyio + + server = build_server() + calls = { + "list_tags": {}, + "list_algorithms": {}, + "list_algorithms_filtered": {"tag": "quantizer"}, + "describe_algorithm": {"name": "torch_compile"}, + "list_benchmarks": {}, + "list_metrics": {}, + "list_datasets": {}, + "validate_config": {"algorithms": ["torch_compile"]}, + } + for label, arguments in calls.items(): + name = "list_algorithms" if label == "list_algorithms_filtered" else label + result = anyio.run(lambda n=name, a=arguments: server.call_tool(n, a)) + assert result.is_error is False, f"{label} reported an error" + assert len(result.content) == 1, f"{label} fragmented its answer into {len(result.content)} blocks" + assert isinstance(result.structured_content, dict), f"{label} did not return an object" + assert "result" not in result.structured_content, f"{label} fell back to the SDK's list wrapper" + + def test_validate_config_reports_its_result(self) -> None: + """The validator tool should return a structured report.""" + import anyio + + result = anyio.run( + lambda: build_server().call_tool("validate_config", {"algorithms": ["deepcache"]}) + ) + assert result.is_error is False + assert result.structured_content["valid"] is True + + def test_describe_unknown_algorithm_keeps_the_error_message(self) -> None: + """ + An unknown algorithm name is a routine mistake, so the guidance must reach the model. + + The SDK turns a plain exception into a generic message and drops the detail, so the + tools raise ToolError deliberately to preserve it. + """ + import anyio + from mcp.server.mcpserver.exceptions import ToolError + + with pytest.raises(ToolError) as caught: + anyio.run(lambda: build_server().call_tool("describe_algorithm", {"name": "nope"})) + + message = str(caught.value) + assert "nope" in message + assert "list_algorithms" in message + + def test_invalid_tag_keeps_the_error_message(self) -> None: + """An unknown category should list the valid ones rather than fail opaquely.""" + import anyio + from mcp.server.mcpserver.exceptions import ToolError + + with pytest.raises(ToolError) as caught: + anyio.run(lambda: build_server().call_tool("list_algorithms", {"tag": "not_a_tag"})) + + assert "not_a_tag" in str(caught.value) + +@pytest.mark.cpu +@pytest.mark.requires_mcp +class TestExecutionTierContracts: + """ + Tests for the execution tools that do not need a model to run. + + These tools call deep into pruna behind function-level imports, so a wrong argument name or + a renamed factory would not surface until a model was downloaded. These tests bind the + arguments to the real signatures instead. + """ + + @staticmethod + def _bind(function, **arguments): + """Return the parameter names the function would actually accept.""" + import inspect + + return set(inspect.signature(function).parameters) + + def test_smash_accepts_the_configuration_the_tool_builds(self) -> None: + """The config the tool builds has to be what smash expects.""" + from pruna.smash import smash + + assert {"model", "smash_config", "experimental"} <= self._bind(smash) + + def test_saving_accepts_the_optimized_model_and_config(self) -> None: + """save_pruna_model takes the model, a path and the config, in that order.""" + from pruna.engine.save import save_pruna_model + + assert self._bind(save_pruna_model) >= {"model", "model_path", "smash_config"} + + def test_loaders_take_a_path(self) -> None: + """Both loaders accept a bare path, which is what the tool passes.""" + from pruna.engine.load import load_diffusers_model, load_transformers_model + + for loader in (load_transformers_model, load_diffusers_model): + assert self._bind(loader) >= {"path"} + + def test_a_named_dataset_resolves_through_the_public_factory(self) -> None: + """ + evaluate_model must build its datamodule through PrunaDataModule.from_string. + + The constructor takes dataset objects and has no data_name argument, so calling it + directly with a name raises a TypeError that the SDK would reduce to a bare + "Error executing tool evaluate_model". + """ + from pruna.data.pruna_datamodule import PrunaDataModule + + assert "data_name" not in self._bind(PrunaDataModule.__init__) + assert "dataset_name" in self._bind(PrunaDataModule.from_string) + + def test_evaluation_agent_accepts_the_tool_keyword_arguments(self) -> None: + """The agent takes request, datamodule and device as keywords.""" + from pruna.evaluation.evaluation_agent import EvaluationAgent + + assert {"request", "datamodule", "device"} <= self._bind(EvaluationAgent.__init__) + assert {"model", "model_name"} <= self._bind(EvaluationAgent.evaluate) + + +@pytest.mark.cpu +@pytest.mark.requires_mcp +class TestMissingExtra: + """Tests for a pruna installed without the mcp extra.""" + + def test_the_package_imports_without_the_extra(self) -> None: + """ + Importing the module must not require the mcp package. + + The console script imports it before main() runs, so an import-time raise would replace + the guidance with a traceback that a client shows to the user verbatim. + """ + result = subprocess.run( + [ + sys.executable, + "-c", + "import pruna.mcp_server.server as s; print(s._MCP_IMPORT_ERROR is None)", + ], + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + assert "True" in result.stdout + + def test_building_without_the_extra_explains_how_to_install_it(self) -> None: + """The failure has to name the extra, not just the missing module.""" + blocker = tempfile.mkdtemp() + Path(blocker, "sitecustomize.py").write_text( + "import sys\n" + "class Blocker:\n" + " def find_spec(self, name, path=None, target=None):\n" + " if name == 'mcp' or name.startswith('mcp.'):\n" + " raise ImportError(\"No module named 'mcp'\")\n" + " return None\n" + "sys.meta_path.insert(0, Blocker())\n", + encoding="utf-8", + ) + environment = {**os.environ, "PYTHONPATH": blocker} + result = subprocess.run( + [sys.executable, "-c", "import pruna.mcp_server.server as s; s.build_server()"], + capture_output=True, + text=True, + env=environment, + timeout=120, + ) + assert result.returncode != 0 + assert "pruna[mcp]" in result.stderr, result.stderr[-500:] + + def test_the_console_script_reports_the_missing_extra_cleanly(self) -> None: + """A user wiring pruna-mcp into a CLI should get one line, not a traceback.""" + blocker = tempfile.mkdtemp() + Path(blocker, "sitecustomize.py").write_text( + "import sys\n" + "class Blocker:\n" + " def find_spec(self, name, path=None, target=None):\n" + " if name == 'mcp' or name.startswith('mcp.'):\n" + " raise ImportError(\"No module named 'mcp'\")\n" + " return None\n" + "sys.meta_path.insert(0, Blocker())\n", + encoding="utf-8", + ) + result = subprocess.run( + [sys.executable, "-m", "pruna.mcp_server.server"], + capture_output=True, + text=True, + env={**os.environ, "PYTHONPATH": blocker}, + timeout=120, + ) + assert result.returncode == 1 + assert "pruna[mcp]" in result.stderr + assert "Traceback" not in result.stderr, result.stderr[-500:] diff --git a/tests/mcp_server/test_validation.py b/tests/mcp_server/test_validation.py new file mode 100644 index 00000000..b1715a90 --- /dev/null +++ b/tests/mcp_server/test_validation.py @@ -0,0 +1,307 @@ +"""Tests for model-free configuration validation used by the MCP server.""" + +import json + +import pytest + +from pruna.mcp_server import introspection, validation + +# Subjects chosen so assertions do not depend on which optional extras are installed. +# Each is a real registered algorithm in pruna's catalogue. +CORE_ALGORITHM = "torch_compile" # needs no tokenizer, processor or dataset, and no pip extra +QUANTIZER_ALGORITHM = "quanto" # core-importable, and exposes a tunable weight-bit hyperparameter +EXTRA_ONLY_ALGORITHM = "stable_fast" # importable only with pruna[stable-fast], absent from the dev extra +TOKENIZER_ALGORITHM = "padding_pruning" # always requires a tokenizer to be attached + + +def _is_importable(algorithm_name: str) -> bool: + """ + Report whether an algorithm's own dependencies can be imported. + + Parameters + ---------- + algorithm_name : str + The algorithm to probe. + + Returns + ------- + bool + True if the algorithm's imports succeed. + """ + from pruna import AlgorithmRegistry + + try: + AlgorithmRegistry[algorithm_name].import_algorithm_packages() + except ImportError: + return False + return True + + +@pytest.mark.cpu +class TestValidateConfig: + """Tests for the GPU-free configuration validator.""" + + def test_core_config_reports_no_errors(self) -> None: + """An algorithm with no unmet requirements should validate cleanly.""" + report = validation.validate_config(algorithms=[CORE_ALGORITHM]) + assert report["errors"] == [] + assert report["valid"] is True + assert report["active_algorithms"] == [CORE_ALGORITHM] + + def test_missing_pip_extra_is_reported(self) -> None: + """An algorithm behind an uninstalled extra should be reported, not swallowed.""" + if _is_importable(EXTRA_ONLY_ALGORITHM): + pytest.skip(f"{EXTRA_ONLY_ALGORITHM} is installed in this environment") + + report = validation.validate_config(algorithms=[EXTRA_ONLY_ALGORITHM]) + assert report["valid"] is False + assert any( + error["check"] == "check_algorithm_packages_availability" for error in report["errors"] + ) + assert EXTRA_ONLY_ALGORITHM in report["requirements"]["pip_install"] or any( + EXTRA_ONLY_ALGORITHM in error["message"] for error in report["errors"] + ) + + def test_missing_tokenizer_is_reported(self) -> None: + """An algorithm that needs a tokenizer should say so when none is attached.""" + report = validation.validate_config(algorithms=[TOKENIZER_ALGORITHM]) + assert report["valid"] is False + assert any(error["check"] == "check_argument_compatibility" for error in report["errors"]) + assert TOKENIZER_ALGORITHM in report["requirements"]["tokenizer"] + + def test_resolves_a_value_for_every_declared_hyperparameter(self) -> None: + """Every declared hyperparameter should appear in the resolved configuration.""" + from pruna import AlgorithmRegistry + + declared = {hp.name for hp in AlgorithmRegistry[CORE_ALGORITHM].hyperparameters} + resolved = validation.validate_config(algorithms=[CORE_ALGORITHM])["resolved_hyperparameters"] + assert set(resolved[CORE_ALGORITHM]) == declared + + def test_resolved_configuration_is_json_serializable(self) -> None: + """The report has to survive the JSON encoding an MCP client applies.""" + json.dumps(validation.validate_config(algorithms=[CORE_ALGORITHM, QUANTIZER_ALGORITHM])) + + def test_untouched_hyperparameter_reports_its_default(self) -> None: + """A hyperparameter that was never set should fall back to its declared default.""" + from pruna import AlgorithmRegistry + + report = validation.validate_config(algorithms=[CORE_ALGORITHM]) + for hyperparameter in AlgorithmRegistry[CORE_ALGORITHM].hyperparameters: + resolved = report["resolved_hyperparameters"][CORE_ALGORITHM][hyperparameter.name] + assert resolved == validation._to_jsonable(hyperparameter.default_value) + + def test_explicit_hyperparameter_is_respected(self) -> None: + """A configured hyperparameter must win over the default.""" + report = validation.validate_config( + algorithms=[CORE_ALGORITHM], + hyperparameters={f"{CORE_ALGORITHM}_max_kv_cache_size": 512}, + ) + resolved = report["resolved_hyperparameters"][CORE_ALGORITHM] + assert resolved[f"{CORE_ALGORITHM}_max_kv_cache_size"] == 512 + + def test_unknown_algorithm_is_reported_not_raised(self) -> None: + """A bad configuration should come back as a report, not an exception.""" + report = validation.validate_config(algorithms=["not_a_real_algorithm"]) + assert report["valid"] is False + assert report["errors"][0]["check"] == "build_config" + + def test_errors_carry_a_hint(self) -> None: + """Every reported error should explain what went wrong.""" + for error in validation.validate_config(algorithms=["not_a_real_algorithm"])["errors"]: + assert set(error) == {"check", "message", "hint"} + assert error["message"] + + def test_report_keeps_its_shape_when_the_config_cannot_be_built(self) -> None: + """ + A rejected hyperparameter value must still produce the full report. + + A client reads a fixed set of keys, so the failure path has to carry every key the + success path does rather than a shorter dict. + """ + from pruna import AlgorithmRegistry + + algorithm = AlgorithmRegistry[QUANTIZER_ALGORITHM] + hyperparameter = algorithm.hyperparameters[0] + + report = validation.validate_config( + algorithms=[QUANTIZER_ALGORITHM], + hyperparameters={f"{QUANTIZER_ALGORITHM}_{hyperparameter.name}": object()}, + ) + success = validation.validate_config(algorithms=[CORE_ALGORITHM]) + + assert set(report) == set(success) + assert report["valid"] is False + assert report["errors"][0]["check"] == "build_config" + assert report["execution_order"] is None + assert report["notes"] + + def test_a_rejected_hyperparameter_value_is_explained(self) -> None: + """The message has to name the offending value so the model can correct it.""" + report = validation.validate_config( + algorithms=[QUANTIZER_ALGORITHM], + hyperparameters={f"{QUANTIZER_ALGORITHM}_weight_bits": "definitely-not-a-value"}, + ) + assert "weight_bits" in report["errors"][0]["message"] + assert "describe_algorithm" in report["errors"][0]["hint"] + + def test_reports_the_execution_order(self) -> None: + """The report should state the order the algorithms will actually be applied in.""" + report = validation.validate_config(algorithms=[CORE_ALGORITHM, QUANTIZER_ALGORITHM]) + order = report["execution_order"] + assert order is not None + assert set(order) == {CORE_ALGORITHM, QUANTIZER_ALGORITHM} + + def test_execution_order_respects_declared_constraints(self) -> None: + """The reported order must match what pruna's own ordering routine produces.""" + from pruna import SmashConfig + from pruna.config.pre_smash_routines import determine_algorithm_order + + algorithms = [CORE_ALGORITHM, QUANTIZER_ALGORITHM] + config = SmashConfig(dict.fromkeys(algorithms, True)) + expected = determine_algorithm_order(None, config) + assert validation.validate_config(algorithms=algorithms)["execution_order"] == expected + + def test_report_notes_what_still_needs_a_model(self) -> None: + """The report should be explicit about the checks it could not perform.""" + report = validation.validate_config(algorithms=[CORE_ALGORITHM]) + assert any("model" in note.lower() for note in report["notes"]) + + def test_reports_requirements_of_the_selected_algorithms(self) -> None: + """The requirements block should reflect only the selected algorithms.""" + requirements = validation.validate_config(algorithms=[CORE_ALGORITHM])["requirements"] + assert set(requirements) == {"tokenizer", "processor", "dataset", "supported_devices", "pip_install"} + assert requirements["tokenizer"] == [] + + +@pytest.mark.cpu +class TestRoundTrip: + """Tests that what the server advertises can actually be configured.""" + + @staticmethod + def _importable_names() -> list[str]: + """Return the registry names whose own dependencies import successfully.""" + from pruna import AlgorithmRegistry + + names = [] + for algorithm in AlgorithmRegistry.get_all_algorithms(): + try: + algorithm.import_algorithm_packages() + except Exception: # an uninstalled extra or a missing optional dependency + continue + names.append(algorithm.algorithm_name) + return names + + def test_every_advertised_default_is_accepted_back(self) -> None: + """ + Feeding describe_algorithm's own defaults into validate_config must succeed. + + This is the contract the model relies on: the server states a default for each + hyperparameter, so that exact value has to be one pruna will accept. + """ + from pruna import SmashConfig + + checked = 0 + for name in self._importable_names(): + described = introspection.describe_algorithm(name) + payload: dict[str, object] = {name: True} + for hyperparameter in described["hyperparameters"]: + payload[hyperparameter["qualified_name"]] = hyperparameter["default"] + try: + SmashConfig(payload) + except Exception as error: # noqa: BLE001 - the assertion reports the detail + pytest.fail(f"{name}: describe_algorithm advertised a default pruna rejected: {error}") + checked += 1 + + assert checked > 10, f"only {checked} algorithms were exercised, too few to be meaningful" + + def test_every_advertised_allowed_value_is_accepted_back(self) -> None: + """ + Every value describe_algorithm lists as allowed must be one pruna accepts. + + Allowed values are read off the ConfigSpace hyperparameter, so this guards the + attribute the reader trusts is the authoritative one. + """ + from pruna import SmashConfig + + exercised = 0 + for name in self._importable_names(): + described = introspection.describe_algorithm(name) + for hyperparameter in described["hyperparameters"]: + allowed = hyperparameter.get("allowed") + if not allowed: + continue + for value in allowed: + payload: dict[str, object] = {name: True, hyperparameter["qualified_name"]: value} + try: + SmashConfig(payload) + except Exception as error: # noqa: BLE001 - the assertion reports the detail + pytest.fail( + f"{name}.{hyperparameter['qualified_name']}: describe_algorithm advertised " + f"{value!r} as allowed but pruna rejected it: {error}" + ) + exercised += 1 + + assert exercised > 20, f"only {exercised} values were exercised, too few to be meaningful" + + +@pytest.mark.cpu +class TestIncompatiblePairs: + """Tests for the static pairwise compatibility check.""" + + @staticmethod + def _registry_names() -> list[str]: + """Return every registered algorithm name.""" + from pruna import AlgorithmRegistry + + return [algorithm.algorithm_name for algorithm in AlgorithmRegistry.get_all_algorithms()] + + def test_single_algorithm_has_no_pairs(self) -> None: + """A lone algorithm cannot conflict with anything.""" + assert validation.find_incompatible_pairs([CORE_ALGORITHM]) == [] + + def test_reported_pairs_are_genuinely_incompatible(self) -> None: + """Every reported pair must be absent from both compatibility lists.""" + from pruna import AlgorithmRegistry + + for first, second in validation.find_incompatible_pairs(self._registry_names()): + other = AlgorithmRegistry[second] + assert first not in other.get_compatible_algorithms() + assert first not in other.get_disjointly_compatible_algorithms() + + def test_compatible_algorithms_are_not_reported(self) -> None: + """Algorithms declared mutually compatible must not be flagged.""" + from pruna import AlgorithmRegistry + + for first, second in validation.find_incompatible_pairs(self._registry_names()): + assert not ( + first in AlgorithmRegistry[second].get_compatible_algorithms() + or second in AlgorithmRegistry[first].get_compatible_algorithms() + ) + + def test_incompatible_pair_surfaces_as_an_error(self) -> None: + """An incompatible combination should make the whole report invalid.""" + found = validation.find_incompatible_pairs(self._registry_names()) + if not found: + pytest.skip("no algorithm pair in the registry is declared incompatible") + + report = validation.validate_config(algorithms=list(found[0])) + assert report["valid"] is False + assert any(error["check"] == "cross_compatibility" for error in report["errors"]) + + +@pytest.mark.cpu +class TestBuildConfig: + """Tests for constructing a SmashConfig from flat tool arguments.""" + + def test_activates_the_named_algorithms(self) -> None: + """Every named algorithm should end up active.""" + config = validation.build_config([CORE_ALGORITHM, QUANTIZER_ALGORITHM]) + assert set(config.get_active_algorithms()) == {CORE_ALGORITHM, QUANTIZER_ALGORITHM} + + def test_records_batch_size(self) -> None: + """The requested batch size should be stored on the configuration.""" + assert validation.build_config([CORE_ALGORITHM], batch_size=4).batch_size == 4 + + def test_records_device_when_given(self) -> None: + """An explicit device should be stored on the configuration.""" + assert str(validation.build_config([CORE_ALGORITHM], device="cpu").device) == "cpu" \ No newline at end of file