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

dayhoff_tools/cli/engine_studio_commands.py

@@ -0,0 +1,323 @@
+"""Engine and Studio management commands for DHT CLI."""
+
+from typing import Optional
+
+import typer
+
+# Initialize Typer apps
+engine_app = typer.Typer(help="Manage compute engines for development.")
+studio_app = typer.Typer(help="Manage persistent development studios.")
+
+# Use lazy loading pattern similar to main.py swarm commands
+# Import functions only when commands are actually called
+
+
+# Engine commands
+@engine_app.command("launch")
+def launch_engine_cmd(
+    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: str = typer.Option(None, "--user", "-u", help="Override username"),
+    boot_disk_size: int = typer.Option(
+        None,
+        "--size",
+        "-s",
+        help="Boot disk size in GB (default: 50GB, min: 20GB, max: 1000GB)",
+    ),
+    availability_zone: 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."""
+    from .engine.launch import launch_engine
+
+    return launch_engine(name, engine_type, user, boot_disk_size, availability_zone)
+
+
+@engine_app.command("list")
+def list_engines_cmd(
+    user: 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)."""
+    from .engine.list import list_engines
+
+    return list_engines(user, running_only, stopped_only, detailed)
+
+
+@engine_app.command("status")
+def engine_status_cmd(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    detailed: bool = typer.Option(
+        False, "--detailed", "-d", help="Show detailed status (slower)"
+    ),
+    show_log: bool = typer.Option(
+        False, "--show-log", help="Show bootstrap log (requires --detailed)"
+    ),
+):
+    """Show engine status and information."""
+    from .engine.status import engine_status
+
+    return engine_status(name_or_id, detailed, show_log)
+
+
+@engine_app.command("start")
+def start_engine_cmd(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+):
+    """Start a stopped engine."""
+    from .engine.lifecycle import start_engine
+
+    return start_engine(name_or_id)
+
+
+@engine_app.command("stop")
+def stop_engine_cmd(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    force: bool = typer.Option(
+        False, "--force", "-f", help="Force stop and detach all studios"
+    ),
+):
+    """Stop an engine."""
+    from .engine.lifecycle import stop_engine
+
+    return stop_engine(name_or_id, force)
+
+
+@engine_app.command("terminate")
+def terminate_engine_cmd(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+):
+    """Permanently terminate an engine."""
+    from .engine.lifecycle import terminate_engine
+
+    return terminate_engine(name_or_id)
+
+
+@engine_app.command("ssh")
+def ssh_engine_cmd(
+    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."""
+    from .engine.ssh import ssh_engine
+
+    return ssh_engine(name_or_id, admin, idle_timeout)
+
+
+@engine_app.command("config-ssh")
+def config_ssh_cmd(
+    clean: bool = typer.Option(False, "--clean", help="Remove all managed entries"),
+    all_engines: bool = typer.Option(
+        False, "--all", "-a", help="Include all engines from all users"
+    ),
+    admin: bool = typer.Option(
+        False,
+        "--admin",
+        help="Generate entries that use ec2-user instead of per-engine owner user",
+    ),
+):
+    """Update SSH config with available engines."""
+    from .engine.config_ssh import config_ssh
+
+    return config_ssh(clean, all_engines, admin)
+
+
+@engine_app.command("resize")
+def resize_engine_cmd(
+    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."""
+    from .engine.resize import resize_engine
+
+    return resize_engine(name_or_id, size, online, force)
+
+
+@engine_app.command("gami")
+def create_ami_cmd(
+    name_or_id: str = typer.Argument(
+        help="Engine name or instance ID to create AMI from"
+    ),
+):
+    """Create a 'Golden AMI' from a running engine."""
+    from .engine.gami import create_ami
+
+    return create_ami(name_or_id)
+
+
+@engine_app.command("coffee")
+def coffee_cmd(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    duration: str = typer.Argument("4h", help="Duration (e.g., 2h, 30m, 2h30m)"),
+    cancel: bool = typer.Option(
+        False, "--cancel", help="Cancel existing coffee lock instead of extending"
+    ),
+):
+    """Pour ☕ for an engine: keeps it awake for the given duration (or cancel)."""
+    from .engine.coffee import coffee
+
+    return coffee(name_or_id, duration, cancel)
+
+
+@engine_app.command("idle")
+def idle_timeout_cmd_wrapper(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    set: Optional[str] = typer.Option(
+        None, "--set", "-s", help="New timeout (e.g., 2h30m, 45m)"
+    ),
+    slack: Optional[str] = typer.Option(
+        None, "--slack", help="Set Slack notifications: none, default, all"
+    ),
+):
+    """Show or set engine idle-detector settings."""
+    from .engine.idle import idle_timeout_cmd
+
+    return idle_timeout_cmd(name_or_id=name_or_id, set=set, slack=slack)
+
+
+@engine_app.command("debug")
+def debug_engine_cmd(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+):
+    """Debug engine bootstrap status and files."""
+    from .engine.debug import debug_engine
+
+    return debug_engine(name_or_id)
+
+
+@engine_app.command("repair")
+def repair_engine_cmd(
+    name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+):
+    """Repair an engine that's stuck in a bad state (e.g., after GAMI creation)."""
+    from .engine.repair import repair_engine
+
+    return repair_engine(name_or_id)
+
+
+# Studio commands
+@studio_app.command("create")
+def create_studio_cmd(
+    size_gb: int = typer.Option(50, "--size", "-s", help="Studio size in GB"),
+):
+    """Create a new studio for the current user."""
+    from .studio.create import create_studio
+
+    return create_studio(size_gb)
+
+
+@studio_app.command("status")
+def studio_status_cmd(
+    user: str = typer.Option(
+        None, "--user", "-u", help="Check status for a different user (admin only)"
+    ),
+):
+    """Show status of your studio."""
+    from .studio.status import studio_status
+
+    return studio_status(user)
+
+
+@studio_app.command("attach")
+def attach_studio_cmd(
+    engine_name_or_id: str = typer.Argument(help="Engine name or instance ID"),
+    user: str = typer.Option(
+        None, "--user", "-u", help="Attach a different user's studio (admin only)"
+    ),
+):
+    """Attach your studio to an engine."""
+    from .studio.attach import attach_studio
+
+    return attach_studio(engine_name_or_id, user)
+
+
+@studio_app.command("detach")
+def detach_studio_cmd(
+    user: str = typer.Option(
+        None, "--user", "-u", help="Detach a different user's studio (admin only)"
+    ),
+):
+    """Detach your studio from its current engine."""
+    from .studio.detach import detach_studio
+
+    return detach_studio(user)
+
+
+@studio_app.command("delete")
+def delete_studio_cmd(
+    user: str = typer.Option(
+        None, "--user", "-u", help="Delete a different user's studio (admin only)"
+    ),
+):
+    """Delete your studio permanently."""
+    from .studio.delete import delete_studio
+
+    return delete_studio(user)
+
+
+@studio_app.command("list")
+def list_studios_cmd(
+    all_users: bool = typer.Option(
+        False, "--all", "-a", help="Show all users' studios"
+    ),
+):
+    """List studios."""
+    from .studio.list import list_studios
+
+    return list_studios(all_users)
+
+
+@studio_app.command("reset")
+def reset_studio_cmd(
+    user: str = typer.Option(
+        None, "--user", "-u", help="Reset a different user's studio"
+    ),
+):
+    """Reset a stuck studio (admin operation)."""
+    from .studio.reset import reset_studio
+
+    return reset_studio(user)
+
+
+@studio_app.command("resize")
+def resize_studio_cmd(
+    size: int = typer.Option(..., "--size", "-s", help="New size in GB"),
+    user: str = typer.Option(
+        None, "--user", "-u", help="Resize a different user's studio (admin only)"
+    ),
+):
+    """Resize your studio volume (requires detachment)."""
+    from .studio.resize import resize_studio
+
+    return resize_studio(size, user)

dayhoff_tools/cli/engine_studio_utils/__init__.py

@@ -0,0 +1 @@
+"""Utility modules for engine and studio commands."""

dayhoff_tools/cli/engine_studio_utils/api_utils.py

@@ -0,0 +1,47 @@
+"""API request utilities for engine and studio commands."""
+
+from typing import Dict, Optional
+
+import requests
+import typer
+
+from .aws_utils import get_api_url
+from .constants import console
+
+
+def make_api_request(
+    method: str,
+    endpoint: str,
+    json_data: Optional[Dict] = None,
+    params: Optional[Dict] = None,
+) -> requests.Response:
+    """Make an API request with error handling."""
+    api_url = get_api_url()
+    url = f"{api_url}{endpoint}"
+
+    try:
+        if method == "GET":
+            response = requests.get(url, params=params)
+        elif method == "POST":
+            response = requests.post(url, json=json_data)
+        elif method == "DELETE":
+            response = requests.delete(url)
+        else:
+            raise ValueError(f"Unsupported HTTP method: {method}")
+
+        return response
+    except requests.exceptions.RequestException as e:
+        console.print(f"[red]❌ API request failed: {e}[/red]")
+        raise typer.Exit(1)
+
+
+def get_user_studio(username: str) -> Optional[Dict]:
+    """Get the current user's studio."""
+    response = make_api_request("GET", "/studios")
+    if response.status_code != 200:
+        return None
+
+    studios = response.json().get("studios", [])
+    user_studios = [s for s in studios if s["user"] == username]
+
+    return user_studios[0] if user_studios else None

dayhoff_tools/cli/engine_studio_utils/aws_utils.py

@@ -0,0 +1,102 @@
+"""AWS-specific utilities for engine and studio commands."""
+
+import subprocess
+from typing import Dict, List
+
+import boto3
+import typer
+from botocore.exceptions import ClientError, NoCredentialsError
+from rich.prompt import Confirm
+
+from .constants import console
+
+
+def check_aws_sso() -> str:
+    """Check AWS SSO status and return username."""
+    try:
+        sts = boto3.client("sts")
+        identity = sts.get_caller_identity()
+        # Parse username from assumed role ARN
+        # Format: arn:aws:sts::123456789012:assumed-role/AWSReservedSSO_DeveloperAccess_xxxx/username
+        arn = identity["Arn"]
+        if "assumed-role" in arn:
+            username = arn.split("/")[-1]
+            return username
+        else:
+            # Fallback for other auth methods
+            return identity["UserId"].split(":")[-1]
+    except (NoCredentialsError, ClientError):
+        console.print("[red]❌ Not logged in to AWS SSO[/red]")
+        console.print("Please run: [cyan]aws sso login[/cyan]")
+        if Confirm.ask("Would you like to login now?"):
+            try:
+                result = subprocess.run(
+                    ["aws", "sso", "login"],
+                    capture_output=True,
+                    text=True,
+                    check=True,
+                )
+                if result.returncode == 0:
+                    console.print("[green]✓ Successfully logged in![/green]")
+                    return check_aws_sso()
+            except subprocess.CalledProcessError as e:
+                console.print(f"[red]Login failed: {e}[/red]")
+        raise typer.Exit(1)
+
+
+def get_api_url() -> str:
+    """Get Studio Manager API URL from SSM Parameter Store."""
+    ssm = boto3.client("ssm", region_name="us-east-1")
+    try:
+        response = ssm.get_parameter(Name="/dev/studio-manager/api-url")
+        return response["Parameter"]["Value"]
+    except ClientError as e:
+        if e.response["Error"]["Code"] == "ParameterNotFound":
+            console.print(
+                "[red]❌ API URL parameter not found in SSM Parameter Store[/red]"
+            )
+            console.print(
+                "Please ensure the Studio Manager infrastructure is deployed."
+            )
+        else:
+            console.print(f"[red]❌ Error retrieving API URL: {e}[/red]")
+        raise typer.Exit(1)
+
+
+def _colour_stage(stage: str) -> str:
+    """Return colourised stage name for table output."""
+    if not stage:
+        return "[dim]-[/dim]"
+    low = stage.lower()
+    if low.startswith("error"):
+        return f"[red]{stage}[/red]"
+    if low == "finished":
+        return f"[green]{stage}[/green]"
+    return f"[yellow]{stage}[/yellow]"
+
+
+def _fetch_init_stages(instance_ids: List[str]) -> Dict[str, str]:
+    """Fetch DayhoffInitStage tag for many instances in one call."""
+    if not instance_ids:
+        return {}
+    ec2 = boto3.client("ec2", region_name="us-east-1")
+    stages: Dict[str, str] = {}
+    try:
+        paginator = ec2.get_paginator("describe_instances")
+        for page in paginator.paginate(InstanceIds=instance_ids):
+            for res in page["Reservations"]:
+                for inst in res["Instances"]:
+                    iid = inst["InstanceId"]
+                    tag_val = next(
+                        (
+                            t["Value"]
+                            for t in inst.get("Tags", [])
+                            if t["Key"] == "DayhoffInitStage"
+                        ),
+                        None,
+                    )
+                    if tag_val:
+                        stages[iid] = tag_val
+    except Exception:
+        pass  # best-effort
+    return stages

dayhoff_tools/cli/engine_studio_utils/constants.py

@@ -0,0 +1,21 @@
+"""Constants used across engine and studio commands."""
+
+from rich.console import Console
+
+console = Console()
+
+# Cost information
+HOURLY_COSTS = {
+    "cpu": 0.50,  # r6i.2xlarge
+    "cpumax": 2.02,  # r7i.8xlarge
+    "t4": 0.75,  # g4dn.2xlarge
+    "a10g": 1.50,  # g5.2xlarge
+    "a100": 21.96,  # p4d.24xlarge
+    "4_t4": 3.91,  # g4dn.12xlarge
+    "8_t4": 7.83,  # g4dn.metal
+    "4_a10g": 6.24,  # g5.12xlarge
+    "8_a10g": 16.29,  # g5.48xlarge
+}
+
+# SSH config management
+SSH_MANAGED_COMMENT = "# Managed by dh engine"