mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev271__py3-none-any.whl

mplang/v1/runtime/server.py

@@ -1,501 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# 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.
-
-"""
-This module implements the HTTP server for the toy backend.
-It uses FastAPI to provide a RESTful API for managing computations.
-"""
-
-import base64
-import logging
-import re
-from typing import Any
-
-from fastapi import (
-    FastAPI,
-    HTTPException,
-    Request,
-)
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-
-from mplang.v1.core import IrReader, TableType, TensorType
-from mplang.v1.core.cluster import ClusterSpec
-from mplang.v1.kernels.base import KernelContext
-from mplang.v1.kernels.value import Value, decode_value, encode_value
-from mplang.v1.protos.v1alpha1 import mpir_pb2
-from mplang.v1.runtime.data_providers import (
-    DataProvider,
-    ResolvedURI,
-    register_provider,
-)
-from mplang.v1.runtime.exceptions import InvalidRequestError, ResourceNotFound
-from mplang.v1.runtime.session import (
-    Computation,
-    Session,
-    Symbol,
-    create_session_from_spec,
-)
-
-logger = logging.getLogger(__name__)
-
-app = FastAPI()
-
-# per-server global state
-_sessions: dict[str, Session] = {}
-_global_symbols: dict[str, Symbol] = {}
-
-
-def register_session(session: Session) -> Session:  # pragma: no cover - test helper
-    existing = _sessions.get(session.name)
-    if existing:
-        return existing
-    _sessions[session.name] = session
-    return session
-
-
-class _SymbolsProvider(DataProvider):
-    """Server-local symbols provider backed by BackendRuntime.state."""
-
-    @staticmethod
-    def _symbol_name(uri: ResolvedURI) -> str:
-        if uri.scheme != "symbols" or uri.parsed is None:
-            raise InvalidRequestError(
-                "symbols provider expects URI in the form symbols://{name}"
-            )
-
-        parsed = uri.parsed
-        if parsed.query or parsed.params or parsed.fragment:
-            raise InvalidRequestError(
-                "symbols:// URI must not contain query or fragment"
-            )
-
-        if parsed.netloc:
-            # e.g. symbols://foo -> name is carried in netloc (path may be empty or "/")
-            if parsed.path not in ("", "/"):
-                raise InvalidRequestError("symbols:// URIs cannot include subpaths")
-            name = parsed.netloc
-        else:
-            # e.g. symbols:///foo -> netloc empty, single path segment is the symbol name
-            path = parsed.path.lstrip("/")
-            if not path or "/" in path:
-                raise InvalidRequestError(
-                    "symbols:// URI must specify a single symbol name"
-                )
-            name = path
-
-        if not name:
-            raise InvalidRequestError("symbols:// URI missing symbol name")
-        return name
-
-    def read(
-        self,
-        uri: ResolvedURI,
-        out_spec: TensorType | TableType,
-        *,
-        ctx: KernelContext,
-    ) -> Any:  # type: ignore[override]
-        name = self._symbol_name(uri)
-        sym = _global_symbols.get(name)
-        if sym is None:
-            raise ResourceNotFound(f"Global symbol '{name}' not found")
-        return sym.data
-
-    def write(
-        self,
-        uri: ResolvedURI,
-        value: Any,
-        *,
-        ctx: KernelContext,
-    ) -> None:  # type: ignore[override]
-        name = self._symbol_name(uri)
-        if not isinstance(value, Value):
-            raise InvalidRequestError(
-                f"symbols:// write expects Value instance, got {type(value)}"
-            )
-        _global_symbols[name] = Symbol(name=name, mptype={}, data=value)
-
-
-# Register symbols provider explicitly for server runtime
-register_provider("symbols", _SymbolsProvider())
-
-
-@app.exception_handler(ResourceNotFound)
-def resource_not_found_handler(request: Request, exc: ResourceNotFound) -> JSONResponse:
-    """Handler for ResourceNotFound exceptions."""
-    logger.warning(f"Resource not found at {request.url}: {exc}")
-    return JSONResponse(
-        status_code=404,
-        content={"detail": str(exc)},
-    )
-
-
-@app.exception_handler(InvalidRequestError)
-def invalid_request_handler(request: Request, exc: InvalidRequestError) -> JSONResponse:
-    """Handler for InvalidRequestError exceptions."""
-    logger.warning(f"Invalid request at {request.url}: {exc}")
-    return JSONResponse(
-        status_code=400,
-        content={"detail": str(exc)},
-    )
-
-
-@app.exception_handler(Exception)
-def general_exception_handler(request: Request, exc: Exception) -> JSONResponse:
-    """Global exception handler for better error reporting."""
-    logger.error(f"Unhandled exception at {request.url}: {exc}", exc_info=True)
-    return JSONResponse(
-        status_code=500,
-        content={
-            "detail": f"Internal server error: {exc!s}",
-            "error_type": type(exc).__name__,
-            "path": str(request.url.path),
-        },
-    )
-
-
-def validate_name(name: str, name_type: str) -> None:
-    """Validate that a name is safe for use in URL paths.
-
-    Args:
-        name: The name to validate
-        name_type: Type of name (for error messages, e.g., "session", "computation")
-
-    Raises:
-        HTTPException: If the name contains invalid characters
-    """
-    if not name:
-        raise HTTPException(status_code=400, detail=f"{name_type} name cannot be empty")
-
-    # Only allow alphanumeric, hyphens, underscores, and dots
-    if not re.match(r"^[a-zA-Z0-9._-]+$", name):
-        raise HTTPException(
-            status_code=400,
-            detail=f"{name_type} name can only contain letters, numbers, dots, hyphens, and underscores",
-        )
-
-
-# Request/Response Models
-class CreateSessionRequest(BaseModel):
-    rank: int
-    cluster_spec: dict
-
-
-class SessionResponse(BaseModel):
-    name: str
-
-
-class CreateComputationRequest(BaseModel):
-    mpprogram: str  # Base64 encoded MPProgram proto
-    input_names: list[str]  # Mandatory input symbol names
-    output_names: list[str]  # Mandatory output symbol names
-
-
-class ComputationResponse(BaseModel):
-    name: str
-
-
-class CreateSymbolRequest(BaseModel):
-    mptype: dict
-    data: str  # Base64 encoded Value data
-
-
-class SymbolResponse(BaseModel):
-    name: str
-    mptype: dict
-    data: str  # Base64 encoded Value data
-
-
-class CommSendRequest(BaseModel):
-    data: str  # Base64 encoded binary data
-    is_raw_bytes: bool = False  # True for SPU channel raw bytes
-
-
-# Response Models for enhanced status
-class SessionListResponse(BaseModel):
-    sessions: list[str]
-
-
-class ComputationListResponse(BaseModel):
-    computations: list[str]
-
-
-class GlobalSymbolResponse(BaseModel):
-    name: str
-    mptype: dict
-    data: str  # Base64 encoded Value data
-
-
-@app.get("/health")
-async def health_check() -> dict[str, str]:
-    """Health check endpoint."""
-    return {"status": "ok"}
-
-
-# List all sessions
-@app.get("/sessions", response_model=SessionListResponse)
-def list_sessions() -> SessionListResponse:
-    """List all session names."""
-    return SessionListResponse(sessions=list(_sessions.keys()))
-
-
-# List all computations in a session
-@app.get(
-    "/sessions/{session_name}/computations", response_model=ComputationListResponse
-)
-def list_session_computations(session_name: str) -> ComputationListResponse:
-    """List all computation names in a session."""
-    sess = _sessions.get(session_name)
-    if not sess:
-        raise ResourceNotFound(f"Session '{session_name}' not found")
-    return ComputationListResponse(computations=sess.list_computations())
-
-
-# Session endpoints
-@app.put("/sessions/{session_name}", response_model=SessionResponse)
-def create_session(session_name: str, request: CreateSessionRequest) -> SessionResponse:
-    validate_name(session_name, "session")
-    # Delegate cluster spec parsing & session construction to resource layer
-
-    if session_name in _sessions:
-        sess = _sessions[session_name]
-    else:
-        spec = ClusterSpec.from_dict(request.cluster_spec)
-        sess = create_session_from_spec(name=session_name, rank=request.rank, spec=spec)
-        _sessions[session_name] = sess
-    return SessionResponse(name=sess.name)
-
-
-@app.get("/sessions/{session_name}", response_model=SessionResponse)
-def get_session(session_name: str) -> SessionResponse:
-    sess = _sessions.get(session_name)
-    if not sess:
-        raise ResourceNotFound(f"Session '{session_name}' not found")
-    return SessionResponse(name=sess.name)
-
-
-@app.delete("/sessions/{session_name}")
-def delete_session(session_name: str) -> dict[str, str]:
-    """Delete a session and all its associated resources."""
-    if session_name in _sessions:
-        del _sessions[session_name]
-        logging.info(f"Session {session_name} deleted successfully")
-        return {"message": f"Session '{session_name}' deleted successfully"}
-    else:
-        raise ResourceNotFound(f"Session '{session_name}' not found")
-
-
-# Computation endpoints
-@app.put(
-    "/sessions/{session_name}/computations/{computation_id}",
-    response_model=ComputationResponse,
-)
-def create_and_execute_computation(
-    session_name: str, computation_id: str, request: CreateComputationRequest
-) -> ComputationResponse:
-    graph_proto = mpir_pb2.GraphProto()
-    try:
-        graph_proto.ParseFromString(base64.b64decode(request.mpprogram))
-    except Exception as e:
-        raise InvalidRequestError(
-            f"Invalid base64 or protobuf for mpprogram: {e!s}"
-        ) from e
-
-    reader = IrReader()
-    expr = reader.loads(graph_proto)
-
-    if expr is None:
-        raise InvalidRequestError("Failed to parse expression from protobuf")
-
-    # Create the computation resource
-    sess = _sessions.get(session_name)
-    if not sess:
-        raise ResourceNotFound(f"Session '{session_name}' not found.")
-    comp = sess.get_computation(computation_id)
-    if not comp:
-        comp = Computation(name=computation_id, expr=expr)
-        sess.add_computation(comp)
-    sess.execute(comp, request.input_names, request.output_names)
-    return ComputationResponse(name=computation_id)
-
-
-@app.delete("/sessions/{session_name}/computations/{computation_id}")
-def delete_computation(session_name: str, computation_id: str) -> dict[str, str]:
-    """Delete a specific computation."""
-    sess = _sessions.get(session_name)
-    if sess and sess.delete_computation(computation_id):
-        logging.info(
-            f"Computation {computation_id} deleted from session {session_name}"
-        )
-        return {"message": f"Computation '{computation_id}' deleted successfully"}
-    else:
-        raise ResourceNotFound(
-            f"Computation '{computation_id}' not found in session '{session_name}'"
-        )
-
-
-# Symbol endpoints
-@app.put(
-    "/sessions/{session_name}/symbols/{symbol_name}", response_model=SymbolResponse
-)
-def create_session_symbol(
-    session_name: str, symbol_name: str, request: CreateSymbolRequest
-) -> SymbolResponse:
-    """Create a symbol in a session."""
-    sess = _sessions.get(session_name)
-    if not sess:
-        raise ResourceNotFound(f"Session '{session_name}' not found.")
-    try:
-        obj = decode_value(base64.b64decode(request.data))
-    except Exception as e:
-        raise InvalidRequestError(f"Invalid symbol data: {e!s}") from e
-    symbol = Symbol(name=symbol_name, mptype=request.mptype, data=obj)
-    sess.add_symbol(symbol)
-    # Return the base64 data back to client; server stores Python object
-    return SymbolResponse(
-        name=symbol.name,
-        mptype=symbol.mptype,
-        data=base64.b64encode(encode_value(symbol.data)).decode("utf-8"),
-    )
-
-
-@app.get(
-    "/sessions/{session_name}/symbols/{symbol_name}", response_model=SymbolResponse
-)
-def get_session_symbol(session_name: str, symbol_name: str) -> SymbolResponse:
-    """Get a symbol from a session."""
-    try:
-        logger.debug(
-            f"Looking for symbol: '{symbol_name}' in session: '{session_name}'"
-        )
-        sess = _sessions.get(session_name)
-        symbol = sess.get_symbol(symbol_name) if sess else None
-        if not symbol:
-            raise HTTPException(
-                status_code=404, detail=f"Symbol {symbol_name} not found"
-            )
-
-        # symbol data is None means this party does not participate the computation
-        # that produced the symbol.
-        if symbol.data is None:
-            raise ResourceNotFound(f"Symbol '{symbol_name}' has no data on this party")
-
-        # Serialize using Value envelope
-        return SymbolResponse(
-            name=symbol.name,
-            mptype=symbol.mptype,
-            data=base64.b64encode(encode_value(symbol.data)).decode("utf-8"),
-        )
-    except ValueError as e:
-        raise HTTPException(status_code=404, detail=str(e)) from e
-
-
-@app.get("/sessions/{session_name}/symbols")
-def list_session_symbols(session_name: str) -> dict[str, list[str]]:
-    """List all symbols in a session."""
-    sess = _sessions.get(session_name)
-    if not sess:
-        raise ResourceNotFound(f"Session '{session_name}' not found.")
-    symbols = sess.list_symbols()
-    return {"symbols": symbols}
-
-
-@app.delete("/sessions/{session_name}/symbols/{symbol_name}")
-def delete_symbol(session_name: str, symbol_name: str) -> dict[str, str]:
-    """Delete a specific symbol."""
-    sess = _sessions.get(session_name)
-    if sess and sess.delete_symbol(symbol_name):
-        logging.info(f"Symbol {symbol_name} deleted from session {session_name}")
-        return {"message": f"Symbol '{symbol_name}' deleted successfully"}
-    else:
-        raise ResourceNotFound(
-            f"Symbol '{symbol_name}' not found in session '{session_name}'"
-        )
-
-
-# Global Symbols endpoints
-@app.put("/api/v1/symbols/{symbol_name}", response_model=GlobalSymbolResponse)
-def create_global_symbol(
-    symbol_name: str, request: CreateSymbolRequest
-) -> GlobalSymbolResponse:
-    validate_name(symbol_name, "symbol")
-    try:
-        obj = decode_value(base64.b64decode(request.data))
-    except Exception as e:
-        raise InvalidRequestError(f"Invalid global symbol data: {e!s}") from e
-    sym = Symbol(name=symbol_name, mptype=request.mptype, data=obj)
-    _global_symbols[symbol_name] = sym
-    return GlobalSymbolResponse(
-        name=sym.name,
-        mptype=sym.mptype,
-        data=base64.b64encode(encode_value(sym.data)).decode("utf-8"),
-    )
-
-
-@app.get("/api/v1/symbols/{symbol_name}", response_model=GlobalSymbolResponse)
-def get_global_symbol(symbol_name: str) -> GlobalSymbolResponse:  # route handler
-    sym = _global_symbols.get(symbol_name)
-    if not sym:
-        raise ResourceNotFound(f"Global symbol '{symbol_name}' not found")
-    # Serialize using Value envelope
-    return GlobalSymbolResponse(
-        name=sym.name,
-        mptype=sym.mptype,
-        data=base64.b64encode(encode_value(sym.data)).decode("utf-8"),
-    )
-
-
-@app.get("/api/v1/symbols")
-def list_global_symbols() -> dict[str, list[str]]:
-    return {"symbols": list(_global_symbols.keys())}
-
-
-@app.delete("/api/v1/symbols/{symbol_name}")
-def delete_global_symbol(symbol_name: str) -> dict[str, str]:  # route handler
-    if symbol_name in _global_symbols:
-        del _global_symbols[symbol_name]
-        return {"message": f"Global symbol '{symbol_name}' deleted successfully"}
-    else:
-        raise ResourceNotFound(f"Global symbol '{symbol_name}' not found")
-
-
-# Communication endpoints
-# TODO(jint) this should be computation level, add multi computation parallel support.
-@app.put("/sessions/{session_name}/comm/{key}/from/{from_rank}")
-def comm_send(
-    session_name: str, key: str, from_rank: int, request: CommSendRequest
-) -> dict[str, str]:
-    """
-    Receive a message from another party and deliver it to the session's communicator.
-    This endpoint runs on the receiver's server.
-    """
-    sess = _sessions.get(session_name)
-    if not sess or not sess.communicator:
-        logger.error(f"Session or communicator not found: session={session_name}")
-        raise HTTPException(status_code=404, detail="Session or communicator not found")
-
-    # The receiver rank should be the rank of the server hosting this endpoint
-    # We don't need to validate to_rank since the request is coming to this server
-
-    # For raw bytes (SPU channel), pass through as dict with flag
-    # For normal data, pass the base64 string directly
-    data_payload: str | dict[str, object]
-    if request.is_raw_bytes:
-        data_payload = {"data": request.data, "is_raw_bytes": True}
-    else:
-        data_payload = request.data
-
-    # Use the proper onSent mechanism from CommunicatorBase
-    sess.communicator.onSent(from_rank, key, data_payload)
-    return {"status": "ok"}