mlspace-plugin 0.2.1__py3-none-any.whl

mlspace_mcp/tools/docker_registry.py

@@ -0,0 +1,156 @@
+"""``mlspace_docker_registry`` — manage the MLSpace Docker registry.
+
+Pure declarative data following the EXEMPLAR ``inference.py``. The registry has
+multiple resource types (registry, repositories, images, tags, credentials), so
+action names suffix the noun. All request mechanics live in the core.
+"""
+
+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] = {
+    "current_registry": Op(
+        "GET", "/public/v2/docker_registry/v2/registries/current",
+        help="Get the current Docker registry for the workspace.",
+    ),
+    "list_repos": Op(
+        "GET", "/public/v2/docker_registry/v2/repositories/",
+        # page/page_size are a working proband lever here. NOTE we deliberately do NOT
+        # paginate list_images/list_tags: there pagination splits ONE image across
+        # pages with partial tags (08#4), so surfacing it would amplify that bug.
+        query_params={"page": "page", "page_size": "page_size"},
+        # registry_repos kind (not plain "list") elides the heavy recent_image.
+        # extra_attrs (~39% of the response) with a pointer to get_image.
+        kind="registry_repos",
+        help="List repositories in the registry.",
+    ),
+    "get_repo": Op(
+        "GET", "/public/v2/docker_registry/v2/repositories/{repository_id}",
+        path_params={"repository_id": "repository_id"},
+        required=("repository_id",),
+        help="Get one repository.",
+    ),
+    "update_repo": Op(
+        "PUT", "/public/v2/docker_registry/v2/repositories/{repository_id}",
+        path_params={"repository_id": "repository_id"},
+        body_field="body",
+        required=("repository_id", "body"),
+        write=True,
+        help="Update a repository; `body` = UpdateRepository: "
+             "{description, labels, labels_to_delete}.",
+    ),
+    "delete_repos": Op(
+        "DELETE", "/public/v2/docker_registry/v2/repositories/",
+        body_field="body",
+        required=("body",),
+        write=True,
+        confirm=True,
+        help="Delete repositories; `body` = DeleteRepositories.",
+    ),
+    "repo_fav": Op(
+        "POST", "/public/v2/docker_registry/v2/repositories/{repository_id}/fav",
+        path_params={"repository_id": "repository_id"},
+        body_field="body",
+        required=("repository_id", "body"),
+        write=True,
+        help="Mark/unmark a repository as favorite; `body` = FavoriteInput.",
+    ),
+    "list_images": Op(
+        "GET", "/public/v2/docker_registry/v2/images/",
+        query_params={"repository_id": "repository_id"},
+        required=("repository_id",),
+        kind="list",
+        help="List images in a repository (requires repository_id).",
+    ),
+    "get_image": Op(
+        "GET", "/public/v2/docker_registry/v2/images/{image_id}",
+        path_params={"image_id": "image_id"},
+        required=("image_id",),
+        help="Get one image.",
+    ),
+    "delete_images": Op(
+        "DELETE", "/public/v2/docker_registry/v2/images/",
+        body_field="body",
+        required=("body",),
+        write=True,
+        confirm=True,
+        help="Delete images; `body` = DeleteImages.",
+    ),
+    "list_tags": Op(
+        "GET", "/public/v2/docker_registry/v2/tags/",
+        query_params={"image_id": "image_id"},
+        required=("image_id",),
+        kind="list",
+        help="List tags for an image (requires image_id).",
+    ),
+    "create_tag": Op(
+        "POST", "/public/v2/docker_registry/v2/tags/",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Create a tag; `body` = CreateTag.",
+    ),
+    "update_tag": Op(
+        "PUT", "/public/v2/docker_registry/v2/tags/{tag_id}",
+        path_params={"tag_id": "tag_id"},
+        body_field="body",
+        required=("tag_id", "body"),
+        write=True,
+        help="Update a tag; `body` = UpdateTag.",
+    ),
+    "delete_tags": Op(
+        "DELETE", "/public/v2/docker_registry/v2/tags/",
+        body_field="body",
+        required=("body",),
+        write=True,
+        confirm=True,
+        help="Delete tags; `body` = DeleteTags.",
+    ),
+    "generate_password": Op(
+        "GET", "/public/v2/docker_registry/v1/users/generate_password",
+        write=True,
+        # irreversible credential rotation: the previous password immediately stops
+        # working, so a single mistaken call would lock out anything using it. Gate it
+        # behind confirm=true like the other destructive ops.
+        confirm=True,
+        help="Rotate and return the registry password for your personal account "
+             "(not service accounts). The previous password stops working. "
+             "Destructive: requires confirm=true.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("repository_id", str, "Repository ID (UUID)."),
+    Param("image_id", str, "Image ID (UUID); also the query filter for list_images/list_tags."),
+    Param("tag_id", str, "Tag ID (UUID)."),
+    Param(
+        "page",
+        int,
+        "Page number (list_repos, 1-based). The response is a bare array with no "
+        "total_count: a short/empty page does NOT prove the end of the list.",
+    ),
+    Param("page_size", int, "Page size (list_repos)."),
+    Param(
+        "body",
+        dict,
+        "JSON body. update_repo=UpdateRepository; delete_repos=DeleteRepositories "
+        "{repository_ids:[..]}; repo_fav=FavoriteInput {favorite}; "
+        "delete_images=DeleteImages {image_ids:[..]}; create_tag=CreateTag "
+        "{image_id, name, color}; update_tag=UpdateTag {color}; "
+        "delete_tags=DeleteTags {tag_ids:[..]}.",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="docker_registry",
+    name="mlspace_docker_registry",
+    title="MLSpace Docker Registry",
+    summary="Manage the MLSpace Docker registry: repositories, images, tags, "
+            "favorites, and registry credentials.",
+    actions=ACTIONS,
+    params=PARAMS,
+    redact_result_keys=("password",),
+)

mlspace_mcp/tools/inference.py

@@ -0,0 +1,88 @@
+"""``mlspace_inference`` — manage and call MLSpace inference services.
+
+EXEMPLAR module: every other domain module follows this exact shape — a single
+``DOMAIN: DomainTool`` made of pure data. No request mechanics live here; the
+core (registry.py + client.py) owns headers, path templating, query encoding,
+bodies, readonly, confirm, formatting. To add a domain, copy this structure and
+fill ``ACTIONS`` / ``PARAMS`` from the domain's slice in ``_planning/slices``.
+"""
+
+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] = {
+    "list": Op(
+        "GET", "/public/v2/inference/v2/",
+        kind="list",
+        help="List inference services in the workspace.",
+    ),
+    "get": Op(
+        "GET", "/public/v2/inference/v2/{inference_name}",
+        path_params={"inference_name": "inference_name"},
+        query_params={"region": "region"},
+        required=("inference_name", "region"),
+        help="Get one inference service.",
+    ),
+    "predict": Op(
+        "POST", "/public/v2/inference/v2/predict/{inference_name}/{model_name}/",
+        path_params={"inference_name": "inference_name", "model_name": "model_name"},
+        body_field="body",
+        required=("inference_name", "model_name"),
+        side_effect=True,  # runs a deployed (paid) model — not a read
+        help="Call a deployed model; put the model's input in `body` (free-form).",
+    ),
+    "create": Op(
+        "POST", "/public/v2/inference/v2/",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Create an inference service; `body` = CreateInferenceServiceRequestV2.",
+    ),
+    "update": Op(
+        "PATCH", "/public/v2/inference/v2/{inference_name}",
+        path_params={"inference_name": "inference_name"},
+        body_field="body",
+        required=("inference_name", "body"),
+        write=True,
+        help="Update an inference service; `body` = PatchInferenceServiceRequestV2.",
+    ),
+    "delete": Op(
+        "DELETE", "/public/v2/inference/v2/{inference_name}",
+        path_params={"inference_name": "inference_name"},
+        query_params={"region": "region"},
+        required=("inference_name", "region"),
+        write=True,
+        confirm=True,
+        help="Delete an inference service.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("inference_name", str, "Name of the inference service."),
+    Param("model_name", str, "Model name within the service (predict only)."),
+    Param(
+        "region",
+        str,
+        "Inference region, e.g. DGX2-INF-001, CCE-INF, INF-002 "
+        "(authoritative list via mlspace_resources configs cluster_type=INF). "
+        "Required for get/delete.",
+    ),
+    Param(
+        "body",
+        dict,
+        "JSON body. create=CreateInferenceServiceRequestV2; update="
+        "PatchInferenceServiceRequestV2; predict=the model's own input payload.",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="inference",
+    name="mlspace_inference",
+    title="MLSpace Inference",
+    summary="Manage and call MLSpace inference services (deploy, inspect, predict, delete).",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/jobs.py

@@ -0,0 +1,155 @@
+"""``mlspace_jobs`` — manage MLSpace training jobs.
+
+Pure declarative data following the exemplar in ``inference.py``. All request
+mechanics (headers, path templating, query encoding, bodies, readonly, confirm,
+formatting) live in the core; this module only declares ``ACTIONS`` / ``PARAMS``.
+"""
+
+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.
+# NOTE: most paths use {job_name}, but /params and /preemptors use {name}.
+ACTIONS: dict[str, Op] = {
+    "list": Op(
+        "GET", "/public/v2/jobs",
+        query_params={
+            "region": "region",
+            "allocation_name": "allocation_name",
+            "status": "status",
+            "job_author": "job_author",
+            "queue_id": "queue_id",
+            "offset": "offset",
+            "limit": "limit",
+            "start_date": "start_date",
+            "end_date": "end_date",
+        },
+        kind="list",
+        help="List one page of training jobs in a target region (filter by status/author/queue/dates). "
+        "For broad connected-workspace reads use mlspace_jobs_overview (dynamic regions/pages). "
+        "Pass `region` (e.g. SR006): without it the API returns count=0 (empty), not all regions.",
+    ),
+    "get": Op(
+        "GET", "/public/v2/jobs/{job_name}",
+        path_params={"job_name": "job_name"},
+        required=("job_name",),
+        # kind='job' triggers _normalise_job: flags the created→'job not found' race
+        # and the failed+empty-error trap, and adds a canonical status.
+        kind="job",
+        help="Get one job's status using resource_ref or its exact target/name.",
+    ),
+    "logs": Op(
+        "GET", "/public/v2/jobs/{job_name}/logs",
+        path_params={"job_name": "job_name"},
+        query_params={"tail": "tail", "verbose": "verbose", "region": "region"},
+        # API rejects logs without region (422) though the spec marks it optional —
+        # both blind-model runs hit this, so require it.
+        required=("job_name", "region"),
+        kind="log",
+        help="Get a job's logs by name (use `tail` to limit lines).",
+    ),
+    "list_nodes": Op(
+        "GET", "/public/v2/jobs/{job_name}/nodes",
+        path_params={"job_name": "job_name"},
+        required=("job_name",),
+        kind="list",
+        help="List nodes that were used for a job's execution.",
+    ),
+    "list_pods": Op(
+        "GET", "/public/v2/jobs/elastic/{job_name}/pods",
+        path_params={"job_name": "job_name"},
+        required=("job_name",),
+        kind="list",
+        help="List Elastic job pods by job name (pytorch_elastic jobs).",
+    ),
+    "get_params": Op(
+        "GET", "/public/v2/jobs/{name}/params",
+        path_params={"name": "job_name"},
+        required=("job_name",),
+        help="Get the parameters a job was run with, by job name.",
+    ),
+    "get_preemptors": Op(
+        "GET", "/public/v2/jobs/{name}/preemptors",
+        path_params={"name": "job_name"},
+        required=("job_name",),
+        kind="list",
+        help="Get the preemptors that gained priority over a job, by job name.",
+    ),
+    "run": Op(
+        "POST", "/public/v2/jobs",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Run a training job; `body` = RunJobRequest: "
+        "{script, base_image, instance_type, region, type, n_workers, ...}. "
+        "NOTE: `region` must be explicit. For type=binary, `script` is a non-empty "
+        "command or executable path and is forwarded unchanged; verify that it is "
+        "available in the selected image. Other job types require an NFS script path "
+        "under /home/jovyan/... in that same region. NFS volumes are region-scoped, "
+        "and the Public API cannot stat the file.",
+    ),
+    "restart": Op(
+        "POST", "/public/v2/jobs/restart",
+        body_field="body",
+        required=("body",),
+        write=True,
+        help="Restart a job by name (new job, same params); `body` = "
+        "RestartJobRequest: {job_name}.",
+    ),
+    "delete": Op(
+        "DELETE", "/public/v2/jobs/{job_name}",
+        path_params={"job_name": "job_name"},
+        query_params={"region": "region"},
+        required=("job_name", "region"),
+        write=True,
+        confirm=True,
+        help="Delete a job by name.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("job_name", str, "Name of the job (path param for get/logs/nodes/pods/"
+          "params/preemptors/delete)."),
+    Param(
+        "region",
+        str,
+        "Model-training region, e.g. DGX2-MT, A100-MT, SR003-SR006, SR008 "
+        "(configs may omit allocated regions; use jobs overview or workspace allocations "
+        "to discover existing-job regions). Required for logs and delete, "
+        "and effectively required for `list` (omitting it returns count=0, not all regions).",
+    ),
+    Param("allocation_name", str, "Filter list by allocation name."),
+    Param(
+        "status",
+        list[str],
+        "Filter list by job statuses (repeatable). Known values: Completed, "
+        "Completing, Deleted, Failed, Pending, Running, Stopped, Succeeded, "
+        "Terminated, Terminating, Restarting.",
+    ),
+    Param("job_author", list[str], "Filter list by job authors (repeatable)."),
+    Param("queue_id", str, "Filter list by queue UUID."),
+    Param("offset", int, "Number of list items to skip.", default=0),
+    Param("limit", int, "Max number of list items to return."),
+    Param("start_date", str, "List start-date filter (RFC3339, e.g. 2026-03-31T00:00:00Z)."),
+    Param("end_date", str, "List end-date filter (RFC3339, e.g. 2026-03-31T23:00:00Z)."),
+    Param("tail", int, "logs: number of recent log lines to display.", default=0),
+    Param("verbose", bool, "logs: if true, include timestamps on each line.", default=True),
+    Param(
+        "body",
+        dict,
+        "JSON body. run=RunJobRequest (requires script, base_image, instance_type); "
+        "for type=binary, script may be a command or executable path; "
+        "restart=RestartJobRequest ({job_name}).",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="jobs",
+    name="mlspace_jobs",
+    title="MLSpace Training Jobs",
+    summary="Manage MLSpace training jobs (run, inspect, logs, nodes, restart, delete).",
+    actions=ACTIONS,
+    params=PARAMS,
+)

mlspace_mcp/tools/notebooks.py

@@ -0,0 +1,174 @@
+"""``mlspace_notebooks`` — manage MLSpace Jupyter Server notebooks.
+
+Pure declarative ``DomainTool`` data; the core owns all request mechanics.
+Mixes v1 and v2 API paths — paths are copied verbatim from the slice.
+"""
+
+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/notebooks/v2/notebooks",
+        query_params={
+            "notebook_type": "notebook_type",
+            "order_by": "order_by",
+            "desc": "desc",
+            "limit": "limit",
+            "offset": "offset",
+            "search": "search",
+            "status": "status",
+            "access_mode": "access_mode",
+        },
+        kind="list",
+        help="List one page of Jupyter servers in the selected target. For the connected-set "
+        "inventory use mlspace_notebooks_overview. Retain paused notebooks with status; "
+        "configured GPUs are not evidence of active use. These are workloads, not compute nodes.",
+    ),
+    "get": Op(
+        "GET", "/public/v2/notebooks/v1/notebook/{notebook_uuid}",
+        path_params={"notebook_uuid": "notebook_uuid"},
+        required=("notebook_uuid",),
+        help="Get one Jupyter server by uuid.",
+    ),
+    "autoshutdown_get": Op(
+        "GET", "/public/v2/notebooks/v2/autoshutdown-rules/workspace/{workspace_id}",
+        path_params={"workspace_id": "workspace_id"},
+        required=("workspace_id",),
+        help="Get the workspace-level autoshutdown rule.",
+    ),
+    "create": Op(
+        "POST", "/public/v2/notebooks/v2/{namespace}/notebook",
+        path_params={"namespace": "namespace"},
+        body_field="body",
+        required=("namespace", "body"),
+        write=True,
+        help="Create a Jupyter server; `body` = CreateNotebookPayload.",
+    ),
+    "delete": Op(
+        "DELETE", "/public/v2/notebooks/v2/notebook/{notebook_uuid}",
+        path_params={"notebook_uuid": "notebook_uuid"},
+        required=("notebook_uuid",),
+        write=True,
+        confirm=True,
+        help="Delete a Jupyter server asynchronously.",
+    ),
+    "users_list": Op(
+        "GET", "/public/v2/notebooks/v2/{namespace}/{notebook_name}/users",
+        path_params={"namespace": "namespace", "notebook_name": "notebook_name"},
+        required=("notebook_name",),
+        kind="list",
+        help="List users a notebook is shared with. Keyed by notebook NAME (not uuid).",
+    ),
+    "users_set": Op(
+        "PUT", "/public/v2/notebooks/v2/{namespace}/{notebook_name}/users",
+        path_params={"namespace": "namespace", "notebook_name": "notebook_name"},
+        body_field="body",
+        required=("notebook_name", "body"),
+        write=True,
+        help="Grant notebook access to users; `body` = user access payload.",
+    ),
+    "users_revoke": Op(
+        "DELETE", "/public/v2/notebooks/v2/{namespace}/{notebook_name}/users",
+        path_params={"namespace": "namespace", "notebook_name": "notebook_name"},
+        body_field="body",
+        required=("notebook_name", "body"),
+        write=True,
+        confirm=True,
+        help="Revoke notebook access from users; `body` = user access payload.",
+    ),
+    "pause": Op(
+        "POST", "/public/v2/notebooks/v2/notebook/{notebook_uuid}/pause",
+        path_params={"notebook_uuid": "notebook_uuid"},
+        required=("notebook_uuid",),
+        write=True,
+        help="Pause a Jupyter server asynchronously.",
+    ),
+    "resume": Op(
+        "POST", "/public/v2/notebooks/v1/{namespace}/notebook/{notebook_uuid}/resume",
+        path_params={"namespace": "namespace", "notebook_uuid": "notebook_uuid"},
+        body_field="body",
+        required=("namespace", "notebook_uuid", "body"),
+        write=True,
+        help="Resume a Jupyter server; `body` = ResumeNotebookPayload.",
+    ),
+    "modify": Op(
+        "POST", "/public/v2/notebooks/v2/notebook/{notebook_uuid}/modify",
+        path_params={"notebook_uuid": "notebook_uuid"},
+        body_field="body",
+        required=("notebook_uuid", "body"),
+        write=True,
+        help="Modify a Jupyter server; `body` = ModifyNotebookV2Request.",
+    ),
+    "autoshutdown_set": Op(
+        "POST", "/public/v2/notebooks/v2/autoshutdown-rules/workspace/{workspace_id}",
+        path_params={"workspace_id": "workspace_id"},
+        body_field="body",
+        required=("workspace_id", "body"),
+        write=True,
+        help="Create the workspace-level autoshutdown rule; `body` = AutoShutdownV2Config.",
+    ),
+    "autoshutdown_delete": Op(
+        "DELETE", "/public/v2/notebooks/v1/autoshutdown-rules/workspace/{workspace_id}",
+        path_params={"workspace_id": "workspace_id"},
+        required=("workspace_id",),
+        write=True,
+        confirm=True,
+        help="Delete the workspace-level autoshutdown rule.",
+    ),
+}
+
+PARAMS: list[Param] = [
+    Param("notebook_uuid", str, "UUID of the Jupyter server (uid in list results), or use resource_ref."),
+    Param("notebook_name", str, "Notebook NAME (users_* actions are keyed by name, "
+          "not by uuid)."),
+    Param("namespace", str, "Namespace the notebook belongs to (create/resume). "
+          "Optional: auto-filled from the selected target if omitted."),
+    Param("workspace_id", str, "Workspace UUID for autoshutdown rules (path parameter). "
+          "Optional: auto-filled from the selected target if omitted."),
+    Param(
+        "notebook_type",
+        Literal["Spark", "Default"],
+        "List filter: notebook type.",
+    ),
+    Param(
+        "order_by",
+        Literal["name", "creationTimestamp"],
+        "List ordering field.",
+    ),
+    Param("desc", bool, "List: sort descending."),
+    Param("limit", int, "List: max results to return."),
+    Param("offset", int, "List: pagination offset."),
+    Param("search", str, "List: search by name."),
+    Param("status", str, "List: filter by status (comma-separated)."),
+    Param("access_mode", str, "List: filter by access mode (comma-separated)."),
+    Param(
+        "body",
+        dict,
+        "JSON body. create=CreateNotebookPayload {name, image, instance_type, ...} — "
+        "`image` is an OBJECT {name, tag, type}, where type is datahub|custom; a bare "
+        "\"repo:tag\" string is rejected with 422. "
+        "resume=ResumeNotebookPayload {region, instance_type} — the body is REQUIRED "
+        "even though both fields have defaults, and the API is known to ignore "
+        "`instance_type`, resuming on the notebook's previous configuration while "
+        "answering 200: check `instanceType` in the response before assuming the "
+        "requested size, or a paid GPU may come back up instead of a free one. "
+        "modify=ModifyNotebookV2Request; autoshutdown_set=AutoShutdownV2Config "
+        "{by_timer | by_load | by_schedule}.",
+    ),
+]
+
+DOMAIN = DomainTool(
+    domain="notebooks",
+    name="mlspace_notebooks",
+    title="MLSpace Notebooks",
+    summary="Manage MLSpace Jupyter Server notebooks (create, inspect, pause, resume, "
+    "modify, delete, autoshutdown rules).",
+    actions=ACTIONS,
+    params=PARAMS,
+)