dayhoff-tools 1.10.0__py3-none-any.whl → 1.10.1__py3-none-any.whl

dayhoff_tools/cli/engine/resize.py

@@ -0,0 +1,195 @@
+"""Engine resize command."""
+
+import time
+
+import boto3
+import typer
+from botocore.exceptions import ClientError
+from rich.prompt import Confirm
+
+from ..engine_studio_utils.api_utils import make_api_request
+from ..engine_studio_utils.aws_utils import check_aws_sso
+from ..engine_studio_utils.constants import console
+from ..engine_studio_utils.formatting import resolve_engine
+
+
+def resize_engine(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    size: int = typer.Option(..., "--size", "-s", help="New size in GB"),
+    online: bool = typer.Option(
+        False,
+        "--online",
+        help="Resize while running (requires manual filesystem expansion)",
+    ),
+    force: bool = typer.Option(
+        False, "--force", "-f", help="Force resize and detach all studios"
+    ),
+):
+    """Resize an engine's boot disk."""
+    check_aws_sso()
+
+    # Get all engines to resolve name
+    response = make_api_request("GET", "/engines")
+    if response.status_code != 200:
+        console.print("[red]❌ Failed to fetch engines[/red]")
+        raise typer.Exit(1)
+
+    engines = response.json().get("engines", [])
+    engine = resolve_engine(name_or_id, engines)
+
+    # Get current volume info to validate size
+    ec2 = boto3.client("ec2", region_name="us-east-1")
+
+    try:
+        # Get instance details to find root volume
+        instance_info = ec2.describe_instances(InstanceIds=[engine["instance_id"]])
+        instance = instance_info["Reservations"][0]["Instances"][0]
+
+        # Find root volume
+        root_device = instance.get("RootDeviceName", "/dev/xvda")
+        root_volume_id = None
+
+        for bdm in instance.get("BlockDeviceMappings", []):
+            if bdm["DeviceName"] == root_device:
+                root_volume_id = bdm["Ebs"]["VolumeId"]
+                break
+
+        if not root_volume_id:
+            console.print("[red]❌ Could not find root volume[/red]")
+            raise typer.Exit(1)
+
+        # Get current volume size
+        volumes = ec2.describe_volumes(VolumeIds=[root_volume_id])
+        current_size = volumes["Volumes"][0]["Size"]
+
+        if size <= current_size:
+            console.print(
+                f"[red]❌ New size ({size}GB) must be larger than current size ({current_size}GB)[/red]"
+            )
+            raise typer.Exit(1)
+
+        console.print(
+            f"[yellow]Resizing engine boot disk from {current_size}GB to {size}GB[/yellow]"
+        )
+
+        # Check if we need to stop the instance
+        if not online and engine["state"].lower() == "running":
+            console.print("Stopping engine for offline resize...")
+            stop_response = make_api_request(
+                "POST",
+                f"/engines/{engine['instance_id']}/stop",
+                json_data={"detach_studios": False},
+            )
+            if stop_response.status_code != 200:
+                console.print("[red]❌ Failed to stop engine[/red]")
+                raise typer.Exit(1)
+
+            # Wait for instance to stop
+            console.print("Waiting for engine to stop...")
+            waiter = ec2.get_waiter("instance_stopped")
+            waiter.wait(InstanceIds=[engine["instance_id"]])
+            console.print("[green]✓ Engine stopped[/green]")
+
+        # Call the resize API
+        console.print("Resizing volume...")
+        resize_response = make_api_request(
+            "POST",
+            f"/engines/{engine['instance_id']}/resize",
+            json_data={"size": size, "detach_studios": force},
+        )
+
+        if resize_response.status_code == 409 and not force:
+            # Engine has attached studios
+            data = resize_response.json()
+            attached_studios = data.get("attached_studios", [])
+
+            console.print("\n[yellow]⚠️  This engine has attached studios:[/yellow]")
+            for studio in attached_studios:
+                console.print(f"  • {studio['user']} ({studio['studio_id']})")
+
+            if Confirm.ask("\nDetach all studios and resize the engine?"):
+                resize_response = make_api_request(
+                    "POST",
+                    f"/engines/{engine['instance_id']}/resize",
+                    json_data={"size": size, "detach_studios": True},
+                )
+            else:
+                console.print("Resize cancelled.")
+                return
+
+        if resize_response.status_code != 200:
+            error = resize_response.json().get("error", "Unknown error")
+            console.print(f"[red]❌ Failed to resize engine: {error}[/red]")
+            raise typer.Exit(1)
+
+        # Check if studios were detached
+        data = resize_response.json()
+        detached_studios = data.get("detached_studios", 0)
+        if detached_studios > 0:
+            console.print(
+                f"[green]✓ Detached {detached_studios} studio(s) before resize[/green]"
+            )
+
+        # Wait for modification to complete
+        console.print("Waiting for volume modification to complete...")
+        while True:
+            mod_state = ec2.describe_volumes_modifications(VolumeIds=[root_volume_id])
+            if not mod_state["VolumesModifications"]:
+                break  # Modification complete
+
+            modification = mod_state["VolumesModifications"][0]
+            state = modification["ModificationState"]
+            progress = modification.get("Progress", 0)
+
+            # Show progress updates only for the resize phase
+            if state == "modifying":
+                console.print(f"[yellow]Progress: {progress}%[/yellow]")
+
+            # Exit as soon as optimization starts (resize is complete)
+            if state == "optimizing":
+                console.print("[green]✓ Volume resized successfully[/green]")
+                console.print(
+                    "[dim]AWS is optimizing the volume in the background (no action needed).[/dim]"
+                )
+                break
+
+            if state == "completed":
+                console.print("[green]✓ Volume resized successfully[/green]")
+                break
+            elif state == "failed":
+                console.print("[red]❌ Volume modification failed[/red]")
+                raise typer.Exit(1)
+
+            time.sleep(2)  # Check more frequently for better UX
+
+        # If offline resize, start the instance back up
+        if not online and engine["state"].lower() == "running":
+            console.print("Starting engine back up...")
+            start_response = make_api_request(
+                "POST", f"/engines/{engine['instance_id']}/start"
+            )
+            if start_response.status_code != 200:
+                console.print(
+                    "[yellow]⚠️  Failed to restart engine automatically[/yellow]"
+                )
+                console.print(
+                    f"Please start it manually: [cyan]dh engine start {engine['name']}[/cyan]"
+                )
+            else:
+                console.print("[green]✓ Engine started[/green]")
+                console.print("The filesystem will be automatically expanded on boot.")
+
+        elif online and engine["state"].lower() == "running":
+            console.print(
+                "\n[yellow]⚠️  Online resize complete. You must now expand the filesystem:[/yellow]"
+            )
+            console.print(f"1. SSH into the engine: [cyan]ssh {engine['name']}[/cyan]")
+            console.print("2. Find the root device: [cyan]lsblk[/cyan]")
+            console.print(
+                "3. Expand the partition: [cyan]sudo growpart /dev/nvme0n1 1[/cyan] (adjust device name as needed)"
+            )
+            console.print("4. Expand the filesystem: [cyan]sudo xfs_growfs /[/cyan]")
+
+    except ClientError as e:
+        console.print(f"[red]❌ Failed to resize engine: {e}[/red]")
+        raise typer.Exit(1)

