[ENH]: Add save_model MCP tool for persisting estimators via sktime MLflow (#48)

#### Reference Issues/PRs
Fixes #31 

#### What does this implement/fix? Explain your changes.
Implements the requested `save_model` MCP tool for `sktime-mcp`.
This PR adds a new MCP tool, `save_model`, to allow persistent storage
of instantiated estimators and pipelines via
`sktime.utils.mlflow_sktime.save_model`.
For `docs`:
updated user-facing docs and examples to document the new tool, clarify
the current local filesystem path behavior of sktime's save_model API,
and mention the MLflow runtime requirement

Changes included:
- added `src/sktime_mcp/tools/save_model.py` with `save_model_tool` .
- resolved estimator handles to the underlying estimator instance
through the handle manager
- registered the tool in `src/sktime_mcp/server.py` with schema and
dispatch wiring
- added unit coverage in `tests/test_core.py`
checked if run locally or not with this command `python -m pytest
tests/test_core.py::TestTools::test_save_model_tool -q`
<img width="1856" height="135" alt="image"
src="https://github.com/user-attachments/assets/fbaf5475-8610-4be6-a09f-43fd7e0297da"
/>
- updated user-facing docs and examples to document the new too

#### Does your contribution introduce a new dependency? If yes, which
one?
No but ,the tool relies on `sktime`'s MLflow integration
(`sktime.utils.mlflow_sktime.save_model`), so MLflow must be available
in the runtime environment for the tool to work.
I documented this behavior, but did not add `mlflow` to project
dependencies .

#### What should reviewers focus on?
I’d really appreciate feedback on a few specific areas:
1. Does `save_model` feel like the right MCP interface for handling
persistence in this project?
2. Should we be more explicit in the contract and say “local filesystem
path” instead of the broader “path or URI”?
3. Does it make sense to keep MLflow as a documented/runtime
requirement, or should it be added as an explicit dependency

#### Any other comments?
The implementation follows the documented
`sktime.utils.mlflow_sktime.save_model` API, which describes saving to a
local path on the filesystem. I aligned the documentation with that
behavior to avoid overstating URI support.

#### PR checklist
##### For all contributions
- [ ] I've added myself to the
`https://github.com/alan-turing-institute/sktime/blob/main/.all-contributorsrc`
.
- [ ] Optionally, I've updated sktime's
`https://github.com/alan-turing-institute/sktime/blob/main/CODEOWNERS`
to receive notifications about future changes to these files.
- [x] I've added unit tests and made sure they pass locally.

##### For new estimators
- [ ] Not applicable, this PR does not add a new estimator.
- [ ] Not applicable, this PR does not add or modify estimator example
notebooks.

ran the test locally they are pssing
<img width="1840" height="176" alt="image"
src="https://github.com/user-attachments/assets/2013fff5-4792-48e3-9ec4-9958aa16ed81"
/>

Author: direkkakkar319-ops

README.md

@@ -22,7 +22,7 @@ This MCP is **not** just documentation or static code analysis. It is a **semant
 
 2. **Registry-First** - Instead of `File → Class → Infer Relationships`, we do `Registry → Semantics → Safe Execution`.
 
-3. **Minimal MCP Surface** - Exposes only what an LLM needs: Discovery, Description, Instantiation, Execution.
+3. **Minimal MCP Surface** - Exposes only what an LLM needs: Discovery, Description, Instantiation, Execution, and model persistence.
 
 ## 🛠️ Installation
 
@@ -235,9 +235,34 @@ Execute a complete workflow: load dataset, fit estimator, and generate predictio
 
 ---
 
+#### 9. `save_model`
+Persist a fitted estimator or pipeline handle to a local filesystem path using `sktime.utils.mlflow_sktime.save_model`.
+
+**Arguments:**
+- `estimator_handle` (required): Handle from `instantiate_estimator` or `instantiate_pipeline`
+- `path` (required): Local filesystem path where the model should be saved
+- `mlflow_params` (optional): Extra keyword arguments forwarded to `sktime.utils.mlflow_sktime.save_model`
+
+**Example:**
+```json
+{
+  "estimator_handle": "est_abc123",
+  "path": "/absolute/path/to/model_dir",
+  "mlflow_params": {
+    "serialization_format": "cloudpickle"
+  }
+}
+```
+
+**Returns:** `{"success": true, "saved_path": "/absolute/path/to/model_dir", "message": "Model saved successfully to '/absolute/path/to/model_dir'"}`
+
+**Note:** This tool requires MLflow to be available in the server environment.
+
+---
+
 ### Datasets
 
-#### 9. `list_datasets`
+#### 10. `list_datasets`
 List all available demo datasets for testing and experimentation.
 
 **Arguments:** None
@@ -248,7 +273,7 @@ List all available demo datasets for testing and experimentation.
 
 ### Handle Management
 
-#### 10. `list_handles`
+#### 11. `list_handles`
 List all active estimator handles and their status.
 
 **Arguments:** None
@@ -257,7 +282,7 @@ List all active estimator handles and their status.
 
 ---
 
-#### 11. `release_handle`
+#### 12. `release_handle`
 Release an estimator handle and free memory.
 
 **Arguments:**

docs/usage-examples.md

@@ -159,3 +159,23 @@ When you are done with an estimator, it's good practice to release it to free up
   }
 }
 ```
+
+### Model Persistence
+
+**Save a Fitted Estimator**
+Use `save_model` after fitting an estimator or pipeline handle. The underlying `sktime.utils.mlflow_sktime.save_model` API saves to a local filesystem path.
+
+```json
+{
+  "name": "save_model",
+  "arguments": {
+    "estimator_handle": "est_abc123",
+    "path": "/absolute/path/to/model_dir",
+    "mlflow_params": {
+      "serialization_format": "pickle"
+    }
+  }
+}
+```
+
+*Returns:* `{"success": true, "saved_path": "/absolute/path/to/model_dir", "message": "Model saved successfully to '/absolute/path/to/model_dir'"}`

docs/user-guide.md

@@ -54,7 +54,7 @@ The `sktime-mcp` server exposes a suite of tools designed for Large Language Mod
 | **Instantiation** | `instantiate_estimator`, `instantiate_pipeline` | Create model instances or complex pipelines. |
 | **Execution** | `fit_predict`, `fit`, `predict` | Train models and generate forecasts. |
 | **Data** | `load_data_source`, `list_datasets` | Load data from Pandas, CSV/Parquet, or SQL. |
-| **Export** | `export_code` | Generate Python code to reproduce your results. |
+| **Export** | `export_code`, `save_model` | Generate Python code or persist fitted estimators to a local path. |
 
 ---
 
@@ -114,6 +114,46 @@ Check if components work together (e.g., Deseasonalizer -> Detrender -> ARIMA).
 }
 ```
 
