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

dayhoff_tools/cli/engine_studio_utils/formatting.py

@@ -0,0 +1,210 @@
+"""Display formatting utilities for engine and studio commands."""
+
+import time
+from datetime import datetime, timedelta, timezone
+from typing import Dict, List, Optional
+
+import boto3
+import typer
+from rich.prompt import IntPrompt
+
+from .constants import HOURLY_COSTS, console
+
+
+def format_duration(duration: timedelta) -> str:
+    """Format a duration as a human-readable string."""
+    total_seconds = int(duration.total_seconds())
+    hours = total_seconds // 3600
+    minutes = (total_seconds % 3600) // 60
+
+    if hours > 0:
+        return f"{hours}h {minutes}m"
+    else:
+        return f"{minutes}m"
+
+
+def parse_launch_time(launch_time_str: str) -> datetime:
+    """Parse launch time from API response."""
+    # Try different datetime formats
+    formats = [
+        "%Y-%m-%dT%H:%M:%S.%fZ",
+        "%Y-%m-%dT%H:%M:%SZ",
+        "%Y-%m-%dT%H:%M:%S%z",  # ISO format with timezone
+        "%Y-%m-%dT%H:%M:%S+00:00",  # Explicit UTC offset
+        "%Y-%m-%d %H:%M:%S",
+    ]
+
+    # First try parsing with fromisoformat for better timezone handling
+    try:
+        # Handle the ISO format properly
+        return datetime.fromisoformat(launch_time_str.replace("Z", "+00:00"))
+    except (ValueError, AttributeError):
+        pass
+
+    # Fallback to manual format parsing
+    for fmt in formats:
+        try:
+            parsed = datetime.strptime(launch_time_str, fmt)
+            # If no timezone info, assume UTC
+            if parsed.tzinfo is None:
+                parsed = parsed.replace(tzinfo=timezone.utc)
+            return parsed
+        except ValueError:
+            continue
+
+    # Fallback: assume it's recent
+    return datetime.now(timezone.utc)
+
+
+def format_status(state: str, ready: Optional[bool]) -> str:
+    """Format engine status with ready indicator."""
+    if state.lower() == "running":
+        if ready is True:
+            return "[green]Running ✓[/green]"
+        elif ready is False:
+            return "[yellow]Running ⚠ (Bootstrapping...)[/yellow]"
+        else:
+            return "[green]Running[/green]"
+    elif state.lower() == "stopped":
+        return "[dim]Stopped[/dim]"
+    elif state.lower() == "stopping":
+        return "[yellow]Stopping...[/yellow]"
+    elif state.lower() == "pending":
+        return "[yellow]Starting...[/yellow]"
+    else:
+        return state
+
+
+def resolve_engine(name_or_id: str, engines: List[Dict]) -> Dict:
+    """Resolve engine by name or ID with interactive selection."""
+    # Exact ID match
+    exact_id = [e for e in engines if e["instance_id"] == name_or_id]
+    if exact_id:
+        return exact_id[0]
+
+    # Exact name match
+    exact_name = [e for e in engines if e["name"] == name_or_id]
+    if len(exact_name) == 1:
+        return exact_name[0]
+
+    # Prefix matches
+    matches = [
+        e
+        for e in engines
+        if e["name"].startswith(name_or_id) or e["instance_id"].startswith(name_or_id)
+    ]
+
+    if len(matches) == 0:
+        console.print(f"[red]❌ No engine found matching '{name_or_id}'[/red]")
+        raise typer.Exit(1)
+    elif len(matches) == 1:
+        return matches[0]
+    else:
+        # Interactive selection
+        console.print(f"Multiple engines match '{name_or_id}':")
+        for i, engine in enumerate(matches, 1):
+            cost = HOURLY_COSTS.get(engine["engine_type"], 0)
+            console.print(
+                f"  {i}. [cyan]{engine['name']}[/cyan] ({engine['instance_id']}) "
+                f"- {engine['engine_type']} - {engine['state']} - ${cost:.2f}/hr"
+            )
+
+        while True:
+            try:
+                choice = IntPrompt.ask(
+                    "Select engine",
+                    default=1,
+                    choices=[str(i) for i in range(1, len(matches) + 1)],
+                )
+                return matches[choice - 1]
+            except (ValueError, IndexError):
+                console.print("[red]Invalid selection, please try again[/red]")
+
+
+def get_disk_usage_via_ssm(instance_id: str) -> Optional[str]:
+    """Get disk usage for an engine via SSM.
+
+    Returns:
+        String like "17/50 GB" or None if failed
+    """
+    try:
+        ssm = boto3.client("ssm", region_name="us-east-1")
+
+        # Run df command to get disk usage
+        response = ssm.send_command(
+            InstanceIds=[instance_id],
+            DocumentName="AWS-RunShellScript",
+            Parameters={
+                "commands": [
+                    # Get root filesystem usage in GB
+                    'df -BG / | tail -1 | awk \'{gsub(/G/, "", $2); gsub(/G/, "", $3); print $3 "/" $2 " GB"}\''
+                ],
+                "executionTimeout": ["10"],
+            },
+        )
+
+        command_id = response["Command"]["CommandId"]
+
+        # Wait for command to complete (with timeout)
+        for _ in range(5):  # 5 second timeout
+            time.sleep(1)
+            result = ssm.get_command_invocation(
+                CommandId=command_id,
+                InstanceId=instance_id,
+            )
+            if result["Status"] in ["Success", "Failed"]:
+                break
+
+        if result["Status"] == "Success":
+            output = result["StandardOutputContent"].strip()
+            return output if output else None
+
+        return None
+
+    except Exception as e:
+        # logger.debug(f"Failed to get disk usage for {instance_id}: {e}") # Original code had this line commented out
+        return None
+
+
+def get_studio_disk_usage_via_ssm(instance_id: str, username: str) -> Optional[str]:
+    """Get disk usage for a studio via SSM.
+
+    Returns:
+        String like "333/500 GB" or None if failed
+    """
+    try:
+        ssm = boto3.client("ssm", region_name="us-east-1")
+
+        # Run df command to get studio disk usage
+        response = ssm.send_command(
+            InstanceIds=[instance_id],
+            DocumentName="AWS-RunShellScript",
+            Parameters={
+                "commands": [
+                    # Get studio filesystem usage in GB
+                    f'df -BG /studios/{username} 2>/dev/null | tail -1 | awk \'{{gsub(/G/, "", $2); gsub(/G/, "", $3); print $3 "/" $2 " GB"}}\''
+                ],
+                "executionTimeout": ["10"],
+            },
+        )
+
+        command_id = response["Command"]["CommandId"]
+
+        # Wait for command to complete (with timeout)
+        for _ in range(5):  # 5 second timeout
+            time.sleep(1)
+            result = ssm.get_command_invocation(
+                CommandId=command_id,
+                InstanceId=instance_id,
+            )
+            if result["Status"] in ["Success", "Failed"]:
+                break
+
+        if result["Status"] == "Success":
+            output = result["StandardOutputContent"].strip()
+            return output if output else None
+
+        return None
+
+    except Exception:
+        return None

