dayhoff-tools 1.3.15__py3-none-any.whl → 1.3.17__py3-none-any.whl

dayhoff_tools/cli/engine_commands.py

@@ -18,6 +18,7 @@ from rich.panel import Panel
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.prompt import Confirm, IntPrompt, Prompt
 from rich.table import Table
+import re
 
 # Initialize Typer apps
 engine_app = typer.Typer(help="Manage compute engines for development.")
@@ -37,6 +38,41 @@ HOURLY_COSTS = {
 # SSH config management
 SSH_MANAGED_COMMENT = "# Managed by dh engine"
 
+# --------------------------------------------------------------------------------
+# Bootstrap stage helpers
+# --------------------------------------------------------------------------------
+
+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
+
 
 def check_aws_sso() -> str:
     """Check AWS SSO status and return username."""
@@ -388,6 +424,12 @@ def launch_engine(
         help="Engine type: cpu, cpumax, t4, a10g, a100",
     ),
     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)",
+    ),
 ):
     """Launch a new engine instance."""
     username = check_aws_sso()
@@ -401,9 +443,19 @@ def launch_engine(
         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}) for ${cost:.2f}/hour..."
+        f"Launching [cyan]{name}[/cyan] ({engine_type}){disk_info} for ${cost:.2f}/hour..."
     )
 
     with Progress(
@@ -413,17 +465,23 @@ def launch_engine(
     ) as progress:
         progress.add_task("Creating engine...", total=None)
 
-        response = make_api_request(
-            "POST",
-            "/engines",
-            json_data={"name": name, "user": username, "engine_type": engine_type},
-        )
+        request_data = {
+            "name": name,
+            "user": username,
+            "engine_type": engine_type,
+        }
+        if boot_disk_size is not None:
+            request_data["boot_disk_size"] = boot_disk_size
+
+        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:
@@ -464,6 +522,9 @@ def list_engines(
             console.print("No engines found.")
             return
 
+        # Fetch bootstrap stages once
+        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")
@@ -471,6 +532,7 @@ def list_engines(
         table.add_column("Type")
         table.add_column("User")
         table.add_column("Status")
+        table.add_column("Stage")
         table.add_column("Disk Usage")
         table.add_column("Uptime/Since")
         table.add_column("$/hour", justify="right")
@@ -493,12 +555,15 @@ def list_engines(
                 time_str = launch_time.strftime("%Y-%m-%d %H:%M")
                 disk_usage = "-"
 
+            stage_display = _colour_stage(stages_map.get(engine["instance_id"], "-"))
+
             table.add_row(
                 engine["name"],
                 engine["instance_id"],
                 engine["engine_type"],
                 engine["user"],
                 format_status(engine["state"], engine.get("ready")),
+                stage_display,
                 disk_usage,
                 time_str,
                 f"${hourly_cost:.2f}",
@@ -544,18 +609,43 @@ def engine_status(
     hourly_cost = HOURLY_COSTS.get(engine["engine_type"], 0)
     total_cost = hourly_cost * (uptime.total_seconds() / 3600)
 
-    # Create status panel
+    stages_map = _fetch_init_stages([engine["instance_id"]])
+    stage_val = stages_map.get(engine["instance_id"], "-")
+
     status_lines = [
         f"[bold]Name:[/bold]        {engine['name']}",
         f"[bold]Instance:[/bold]    {engine['instance_id']}",
         f"[bold]Type:[/bold]        {engine['engine_type']} ({engine['instance_type']})",
         f"[bold]Status:[/bold]      {format_status(engine['state'], engine.get('ready'))}",
+        f"[bold]Bootstrap:[/bold]   {_colour_stage(stage_val)}",
         f"[bold]User:[/bold]        {engine['user']}",
         f"[bold]IP:[/bold]          {engine.get('public_ip', 'N/A')}",
         f"[bold]Launched:[/bold]    {launch_time.strftime('%Y-%m-%d %H:%M:%S')} ({format_duration(uptime)} ago)",
         f"[bold]Cost:[/bold]        ${hourly_cost:.2f}/hour (${total_cost:.2f} total)",
     ]
 
+    # Health report (only if bootstrap finished)
+    if stage_val == "finished":
+        try:
+            ssm = boto3.client("ssm", region_name="us-east-1")
+            res = ssm.send_command(
+                InstanceIds=[engine["instance_id"]],
+                DocumentName="AWS-RunShellScript",
+                Parameters={"commands": ["cat /var/run/engine-health.json || true"], "executionTimeout": ["10"]},
+            )
+            cid = res["Command"]["CommandId"]
+            time.sleep(1)
+            inv = ssm.get_command_invocation(CommandId=cid, InstanceId=engine["instance_id"])
+            if inv["Status"] == "Success":
+                import json as _json
+                health = _json.loads(inv["StandardOutputContent"].strip() or "{}")
+                status_lines.append("")
+                status_lines.append("[bold]Health:[/bold]")
+                status_lines.append(f"  • GPU Drivers: {'OK' if health.get('drivers_ok') else 'MISSING'}")
+                status_lines.append(f"  • Idle Detector: {health.get('idle_detector_timer', 'unknown')}")
+        except Exception:
+            pass
+
     if attached_studios:
         status_lines.append("")
         status_lines.append("[bold]Attached Studios:[/bold]")
@@ -1887,3 +1977,72 @@ def resize_studio(
     
     console.print("\n[dim]The filesystem will be automatically expanded when you next attach the studio.[/dim]")
     console.print(f"To attach: [cyan]dh studio attach <engine-name>[/cyan]")
+
+# ================= Idle timeout command =================
+
+
+@engine_app.command("idle-timeout")
+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)")
+):
+    """Show or set the engine idle-detector timeout."""
+    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")
+
+    if set is None:
+        # Show current
+        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
+
+    # ----- set new value -----
+    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.timer"
+    )
+
+    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]")

{dayhoff_tools-1.3.15.dist-info → dayhoff_tools-1.3.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dayhoff-tools
-Version: 1.3.15
+Version: 1.3.17
 Summary: Common tools for all the repos at Dayhoff Labs
 Author: Daniel Martin-Alarcon
 Author-email: dma@dayhofflabs.com

{dayhoff_tools-1.3.15.dist-info → dayhoff_tools-1.3.17.dist-info}/RECORD

@@ -3,7 +3,7 @@ dayhoff_tools/chemistry/standardizer.py,sha256=uMn7VwHnx02nc404eO6fRuS4rsl4dvSPf
 dayhoff_tools/chemistry/utils.py,sha256=jt-7JgF-GeeVC421acX-bobKbLU_X94KNOW24p_P-_M,2257
 dayhoff_tools/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dayhoff_tools/cli/cloud_commands.py,sha256=33qcWLmq-FwEXMdL3F0OHm-5Stlh2r65CldyEZgQ1no,40904
-dayhoff_tools/cli/engine_commands.py,sha256=DPD0rsctEz4iI4KXyPdMwS5XzlXaTv0DKUiYAtlsT1w,69414
+dayhoff_tools/cli/engine_commands.py,sha256=DlWw8oUzOJqN6uUgsxF5whfY8jfwvBko8hvRm3ZHxVc,75814
 dayhoff_tools/cli/main.py,sha256=rgeEHD9lJ8SBCR34BTLb7gVInHUUdmEBNXAJnq5yEU4,4795
 dayhoff_tools/cli/swarm_commands.py,sha256=5EyKj8yietvT5lfoz8Zx0iQvVaNgc3SJX1z2zQR6o6M,5614
 dayhoff_tools/cli/utility_commands.py,sha256=qs8vH9TBFHsOPC3X8cU3qZigM3dDn-2Ytq4o_F2WubU,27874
@@ -27,7 +27,7 @@ dayhoff_tools/intake/uniprot.py,sha256=BZYJQF63OtPcBBnQ7_P9gulxzJtqyorgyuDiPeOJq
 dayhoff_tools/logs.py,sha256=DKdeP0k0kliRcilwvX0mUB2eipO5BdWUeHwh-VnsICs,838
 dayhoff_tools/sqlite.py,sha256=jV55ikF8VpTfeQqqlHSbY8OgfyfHj8zgHNpZjBLos_E,18672
 dayhoff_tools/warehouse.py,sha256=8YbnQ--usrEgDQGfvpV4MrMji55A0rq2hZaOgFGh6ag,15896
-dayhoff_tools-1.3.15.dist-info/METADATA,sha256=2JVTqkvXrTEVUfMuItJ9WEr9xQtmd2TLqHadTZlzUgc,2825
-dayhoff_tools-1.3.15.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-dayhoff_tools-1.3.15.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.3.15.dist-info/RECORD,,
+dayhoff_tools-1.3.17.dist-info/METADATA,sha256=69jNn8FqivgUWrwrLZ9S-ZEherpCtEHUrBABW5dewWs,2825
+dayhoff_tools-1.3.17.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+dayhoff_tools-1.3.17.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.3.17.dist-info/RECORD,,