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

dayhoff_tools/cli/engine/gami.py

@@ -0,0 +1,160 @@
+"""Engine GAMI (Golden AMI) creation command."""
+
+from datetime import datetime
+
+import boto3
+import typer
+from botocore.exceptions import ClientError
+from rich.progress import Progress, SpinnerColumn, TextColumn
+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 create_ami(
+    name_or_id: str = typer.Argument(
+        help="Engine name or instance ID to create AMI from"
+    ),
+):
+    """Create a 'Golden AMI' from a running engine.
+
+    This process is for creating a pre-warmed, standardized machine image
+    that can be used to launch new engines more quickly.
+
+    IMPORTANT:
+    - The engine MUST have all studios detached before running this command.
+    - This process will make the source engine unusable. You should
+      plan to TERMINATE the engine after the AMI is created.
+    """
+    check_aws_sso()
+
+    # Get all engines to resolve name and check status
+    # We pass check_ready=True to get attached studio info
+    response = make_api_request("GET", "/engines", params={"check_ready": "true"})
+    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)
+
+    # --- Pre-flight checks ---
+
+    # 1. Check if engine is running
+    if engine["state"].lower() != "running":
+        console.print(f"[red]❌ Engine '{engine['name']}' is not running.[/red]")
+        console.print("Please start it before creating an AMI.")
+        raise typer.Exit(1)
+
+    # 2. Check for attached studios from the detailed API response
+    attached_studios = engine.get("studios", [])
+    if attached_studios:
+        console.print(
+            f"[bold red]❌ Engine '{engine['name']}' has studios attached.[/bold red]"
+        )
+        console.print("Please detach all studios before creating an AMI:")
+        for studio in attached_studios:
+            console.print(f"  - {studio['user']} ({studio['studio_id']})")
+        console.print("\nTo detach, run [bold]dh studio detach[/bold]")
+        raise typer.Exit(1)
+
+    # Construct AMI name and description
+    ami_name = (
+        f"prewarmed-engine-{engine['engine_type']}-{datetime.now().strftime('%Y%m%d')}"
+    )
+    description = (
+        f"Amazon Linux 2023 with NVIDIA drivers, Docker, and pre-pulled "
+        f"dev container image for {engine['engine_type']} engines"
+    )
+
+    console.print(f"Creating AMI from engine [cyan]{engine['name']}[/cyan]...")
+    console.print(f"[bold]AMI Name:[/] {ami_name}")
+    console.print(f"[bold]Description:[/] {description}")
+
+    console.print(
+        "\n[bold yellow]⚠️  Important:[/bold yellow]\n"
+        "1. This process will run cleanup scripts on the engine.\n"
+        "2. The source engine should be [bold]terminated[/bold] after the AMI is created.\n"
+    )
+
+    if not Confirm.ask("Continue with AMI creation?"):
+        raise typer.Exit()
+
+    # Create AMI using EC2 client directly, as the backend logic is too complex
+    ec2 = boto3.client("ec2", region_name="us-east-1")
+    ssm = boto3.client("ssm", region_name="us-east-1")
+
+    try:
+        # Clean up instance state before snapshotting
+        console.print("Cleaning up instance for AMI creation...")
+        cleanup_commands = [
+            "sudo rm -f /opt/dayhoff/first_boot_complete.sentinel",
+            "history -c",
+            "sudo rm -rf /tmp/* /var/log/messages /var/log/cloud-init.log",
+            "sudo rm -rf /var/lib/amazon/ssm/* /etc/amazon/ssm/*",
+            "sleep 2 && sudo systemctl stop amazon-ssm-agent &",  # Stop agent last
+        ]
+
+        cleanup_response = ssm.send_command(
+            InstanceIds=[engine["instance_id"]],
+            DocumentName="AWS-RunShellScript",
+            Parameters={"commands": cleanup_commands, "executionTimeout": ["120"]},
+        )
+
+        # Acknowledge that the SSM command might be in progress as the agent shuts down
+        console.print(
+            "[dim]ℹ️  Cleanup command sent (status may show 'InProgress' as SSM agent stops)[/dim]"
+        )
+
+        # Create the AMI
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            transient=True,
+        ) as progress:
+            task = progress.add_task(
+                "Creating AMI (this will take several minutes)...", total=None
+            )
+
+            response = ec2.create_image(
+                InstanceId=engine["instance_id"],
+                Name=ami_name,
+                Description=description,
+                NoReboot=False,
+                TagSpecifications=[
+                    {
+                        "ResourceType": "image",
+                        "Tags": [
+                            {"Key": "Environment", "Value": "dev"},
+                            {"Key": "Type", "Value": "golden-ami"},
+                            {"Key": "EngineType", "Value": engine["engine_type"]},
+                            {"Key": "Name", "Value": ami_name},
+                        ],
+                    }
+                ],
+            )
+
+            ami_id = response["ImageId"]
+            progress.update(
+                task,
+                completed=True,
+                description=f"[green]✓ AMI creation initiated![/green]",
+            )
+
+        console.print(f"  [bold]AMI ID:[/] {ami_id}")
+        console.print("\nThe AMI creation process will continue in the background.")
+        console.print("You can monitor progress in the EC2 Console under 'AMIs'.")
+        console.print(
+            "\nOnce complete, update the AMI ID in [bold]terraform/environments/dev/variables.tf[/bold] "
+            "and run [bold]terraform apply[/bold]."
+        )
+        console.print(
+            f"\nRemember to [bold red]terminate the source engine '{engine['name']}'[/bold red] to save costs."
+        )
+
+    except ClientError as e:
+        console.print(f"[red]❌ Failed to create AMI: {e}[/red]")
+        raise typer.Exit(1)

