uxarray-mcp 0.1.0__py3-none-any.whl

uxarray_mcp/remote/config.py

@@ -0,0 +1,322 @@
+"""Configuration management for remote execution."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+from uuid import UUID
+
+import yaml
+
+USER_CONFIG_PATH = Path.home() / ".config" / "uxarray-mcp" / "config.yaml"
+
+
+def discover_config_path() -> Path | None:
+    """Return the first existing config file in the discovery order.
+
+    Order:
+      1. ``$UXARRAY_MCP_CONFIG`` (explicit override)
+      2. ``./config.yaml`` (current working directory — project-local)
+      3. ``~/.config/uxarray-mcp/config.yaml`` (user install)
+      4. ``<repo_root>/config.yaml`` (editable install fallback)
+
+    Project-local (cwd) wins over the user config so that running the CLI
+    from inside a checkout uses the repo's config, even when an empty user
+    config was previously written by ``uxarray-mcp setup``.
+
+    Returns ``None`` when no config file is found.
+    """
+    env_path = os.environ.get("UXARRAY_MCP_CONFIG")
+    if env_path:
+        candidate = Path(env_path).expanduser()
+        if candidate.exists():
+            return candidate
+
+    cwd_config = Path.cwd() / "config.yaml"
+    if cwd_config.exists():
+        return cwd_config
+
+    if USER_CONFIG_PATH.exists():
+        return USER_CONFIG_PATH
+
+    repo_config = Path(__file__).resolve().parent.parent.parent.parent / "config.yaml"
+    if repo_config.exists():
+        return repo_config
+
+    return None
+
+
+def discover_config_search_paths() -> list[Path]:
+    """Return the ordered list of paths discover_config_path inspects.
+
+    Useful for diagnostics — ``endpoints list`` prints this when no
+    endpoints are configured so the user can see exactly which file was
+    used or which paths were searched.
+    """
+    paths: list[Path] = []
+    env_path = os.environ.get("UXARRAY_MCP_CONFIG")
+    if env_path:
+        paths.append(Path(env_path).expanduser())
+    paths.append(Path.cwd() / "config.yaml")
+    paths.append(USER_CONFIG_PATH)
+    paths.append(Path(__file__).resolve().parent.parent.parent.parent / "config.yaml")
+    return paths
+
+
+_VALID_EXECUTION_MODES = {"local", "hpc", "auto"}
+_EXECUTION_MODE_ALIASES = {"remote": "hpc"}
+
+
+def _is_uuid(value: str) -> bool:
+    try:
+        UUID(value)
+    except ValueError:
+        return False
+    return True
+
+
+@dataclass(frozen=True)
+class EndpointProfile:
+    """Named Globus Compute endpoint profile."""
+
+    name: str
+    endpoint_id: str
+    path_prefixes: tuple[str, ...] = ()
+    timeout_seconds: int | None = None
+
+
+def normalize_execution_mode(execution_mode: str) -> str:
+    """Return the canonical execution mode name.
+
+    The repository previously used ``remote`` for the HPC-only mode. Accept it
+    as a backwards-compatible alias so older configs and tests keep working.
+    """
+    canonical_mode = _EXECUTION_MODE_ALIASES.get(execution_mode, execution_mode)
+    if canonical_mode not in _VALID_EXECUTION_MODES:
+        raise ValueError(
+            f"Invalid execution mode {execution_mode!r}. "
+            f"Must be one of: {', '.join(sorted(_VALID_EXECUTION_MODES))}"
+        )
+    return canonical_mode
+
+
+class HPCConfig:
+    """HPC execution configuration.
+
+    Parameters
+    ----------
+    endpoint_id : str | None
+        Globus Compute endpoint UUID
+    execution_mode : str
+        Execution mode: "local", "hpc", or "auto"
+    timeout_seconds : int
+        Timeout for remote execution in seconds
+
+    Examples
+    --------
+    >>> config = HPCConfig(endpoint_id="abc-123", execution_mode="hpc")
+    >>> config.has_endpoint
+    True
+    """
+
+    def __init__(
+        self,
+        endpoint_id: Optional[str] = None,
+        execution_mode: str = "local",
+        timeout_seconds: int = 300,
+        endpoints: dict[str, EndpointProfile] | None = None,
+        default_endpoint: str | None = None,
+        endpoint_name: str | None = None,
+    ):
+        self.endpoints = endpoints or {}
+        self.default_endpoint = default_endpoint
+        self.endpoint_name = endpoint_name
+        self.endpoint_id = endpoint_id
+        self.execution_mode = normalize_execution_mode(execution_mode)
+        self.timeout_seconds = timeout_seconds
+
+    @property
+    def has_endpoint(self) -> bool:
+        """Check if Globus Compute endpoint is configured."""
+        return self.endpoint_id is not None or bool(self.endpoints)
+
+    @property
+    def should_use_remote(self) -> bool:
+        """Determine if remote execution should be used."""
+        if self.execution_mode == "local":
+            return False
+        elif self.execution_mode == "hpc":
+            return self.has_endpoint
+        elif self.execution_mode == "auto":
+            return self.has_endpoint
+        else:
+            return False
+
+    @property
+    def endpoint_names(self) -> list[str]:
+        """Return configured endpoint profile names."""
+        return sorted(self.endpoints)
+
+    def resolve_endpoint(
+        self, endpoint: str | None = None, path: str | None = None
+    ) -> EndpointProfile | None:
+        """Resolve an explicit endpoint name, default endpoint, or raw UUID."""
+        if endpoint:
+            if endpoint in self.endpoints:
+                return self.endpoints[endpoint]
+            if endpoint == self.endpoint_id:
+                return EndpointProfile(
+                    name=self.endpoint_name or endpoint,
+                    endpoint_id=endpoint,
+                    timeout_seconds=self.timeout_seconds,
+                )
+            if _is_uuid(endpoint):
+                return EndpointProfile(
+                    name=endpoint,
+                    endpoint_id=endpoint,
+                    timeout_seconds=self.timeout_seconds,
+                )
+            configured = ", ".join(self.endpoint_names) or "none"
+            raise ValueError(
+                f"Unknown endpoint {endpoint!r}. "
+                f"Configured endpoint names: {configured}. "
+                "Pass a configured endpoint name or a Globus Compute endpoint UUID."
+            )
+
+        if self.default_endpoint and self.default_endpoint in self.endpoints:
+            return self.endpoints[self.default_endpoint]
+
+        if self.endpoint_id is not None:
+            return EndpointProfile(
+                name=self.endpoint_name or "default",
+                endpoint_id=self.endpoint_id,
+                timeout_seconds=self.timeout_seconds,
+            )
+
+        if len(self.endpoints) == 1:
+            return next(iter(self.endpoints.values()))
+
+        return None
+
+    def for_endpoint(
+        self, endpoint: str | None = None, path: str | None = None
+    ) -> "HPCConfig":
+        """Return a copy configured for the selected endpoint profile."""
+        profile = self.resolve_endpoint(endpoint=endpoint, path=path)
+        if profile is None:
+            return HPCConfig(
+                endpoint_id=None,
+                execution_mode=self.execution_mode,
+                timeout_seconds=self.timeout_seconds,
+                endpoints=self.endpoints,
+                default_endpoint=self.default_endpoint,
+            )
+
+        return HPCConfig(
+            endpoint_id=profile.endpoint_id,
+            execution_mode=self.execution_mode,
+            timeout_seconds=profile.timeout_seconds or self.timeout_seconds,
+            endpoints=self.endpoints,
+            default_endpoint=self.default_endpoint,
+            endpoint_name=profile.name,
+        )
+
+
+def _coerce_prefixes(value: Any) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, str):
+        return (value,)
+    if isinstance(value, list):
+        return tuple(str(item) for item in value if item)
+    return ()
+
+
+def _parse_endpoint_profiles(raw_endpoints: Any) -> dict[str, EndpointProfile]:
+    if not isinstance(raw_endpoints, dict):
+        return {}
+
+    profiles: dict[str, EndpointProfile] = {}
+    for name, raw_profile in raw_endpoints.items():
+        if not isinstance(raw_profile, dict):
+            continue
+        endpoint_id = raw_profile.get("endpoint_id")
+        if not endpoint_id:
+            continue
+        timeout = raw_profile.get("timeout_seconds")
+        profiles[str(name)] = EndpointProfile(
+            name=str(name),
+            endpoint_id=str(endpoint_id),
+            path_prefixes=_coerce_prefixes(raw_profile.get("path_prefixes")),
+            timeout_seconds=int(timeout) if timeout is not None else None,
+        )
+    return profiles
+
+
+def load_config(config_path: Optional[Path] = None) -> HPCConfig:
+    """Load HPC configuration from YAML file.
+
+    Parameters
+    ----------
+    config_path : Path | None
+        Path to config.yaml. If None, uses default location.
+
+    Returns
+    -------
+    HPCConfig
+        Loaded configuration object
+
+    Examples
+    --------
+    >>> config = load_config()
+    >>> config.execution_mode
+    'local'
+    """
+    if config_path is None:
+        config_path = discover_config_path()
+
+    if config_path is None or not config_path.exists():
+        return HPCConfig()
+
+    with open(config_path, "r", encoding="utf-8") as f:
+        data = yaml.safe_load(f) or {}
+
+    if not isinstance(data, dict):
+        return HPCConfig()
+
+    hpc_config = data.get("hpc", {})
+    if not isinstance(hpc_config, dict):
+        hpc_config = {}
+
+    globus_config = hpc_config.get("globus_compute", {})
+    if not isinstance(globus_config, dict):
+        globus_config = {}
+
+    endpoints = _parse_endpoint_profiles(hpc_config.get("endpoints"))
+    if not endpoints:
+        endpoints = _parse_endpoint_profiles(globus_config.get("endpoints"))
+
+    default_endpoint = hpc_config.get("default_endpoint") or globus_config.get(
+        "default_endpoint"
+    )
+    endpoint_id = globus_config.get("endpoint_id")
+
+    if endpoint_id is None and default_endpoint in endpoints:
+        endpoint_id = endpoints[default_endpoint].endpoint_id
+    endpoint_name = (
+        default_endpoint
+        if default_endpoint in endpoints
+        and endpoint_id == endpoints[default_endpoint].endpoint_id
+        else None
+    )
+
+    return HPCConfig(
+        endpoint_id=endpoint_id,
+        execution_mode=hpc_config.get("execution_mode", "local"),
+        timeout_seconds=hpc_config.get("timeout_seconds", 300),
+        endpoints=endpoints,
+        default_endpoint=default_endpoint,
+        endpoint_name=endpoint_name,
+    )

uxarray_mcp/remote/health.py

@@ -0,0 +1,372 @@
+"""Endpoint manager status checking for HPC execution.
+
+## Status vocabulary
+
+Every function in this module returns one of four manager-level statuses:
+
+``"registered"``
+    The endpoint manager process is running and visible to Globus Compute.
+    Slurm/PBS jobs will be submitted when tasks arrive. Workers are not
+    necessarily running yet — the scheduler allocates them on demand.
+    This was formerly reported as ``"online"`` by the Globus SDK.
+
+``"active"``
+    Manager is registered AND a lightweight probe task ran successfully on a
+    real compute node. Use :func:`probe_endpoint_worker` to obtain this status.
+
+``"offline"``
+    The endpoint manager is not running. Someone must SSH in and run
+    ``globus-compute-endpoint start <name>``.
+
+``"unreachable"``
+    The Globus Compute service cannot be contacted (auth error, network issue).
+
+``"no_endpoint"``
+    No endpoint UUID is configured for this name.
+
+## Why "registered" not "online"
+
+Globus reports ``"online"`` when the manager process is registered. This is
+purely a registration state — it says nothing about whether Slurm has idle
+nodes, whether the worker environment is healthy, or whether submitted tasks
+will actually run. Using ``"online"`` led to misleading status messages.
+``"registered"`` is honest: the manager is up and accepting submissions.
+
+## Caching
+
+The Globus Compute ``Client()`` constructor and ``get_endpoint_status()`` are
+both network round-trips that take 1–5 s. This module keeps a process-level
+Client cache plus a short-TTL per-endpoint status cache so back-to-back checks
+are effectively free.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+import warnings
+from typing import Any
+
+# ── Process-wide cached Client ────────────────────────────────────────────────
+_CLIENT: Any = None
+_CLIENT_LOCK = threading.Lock()
+
+# ── Per-endpoint status cache ─────────────────────────────────────────────────
+# Keyed by endpoint_id; value is (monotonic_ts, status_dict).
+# The stored dict never contains "cached"/"cache_age_seconds" — those are
+# added on the way out by _cached_entry().
+_STATUS_CACHE: dict[str, tuple[float, dict[str, Any]]] = {}
+_STATUS_CACHE_LOCK = threading.Lock()
+
+# Registered endpoints stay valid longer; offline/unreachable are re-checked sooner.
+_TTL_REGISTERED_SECONDS = 10.0
+_TTL_OTHER_SECONDS = 3.0
+
+# ── Globus → our vocabulary mapping ──────────────────────────────────────────
+_GLOBUS_TO_STATUS: dict[str, str] = {
+    "online": "registered",
+    "healthy": "registered",
+    "offline": "offline",
+    "stopped": "offline",
+}
+
+
+def _translate_globus_status(raw: str) -> str:
+    """Map a raw Globus status string to our vocabulary."""
+    return _GLOBUS_TO_STATUS.get(raw.lower(), "unreachable")
+
+
+def _endpoint_public_fields(config: Any) -> dict[str, Any]:
+    """Return non-sensitive endpoint metadata for public tool payloads."""
+    return {
+        "endpoint_name": getattr(config, "endpoint_name", None) or "configured",
+        "endpoint_configured": bool(getattr(config, "endpoint_id", None)),
+    }
+
+
+def _get_client() -> Any:
+    """Return a process-wide cached Globus Compute ``Client``."""
+    global _CLIENT
+    if _CLIENT is None:
+        with _CLIENT_LOCK:
+            if _CLIENT is None:
+                from globus_compute_sdk import Client
+
+                _CLIENT = Client()
+    return _CLIENT
+
+
+def invalidate_cache(endpoint_id: str | None = None) -> None:
+    """Drop cached status entries.
+
+    Parameters
+    ----------
+    endpoint_id : str | None
+        If ``None``, clear the entire cache. The cached Client is preserved.
+    """
+    with _STATUS_CACHE_LOCK:
+        if endpoint_id is None:
+            _STATUS_CACHE.clear()
+        else:
+            _STATUS_CACHE.pop(endpoint_id, None)
+
+
+def _cached_entry(endpoint_id: str) -> dict[str, Any] | None:
+    with _STATUS_CACHE_LOCK:
+        entry = _STATUS_CACHE.get(endpoint_id)
+    if entry is None:
+        return None
+    ts, payload = entry
+    age = time.monotonic() - ts
+    ttl = (
+        _TTL_REGISTERED_SECONDS
+        if payload.get("status") == "registered"
+        else _TTL_OTHER_SECONDS
+    )
+    if age > ttl:
+        return None
+    cached = dict(payload)
+    cached["cached"] = True
+    cached["cache_age_seconds"] = round(age, 3)
+    return cached
+
+
+def _store(endpoint_id: str, payload: dict[str, Any]) -> dict[str, Any]:
+    with _STATUS_CACHE_LOCK:
+        _STATUS_CACHE[endpoint_id] = (time.monotonic(), payload)
+    fresh = dict(payload)
+    fresh["cached"] = False
+    return fresh
+
+
+def check_endpoint_manager_status(
+    config: Any, *, force: bool = False
+) -> dict[str, Any]:
+    """Check whether the Globus Compute endpoint manager is registered.
+
+    This is a fast, cached check that queries the Globus cloud about the
+    manager process. It does **not** verify that workers are running or that
+    submitted tasks will succeed.
+
+    Parameters
+    ----------
+    config : HPCConfig
+        HPC configuration with ``endpoint_id``.
+    force : bool, default False
+        Bypass the short-TTL cache and re-query the SDK.
+
+    Returns
+    -------
+    dict
+        Keys:
+
+        ``status``
+            One of ``"registered"``, ``"offline"``, ``"unreachable"``,
+            ``"no_endpoint"``.
+        ``endpoint_name`` / ``endpoint_configured``
+            Non-sensitive endpoint metadata. Raw UUIDs stay in private config.
+        ``error``
+            Error message when ``status`` is ``"unreachable"``.
+        ``cached``
+            ``True`` if this result was served from the in-process cache.
+        ``cache_age_seconds``
+            Age of the cached entry in seconds (only when ``cached=True``).
+
+    Examples
+    --------
+    >>> from uxarray_mcp.remote.config import load_config
+    >>> config = load_config()
+    >>> check_endpoint_manager_status(config)
+    {"status": "registered", "endpoint_name": "improv", "cached": False}
+    """
+    endpoint_id = getattr(config, "endpoint_id", None)
+    if not endpoint_id:
+        return {"status": "no_endpoint", "message": "No endpoint configured"}
+
+    if not force:
+        cached = _cached_entry(endpoint_id)
+        if cached is not None:
+            return cached
+
+    try:
+        client = _get_client()
+        raw = client.get_endpoint_status(endpoint_id)
+        raw_status = raw.get("status", "unknown") if isinstance(raw, dict) else str(raw)
+        status = _translate_globus_status(raw_status)
+        payload: dict[str, Any] = {"status": status, **_endpoint_public_fields(config)}
+        return _store(endpoint_id, payload)
+    except Exception as exc:
+        payload = {
+            "status": "unreachable",
+            **_endpoint_public_fields(config),
+            "error": str(exc),
+        }
+        return _store(endpoint_id, payload)
+
+
+def probe_endpoint_worker(
+    config: Any,
+    *,
+    timeout_seconds: int = 60,
+) -> dict[str, Any]:
+    """Submit a lightweight task to confirm a real worker responds.
+
+    Unlike :func:`check_endpoint_manager_status`, this function actually
+    submits a task via Globus Compute and waits for it to run on a scheduler
+    node. This is the only way to confirm that Slurm/PBS allocated a worker
+    and the Python environment is intact.
+
+    Parameters
+    ----------
+    config : HPCConfig
+        HPC configuration with ``endpoint_id``.
+    timeout_seconds : int, default 60
+        How long to wait for the probe task to complete.
+
+    Returns
+    -------
+    dict
+        Keys:
+
+        ``status``
+            ``"active"`` if a worker responded, ``"registered"`` if the
+            manager is up but the probe timed out, ``"offline"`` or
+            ``"unreachable"`` if the manager itself is down.
+        ``node``
+            Hostname of the compute node that ran the task (when active).
+        ``python``
+            Python version on the worker (when active).
+        ``slurm_job_id``
+            Slurm job ID (when active and running under Slurm).
+        ``elapsed_seconds``
+            Wall-clock time the probe took.
+        ``error``
+            Error message on failure.
+
+    Examples
+    --------
+    >>> probe_endpoint_worker(config, timeout_seconds=90)
+    {"status": "active", "node": "chr-0497", "python": "3.13.13",
+     "slurm_job_id": "1228500", "elapsed_seconds": 28.4}
+    """
+    endpoint_id = getattr(config, "endpoint_id", None)
+    if not endpoint_id:
+        return {"status": "no_endpoint", "message": "No endpoint configured"}
+
+    # Fast manager check first — no point probing an offline endpoint
+    manager = check_endpoint_manager_status(config, force=True)
+    if manager["status"] not in ("registered", "active"):
+        return manager
+
+    try:
+        from globus_compute_sdk import Executor
+        from globus_compute_sdk.serialize import AllCodeStrategies, ComputeSerializer
+    except ImportError:
+        return {
+            "status": "unreachable",
+            **_endpoint_public_fields(config),
+            "error": "globus-compute-sdk not installed. Run: uv sync --extra hpc",
+        }
+
+    def _worker_probe() -> dict:
+        import os
+        import platform
+
+        return {
+            "node": platform.node(),
+            "python": platform.python_version(),
+            "slurm_job_id": os.environ.get("SLURM_JOB_ID", ""),
+            "pbs_job_id": os.environ.get("PBS_JOBID", ""),
+            "pythonpath": os.environ.get("PYTHONPATH", ""),
+        }
+
+    t0 = time.monotonic()
+    try:
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                message=r"(?s).*Environment differences detected between local SDK and endpoint.*",
+                category=UserWarning,
+            )
+            with Executor(
+                endpoint_id=endpoint_id,
+                serializer=ComputeSerializer(strategy_code=AllCodeStrategies()),
+            ) as ex:
+                fut = ex.submit(_worker_probe)
+                result = fut.result(timeout=timeout_seconds)
+
+        elapsed = round(time.monotonic() - t0, 1)
+        payload = {
+            "status": "active",
+            **_endpoint_public_fields(config),
+            "node": result.get("node", ""),
+            "python": result.get("python", ""),
+            "slurm_job_id": result.get("slurm_job_id") or None,
+            "pbs_job_id": result.get("pbs_job_id") or None,
+            "pythonpath_set": bool(result.get("pythonpath")),
+            "elapsed_seconds": elapsed,
+        }
+        # Warn if PYTHONPATH is set — this is the root cause of most worker crashes
+        if result.get("pythonpath"):
+            payload["warning"] = (
+                "PYTHONPATH is set on the worker. This can cause pydantic/dill "
+                "conflicts. Add 'unset PYTHONPATH' to worker_init in the endpoint "
+                "config, and set PYTHONPATH: '' in user_environment.yaml."
+            )
+        return payload
+
+    except TimeoutError:
+        return {
+            "status": "registered",
+            **_endpoint_public_fields(config),
+            "error": f"Probe timed out after {timeout_seconds}s. Manager is up but "
+            "no worker responded. The scheduler may be busy or the worker "
+            "environment may be broken.",
+            "elapsed_seconds": round(time.monotonic() - t0, 1),
+        }
+    except Exception as exc:
+        return {
+            "status": "unreachable",
+            **_endpoint_public_fields(config),
+            "error": str(exc),
+            "elapsed_seconds": round(time.monotonic() - t0, 1),
+        }
+
+
+def check_all_endpoints_manager_status(
+    config: Any, *, force: bool = False
+) -> list[dict[str, Any]]:
+    """Return a status row for every configured endpoint.
+
+    Each row contains ``name`` and the same fields produced
+    by :func:`check_endpoint_manager_status`. Uses the cache by default.
+    """
+    rows: list[dict[str, Any]] = []
+    names = list(getattr(config, "endpoint_names", []) or [])
+    if not names:
+        if getattr(config, "endpoint_id", None):
+            status = check_endpoint_manager_status(config, force=force)
+            rows.append({"name": config.endpoint_name or "default", **status})
+        return rows
+
+    for name in names:
+        endpoint_cfg = config.for_endpoint(endpoint=name)
+        status = check_endpoint_manager_status(endpoint_cfg, force=force)
+        rows.append({"name": name, **status})
+    return rows
+
+
+# ---------------------------------------------------------------------------
+# Backwards-compatibility aliases — kept so existing callers don't break
+# while we migrate. Will be removed in a future release.
+# ---------------------------------------------------------------------------
+def check_endpoint_health(config: Any, *, force: bool = False) -> dict[str, Any]:
+    """Deprecated alias for :func:`check_endpoint_manager_status`."""
+    return check_endpoint_manager_status(config, force=force)
+
+
+def check_all_endpoints_health(
+    config: Any, *, force: bool = False
+) -> list[dict[str, Any]]:
+    """Deprecated alias for :func:`check_all_endpoints_manager_status`."""
+    return check_all_endpoints_manager_status(config, force=force)