dayhoff-tools 1.3.16__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."""
@@ -486,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")
@@ -493,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")
@@ -515,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}",
@@ -566,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]")
@@ -1909,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.16.dist-info → dayhoff_tools-1.3.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dayhoff-tools
-Version: 1.3.16
+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.16.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=TX9IwHkpb-x3OvvydwwhsXqxCpXnZ9TCNiOvYXYGP94,70265
+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.16.dist-info/METADATA,sha256=Ylw3uOqRFudCtgZdphUqKDLMBn0bH7O6-Ns8ZGTQ5R4,2825
-dayhoff_tools-1.3.16.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-dayhoff_tools-1.3.16.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.3.16.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,,