+### 3. Save a Fitted Model
+
+Persist a trained estimator to a local filesystem path using sktime's MLflow integration.
+
+**Step 1: Fit the estimator**
+```json
+{
+  "tool": "fit_predict",
+  "arguments": {
+    "estimator_handle": "est_abc123",
+    "dataset": "airline",
+    "horizon": 12
+  }
+}
+```
+
+**Step 2: Save the fitted model**
+```json
+{
+  "tool": "save_model",
+  "arguments": {
+    "estimator_handle": "est_abc123",
+    "path": "/absolute/path/to/model_dir",
+    "mlflow_params": {
+      "serialization_format": "cloudpickle"
+    }
+  }
+}
+```
+
+**Typical response**
+```json
+{
+  "success": true,
+  "estimator_handle": "est_abc123",
+  "saved_path": "/absolute/path/to/model_dir",
+  "message": "Model saved successfully to '/absolute/path/to/model_dir'"
+}
+```
+
 ---
 
 ## 💾 Data Management
@@ -150,6 +190,7 @@ Bring your own data into the MCP server.
 
 - **Resource Management**: Explicitly release handles (`release_handle`, `release_data_handle`) when done to free up memory.
 - **Reproducibility**: Always use `export_code` after a successful experiment to save your work.
+- **Persistence**: Use `save_model` after fitting if you need the estimator to survive server restarts.
 - **Data Hygiene**: Use `auto_format_on_load` for messy real-world data to avoid frequent validation errors.
 
 ---