dayhoff_tools/cli/engine_studio_utils/ssh_utils.py

@@ -0,0 +1,141 @@
+"""SSH-related utilities for engine and studio commands."""
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+from .constants import SSH_MANAGED_COMMENT
+
+
+def get_ssh_public_key() -> str:
+    """Get the user's SSH public key.
+
+    Discovery order (container-friendly):
+    1) DHT_SSH_PUBLIC_KEY env var (direct key content)
+    2) DHT_SSH_PUBLIC_KEY_PATH env var (path to a .pub file)
+    3) ssh-agent via `ssh-add -L` (requires SSH_AUTH_SOCK)
+    4) Conventional files: ~/.ssh/id_ed25519.pub, ~/.ssh/id_rsa.pub
+
+    Raises:
+        FileNotFoundError: If no public key can be discovered.
+    """
+    # 1) Direct env var content
+    env_key = os.environ.get("DHT_SSH_PUBLIC_KEY")
+    if env_key and env_key.strip():
+        return env_key.strip()
+
+    # 2) Env var path
+    env_path = os.environ.get("DHT_SSH_PUBLIC_KEY_PATH")
+    if env_path:
+        p = Path(env_path).expanduser()
+        if p.is_file():
+            try:
+                return p.read_text().strip()
+            except Exception:
+                pass
+
+    # 3) Agent lookup (ssh-add -L)
+    try:
+        if shutil.which("ssh-add") is not None:
+            proc = subprocess.run(["ssh-add", "-L"], capture_output=True, text=True)
+            if proc.returncode == 0 and proc.stdout:
+                keys = [
+                    line.strip() for line in proc.stdout.splitlines() if line.strip()
+                ]
+                # Prefer ed25519, then rsa
+                for pref in ("ssh-ed25519", "ssh-rsa", "ecdsa-sha2-nistp256"):
+                    for k in keys:
+                        if k.startswith(pref + " "):
+                            return k
+                # Fallback to first key if types not matched
+                if keys:
+                    return keys[0]
+    except Exception:
+        pass
+
+    # 4) Conventional files
+    home = Path.home()
+    key_paths = [home / ".ssh" / "id_ed25519.pub", home / ".ssh" / "id_rsa.pub"]
+    for key_path in key_paths:
+        if key_path.is_file():
+            try:
+                return key_path.read_text().strip()
+            except Exception:
+                continue
+
+    raise FileNotFoundError(
+        "No SSH public key found. Please create one with 'ssh-keygen' first."
+    )
+
+
+def check_session_manager_plugin():
+    """Check if AWS Session Manager Plugin is available and warn if not."""
+    from .constants import console
+
+    if shutil.which("session-manager-plugin") is None:
+        console.print(
+            "[bold red]⚠️  AWS Session Manager Plugin not found![/bold red]\n"
+            "SSH connections to engines require the Session Manager Plugin.\n"
+            "Please install it following the setup guide:\n"
+            "[link]https://github.com/dayhofflabs/nutshell/blob/main/REFERENCE/setup_guides/new-laptop.md[/link]"
+        )
+        return False
+    return True
+
+
+def update_ssh_config_entry(
+    engine_name: str, instance_id: str, ssh_user: str, idle_timeout: int = 600
+):
+    """Add or update a single SSH config entry for the given SSH user.
+
+    Args:
+        engine_name:  Host alias to write into ~/.ssh/config
+        instance_id:  EC2 instance-id (used by the proxy command)
+        ssh_user:     Username to place into the SSH stanza
+        idle_timeout: Idle timeout **in seconds** to pass to the SSM port-forward. 600 = 10 min.
+    """
+    config_path = Path.home() / ".ssh" / "config"
+    config_path.parent.mkdir(mode=0o700, exist_ok=True)
+
+    # Touch the file if it doesn't exist
+    if not config_path.exists():
+        config_path.touch(mode=0o600)
+
+    # Read existing config
+    content = config_path.read_text()
+    lines = content.splitlines() if content else []
+
+    # Remove any existing entry for this engine
+    new_lines = []
+    skip_until_next_host = False
+    for line in lines:
+        # Check if this is our managed host
+        if (
+            line.strip().startswith(f"Host {engine_name}")
+            and SSH_MANAGED_COMMENT in line
+        ):
+            skip_until_next_host = True
+        elif line.strip().startswith("Host ") and skip_until_next_host:
+            skip_until_next_host = False
+            # This is a different host entry, keep it
+            new_lines.append(line)
+        elif not skip_until_next_host:
+            new_lines.append(line)
+
+    # Add the new entry
+    if new_lines and new_lines[-1].strip():  # Add blank line if needed
+        new_lines.append("")
+
+    new_lines.extend(
+        [
+            f"Host {engine_name} {SSH_MANAGED_COMMENT}",
+            f"    HostName {instance_id}",
+            f"    User {ssh_user}",
+            f"    ProxyCommand sh -c \"AWS_SSM_IDLE_TIMEOUT={idle_timeout} aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'\"",
+        ]
+    )
+
+    # Write back
+    config_path.write_text("\n".join(new_lines))
+    config_path.chmod(0o600)