dayhoff_tools/cli/engine/ssh.py

@@ -0,0 +1,62 @@
+"""Engine SSH command."""
+
+import subprocess
+
+import typer
+
+from ..engine_studio_utils.api_utils import make_api_request
+from ..engine_studio_utils.aws_utils import check_aws_sso
+from ..engine_studio_utils.constants import console
+from ..engine_studio_utils.formatting import resolve_engine
+from ..engine_studio_utils.ssh_utils import check_session_manager_plugin, update_ssh_config_entry
+
+
+def ssh_engine(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    admin: bool = typer.Option(
+        False, "--admin", help="Connect as ec2-user instead of the engine owner user"
+    ),
+    idle_timeout: int = typer.Option(
+        600,
+        "--idle-timeout",
+        help="Idle timeout (seconds) for the SSM port-forward (0 = disable)",
+    ),
+):
+    """Connect to an engine via SSH.
+
+    By default the CLI connects using the engine's owner username (the same one stored in the `User` tag).
+    Pass `--admin` to connect with the underlying [`ec2-user`] account for break-glass or debugging.
+    """
+    username = check_aws_sso()
+
+    # Check for Session Manager Plugin
+    if not check_session_manager_plugin():
+        raise typer.Exit(1)
+
+    # Get all engines to resolve name
+    response = make_api_request("GET", "/engines")
+    if response.status_code != 200:
+        console.print("[red]❌ Failed to fetch engines[/red]")
+        raise typer.Exit(1)
+
+    engines = response.json().get("engines", [])
+    engine = resolve_engine(name_or_id, engines)
+
+    if engine["state"].lower() != "running":
+        console.print(f"[red]❌ Engine is not running (state: {engine['state']})[/red]")
+        raise typer.Exit(1)
+
+    # Choose SSH user
+    ssh_user = "ec2-user" if admin else username
+
+    # Update SSH config
+    console.print(
+        f"Updating SSH config for [cyan]{engine['name']}[/cyan] (user: {ssh_user})..."
+    )
+    update_ssh_config_entry(
+        engine["name"], engine["instance_id"], ssh_user, idle_timeout
+    )
+
+    # Connect
+    console.print(f"[green]✓ Connecting to {engine['name']}...[/green]")
+    subprocess.run(["ssh", engine["name"]])

