mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-05-07 18:18:42 -05:00
1eb30f5f86
Align the MLSys·im code, docs, paper, website, workflows, and lab wheel for the 0.1.1 release. This also fixes runtime/API issues found during release review and prepares the paper PDF plus archive package.
106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
# mlsysim/examples/mcp_server.py
|
|
"""
|
|
MLSys·im Model Context Protocol (MCP) Server
|
|
============================================
|
|
This script runs a fast, native MCP server that exposes the MLSys·im engine
|
|
to any compatible AI Agent (e.g., Claude Desktop, Cursor, or custom agents).
|
|
|
|
It exposes two primary tools:
|
|
1. `get_schemas`: Returns the exact JSON schema required to formulate an MLSysPlan.
|
|
2. `evaluate_cluster_yaml`: Takes a YAML string, evaluates it, and returns the strict JSON physics profile.
|
|
|
|
To use with Claude Desktop, add to your config:
|
|
```json
|
|
"mcpServers": {
|
|
"mlsysim": {
|
|
"command": "python3",
|
|
"args": ["/path/to/MLSysBook/mlsysim/examples/mcp_server.py"]
|
|
}
|
|
}
|
|
```
|
|
"""
|
|
|
|
import sys
|
|
import json
|
|
import yaml
|
|
from pathlib import Path
|
|
from pydantic import ValidationError
|
|
|
|
try:
|
|
from mcp.server import FastMCP
|
|
except ImportError:
|
|
print("Error: The 'mcp' library is required to run the MCP server.", file=sys.stderr)
|
|
print("Install it with: pip install mcp", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
from mlsysim.cli.schemas import MlsysPlanSchema
|
|
from mlsysim.core.evaluation import SystemEvaluator
|
|
|
|
# Initialize the MCP Server
|
|
mcp = FastMCP("MLSys·im")
|
|
|
|
@mcp.tool()
|
|
def get_schemas() -> str:
|
|
"""
|
|
Returns the strict JSON schemas required to interact with the MLSys·im engine.
|
|
Always read these schemas before attempting to generate a cluster YAML.
|
|
"""
|
|
return json.dumps({
|
|
"MlsysPlanSchema": MlsysPlanSchema.model_json_schema()
|
|
}, indent=2)
|
|
|
|
@mcp.tool()
|
|
def evaluate_cluster_yaml(yaml_config_string: str) -> str:
|
|
"""
|
|
Evaluates an MLSys·im cluster configuration against the Iron Law of ML Performance.
|
|
|
|
Args:
|
|
yaml_config_string: A complete YAML string adhering to MlsysPlanSchema.
|
|
|
|
Returns:
|
|
A JSON string containing the physics-based evaluation (feasibility, performance, economics).
|
|
If the configuration violates physical laws (e.g., OOM), the JSON will detail the failure.
|
|
"""
|
|
try:
|
|
raw_data = yaml.safe_load(yaml_config_string)
|
|
if not raw_data:
|
|
return json.dumps({"error": "Empty YAML provided."})
|
|
|
|
# 1. Gate: Schema Validation
|
|
try:
|
|
schema = MlsysPlanSchema(**raw_data)
|
|
except ValidationError as e:
|
|
return json.dumps({
|
|
"status": "VALIDATION_FAILED",
|
|
"message": "The provided YAML did not match the required schema.",
|
|
"errors": e.errors()
|
|
})
|
|
|
|
# 2. Gate: Physics Evaluation
|
|
eval_obj = SystemEvaluator.evaluate(
|
|
scenario_name=schema.name,
|
|
model_obj=schema.model_obj,
|
|
hardware_obj=schema.hardware_obj,
|
|
batch_size=schema.workload.batch_size,
|
|
precision=schema.hardware.precision,
|
|
efficiency=schema.hardware.efficiency,
|
|
fleet_obj=schema.fleet_obj,
|
|
nodes=schema.hardware.nodes,
|
|
duration_days=schema.ops.duration_days if schema.ops else None
|
|
)
|
|
|
|
# 3. Output standard JSON payload
|
|
return json.dumps(eval_obj.to_dict(), indent=2)
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
return json.dumps({
|
|
"status": "CRITICAL_ERROR",
|
|
"error_type": type(e).__name__,
|
|
"message": str(e),
|
|
"traceback": traceback.format_exc()
|
|
})
|
|
|
|
if __name__ == "__main__":
|
|
mcp.run()
|