dayhoff_tools/cli/main.py

@@ -4,9 +4,8 @@ import sys
 from importlib.metadata import PackageNotFoundError, version
 
 import typer
-
 from dayhoff_tools.cli.cloud_commands import aws_app, gcp_app
-from dayhoff_tools.cli.engine import engine_app, studio_app
+from dayhoff_tools.cli.engine_studio_commands import engine_app, studio_app
 from dayhoff_tools.cli.utility_commands import (
     build_and_upload_wheel,
     delete_local_branch,

dayhoff_tools/cli/studio/__init__.py

@@ -0,0 +1 @@
+"""Studio management commands."""

dayhoff_tools/cli/studio/attach.py

@@ -0,0 +1,314 @@
+"""Studio attach command."""
+
+import time
+from typing import Optional
+
+import typer
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from rich.prompt import Confirm, IntPrompt
+
+from ..engine_studio_utils.api_utils import get_user_studio, 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,
+    get_ssh_public_key,
+    update_ssh_config_entry,
+)
+
+
+def attach_studio(
+    engine_name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    user: Optional[str] = typer.Option(
+        None, "--user", "-u", help="Attach a different user's studio (admin only)"
+    ),
+):
+    """Attach your studio to an engine."""
+    username = check_aws_sso()
+
+    # Check for Session Manager Plugin since we'll update SSH config
+    if not check_session_manager_plugin():
+        raise typer.Exit(1)
+
+    # Use specified user if provided, otherwise use current user
+    target_user = user if user else username
+
+    # Add confirmation when attaching another user's studio
+    if target_user != username:
+        console.print(f"[yellow]⚠️  Managing studio for user: {target_user}[/yellow]")
+        if not Confirm.ask(f"Are you sure you want to attach {target_user}'s studio?"):
+            console.print("Operation cancelled.")
+            return
+
+    # Get user's studio
+    studio = get_user_studio(target_user)
+    if not studio:
+        if target_user == username:
+            console.print("[yellow]You don't have a studio yet.[/yellow]")
+            if Confirm.ask("Would you like to create one now?"):
+                size = IntPrompt.ask("Studio size (GB)", default=50)
+                response = make_api_request(
+                    "POST",
+                    "/studios",
+                    json_data={"user": username, "size_gb": size},
+                )
+                if response.status_code != 201:
+                    console.print("[red]❌ Failed to create studio[/red]")
+                    raise typer.Exit(1)
+                studio = response.json()
+                studio["studio_id"] = studio["studio_id"]  # Normalize key
+            else:
+                raise typer.Exit(0)
+        else:
+            console.print(f"[red]❌ User {target_user} doesn't have a studio.[/red]")
+            raise typer.Exit(1)
+
+    # Check if already attached
+    if studio.get("status") == "in-use":
+        console.print(
+            f"[yellow]Studio is already attached to {studio.get('attached_vm_id')}[/yellow]"
+        )
+        if not Confirm.ask("Detach and reattach to new engine?"):
+            return
+        # Detach first
+        response = make_api_request("POST", f"/studios/{studio['studio_id']}/detach")
+        if response.status_code != 200:
+            console.print("[red]❌ Failed to detach studio[/red]")
+            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(engine_name_or_id, engines)
+
+    # Flag to track if we started the engine in this command (affects retry length)
+    engine_started_now: bool = False
+
+    if engine["state"].lower() != "running":
+        console.print(f"[yellow]⚠️  Engine is {engine['state']}[/yellow]")
+        if engine["state"].lower() == "stopped" and Confirm.ask(
+            "Start the engine first?"
+        ):
+            response = make_api_request(
+                "POST", f"/engines/{engine['instance_id']}/start"
+            )
+            if response.status_code != 200:
+                console.print("[red]❌ Failed to start engine[/red]")
+                raise typer.Exit(1)
+            console.print("[green]✓ Engine started[/green]")
+            # Mark that we booted the engine so attach loop gets extended retries
+            engine_started_now = True
+            # No further waiting here – attachment attempts below handle retry logic while the
+            # engine finishes booting.
+        else:
+            raise typer.Exit(1)
+
+    # Retrieve SSH public key (required for authorised_keys provisioning)
+    try:
+        public_key = get_ssh_public_key()
+    except FileNotFoundError as e:
+        console.print(f"[red]❌ {e}[/red]")
+        raise typer.Exit(1)
+
+    console.print(f"Attaching studio to engine [cyan]{engine['name']}[/cyan]...")
+
+    # Determine retry strategy based on whether we just started the engine
+    if engine_started_now:
+        max_attempts = 40  # About 7 minutes total with exponential backoff
+        base_delay = 8
+        max_delay = 20
+    else:
+        max_attempts = 15  # About 2 minutes total with exponential backoff
+        base_delay = 5
+        max_delay = 10
+
+    # Unified retry loop with exponential backoff
+    with Progress(
+        SpinnerColumn(),
+        TimeElapsedColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        transient=True,
+    ) as prog:
+        desc = (
+            "Attaching studio (engine is still booting)…"
+            if engine_started_now
+            else "Attaching studio…"
+        )
+        task = prog.add_task(desc, total=None)
+
+        consecutive_not_ready = 0
+        last_error = None
+
+        for attempt in range(max_attempts):
+            # Check if the attach already completed
+            if _is_studio_attached(studio["studio_id"], engine["instance_id"]):
+                success = True
+                break
+
+            success, error_msg = _attempt_studio_attach(
+                studio, engine, target_user, public_key
+            )
+
+            if success:
+                break  # success!
+
+            if error_msg:
+                # Fatal error – bubble up immediately
+                console.print(f"[red]❌ Failed to attach studio: {error_msg}[/red]")
+
+                # Suggest repair command if engine seems broken
+                if "not ready" in error_msg.lower() and attempt > 5:
+                    console.print(
+                        f"\n[yellow]Engine may be in a bad state. Try:[/yellow]"
+                    )
+                    console.print(f"[dim]  dh engine repair {engine['name']}[/dim]")
+                return
+
+            # Track consecutive "not ready" responses
+            consecutive_not_ready += 1
+            last_error = "Engine not ready"
+
+            # Update progress display
+            if attempt % 3 == 0:
+                prog.update(
+                    task,
+                    description=f"{desc} attempt {attempt+1}/{max_attempts}",
+                )
+
+            # If engine seems stuck after many attempts, show a hint
+            if consecutive_not_ready > 10 and attempt == 10:
+                console.print(
+                    "[yellow]Engine is taking longer than expected to become ready.[/yellow]"
+                )
+                console.print(
+                    "[dim]This can happen after GAMI creation or if the engine is still bootstrapping.[/dim]"
+                )
+
+            # Exponential backoff with jitter
+            delay = min(base_delay * (1.5 ** min(attempt, 5)), max_delay)
+            delay += time.time() % 2  # Add 0-2 seconds of jitter
+            time.sleep(delay)
+
+        else:
+            # All attempts exhausted
+            console.print(
+                f"[yellow]Engine is not becoming ready after {max_attempts} attempts.[/yellow]"
+            )
+            if last_error:
+                console.print(f"[dim]Last issue: {last_error}[/dim]")
+            console.print("\n[yellow]You can try:[/yellow]")
+            console.print(
+                f"  1. Wait a minute and retry: [cyan]dh studio attach {engine['name']}[/cyan]"
+            )
+            console.print(
+                f"  2. Check engine status: [cyan]dh engine status {engine['name']}[/cyan]"
+            )
+            console.print(
+                f"  3. Repair the engine: [cyan]dh engine repair {engine['name']}[/cyan]"
+            )
+            return
+
+    # Successful attach path
+    console.print(f"[green]✓ Studio attached successfully![/green]")
+
+    # Update SSH config - use target_user for the connection
+    update_ssh_config_entry(engine["name"], engine["instance_id"], target_user)
+    console.print(f"[green]✓ SSH config updated[/green]")
+    console.print(f"\nConnect with: [cyan]ssh {engine['name']}[/cyan]")
+    console.print(f"Files are at: [cyan]/studios/{target_user}[/cyan]")
+
+
+def _is_studio_attached(target_studio_id: str, target_vm_id: str) -> bool:
+    """Check if a studio is attached to a specific VM."""
+    response = make_api_request("GET", "/studios")
+    if response.status_code != 200:
+        return False
+
+    studios = response.json().get("studios", [])
+    for studio in studios:
+        if (
+            studio["studio_id"] == target_studio_id
+            and studio.get("attached_vm_id") == target_vm_id
+            and studio.get("status") == "in-use"
+        ):
+            return True
+    return False
+
+
+def _attempt_studio_attach(studio, engine, target_user, public_key):
+    response = make_api_request(
+        "POST",
+        f"/studios/{studio['studio_id']}/attach",
+        json_data={
+            "vm_id": engine["instance_id"],
+            "user": target_user,
+            "public_key": public_key,
+        },
+    )
+
+    # Fast-path success
+    if response.status_code == 200:
+        return True, None
+
+    # Asynchronous path – API returned 202 Accepted and operation tracking ID
+    if response.status_code == 202:
+        # The operation status polling is broken in the Lambda, so we just
+        # wait and check if the studio is actually attached
+        time.sleep(5)  # Give the async operation a moment to start
+
+        # Check periodically if the studio is attached
+        for check in range(20):  # Check for up to 60 seconds
+            if _is_studio_attached(studio["studio_id"], engine["instance_id"]):
+                return True, None
+            time.sleep(3)
+
+        # If we get here, attachment didn't complete in reasonable time
+        return False, None  # Return None to trigger retry
+
+    # --- determine if we should retry ---
+    recoverable = False
+    error_text = response.json().get("error", "Unknown error")
+    err_msg = error_text.lower()
+
+    # Check for "Studio is not available (status: in-use)" which means it's already attached
+    if (
+        response.status_code == 400
+        and "not available" in err_msg
+        and "in-use" in err_msg
+    ):
+        # Studio is already attached somewhere - check if it's to THIS engine
+        if _is_studio_attached(studio["studio_id"], engine["instance_id"]):
+            return True, None  # It's attached to our target engine - success!
+        else:
+            return False, error_text  # It's attached elsewhere - fatal error
+
+    if response.status_code in (409, 503):
+        recoverable = True
+    else:
+        RECOVERABLE_PATTERNS = [
+            "not ready",
+            "still starting",
+            "initializing",
+            "failed to mount",
+            "device busy",
+            "pending",  # VM state pending
+        ]
+        FATAL_PATTERNS = [
+            "permission",
+        ]
+        if any(p in err_msg for p in FATAL_PATTERNS):
+            recoverable = False
+        elif any(p in err_msg for p in RECOVERABLE_PATTERNS):
+            recoverable = True
+
+    if not recoverable:
+        # fatal – abort immediately
+        return False, error_text
+
+    # recoverable – signal caller to retry without treating as error
+    return False, None