dayhoff_tools/cli/engine/{engine_core.py → status.py}

@@ -1,4 +1,4 @@
-"""Core engine commands: launch, list, and status."""
+"""Engine status command."""
 
 import json
 import time
@@ -7,214 +7,19 @@ from typing import Any, Dict, Optional
 
 import boto3
 import typer
-from rich import box
 from rich.panel import Panel
-from rich.progress import Progress, SpinnerColumn, TextColumn
-from rich.table import Table
-
-from .shared import (
-    HOURLY_COSTS,
-    _fetch_init_stages,
-    check_aws_sso,
-    console,
+
+from ..engine_studio_utils.api_utils import make_api_request
+from ..engine_studio_utils.aws_utils import _fetch_init_stages, check_aws_sso
+from ..engine_studio_utils.constants import HOURLY_COSTS, console
+from ..engine_studio_utils.formatting import (
     format_duration,
-    format_status,
     get_disk_usage_via_ssm,
-    make_api_request,
     parse_launch_time,
     resolve_engine,
 )
 
 
-def launch_engine(
-    name: str = typer.Argument(help="Name for the new engine"),
-    engine_type: str = typer.Option(
-        "cpu",
-        "--type",
-        "-t",
-        help="Engine type: cpu, cpumax, t4, a10g, a100, 4_t4, 8_t4, 4_a10g, 8_a10g",
-    ),
-    user: Optional[str] = typer.Option(None, "--user", "-u", help="Override username"),
-    boot_disk_size: Optional[int] = typer.Option(
-        None,
-        "--size",
-        "-s",
-        help="Boot disk size in GB (default: 50GB, min: 20GB, max: 1000GB)",
-    ),
-    availability_zone: Optional[str] = typer.Option(
-        None,
-        "--az",
-        help="Prefer a specific Availability Zone (e.g., us-east-1b). If omitted the service will try all public subnets.",
-    ),
-):
-    """Launch a new engine instance."""
-    username = check_aws_sso()
-    if user:
-        username = user
-
-    # Validate engine type
-    valid_types = [
-        "cpu",
-        "cpumax",
-        "t4",
-        "a10g",
-        "a100",
-        "4_t4",
-        "8_t4",
-        "4_a10g",
-        "8_a10g",
-    ]
-    if engine_type not in valid_types:
-        console.print(f"[red]❌ Invalid engine type: {engine_type}[/red]")
-        console.print(f"Valid types: {', '.join(valid_types)}")
-        raise typer.Exit(1)
-
-    # Validate boot disk size
-    if boot_disk_size is not None:
-        if boot_disk_size < 20:
-            console.print("[red]❌ Boot disk size must be at least 20GB[/red]")
-            raise typer.Exit(1)
-        if boot_disk_size > 1000:
-            console.print("[red]❌ Boot disk size cannot exceed 1000GB[/red]")
-            raise typer.Exit(1)
-
-    cost = HOURLY_COSTS.get(engine_type, 0)
-    disk_info = f" with {boot_disk_size}GB boot disk" if boot_disk_size else ""
-    console.print(
-        f"Launching [cyan]{name}[/cyan] ({engine_type}){disk_info} for ${cost:.2f}/hour..."
-    )
-
-    with Progress(
-        SpinnerColumn(),
-        TextColumn("[progress.description]{task.description}"),
-        transient=True,
-    ) as progress:
-        progress.add_task("Creating engine...", total=None)
-
-        request_data: Dict[str, Any] = {
-            "name": name,
-            "user": username,
-            "engine_type": engine_type,
-        }
-        if boot_disk_size is not None:
-            request_data["boot_disk_size"] = boot_disk_size
-        if availability_zone:
-            request_data["availability_zone"] = availability_zone
-
-        response = make_api_request("POST", "/engines", json_data=request_data)
-
-    if response.status_code == 201:
-        data = response.json()
-        console.print(f"[green]✓ Engine launched successfully![/green]")
-        console.print(f"Instance ID: [cyan]{data['instance_id']}[/cyan]")
-        console.print(f"Type: {data['instance_type']} (${cost:.2f}/hour)")
-        if boot_disk_size:
-            console.print(f"Boot disk: {boot_disk_size}GB")
-        console.print("\nThe engine is initializing. This may take a few minutes.")
-        console.print(f"Check status with: [cyan]dh engine status {name}[/cyan]")
-    else:
-        error = response.json().get("error", "Unknown error")
-        console.print(f"[red]❌ Failed to launch engine: {error}[/red]")
-
-
-def list_engines(
-    user: Optional[str] = typer.Option(None, "--user", "-u", help="Filter by user"),
-    running_only: bool = typer.Option(
-        False, "--running", help="Show only running engines"
-    ),
-    stopped_only: bool = typer.Option(
-        False, "--stopped", help="Show only stopped engines"
-    ),
-    detailed: bool = typer.Option(
-        False, "--detailed", "-d", help="Show detailed status (slower)"
-    ),
-):
-    """List engines (shows all engines by default)."""
-    current_user = check_aws_sso()
-
-    params = {}
-    if user:
-        params["user"] = user
-    if detailed:
-        params["check_ready"] = "true"
-
-    response = make_api_request("GET", "/engines", params=params)
-
-    if response.status_code == 200:
-        data = response.json()
-        engines = data.get("engines", [])
-
-        # Filter by state if requested
-        if running_only:
-            engines = [e for e in engines if e["state"].lower() == "running"]
-        elif stopped_only:
-            engines = [e for e in engines if e["state"].lower() == "stopped"]
-
-        if not engines:
-            console.print("No engines found.")
-            return
-
-        # Only fetch detailed info if requested (slow)
-        stages_map = {}
-        if detailed:
-            stages_map = _fetch_init_stages([e["instance_id"] for e in engines])
-
-        # Create table
-        table = Table(title="Engines", box=box.ROUNDED)
-        table.add_column("Name", style="cyan")
-        table.add_column("Instance ID", style="dim")
-        table.add_column("Type")
-        table.add_column("User")
-        table.add_column("Status")
-        if detailed:
-            table.add_column("Disk Usage")
-        table.add_column("Uptime/Since")
-        table.add_column("$/hour", justify="right")
-
-        for engine in engines:
-            launch_time = parse_launch_time(engine["launch_time"])
-            uptime = datetime.now(timezone.utc) - launch_time
-            hourly_cost = HOURLY_COSTS.get(engine["engine_type"], 0)
-
-            if engine["state"].lower() == "running":
-                time_str = format_duration(uptime)
-                # Only get disk usage if detailed mode
-                if detailed:
-                    disk_usage = get_disk_usage_via_ssm(engine["instance_id"]) or "-"
-                else:
-                    disk_usage = None
-            else:
-                time_str = launch_time.strftime("%Y-%m-%d %H:%M")
-                disk_usage = "-" if detailed else None
-
-            row_data = [
-                engine["name"],
-                engine["instance_id"],
-                engine["engine_type"],
-                engine["user"],
-                format_status(engine["state"], engine.get("ready")),
-            ]
-            if detailed:
-                row_data.append(disk_usage)
-            row_data.extend(
-                [
-                    time_str,
-                    f"${hourly_cost:.2f}",
-                ]
-            )
-
-            table.add_row(*row_data)
-
-        console.print(table)
-        if not detailed and any(e["state"].lower() == "running" for e in engines):
-            console.print(
-                "\n[dim]Tip: Use --detailed to see disk usage and bootstrap status (slower)[/dim]"
-            )
-    else:
-        error = response.json().get("error", "Unknown error")
-        console.print(f"[red]❌ Failed to list engines: {error}[/red]")
-
-
 def engine_status(
     name_or_id: str = typer.Argument(help="Engine name or instance ID"),
     detailed: bool = typer.Option(