mas-lab-component-controller 0.2.0__py3-none-any.whl

mas/lab/controller/client.py

@@ -0,0 +1,271 @@
+#  Copyright (c) 2026 Cisco Systems, Inc. and its affiliates
+#  SPDX-License-Identifier: Apache-2.0
+"""Unix-socket JSON-RPC client for the controller daemon."""
+from __future__ import annotations
+
+import atexit
+import json
+import logging
+import os
+import socket
+import subprocess
+import sys
+import time
+import uuid
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from mas.lab.controller import config as cfg
+from mas.lab.controller.config import DEFAULT_HTTP_PORT, ensure_mas_dirs
+
+logger = logging.getLogger(__name__)
+
+_process_session_id: str | None = None
+_process_session_registered = False
+
+
+def _release_process_session() -> None:
+    global _process_session_id
+    if not _process_session_id:
+        return
+    try:
+        ControllerClient().call(
+            "release_session",
+            {"session_id": _process_session_id},
+            timeout=5.0,
+        )
+    except Exception:
+        pass
+    _process_session_id = None
+
+
+def _register_process_session(client: ControllerClient) -> str:
+    """Acquire one daemon session per OS process (CLI or test subprocess)."""
+    global _process_session_id, _process_session_registered
+    if _process_session_id is None:
+        _process_session_id = f"cli-{os.getpid()}-{uuid.uuid4().hex[:8]}"
+    if not _process_session_registered:
+        atexit.register(_release_process_session)
+        _process_session_registered = True
+    client.call("acquire_session", {"session_id": _process_session_id}, timeout=5.0)
+    return _process_session_id
+
+
+class ControllerClient:
+    def __init__(self, socket_path: Path | None = None) -> None:
+        self.socket_path = Path(socket_path or cfg.socket_path())
+
+    def is_running(self) -> bool:
+        if not self.socket_path.exists():
+            return False
+        try:
+            self.call("ping")
+            return True
+        except OSError:
+            return False
+
+    def call(self, method: str, params: Optional[dict] = None, timeout: float = 30.0) -> Any:
+        payload = json.dumps({"method": method, "params": params or {}}).encode("utf-8")
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        sock.settimeout(timeout)
+        try:
+            sock.connect(str(self.socket_path))
+            sock.sendall(payload + b"\n")
+            chunks: list[bytes] = []
+            while True:
+                part = sock.recv(65536)
+                if not part:
+                    break
+                chunks.append(part)
+                if b"\n" in part:
+                    break
+            raw = b"".join(chunks).split(b"\n", 1)[0]
+            data = json.loads(raw.decode("utf-8"))
+            if "error" in data:
+                raise RuntimeError(data["error"])
+            return data.get("result")
+        finally:
+            sock.close()
+
+    def ensure_running(
+        self,
+        *,
+        port: int = DEFAULT_HTTP_PORT,
+        auto_start: bool = True,
+        restart_stale: bool = False,
+    ) -> None:
+        """Connect to the daemon, starting it if needed.
+
+        A daemon started from other ``mas.*`` sources or another environment
+        than this process is reported. With *restart_stale* an idle one is
+        restarted so new work runs with the caller's code and environment.
+        """
+        if self.is_running():
+            drift = self.drift()
+            if not drift:
+                _register_process_session(self)
+                return
+            if not restart_stale:
+                logger.warning(drift_message(drift))
+                _register_process_session(self)
+                return
+            self._stop_stale(drift)
+        if not auto_start:
+            raise RuntimeError(
+                f"Controller daemon is not running (socket {self.socket_path}). "
+                "Start it with: mas-lab control start"
+            )
+        start_daemon(port=port, detach=True)
+        for _ in range(40):
+            if self.is_running():
+                _register_process_session(self)
+                return
+            time.sleep(0.25)
+        raise RuntimeError("Controller daemon failed to start")
+
+    def drift(self) -> list[str]:
+        """Differences between the running daemon and this process (empty when in sync)."""
+        info = self.call("ping") or {}
+        issues: list[str] = []
+        if info.get("code") != cfg.code_fingerprint():
+            issues.append("it runs another mas-lab installation or older sources")
+        theirs = info.get("env")
+        if isinstance(theirs, dict):
+            ours = cfg.env_fingerprint()
+            changed = sorted(k for k in set(theirs) | set(ours) if theirs.get(k) != ours.get(k))
+            if changed:
+                issues.append("its environment differs for " + ", ".join(changed))
+        return issues
+
+    def _stop_stale(self, drift: list[str]) -> None:
+        running = (self.call("status") or {}).get("running", 0)
+        if running:
+            raise RuntimeError(
+                f"{drift_message(drift)} It has {running} running worker(s); wait for "
+                "them to finish, then run: mas-lab control stop"
+            )
+        logger.warning("Restarting controller daemon: %s.", "; ".join(drift))
+        pid_file = cfg.pid_path()
+        pid = pid_file.read_text(encoding="utf-8").strip() if pid_file.exists() else ""
+        stop_daemon()
+        # The old daemon unlinks the socket and pid file on exit; wait so it
+        # cannot remove the ones the new daemon creates.
+        if pid.isdigit():
+            for _ in range(40):
+                try:
+                    os.kill(int(pid), 0)
+                except ProcessLookupError:
+                    break
+                time.sleep(0.25)
+
+
+def drift_message(drift: list[str]) -> str:
+    return (
+        f"The controller daemon ({cfg.socket_path()}) does not match this shell: "
+        + "; ".join(drift)
+        + ". Work it runs may behave differently from what you expect. "
+        "Run `mas-lab control stop`; the next command starts a fresh daemon from this shell."
+    )
+
+
+def controller_session(*, port: int = DEFAULT_HTTP_PORT) -> ControllerClient:
+    """Context manager: ensure daemon + acquire session; release on exit."""
+    client = ControllerClient()
+    client.ensure_running(port=port, auto_start=True)
+    return client
+
+
+def start_daemon(*, port: int = DEFAULT_HTTP_PORT, detach: bool = True, foreground: bool = False) -> None:
+    ensure_mas_dirs()
+    cmd = [
+        sys.executable,
+        "-m",
+        "mas.lab.controller.daemon",
+        "--port",
+        str(port),
+    ]
+    if os.environ.get("MAS_CONTROLLER_NO_HTTP", "1").strip().lower() in ("1", "true", "yes"):
+        cmd.append("--no-http")
+    if foreground:
+        subprocess.run(cmd, check=False)
+        return
+    env = os.environ.copy()
+    env.setdefault("MAS_CONTROLLER_PORT", str(port))
+    log_path = Path(os.environ.get("MAS_CONTROLLER_LOG", cfg.controller_dir() / "daemon.log"))
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    log_file = open(log_path, "a", encoding="utf-8")
+    subprocess.Popen(
+        cmd,
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        start_new_session=detach,
+        env=env,
+    )
+
+
+def stop_daemon() -> bool:
+    client = ControllerClient()
+    if not client.is_running():
+        cfg.pid_path().unlink(missing_ok=True)
+        cfg.socket_path().unlink(missing_ok=True)
+        return True
+    try:
+        client.call("shutdown")
+    except Exception:
+        pass
+    cfg.pid_path().unlink(missing_ok=True)
+    cfg.socket_path().unlink(missing_ok=True)
+    return True
+
+
+def follow_worker(
+    worker_id: str,
+    *,
+    timeout: float = 86400.0,
+    poll: float = 0.5,
+    stream: bool = True,
+) -> Dict[str, Any]:
+    """Poll worker state until terminal; optionally stream stdout/stderr deltas to the terminal."""
+    import sys
+
+    client = ControllerClient()
+    client.ensure_running()
+    last_stdout = ""
+    last_stderr = ""
+    deadline = time.time() + timeout
+    terminal = {"completed", "failed", "cancelled"}
+
+    while time.time() < deadline:
+        detail = client.call("get_worker", {"worker_id": worker_id})
+        if detail is None:
+            raise RuntimeError(f"Worker {worker_id!r} not found")
+
+        if stream:
+            stdout = detail.get("stdout") or ""
+            stderr = detail.get("stderr") or ""
+            if len(stdout) > len(last_stdout):
+                sys.stdout.write(stdout[len(last_stdout) :])
+                sys.stdout.flush()
+                last_stdout = stdout
+            if len(stderr) > len(last_stderr):
+                sys.stderr.write(stderr[len(last_stderr) :])
+                sys.stderr.flush()
+                last_stderr = stderr
+
+        status = detail.get("status")
+        if status in terminal:
+            return detail
+        time.sleep(poll)
+
+    raise TimeoutError(f"Worker {worker_id!r} did not finish within {timeout}s")
+
+
+def wait_for_worker(
+    worker_id: str,
+    *,
+    timeout: float = 86400.0,
+    poll: float = 1.0,
+    stream: bool = False,
+) -> Dict[str, Any]:
+    """Wait for a worker to finish (optional streaming via :func:`follow_worker`)."""
+    return follow_worker(worker_id, timeout=timeout, poll=poll, stream=stream)

