mlspace-plugin 0.2.1__py3-none-any.whl

mlspace_mcp/tools/queues.py

@@ -0,0 +1,202 @@
+"""``mlspace_queues`` — manage MLSpace allocation queues and shared-cluster queues.
+
+Pure declarative data following the EXEMPLAR (``inference.py``). The core owns
+all request mechanics; this module only declares ``ACTIONS`` / ``PARAMS`` and the
+``DomainTool`` metadata for the ``queues`` domain.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from ..registry import DomainTool, Op, Param
+
+# action -> Op. (method, path) pairs are the single source of truth for drift.
+ACTIONS: dict[str, Op] = {
+    "list": Op(
+        "GET", "/public/v2/queues/",
+        query_params={"allocation_id": "allocation_id", "workspace_id": "workspace_id"},
+        required=("allocation_id",),
+        kind="list",
+        help="List queues for an allocation.",
+    ),
+    "get": Op(
+        "GET", "/public/v2/queues/{queue_id}",
+        path_params={"queue_id": "queue_id"},
+        required=("queue_id",),
+        help="Get detailed information about one queue.",
+    ),
+    "defaults": Op(
+        "GET", "/public/v2/queues/defaults",
+        query_params={"allocation_id": "allocation_id"},
+        required=("allocation_id",),
+        kind="list",
+        help="List default queues (with types) for an allocation.",
+    ),
+    "instance_types": Op(
+        "GET", "/public/v2/queues/{queue_id}/instance-types",
+        path_params={"queue_id": "queue_id"},
+        required=("queue_id",),
+        kind="list",
+        help="List instance types available for a queue.",
+    ),
+    "jobs": Op(
+        "GET", "/public/v2/queues/{queue_id}/jobs",
+        path_params={"queue_id": "queue_id"},
+        required=("queue_id",),
+        kind="list",
+        help="List jobs in a queue.",
+    ),
+    "notebooks": Op(
+        "GET", "/public/v2/queues/{queue_id}/notebooks",
+        path_params={"queue_id": "queue_id"},
+        required=("queue_id",),
+        kind="list",
+        help="List notebooks in a queue.",
+    ),
+    "pods": Op(
+        "GET", "/public/v2/queues/{queue_id}/pods",
+        path_params={"queue_id": "queue_id"},
+        required=("queue_id",),
+        kind="list",
+        help="List pods in a queue.",
+    ),
+    "awaiting_launch_resources": Op(
+        "GET", "/public/v2/queues/{queue_id}/awaiting-launch-resources",
+        path_params={"queue_id": "queue_id"},
+        query_params={"priority": "priority", "limit": "limit", "offset": "offset"},
+        required=("queue_id", "priority", "limit", "offset"),
+        kind="list",
+        help="List one priority page; priority, limit and offset are required. Use mlspace_queue_inspect for all-priority coverage. List the resources a queue is waiting on. THE way to answer 'why is my "
+        "job still pending?' — `jobs get` only reports the bare status `pending` and "
+        "`jobs list_nodes` answers 404 while the job waits.",
+    ),
+    "queue_defaults": Op(
+        "GET", "/public/v2/queues/{queue_id}/defaults",
+        path_params={"queue_id": "queue_id"},
+        required=("queue_id",),
+        kind="list",
+        help="List default settings of ONE queue (the `defaults` action lists default "
+        "queues of an allocation — different question, similar name).",
+    ),
+    "assign_workspace": Op(
+        "POST", "/public/v2/queues/{queue_id}/workspaces/assign",
+        path_params={"queue_id": "queue_id"},
+        body_field="body",
+        required=("queue_id", "body"),
+        write=True,
+        help="Assign workspaces to a queue; `body` = workspace assignment request.",
+    ),
+    "delete": Op(
+        "DELETE", "/public/v2/queues/{queue_id}",
+        path_params={"queue_id": "queue_id"},
+        required=("queue_id",),
+        write=True,
+        confirm=True,
+        help="Delete a queue. Irreversible; jobs and notebooks bound to it lose their "
+        "scheduling target.",
+    ),
+    "create": Op(
+        "POST", "/public/v2/queues/",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Create a queue; `body` = CreateQueueRequest.",
+    ),
+    "update": Op(
+        "PUT", "/public/v2/queues/{queue_id}",
+        path_params={"queue_id": "queue_id"},
+        body_field="body",
+        required=("queue_id", "body"),
+        write=True,
+        help="Update a queue; `body` = UpdateQueueRequest.",
+    ),
+    "assign": Op(
+        "POST", "/public/v2/queues/{queue_id}/assign",
+        path_params={"queue_id": "queue_id"},
+        body_field="body",
+        required=("queue_id", "body"),
+        write=True,
+        help="Assign a queue to a workspace; `body` = AssignmentRequest.",
+    ),
+    "unassign": Op(
+        "DELETE", "/public/v2/queues/{queue_id}/assign",
+        path_params={"queue_id": "queue_id"},
+        body_field="body",
+        required=("queue_id", "body"),
+        write=True,
+        help="Unassign a queue from a workspace; `body` = AssignmentRequest.",
+    ),
+    "set_default": Op(
+        "POST", "/public/v2/queues/{queue_id}/defaults",
+        path_params={"queue_id": "queue_id"},
+        body_field="body",
+        required=("queue_id", "body"),
+        write=True,
+        help="Set the queue as default for a job/notebook; "
+             "`body` = DefaultQueueRequest: {org_structure, default_for}.",
+    ),
+    "unset_default": Op(
+        "DELETE", "/public/v2/queues/{queue_id}/defaults",
+        path_params={"queue_id": "queue_id"},
+        body_field="body",
+        required=("queue_id", "body"),
+        write=True,
+        help="Unset the queue as default for a job/notebook; "
+             "`body` = DefaultQueueRequest: {org_structure, default_for}.",
+    ),
+    "add_nodes": Op(
+        "POST", "/public/v2/queues/{queue_id}/nodes",
+        path_params={"queue_id": "queue_id"},
+        body_field="body",
+        required=("queue_id", "body"),
+        write=True,
+        help="Add compute nodes to the destination queue; `body` = AddNodesRequest "
+        "{nodes, force_withdrawal}. Inspect source workload with mlspace_queue_inspect first. "
+        "force_withdrawal=true forces withdrawal of all load; never escalate silently.",
+    ),
+    "assign_shared": Op(
+        "POST", "/public/v2/shared-cluster-queues/assign",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Assign a shared cluster queue to a workspace; "
+             "`body` = AssignmentRequest: {region_key, org_structure}.",
+    ),
+    "unassign_shared": Op(
+        "DELETE", "/public/v2/shared-cluster-queues/assign",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Unassign a shared cluster queue from a workspace; "
+             "`body` = AssignmentRequest: {region_key, org_structure}.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("priority", Literal["shared-low", "shared-medium", "shared-cluster", "low", "medium", "high"], "Required priority for awaiting_launch_resources; one priority is not complete demand."),
+    Param("limit", int, "Required pending page size."),
+    Param("offset", int, "Required pending page offset (start at zero)."),
+    Param("queue_id", str, "Queue unique ID (UUID); required for queue-scoped actions."),
+    Param("allocation_id", str, "Allocation unique ID (UUID); required for list/defaults."),
+    Param("workspace_id", str, "Optional workspace UUID to filter the queue list by."),
+    Param(
+        "body",
+        dict,
+        "JSON body. create=CreateQueueRequest; update=UpdateQueueRequest; "
+        "assign/unassign=AssignmentRequest {org_structure}; "
+        "set_default/unset_default=DefaultQueueRequest {org_structure, default_for}; "
+        "add_nodes=AddNodesRequest {nodes, force_withdrawal}; "
+        "assign_shared/unassign_shared=AssignmentRequest {region_key, org_structure}.",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="queues",
+    name="mlspace_queues",
+    title="MLSpace Queues",
+    summary="Manage MLSpace allocation queues and shared-cluster queues "
+            "(list, inspect, create, update, assign, set defaults, add nodes).",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/resources.py

@@ -0,0 +1,86 @@
+"""``mlspace_resources`` — read MLSpace compute catalog and node load.
+
+A merged read-only catalog over configs, available instance types, nodes, and
+node load. Pure declarative data; the core owns all request mechanics.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from ..registry import DomainTool, Op, Param
+
+# action -> Op. The (method, path) pairs are the single source of truth for the
+# drift test; paths are copied verbatim from the spec, placeholders included.
+ACTIONS: dict[str, Op] = {
+    "configs": Op(
+        "GET", "/public/v2/configs",
+        query_params={"cluster_type": "cluster_type"},
+        required=("cluster_type",),
+        kind="catalog",
+        help="Get region configs: available instance_types and images for Jobs/Notebooks.",
+    ),
+    "instance_types": Op(
+        "GET", "/public/v2/instance_types/{region}/available",
+        path_params={"region": "region"},
+        query_params={"allocation_name": "allocation_name", "queue_name": "queue_name"},
+        required=("region",),
+        kind="list",
+        help="List available instance types by connected allocation in a region.",
+    ),
+    "nodes": Op(
+        "GET", "/public/v2/nodes/{region}/nodes",
+        path_params={"region": "region"},
+        required=("region",),
+        kind="list",
+        help="List available nodes to allocation in a region.",
+    ),
+    "nodes_load": Op(
+        "GET", "/public/v2/nodes/load",
+        query_params={
+            "queue_id": "queue_id",
+            "allocation_id": "allocation_id",
+            "node_names": "node_names",
+        },
+        required=("node_names",),
+        kind="list",
+        help="Get load on nodes; pass node_names (and optionally queue_id/allocation_id).",
+    ),
+    "rate_limits": Op(
+        "GET", "/public/v2/limiter/",
+        kind="list",
+        help="List this account's API rate limits and current usage, per endpoint and "
+        "period. Call this BEFORE a burst of writes: the contract documents 429 on "
+        "only ~45% of operations, so this is the only reliable way to know the budget.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param(
+        "cluster_type",
+        Literal["MT", "INF"],
+        "Cluster type to fetch configs for (required for configs).",
+    ),
+    Param(
+        "region",
+        str,
+        "Model-training region, e.g. DGX2-MT, A100-MT, SR003-SR006, SR008 "
+        "(authoritative list via the configs action). Required for instance_types/nodes. "
+        "Optional for configs: filters the catalog client-side to this one region "
+        "(case-insensitive); omit to get all regions.",
+    ),
+    Param("allocation_name", str, "Filter instance_types by allocation name."),
+    Param("queue_name", str, "Filter instance_types by queue name."),
+    Param("queue_id", str, "Filter node load by queue UUID (nodes_load)."),
+    Param("allocation_id", str, "Filter node load by allocation UUID (nodes_load)."),
+    Param("node_names", list[str], "Node names to query load for (nodes_load)."),
+]
+
+DOMAIN = DomainTool(
+    domain="resources",
+    name="mlspace_resources",
+    title="MLSpace Resources",
+    summary="Read MLSpace compute catalog: configs, available instance types, nodes, node load.",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/tensorboards.py

@@ -0,0 +1,112 @@
+"""``mlspace_tensorboards`` — manage MLSpace TensorBoard instances.
+
+Pure declarative data following the inference.py exemplar. No request
+mechanics here; the core (registry.py + client.py) owns headers, path
+templating, query encoding, bodies, readonly, confirm, and formatting.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from ..registry import DomainTool, Op, Param
+
+# action -> Op. The (method, path) pairs are the single source of truth for the
+# drift test; paths are copied verbatim from the spec, placeholders included.
+ACTIONS: dict[str, Op] = {
+    "list": Op(
+        "GET", "/public/v2/tensorboards/v2/tensorboards",
+        query_params={
+            "order_by": "order_by",
+            "desc": "desc",
+            "limit": "limit",
+            "offset": "offset",
+            "search": "search",
+            "status": "status",
+        },
+        kind="list",
+        help="List TensorBoards in the workspace (filters: order_by, desc, "
+        "limit, offset, search, status).",
+    ),
+    "get": Op(
+        "GET", "/public/v2/tensorboards/v2/tensorboard/{tensorboard_uuid}",
+        path_params={"tensorboard_uuid": "tensorboard_uuid"},
+        required=("tensorboard_uuid",),
+        help="Get one TensorBoard by uuid.",
+    ),
+    "create": Op(
+        "POST", "/public/v2/tensorboards/v2/{namespace}/tensorboard",
+        path_params={"namespace": "namespace"},
+        body_field="body",
+        required=("namespace", "body"),
+        write=True,
+        help="Create a TensorBoard; `body` = CreateTensorboardPayloadV2 "
+        "({name, image, instance_type, logdir, ...}).",
+    ),
+    "modify": Op(
+        "POST", "/public/v2/tensorboards/v2/tensorboard/{tensorboard_uuid}/modify",
+        path_params={"tensorboard_uuid": "tensorboard_uuid"},
+        body_field="body",
+        required=("tensorboard_uuid", "body"),
+        write=True,
+        help="Modify a TensorBoard; `body` = ModifyTensorboardRequest "
+        "({autoshutdown_config, description, s3_credentials, s3_buckets}).",
+    ),
+    "pause": Op(
+        "POST", "/public/v2/tensorboards/v2/tensorboard/{tensorboard_uuid}/pause",
+        path_params={"tensorboard_uuid": "tensorboard_uuid"},
+        required=("tensorboard_uuid",),
+        write=True,
+        help="Pause a TensorBoard.",
+    ),
+    "resume": Op(
+        "POST",
+        "/public/v2/tensorboards/v2/{namespace}/tensorboard/{tensorboard_uuid}/resume",
+        path_params={"namespace": "namespace", "tensorboard_uuid": "tensorboard_uuid"},
+        body_field="body",
+        required=("namespace", "tensorboard_uuid", "body"),
+        write=True,
+        help="Resume a TensorBoard; `body` = ResumeTensorboardPayload "
+        "({region, instance_type}).",
+    ),
+    "delete": Op(
+        "DELETE", "/public/v2/tensorboards/v2/tensorboard/{tensorboard_uuid}",
+        path_params={"tensorboard_uuid": "tensorboard_uuid"},
+        required=("tensorboard_uuid",),
+        write=True,
+        confirm=True,
+        help="Delete a TensorBoard.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("tensorboard_uuid", str, "TensorBoard UUID (get/modify/pause/resume/delete)."),
+    Param("namespace", str, "Namespace the TensorBoard belongs to (create/resume). "
+          "Optional: auto-filled from the server's workspace if omitted."),
+    Param(
+        "order_by",
+        Literal["name", "creationTimestamp"],
+        "List ordering field (default creationTimestamp).",
+    ),
+    Param("desc", bool, "Sort descending (list)."),
+    Param("limit", int, "Max results to return (list)."),
+    Param("offset", int, "Pagination offset (list)."),
+    Param("search", str, "Filter by name substring (list)."),
+    Param("status", str, "Filter by comma-separated statuses, e.g. 'creating,running' (list)."),
+    Param(
+        "body",
+        dict,
+        "JSON body. create=CreateTensorboardPayloadV2; modify="
+        "ModifyTensorboardRequest; resume=ResumeTensorboardPayload.",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="tensorboards",
+    name="mlspace_tensorboards",
+    title="MLSpace TensorBoards",
+    summary="Manage MLSpace TensorBoard instances (list, inspect, create, "
+    "modify, pause, resume, delete).",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/workspaces.py

@@ -0,0 +1,92 @@
+"""``mlspace_workspaces`` — inspect MLSpace workspaces (read-only).
+
+Pure declarative data following the ``inference`` exemplar: the core owns all
+request mechanics (headers, path templating, query encoding, bodies, readonly,
+confirm, formatting). This module only declares ``ACTIONS`` / ``PARAMS``.
+
+SAFETY: workspace bootstrap (POST /workspaces/v3/) and debootstrap (DELETE
+/workspaces/v3/) are DELIBERATELY NOT exposed — creating/deleting an entire
+workspace is a catastrophic blast radius for a shared, company-wide server, so
+those ops are intentionally omitted (and listed in the drift test's exclusion
+set so the bijection check stays green). Do workspace lifecycle in the console.
+"""
+
+from __future__ import annotations
+
+from ..registry import DomainTool, Op, Param
+
+# action -> Op. The (method, path) pairs are the single source of truth for the
+# drift test; paths are copied verbatim from the spec, placeholders included.
+ACTIONS: dict[str, Op] = {
+    "get_api_key": Op(
+        "GET", "/public/v2/workspaces/v1/x_api_key",
+        query_params={"workspace_id": "workspace_id"},
+        required=("workspace_id",),
+        help="Get the workspace X-API-KEY (secret, redacted in output).",
+    ),
+    "status": Op(
+        "GET", "/public/v2/workspaces/v1/{workspace_id}/status",
+        path_params={"workspace_id": "workspace_id"},
+        required=("workspace_id",),
+        help="Get a workspace's status.",
+    ),
+    "users": Op(
+        "GET", "/public/v2/workspaces/v3/{workspace_id}/users",
+        path_params={"workspace_id": "workspace_id"},
+        required=("workspace_id",),
+        kind="list",
+        help="List the users who have access to a workspace.",
+    ),
+    "list": Op(
+        "GET", "/public/v2/workspaces/v3/",
+        query_params={"customer_id": "customer_id"},
+        kind="list",
+        help="List user workspaces (optionally filtered by customer_id).",
+    ),
+    "get": Op(
+        "GET", "/public/v2/workspaces/v3/{workspace_id}",
+        path_params={"workspace_id": "workspace_id"},
+        required=("workspace_id",),
+        help="Get workspace details.",
+    ),
+    "allocations": Op(
+        "GET", "/public/v2/workspaces/v3/{workspace_id}/allocations",
+        path_params={"workspace_id": "workspace_id"},
+        required=("workspace_id",),
+        kind="list",
+        help="List a workspace's allocations.",
+    ),
+    "allocation_queues": Op(
+        "GET",
+        "/public/v2/workspaces/v3/{workspace_id}/allocations/{allocation_id}/queues",
+        path_params={"workspace_id": "workspace_id", "allocation_id": "allocation_id"},
+        required=("workspace_id", "allocation_id"),
+        kind="list",
+        help="List queues for one allocation in a workspace.",
+    ),
+    # NOTE: bootstrap (POST /workspaces/v3/) and debootstrap (DELETE /workspaces/v3/)
+    # are intentionally NOT exposed — see the SAFETY note in the module docstring.
+}
+
+PARAMS: list[Param] = [
+    Param(
+        "workspace_id",
+        str,
+        "Workspace UUID (path/query param for get_api_key/status/get/"
+        "allocations/allocation_queues). Distinct from the x-workspace-id header. "
+        "Optional: defaults to the server's configured workspace if omitted; pass only to target another.",
+    ),
+    Param("allocation_id", str, "Allocation UUID (allocation_queues only)."),
+    Param("customer_id", str, "Filter workspaces by customer (list only)."),
+]
+
+DOMAIN = DomainTool(
+    domain="workspaces",
+    name="mlspace_workspaces",
+    title="MLSpace Workspaces",
+    summary="Inspect MLSpace workspaces (status, details, allocations, queues, "
+    "api-key). Read-only: workspace create/delete is intentionally not exposed.",
+    actions=ACTIONS,
+    params=PARAMS,
+    redact_result_keys=("x-api-key",),
+)

mlspace_mcp/workspace_context.py

@@ -0,0 +1,107 @@
+"""A fixed selected catalogue and independent, lazily initialized HTTP clients."""
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from dataclasses import dataclass, field
+from urllib.parse import urlsplit, urlunsplit
+
+import httpx
+
+from .auth import TokenManager
+from .client import MLSpaceClient
+from .config import Settings
+from .errors import MLSpaceError
+
+
+def normalized_endpoint(endpoint: str) -> str:
+    parsed = urlsplit(endpoint)
+    host = (parsed.hostname or "").lower()
+    if ":" in host:
+        host = f"[{host}]"
+    port = parsed.port
+    if port and (parsed.scheme.lower(), port) not in {("https", 443), ("http", 80)}:
+        host += f":{port}"
+    return urlunsplit((parsed.scheme.lower(), host, parsed.path.rstrip("/"), "", ""))
+
+
+def environment_id(endpoint: str) -> str:
+    return "mlspace:" + hashlib.sha256(normalized_endpoint(endpoint).encode()).hexdigest()
+
+
+@dataclass(frozen=True)
+class WorkspaceSelection:
+    id: str
+    name: str
+    project_name: str
+    api_key: str = field(default="", repr=False)
+    namespace: str = ""
+
+
+def selected_workspaces(settings: Settings) -> tuple[WorkspaceSelection, ...]:
+    if settings.workspaces:
+        return tuple(WorkspaceSelection(
+            item.id, item.name, item.project_name,
+            item.api_key.get_secret_value() if item.api_key else "",
+            settings.namespace if len(settings.workspaces) == 1
+            and item.id == settings.workspace_id else "",
+        ) for item in settings.workspaces)
+    return (WorkspaceSelection(settings.workspace_id, settings.workspace_id, "",
+                               settings.api_key.get_secret_value(), settings.namespace),)
+
+
+def select_target(selected: tuple[WorkspaceSelection, ...], target: str | None) -> WorkspaceSelection:
+    if target is None:
+        if len(selected) == 1:
+            return selected[0]
+        raise MLSpaceError("Choose target from mlspace_contexts; several workspaces are configured.")
+    by_id = [item for item in selected if item.id == target]
+    matches = by_id or [item for item in selected if item.name == target]
+    if len(matches) > 1:
+        raise MLSpaceError("Workspace target is ambiguous; use a configured workspace ID.")
+    if not matches:
+        raise MLSpaceError("Workspace target is not configured; see mlspace_contexts.")
+    return matches[0]
+
+
+class WorkspaceClients:
+    """Never discover membership or switch an existing client's credentials."""
+
+    def __init__(self, settings: Settings, http: httpx.AsyncClient, tokens: TokenManager):
+        self.selected = selected_workspaces(settings)
+        if len({item.id for item in self.selected}) != len(self.selected):
+            raise MLSpaceError("Duplicate configured workspace IDs.")
+        self.endpoint = normalized_endpoint(settings.base_url)
+        self.environment = environment_id(self.endpoint)
+        self._http = http
+        self._tokens = tokens
+        self._clients: dict[str, MLSpaceClient] = {}
+        self._locks = {item.id: asyncio.Lock() for item in self.selected}
+
+    def select(self, target: str | None) -> WorkspaceSelection:
+        """Resolve membership without HTTP; use before validating operation addresses."""
+        return select_target(self.selected, target)
+
+    def catalogue(self) -> dict:
+        return {"environment": self.environment, "endpoint": self.endpoint,
+                "workspaces": [{"id": item.id, "name": item.name,
+                                "project_name": item.project_name} for item in self.selected]}
+
+    async def resolve(self, target: str | None) -> MLSpaceClient:
+        selection = self.select(target)
+        async with self._locks[selection.id]:
+            if selection.id not in self._clients:
+                key = selection.api_key
+                if not key:
+                    bootstrap = MLSpaceClient(self._http, self._tokens, base_url=self.endpoint,
+                                              api_key="", workspace_id=selection.id)
+                    data = await bootstrap.request("GET", "/public/v2/workspaces/v1/x_api_key")
+                    resolved_key = data.get("x-api-key") if isinstance(data, dict) else None
+                    if not isinstance(resolved_key, str) or not resolved_key.strip():
+                        raise MLSpaceError("API key unavailable for the selected workspace.")
+                    key = resolved_key
+                self._clients[selection.id] = MLSpaceClient(
+                    self._http, self._tokens, base_url=self.endpoint, api_key=key,
+                    workspace_id=selection.id, namespace=selection.namespace,
+                )
+            return self._clients[selection.id]