Add beta cluster remediation CLI commands

src/together/lib/cli/__init__.py

@@ -61,9 +61,11 @@
     FINE_TUNING_DOWNLOAD_HELP_EXAMPLES,
     BETA_CLUSTERS_STORAGE_HELP_EXAMPLES,
     FILES_RETRIEVE_CONTENT_HELP_EXAMPLES,
+    BETA_CLUSTERS_REMEDIATIONS_HELP_EXAMPLES,
     BETA_CLUSTERS_STORAGE_CREATE_HELP_EXAMPLES,
     BETA_CLUSTERS_STORAGE_UPDATE_HELP_EXAMPLES,
     BETA_CLUSTERS_GET_CREDENTIALS_HELP_EXAMPLES,
+    BETA_CLUSTERS_REMEDIATIONS_CREATE_HELP_EXAMPLES,
 )
 from together.lib.cli.utils._help_formatter import help_formatter
 from together.lib.cli.utils._preparse_tokens import preparse_tokens
@@ -486,6 +488,44 @@ async def run_command() -> None:
 )
 storage_app.command((f"{_CLI}.beta.clusters.storage.delete:delete"), help="Delete a storage volume", alias="-d")
 
+### Clusters > Remediations API commands
+remediations_app = clusters_app.command(
+    App(
+        name="remediations",
+        help="Manage node remediations",
+        group="Subcommands",
+        help_epilogue=BETA_CLUSTERS_REMEDIATIONS_HELP_EXAMPLES,
+    )
+)
+remediations_app.command(
+    (f"{_CLI}.beta.clusters.remediations.create:create"),
+    alias="-c",
+    help="Create a node remediation",
+    help_epilogue=BETA_CLUSTERS_REMEDIATIONS_CREATE_HELP_EXAMPLES,
+)
+remediations_app.command(
+    (f"{_CLI}.beta.clusters.remediations.list:list"),
+    alias="ls",
+    help="List node remediations",
+)
+remediations_app.command(
+    (f"{_CLI}.beta.clusters.remediations.retrieve:retrieve"),
+    alias="get",
+    help="Get remediation details",
+)
+remediations_app.command(
+    (f"{_CLI}.beta.clusters.remediations.approve:approve"),
+    help="Approve a pending remediation",
+)
+remediations_app.command(
+    (f"{_CLI}.beta.clusters.remediations.cancel:cancel"),
+    help="Cancel a pending remediation",
+)
+remediations_app.command(
+    (f"{_CLI}.beta.clusters.remediations.reject:reject"),
+    help="Reject a pending remediation",
+)
+
 ### Jig commands
 jig_app = beta_app.command(
     App(name="jig", help="Build, deploy, and manage custom containers", help_epilogue=JIG_HELP_EXAMPLES)

src/together/lib/cli/api/beta/clusters/remediations/_resolve_remediation.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import sys
+
+from together import omit
+from together._types import Omit
+from together.lib.cli.utils.config import CLIConfigParameter
+from together.lib.cli.utils._console import console
+from together.types.beta.clusters.remediation import Remediation
+
+
+async def resolve_remediation(config: CLIConfigParameter, remediation_id: str) -> Remediation:
+    clusters = await config.client.beta.clusters.list()
+
+    for cluster in clusters.clusters:
+        page_token: str | Omit = omit
+        while True:
+            response = await config.client.beta.clusters.remediations.list(
+                "-",
+                cluster_id=cluster.cluster_id,
+                page_size=100,
+                page_token=page_token,
+            )
+            for remediation in response.remediations:
+                if remediation.id == remediation_id:
+                    return remediation
+
+            if not response.has_next or not response.next_page_token:
+                break
+            page_token = response.next_page_token
+
+    console.print(f"[red]Error:[/red] Remediation not found: {remediation_id}")
+    sys.exit(1)

src/together/lib/cli/api/beta/clusters/remediations/approve.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import Optional, Annotated
+
+from cyclopts import Parameter
+
+from together import omit
+from together._utils._json import openapi_dumps
+from together.lib.cli.utils.config import CLIConfigParameter
+from together.lib.cli.utils._console import console
+from together.lib.cli.components.loader import show_loading_status
+from together.lib.cli.api.beta.clusters.remediations._resolve_remediation import resolve_remediation
+
+
+async def approve(
+    remediation_id: str,
+    comment: Annotated[Optional[str], Parameter(help="Comment explaining the approval")] = None,
+    *,
+    config: CLIConfigParameter,
+) -> None:
+    """Approve a pending remediation."""
+    remediation = await show_loading_status("Finding remediation...", resolve_remediation(config, remediation_id))
+    response = await show_loading_status(
+        "Approving remediation...",
+        config.client.beta.clusters.remediations.approve(
+            remediation_id,
+            cluster_id=remediation.cluster_id,
+            instance_id=remediation.instance_id,
+            comment=comment or omit,
+        ),
+    )
+
+    if config.json:
+        console.print_json(openapi_dumps(response).decode("utf-8"))
+        return
+
+    console.print(f"[blue]Remediation approved.[/blue] ({response.id})")

src/together/lib/cli/api/beta/clusters/remediations/cancel.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from together._utils._json import openapi_dumps
+from together.lib.cli.utils.config import CLIConfigParameter
+from together.lib.cli.utils._console import console
+from together.lib.cli.components.loader import show_loading_status
+from together.lib.cli.api.beta.clusters.remediations._resolve_remediation import resolve_remediation
+
+
+async def cancel(
+    remediation_id: str,
+    *,
+    config: CLIConfigParameter,
+) -> None:
+    """Cancel a pending remediation."""
+    remediation = await show_loading_status("Finding remediation...", resolve_remediation(config, remediation_id))
+    response = await show_loading_status(
+        "Cancelling remediation...",
+        config.client.beta.clusters.remediations.cancel(
+            remediation_id,
+            cluster_id=remediation.cluster_id,
+            instance_id=remediation.instance_id,
+        ),
+    )
+
+    if config.json:
+        console.print_json(openapi_dumps(response).decode("utf-8"))
+        return
+
+    console.print(f"[blue]Remediation cancelled.[/blue] ({response.id})")

src/together/lib/cli/api/beta/clusters/remediations/create.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from typing import Literal, Optional, Annotated, cast
+
+from cyclopts import Parameter
+
+from together import omit
+from together._utils._json import openapi_dumps
+from together.lib.cli.utils.config import CLIConfigParameter
+from together.lib.cli.utils._console import console
+from together.lib.cli.components.loader import show_loading_status
+
+RemediationModeParameter = Annotated[
+    Literal[
+        "VM_ONLY",
+        "HOST_AWARE",
+        "EVICT_WITHOUT_REPLACEMENT",
+        "REBOOT_VM",
+    ],
+    Parameter(help="The type of remediation to perform"),
+]
+
+
+async def create(
+    cluster_id: Annotated[str, Parameter(help="The ID of the cluster")],
+    instance_id: Annotated[str, Parameter(help="The ID of the node within the cluster to remediate")],
+    *,
+    mode: RemediationModeParameter,
+    remediation_id: Annotated[Optional[str], Parameter(help="Client-specified ID for idempotency")] = None,
+    reason: Annotated[Optional[str], Parameter(help="Reason for the remediation")] = None,
+    config: CLIConfigParameter,
+) -> None:
+    """Create a node remediation for an instance."""
+    safe_mode = cast(
+        Literal[
+            "REMEDIATION_MODE_VM_ONLY",
+            "REMEDIATION_MODE_HOST_AWARE",
+            "REMEDIATION_MODE_EVICT_WITHOUT_REPLACEMENT",
+            "REMEDIATION_MODE_REBOOT_VM",
+        ],
+        f"REMEDIATION_MODE_{mode}",
+    )
+
+    response = await show_loading_status(
+        "Creating remediation...",
+        config.client.beta.clusters.remediations.create(
+            instance_id,
+            cluster_id=cluster_id,
+            mode=safe_mode,
+            remediation_id=remediation_id or omit,
+            reason=reason or omit,
+        ),
+    )
+
+    if config.json:
+        console.print_json(openapi_dumps(response).decode("utf-8"))
+        return
+
+    console.print(f"[green]√ Remediation created[/green] [dim]({response.id})[/dim]")
+    console.print(f"  Remediations may take some time to complete.\n")
+    console.print(f"  To retrieve the status:")
+    console.print(f"    [dim]-[/dim] [primary]tg beta clusters remediations {response.id}[/primary]")

src/together/lib/cli/api/beta/clusters/remediations/list.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from typing import Optional, Annotated
+
+from cyclopts import Parameter
+
+from together import omit
+from together._utils._json import openapi_dumps
+from together.lib.utils.tools import format_datetime
+from together.lib.cli.utils.config import CLIConfigParameter
+from together.lib.cli.utils._console import console
+from together.lib.cli.components.list import ListTable
+from together.lib.cli.components.loader import show_loading_status
+
+
+async def list(
+    cluster_id: str,
+    instance_id: Annotated[Optional[str], Parameter(help="Instance ID to list remediations for")] = None,
+    after: Annotated[Optional[str], Parameter(help="Pagination token from a previous request")] = None,
+    *,
+    config: CLIConfigParameter,
+) -> None:
+    """List node remediations for a cluster or instance."""
+    response = await show_loading_status(
+        "Loading remediations...",
+        config.client.beta.clusters.remediations.list(
+            instance_id or "-",
+            cluster_id=cluster_id,
+            page_token=after or omit,
+        ),
+    )
+
+    if config.json:
+        console.print_json(openapi_dumps(response).decode("utf-8"))
+        return
+
+    table = ListTable(title="Cluster Remediations", empty_message="No remediations found for this cluster.")
+    table.add_column("Created")
+    table.add_primary_column("Instance", ratio=3)
+    table.add_column("Mode")
+    table.add_column("State")
+    table.add_column("Remediation ID", ratio=3)
+
+    for remediation in response.remediations:
+        table.add_row(
+            format_datetime(remediation.create_time) if remediation.create_time else "-",
+            remediation.instance_id,
+            remediation.mode.replace("REMEDIATION_MODE_", ""),
+            _colorize(remediation.state),
+            remediation.id,
+        )
+
+    console.print(table)
+    if response.has_next and response.next_page_token:
+        command = f"tg beta clusters remediations ls {cluster_id}"
+        if instance_id:
+            command += f" {instance_id}"
+        console.print("\n[blue dim]To display the next page, run:[/blue dim]")
+        console.print(f"  [dim]-[/dim] [white]{command} --after {response.next_page_token}[/white]")
+
+
+def _colorize(state: str) -> str:
+    state_colors = {
+        "PENDING_APPROVAL": "yellow",
+        "PENDING": "yellow",
+        "RUNNING": "yellow",
+        "SUCCEEDED": "green",
+        "FAILED": "red",
+        "CANCELLED": "dim",
+        "AUTO_RESOLVED": "green",
+    }
+    color = state_colors[state] if state in state_colors else "white"
+    return f"[{color}]{state}[/{color}]"

src/together/lib/cli/api/beta/clusters/remediations/reject.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import Optional, Annotated
+
+from cyclopts import Parameter
+
+from together import omit
+from together._utils._json import openapi_dumps
+from together.lib.cli.utils.config import CLIConfigParameter
+from together.lib.cli.utils._console import console
+from together.lib.cli.components.loader import show_loading_status
+from together.lib.cli.api.beta.clusters.remediations._resolve_remediation import resolve_remediation
+
+
+async def reject(
+    remediation_id: str,
+    comment: Annotated[Optional[str], Parameter(help="Comment explaining the rejection")] = None,
+    *,
+    config: CLIConfigParameter,
+) -> None:
+    """Reject a pending remediation."""
+    remediation = await show_loading_status("Finding remediation...", resolve_remediation(config, remediation_id))
+    response = await show_loading_status(
+        "Rejecting remediation...",
+        config.client.beta.clusters.remediations.reject(
+            remediation_id,
+            cluster_id=remediation.cluster_id,
+            instance_id=remediation.instance_id,
+            comment=comment or omit,
+        ),
+    )
+
+    if config.json:
+        console.print_json(openapi_dumps(response).decode("utf-8"))
+        return
+
+    console.print(f"[blue]Remediation rejected.[/blue] ({response.id})")

src/together/lib/cli/api/beta/clusters/remediations/retrieve.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from together._utils._json import openapi_dumps
+from together.lib.cli.utils.config import CLIConfigParameter
+from together.lib.cli.utils._console import console
+from together.lib.cli.components.loader import show_loading_status
+from together.lib.cli.components.model_dump import print_model_dump
+from together.lib.cli.api.beta.clusters.remediations._resolve_remediation import resolve_remediation
+
+
+async def retrieve(
+    remediation_id: str,
+    *,
+    config: CLIConfigParameter,
+) -> None:
+    """Retrieve remediation details."""
+    remediation = await show_loading_status("Finding remediation...", resolve_remediation(config, remediation_id))
+    response = await show_loading_status(
+        "Retrieving remediation...",
+        config.client.beta.clusters.remediations.retrieve(
+            remediation_id,
+            cluster_id=remediation.cluster_id,
+            instance_id=remediation.instance_id,
+        ),
+    )
+
+    if config.json:
+        console.print_json(openapi_dumps(response).decode("utf-8"))
+        return
+
+    print_model_dump(response, show_nulls=False, only_set_fields=True)