mas/lab/controller/config.py

@@ -0,0 +1,75 @@
+#  Copyright (c) 2026 Cisco Systems, Inc. and its affiliates
+#  SPDX-License-Identifier: Apache-2.0
+"""Controller daemon paths and defaults."""
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from pathlib import Path
+
+from mas.runtime.xdg import mas_data_root
+
+DEFAULT_HTTP_PORT = int(os.environ.get("MAS_CONTROLLER_PORT", "9000"))
+
+
+def mas_dir() -> Path:
+    return Path(os.environ.get("MAS_HOME", mas_data_root()))
+
+
+def controller_dir() -> Path:
+    return Path(os.environ.get("MAS_CONTROLLER_DIR", mas_dir() / "controller"))
+
+
+def socket_path() -> Path:
+    return Path(os.environ.get("MAS_CONTROLLER_SOCKET", mas_dir() / "controller.sock"))
+
+
+def pid_path() -> Path:
+    return Path(os.environ.get("MAS_CONTROLLER_PID", controller_dir() / "controller.pid"))
+
+
+def ensure_mas_dirs() -> None:
+    mas_dir().mkdir(parents=True, exist_ok=True)
+    controller_dir().mkdir(parents=True, exist_ok=True)
+
+def code_fingerprint() -> str:
+    """Identify the interpreter and installed ``mas.*`` sources.
+
+    The daemon keeps the modules it imported at startup, so a client compares
+    this value with the daemon's to detect an upgrade (wheel or editable).
+    """
+    import hashlib
+    import sys
+
+    import mas
+
+    digest = hashlib.sha256(sys.executable.encode("utf-8"))
+    for root in sorted(mas.__path__):
+        for dirpath, dirnames, filenames in os.walk(root):
+            dirnames[:] = sorted(d for d in dirnames if d != "__pycache__")
+            for name in sorted(filenames):
+                if not name.endswith(".py"):
+                    continue
+                path = os.path.join(dirpath, name)
+                st = os.stat(path)
+                digest.update(f"{path}:{st.st_size}:{st.st_mtime_ns}\n".encode("utf-8"))
+    return digest.hexdigest()[:16]
+
+
+_RUN_ENV_PREFIXES = ("MAS_", "OPENAI_", "AZURE_", "ANTHROPIC_", "LLM_", "LITELLM_", "OTEL_")
+_RUN_ENV_NAMES = frozenset({"SSL_CERT_FILE", "REQUESTS_CA_BUNDLE", "XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME"})
+# Set by the client/daemon machinery itself, not by the user.
+_DAEMON_ENV_PREFIXES = ("MAS_CONTROLLER_", "MAS_DOTENV_LOADED")
+
+
+def env_fingerprint(environ: Mapping[str, str] | None = None) -> dict[str, str]:
+    """Hash of each environment variable that can change a run; values never leave the process."""
+    import hashlib
+
+    env = os.environ if environ is None else environ
+    return {
+        name: hashlib.sha256(value.encode("utf-8")).hexdigest()[:12]
+        for name, value in env.items()
+        if (name.startswith(_RUN_ENV_PREFIXES) or name in _RUN_ENV_NAMES)
+        and not name.startswith(_DAEMON_ENV_PREFIXES)
+    }