@@ -158,9 +199,9 @@ Bring your own data into the MCP server.
 
 While `sktime-mcp` is a powerful tool for prototyping, please be aware of the current architectural limitations.
 
-#### 1. In-Memory "Amnesia" (No Persistence)
-The server stores state in standard Python dictionaries.
-> **Impact**: If the server restarts or connection drops, all loaded data and trained models are lost. There is no disk-backed checkpointing.
+#### 1. In-Memory Handles (Explicit Persistence Required)
+The server stores active handles in standard Python dictionaries.
+> **Impact**: If the server restarts or connection drops, in-memory handles are lost. Use `save_model` to persist fitted estimators to a local filesystem path when needed.
 
 #### 2. Synchronous Execution (GIL Blocking)
 Heavy operations (like `AutoARIMA` fitting) run on the main thread.
@@ -197,5 +238,6 @@ Complex sktime types (Periods, Intervals) are converted to strings for LLM consu
 |-------|----------|
 | **"Unknown estimator"** | Use `search_estimators` to find the exact case-sensitive name. |
 | **"Missing dependencies"** | Run `pip install -e ".[all]"` to ensure all extras are present. |
+| **`save_model` import/runtime errors** | Install MLflow in the environment used by the server. The tool relies on `sktime.utils.mlflow_sktime.save_model` and saves to a local filesystem path. |
 | **Validation Failures** | Enable `auto_format_on_load` or use `format_time_series` to clean your data. |
 | **Server Timeout** | Heavy models take time. Be patient or try a simpler model (e.g., `NaiveForecaster`) first. |

src/sktime_mcp/server.py

@@ -56,6 +56,7 @@
     delete_job_tool,
     cleanup_old_jobs_tool,
 )
+from sktime_mcp.tools.save_model import save_model_tool
 from sktime_mcp.composition.validator import get_composition_validator
 
 # Configure logging to stderr with detailed format
@@ -486,6 +487,28 @@ async def list_tools() -> List[Tool]:
                 },
             },
         ),
+        Tool(
+            name="save_model",
+            description="Save an estimator/pipeline handle using sktime MLflow integration",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "estimator_handle": {
+                        "type": "string",
+                        "description": "Handle ID of the estimator to save",
+                    },
+                    "path": {
+                        "type": "string",
+                        "description": "Local directory or URI where the model will be saved",
+                    },
+                    "mlflow_params": {
+                        "type": "object",
+                        "description": "Optional extra parameters for sktime.utils.mlflow_sktime.save_model",
+                    },
+                },
+                "required": ["estimator_handle", "path"],
+            },
+        ),
     ]
 
 
@@ -589,6 +612,12 @@ async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
             result = delete_job_tool(arguments["job_id"])
         elif name == "cleanup_old_jobs":
             result = cleanup_old_jobs_tool(arguments.get("max_age_hours", 24))
+        elif name == "save_model":
+            result = save_model_tool(
+                arguments["estimator_handle"],
+                arguments["path"],
+                arguments.get("mlflow_params"),
+            )
         else:
             result = {"error": f"Unknown tool: {name}"}
 

src/sktime_mcp/tools/__init__.py

@@ -5,6 +5,7 @@
 from sktime_mcp.tools.instantiate import instantiate_estimator_tool
 from sktime_mcp.tools.fit_predict import fit_predict_tool
 from sktime_mcp.tools.codegen import export_code_tool
