Spaces:
Sleeping
Sleeping
File size: 4,191 Bytes
c509967 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import argparse
import logging
import os
import sys
import uvicorn
from mcp.server.fastmcp import FastMCP
def run_mcp_server(
mcp: FastMCP, default_port: int = 8000, default_host: str = "127.0.0.1"
):
"""Standardize the startup process for ChemGraph MCP servers.
Supports 'stdio' (default) and 'sse' (HTTP) transports.
Ensures logging is correctly routed to stderr to avoid corrupting stdio transport.
Parameters
----------
mcp : FastMCP
MCP server instance to run.
default_port : int, optional
Default port for streamable HTTP transport.
default_host : str, optional
Default host for streamable HTTP transport.
"""
parser = argparse.ArgumentParser(description=f"Run {mcp.name} MCP Server")
parser.add_argument(
"--transport",
choices=["stdio", "streamable_http"],
default="stdio",
help="Transport protocol to use: 'stdio' (default) or 'streamable_http' (HTTP/SSE).",
)
parser.add_argument(
"--port",
type=int,
default=default_port,
help=f"Port for streamable_http transport (default: {default_port})",
)
parser.add_argument(
"--host",
type=str,
default=default_host,
help=f"Host for streamable_http transport (default: {default_host})",
)
args = parser.parse_args()
# Configure logging to write to stderr.
# This is CRITICAL for stdio transport mode, as stdout is used for communication.
logger = logging.getLogger()
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
"%(asctime)s - %(levelname)s - %(name)s - %(message)s"
)
# Always log to stderr
stderr_handler = logging.StreamHandler(sys.stderr)
stderr_handler.setFormatter(formatter)
logger.addHandler(stderr_handler)
# If CHEMGRAPH_LOG_DIR is set, also log to a file
log_dir = os.environ.get("CHEMGRAPH_LOG_DIR")
if log_dir:
try:
os.makedirs(log_dir, exist_ok=True)
log_file = os.path.join(log_dir, f"mcp_server_{mcp.name}.log")
file_handler = logging.FileHandler(log_file)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
logging.info(f"Logging to file: {log_file}")
except Exception as e:
logging.error(f"Failed to setup file logging to {log_dir}: {e}")
if args.transport == "streamable_http":
logging.info(
"Starting %s on %s:%s via streamable_http transport...",
mcp.name,
args.host,
args.port,
)
# FastMCP.streamable_http_app() returns a Starlette/FastAPI-compatible app
app = mcp.streamable_http_app()
uvicorn.run(app, host=args.host, port=args.port)
else:
logging.info("Starting %s via stdio transport...", mcp.name)
# FastMCP.run(transport='stdio') handles the stdio loop
mcp.run(transport="stdio")
# ---------------------------------------------------------------------------
# Parsl HPC configuration loader
# ---------------------------------------------------------------------------
def load_parsl_config(system_name: str):
"""Dynamically import and return a Parsl config for the given HPC system.
Parameters
----------
system_name : str
Target system name. Supported: ``polaris``, ``aurora``.
Returns
-------
parsl.config.Config
Parsl configuration object for the requested system.
Raises
------
ValueError
If the system name is not recognised.
"""
system_name = system_name.lower()
run_dir = os.getcwd()
logging.info("Initializing Parsl for system: %s", system_name)
if system_name == "polaris":
from chemgraph.hpc_configs.polaris_parsl import get_polaris_config
return get_polaris_config(run_dir=run_dir)
elif system_name == "aurora":
from chemgraph.hpc_configs.aurora_parsl import get_aurora_config
return get_aurora_config(run_dir=run_dir)
else:
raise ValueError(
f"Unknown system specified: '{system_name}'. Supported: polaris, aurora"
)
|