dayhoff_tools/cli/engine/idle.py

@@ -0,0 +1,148 @@
+"""Engine idle timeout command."""
+
+import re
+import time
+from typing import Optional
+
+import boto3
+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
+
+
+def idle_timeout_cmd(
+    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."""
+    check_aws_sso()
+
+    # Resolve engine
+    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)
+
+    ssm = boto3.client("ssm", region_name="us-east-1")
+
+    # Handle slack notifications change
+    if slack:
+        slack = slack.lower()
+        if slack not in ["none", "default", "all"]:
+            console.print("[red]❌ Invalid slack option. Use: none, default, all[/red]")
+            raise typer.Exit(1)
+
+        console.print(f"Setting Slack notifications to [bold]{slack}[/bold]...")
+
+        if slack == "none":
+            settings = {
+                "SLACK_NOTIFY_WARNINGS": "false",
+                "SLACK_NOTIFY_IDLE_START": "false",
+                "SLACK_NOTIFY_IDLE_END": "false",
+                "SLACK_NOTIFY_SHUTDOWN": "false",
+            }
+        elif slack == "default":
+            settings = {
+                "SLACK_NOTIFY_WARNINGS": "true",
+                "SLACK_NOTIFY_IDLE_START": "false",
+                "SLACK_NOTIFY_IDLE_END": "false",
+                "SLACK_NOTIFY_SHUTDOWN": "true",
+            }
+        else:  # all
+            settings = {
+                "SLACK_NOTIFY_WARNINGS": "true",
+                "SLACK_NOTIFY_IDLE_START": "true",
+                "SLACK_NOTIFY_IDLE_END": "true",
+                "SLACK_NOTIFY_SHUTDOWN": "true",
+            }
+
+        commands = []
+        for key, value in settings.items():
+            # Use a robust sed command that adds the line if it doesn't exist
+            commands.append(
+                f"grep -q '^{key}=' /etc/engine.env && sudo sed -i 's|^{key}=.*|{key}={value}|' /etc/engine.env || echo '{key}={value}' | sudo tee -a /etc/engine.env > /dev/null"
+            )
+
+        # Instead of restarting service, send SIGHUP to reload config
+        commands.append(
+            "sudo pkill -HUP -f engine-idle-detector.py || sudo systemctl restart engine-idle-detector.service"
+        )
+
+        resp = ssm.send_command(
+            InstanceIds=[engine["instance_id"]],
+            DocumentName="AWS-RunShellScript",
+            Parameters={"commands": commands, "executionTimeout": ["60"]},
+        )
+        cid = resp["Command"]["CommandId"]
+        time.sleep(2)  # Give it a moment to process
+        console.print(f"[green]✓ Slack notifications updated to '{slack}'[/green]")
+        console.print("[dim]Note: Settings updated without resetting idle timer[/dim]")
+
+    # Handle setting new timeout value
+    if set is not None:
+        m = re.match(r"^(?:(\d+)h)?(?:(\d+)m)?$", set)
+        if not m:
+            console.print(
+                "[red]❌ Invalid duration format. Use e.g. 2h, 45m, 1h30m[/red]"
+            )
+            raise typer.Exit(1)
+        hours = int(m.group(1) or 0)
+        minutes = int(m.group(2) or 0)
+        seconds = hours * 3600 + minutes * 60
+        if seconds == 0:
+            console.print("[red]❌ Duration must be greater than zero[/red]")
+            raise typer.Exit(1)
+
+        console.print(f"Setting idle timeout to {set} ({seconds} seconds)…")
+
+        cmd = (
+            "sudo sed -i '/^IDLE_TIMEOUT_SECONDS=/d' /etc/engine.env && "
+            f"echo 'IDLE_TIMEOUT_SECONDS={seconds}' | sudo tee -a /etc/engine.env >/dev/null && "
+            "sudo systemctl restart engine-idle-detector.service"
+        )
+
+        resp = ssm.send_command(
+            InstanceIds=[engine["instance_id"]],
+            DocumentName="AWS-RunShellScript",
+            Parameters={"commands": [cmd], "executionTimeout": ["60"]},
+        )
+        cid = resp["Command"]["CommandId"]
+        time.sleep(2)
+        console.print(f"[green]✓ Idle timeout updated to {set}[/green]")
+
+    # If no action was specified, show current timeout
+    if set is None and slack is None:
+        # Show current timeout setting
+        resp = ssm.send_command(
+            InstanceIds=[engine["instance_id"]],
+            DocumentName="AWS-RunShellScript",
+            Parameters={
+                "commands": [
+                    "grep -E '^IDLE_TIMEOUT_SECONDS=' /etc/engine.env || echo 'IDLE_TIMEOUT_SECONDS=1800'"
+                ],
+                "executionTimeout": ["10"],
+            },
+        )
+        cid = resp["Command"]["CommandId"]
+        time.sleep(1)
+        inv = ssm.get_command_invocation(
+            CommandId=cid, InstanceId=engine["instance_id"]
+        )
+        if inv["Status"] == "Success":
+            line = inv["StandardOutputContent"].strip()
+            secs = int(line.split("=")[1]) if "=" in line else 1800
+            console.print(f"Current idle timeout: {secs//60}m ({secs} seconds)")
+        else:
+            console.print("[red]❌ Could not retrieve idle timeout[/red]")
+        return

dayhoff_tools/cli/engine/launch.py

@@ -0,0 +1,101 @@
+"""Engine launch command."""
+
+from typing import Any, Dict, Optional
+
+import typer
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+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 HOURLY_COSTS, console
+
+
+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]")

dayhoff_tools/cli/engine/list.py

@@ -0,0 +1,116 @@
+"""Engine list command."""
+
+from datetime import datetime, timezone
+from typing import Optional
+
+import typer
+from rich import box
+from rich.table import Table
+
+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,
+    parse_launch_time,
+)
+
+
+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]")

dayhoff_tools/cli/engine/repair.py

@@ -0,0 +1,128 @@
+"""Engine repair command."""
+
+import time
+
+import boto3
+import typer
+from rich.progress import Progress, SpinnerColumn, TextColumn
+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 repair_engine(
+    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)."""
+    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)
+
+    if engine["state"].lower() != "running":
+        console.print(
+            f"[yellow]⚠️  Engine is {engine['state']}. Must be running to repair.[/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]")
+            console.print("Waiting for engine to become ready...")
+            time.sleep(30)  # Give it time to boot
+        else:
+            raise typer.Exit(1)
+
+    console.print(f"[bold]Repairing engine [cyan]{engine['name']}[/cyan][/bold]")
+    console.print(
+        "[dim]This will restore bootstrap state and ensure all services are running[/dim]\n"
+    )
+
+    ssm = boto3.client("ssm", region_name="us-east-1")
+
+    # Repair commands
+    repair_commands = [
+        # Create necessary directories
+        "sudo mkdir -p /opt/dayhoff /opt/dayhoff/state /opt/dayhoff/scripts",
+        # Download scripts from S3 if missing
+        "source /etc/engine.env && sudo aws s3 sync s3://${VM_SCRIPTS_BUCKET}/ /opt/dayhoff/scripts/ --exclude '*' --include '*.sh' --quiet",
+        "sudo chmod +x /opt/dayhoff/scripts/*.sh 2>/dev/null || true",
+        # Restore bootstrap state
+        "sudo touch /opt/dayhoff/first_boot_complete.sentinel",
+        "echo 'finished' | sudo tee /opt/dayhoff/state/engine-init.stage > /dev/null",
+        # Ensure SSM agent is running
+        "sudo systemctl restart amazon-ssm-agent 2>/dev/null || true",
+        # Restart idle detector (service only)
+        "sudo systemctl restart engine-idle-detector.service 2>/dev/null || true",
+        # Report status
+        "echo '=== Repair Complete ===' && echo 'Sentinel: ' && ls -la /opt/dayhoff/first_boot_complete.sentinel",
+        "echo 'Stage: ' && cat /opt/dayhoff/state/engine-init.stage",
+        "echo 'Scripts: ' && ls /opt/dayhoff/scripts/*.sh 2>/dev/null | wc -l",
+    ]
+
+    try:
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            transient=True,
+        ) as progress:
+            task = progress.add_task("Repairing engine...", total=None)
+
+            response = ssm.send_command(
+                InstanceIds=[engine["instance_id"]],
+                DocumentName="AWS-RunShellScript",
+                Parameters={
+                    "commands": repair_commands,
+                    "executionTimeout": ["60"],
+                },
+            )
+
+            command_id = response["Command"]["CommandId"]
+
+            # Wait for command
+            for _ in range(60):
+                time.sleep(1)
+                result = ssm.get_command_invocation(
+                    CommandId=command_id,
+                    InstanceId=engine["instance_id"],
+                )
+                if result["Status"] in ["Success", "Failed"]:
+                    break
+
+        if result["Status"] == "Success":
+            output = result["StandardOutputContent"]
+            console.print("[green]✓ Engine repaired successfully![/green]\n")
+
+            # Show repair results
+            if "=== Repair Complete ===" in output:
+                repair_section = output.split("=== Repair Complete ===")[1].strip()
+                console.print("[bold]Repair Results:[/bold]")
+                console.print(repair_section)
+
+            console.print(
+                "\n[dim]You should now be able to attach studios to this engine.[/dim]"
+            )
+        else:
+            console.print(
+                f"[red]❌ Repair failed: {result.get('StandardErrorContent', 'Unknown error')}[/red]"
+            )
+            console.print(
+                "\n[yellow]Try running 'dh engine debug' for more information.[/yellow]"
+            )
+
+    except Exception as e:
+        console.print(f"[red]❌ Failed to repair engine: {e}[/red]")