mas/lab/controller/constants.py

@@ -0,0 +1,30 @@
+#  Copyright (c) 2026 Cisco Systems, Inc. and its affiliates
+#  SPDX-License-Identifier: Apache-2.0
+"""Shared constants for the FastAPI controller API."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from mas.runtime.xdg import mas_cache_root, mas_data_root
+
+LIBRARIES_DIR = Path(
+    os.getenv(
+        "MAS_LIBRARIES_DIR",
+        str(Path.home() / "mas-lab"),
+    )
+)
+
+MAX_TIMEOUT = 1800
+MAX_UPLOAD_BYTES = 10 * 1024 * 1024
+
+MAS_LAB_ROOT = Path(os.environ.get("MAS_LAB_ROOT", mas_data_root()))
+
+WEB_SEARCH_CACHE_DIR = mas_cache_root() / "web_search"
+
+HIDDEN_FILES = {".DS_Store", ".run_ref", ".gitkeep"}
+
+SCHEMAS_DIR = Path(__file__).parent / "schemas"
+PIPELINE_STEP_TYPES_PRE_PATH = SCHEMAS_DIR / "pipeline-step-types-pre.json"
+PIPELINE_STEP_TYPES_POST_PATH = SCHEMAS_DIR / "pipeline-step-types-post.json"

mas/lab/controller/daemon.py

@@ -0,0 +1,190 @@
+#  Copyright (c) 2026 Cisco Systems, Inc. and its affiliates
+#  SPDX-License-Identifier: Apache-2.0
+"""Controller daemon — Unix socket IPC + optional HTTP."""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+import signal
+import socket
+import threading
+from pathlib import Path
+from typing import Any
+
+from mas.lab.controller import config as cfg
+from mas.lab.controller.api import ControllerAPI, set_controller_api
+from mas.lab.controller.config import DEFAULT_HTTP_PORT, ensure_mas_dirs
+from mas.lab.controller.sessions import SessionRegistry
+
+logger = logging.getLogger(__name__)
+
+_shutdown = threading.Event()
+_api: ControllerAPI | None = None
+_sessions = SessionRegistry()
+_code_fingerprint: str | None = None
+_env_fingerprint: dict[str, str] | None = None
+
+
+def _serve_http(*, host: str, port: int, settings=None) -> None:
+    import uvicorn
+    from mas.lab.controller.fastapi_app import create_app
+
+    app = create_app(settings) if settings is not None else create_app()
+    uvicorn.run(app, host=host, port=port, log_level="info")
+
+
+def _handle_request(payload: dict) -> dict:
+    global _api
+    assert _api is not None
+    method = payload.get("method")
+    params = payload.get("params") or {}
+
+    dispatch = {
+        "ping": lambda _p: {"status": "ok", "code": _code_fingerprint, "env": _env_fingerprint},
+        "status": lambda _p: _api.status(),
+        "shutdown": lambda _p: _shutdown.set() or {"status": "stopping"},
+        "acquire_session": lambda p: _sessions.acquire(p["session_id"]) or {"ok": True},
+        "release_session": lambda p: _sessions.release(p["session_id"]) or {"ok": True},
+        "list_workers": lambda p: _api.list_workers(kind=p.get("kind"), status=p.get("status")),
+        "get_worker": lambda p: _api.get_worker(p["worker_id"]),
+        "cancel_worker": lambda p: {"ok": _api.cancel_worker(p["worker_id"])},
+        "submit_benchmark": lambda p: _api.submit_benchmark(p),
+        "submit_application": lambda p: _api.submit_application(p),
+        "submit_pipeline": lambda p: _api.submit_pipeline(p),
+        "list_libraries": lambda _p: _api.list_libraries(),
+        "list_runtime_runners": lambda _p: _api.list_runtime_runners(),
+    }
+    if method not in dispatch:
+        return {"error": f"unknown method: {method}"}
+    try:
+        return {"result": dispatch[method](params)}
+    except Exception as exc:
+        logger.exception("RPC %s failed", method)
+        return {"error": str(exc)}
+
+
+def _serve_socket(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.exists():
+        path.unlink()
+    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    server.bind(str(path))
+    server.listen(5)
+    logger.info("Controller IPC listening on %s", path)
+
+    while not _shutdown.is_set():
+        server.settimeout(1.0)
+        try:
+            conn, _addr = server.accept()
+        except TimeoutError:
+            continue
+        except OSError:
+            break
+        with conn:
+            data = b""
+            while b"\n" not in data:
+                chunk = conn.recv(65536)
+                if not chunk:
+                    break
+                data += chunk
+            if not data:
+                continue
+            line = data.split(b"\n", 1)[0]
+            try:
+                payload = json.loads(line.decode("utf-8"))
+            except json.JSONDecodeError:
+                try:
+                    conn.sendall(json.dumps({"error": "invalid json"}).encode("utf-8") + b"\n")
+                except OSError:
+                    pass
+                continue
+            response = _handle_request(payload)
+            conn.sendall(json.dumps(response).encode("utf-8") + b"\n")
+    server.close()
+    path.unlink(missing_ok=True)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="MAS Lab controller daemon")
+    parser.add_argument("--port", type=int, default=int(os.environ.get("MAS_CONTROLLER_PORT", DEFAULT_HTTP_PORT)))
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--socket", type=Path, default=None)
+    parser.add_argument("--no-http", action="store_true")
+    parser.add_argument("--cors-origin", action="append", default=[])
+    parser.add_argument("--token-file", type=Path, default=None)
+    parser.add_argument("--state-dir", type=Path, default=None)
+    parser.add_argument("--workspace", type=Path, default=None)
+    args = parser.parse_args(argv)
+    if args.workspace is not None:
+        os.environ["MAS_WORKSPACE_ROOT"] = str(Path(args.workspace).expanduser().resolve())
+    socket = args.socket or cfg.socket_path()
+
+    from mas.lab.controller.settings import ControllerSettings, load_token_file
+
+    token = load_token_file(args.token_file) if args.token_file else None
+    http_settings = ControllerSettings(
+        host=args.host,
+        port=args.port,
+        cors_origins=list(args.cors_origin or []),
+        token=token,
+        state_dir=args.state_dir or (cfg.controller_dir() / "state"),
+    )
+
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+    ensure_mas_dirs()
+
+    global _api, _code_fingerprint, _env_fingerprint
+    _code_fingerprint = cfg.code_fingerprint()
+    _env_fingerprint = cfg.env_fingerprint()
+    try:
+        from mas.lab.workspace import WorkspaceConfig
+
+        workspace = WorkspaceConfig.load()
+    except Exception:
+        workspace = None
+    _api = ControllerAPI(workspace)
+    set_controller_api(_api)
+
+    cfg.pid_path().write_text(str(os.getpid()), encoding="utf-8")
+
+    def _handle_signal(signum: int, _frame: Any) -> None:
+        logger.info("Received signal %s — shutting down", signum)
+        _shutdown.set()
+
+    signal.signal(signal.SIGTERM, _handle_signal)
+    signal.signal(signal.SIGINT, _handle_signal)
+
+    threads = [
+        threading.Thread(target=_serve_socket, args=(socket,), daemon=True, name="ipc"),
+    ]
+    if not args.no_http:
+        threads.append(
+            threading.Thread(
+                target=_serve_http,
+                kwargs={"host": args.host, "port": args.port, "settings": http_settings},
+                daemon=True,
+                name="http",
+            )
+        )
+    for thread in threads:
+        thread.start()
+
+    logger.info("MAS Lab controller running (HTTP :%s)", args.port)
+    try:
+        while not _shutdown.is_set():
+            if _sessions.should_shutdown(_api):
+                logger.info(
+                    "Idle shutdown — no client sessions and no active workers"
+                )
+                _shutdown.set()
+                break
+            _shutdown.wait(timeout=5.0)
+    finally:
+        cfg.pid_path().unlink(missing_ok=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())