+from sktime_mcp.tools.save_model import save_model_tool
 from sktime_mcp.tools.format_tools import (
     format_time_series_tool,
     auto_format_on_load_tool,
@@ -16,6 +17,7 @@
     "instantiate_estimator_tool",
     "fit_predict_tool",
     "export_code_tool",
+    "save_model_tool",
     "format_time_series_tool",
     "auto_format_on_load_tool",
 ]

src/sktime_mcp/tools/save_model.py

@@ -0,0 +1,60 @@
+"""
+save_model tool for sktime MCP.
+
+Saves estimator instances via sktime's MLflow integration.
+"""
+
+from typing import Any, Callable, Dict, Optional
+
+from sktime_mcp.runtime.handles import get_handle_manager
+
+
+def _get_mlflow_save_model() -> Callable[..., Any]:
+    """Resolve sktime MLflow save utility lazily for better runtime compatibility."""
+    try:
+        from sktime.utils.mlflow_sktime import save_model as mlflow_save_model
+    except Exception as exc:
+        raise ImportError(
+            "Unable to import sktime MLflow save_model utility. "
+            "Ensure sktime and MLflow dependencies are installed."
+        ) from exc
+    return mlflow_save_model
+
+
+def save_model_tool(
+    estimator_handle: str,
+    path: str,
+    mlflow_params: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Save an instantiated estimator to a local path or URI using sktime+MLflow.
+
+    Args:
+        estimator_handle: Handle ID from instantiate_estimator / instantiate_pipeline
+        path: Local directory or URI where the model should be saved
+        mlflow_params: Optional extra keyword arguments for sktime MLflow save_model
+
+    Returns:
+        Dictionary with success status and confirmation message/path.
+    """
+    handle_manager = get_handle_manager()
+
+    try:
+        estimator = handle_manager.get_instance(estimator_handle)
+    except KeyError:
+        return {"success": False, "error": f"Handle not found: {estimator_handle}"}
+
+    if mlflow_params is not None and not isinstance(mlflow_params, dict):
+        return {"success": False, "error": "mlflow_params must be a dictionary"}
+
+    try:
+        save_model = _get_mlflow_save_model()
+        save_model(sktime_model=estimator, path=path, **(mlflow_params or {}))
+        return {
+            "success": True,
+            "estimator_handle": estimator_handle,
+            "saved_path": path,
+            "message": f"Model saved successfully to '{path}'",
+        }
+    except Exception as exc:
+        return {"success": False, "error": str(exc)}

tests/test_core.py

@@ -153,6 +153,36 @@ def test_describe_unknown_estimator(self):
         assert not result["success"]
         assert "error" in result
 
+    def test_save_model_tool(self, monkeypatch, tmp_path):
+        """Test save_model tool resolves handle and forwards parameters."""
+        from sktime_mcp.runtime.handles import get_handle_manager
+        from sktime_mcp.tools.save_model import save_model_tool
+        import sktime_mcp.tools.save_model as save_model_module
+
+        calls = {}
+
+        def fake_save_model(**kwargs):
+            calls.update(kwargs)
+
+        monkeypatch.setattr(save_model_module, "_get_mlflow_save_model", lambda: fake_save_model)
+
+        handle_manager = get_handle_manager()
+        handle = handle_manager.create_handle("DummyEstimator", object())
+
+        try:
+            result = save_model_tool(
+                estimator_handle=handle,
+                path=str(tmp_path / "model_dir"),
+                mlflow_params={"serialization_format": "pickle"},
+            )
+        finally:
+            handle_manager.release_handle(handle)
+
+        assert result["success"]
+        assert result["saved_path"] == str(tmp_path / "model_dir")
+        assert calls["path"] == str(tmp_path / "model_dir")
+        assert calls["serialization_format"] == "pickle"
+
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])