dayhoff-tools 1.3.25__py3-none-any.whl → 1.4.10__py3-none-any.whl

dayhoff_tools/cli/engine_commands.py

@@ -15,7 +15,7 @@ from botocore.exceptions import ClientError, NoCredentialsError
 from rich import box
 from rich.console import Console
 from rich.panel import Panel
-from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.prompt import Confirm, IntPrompt, Prompt
 from rich.table import Table
 import re
@@ -382,8 +382,7 @@ def update_ssh_config_entry(engine_name: str, instance_id: str, ssh_user: str):
 Host {engine_name} {SSH_MANAGED_COMMENT}
     HostName {instance_id}
     User {ssh_user}
-    ProxyCommand sh -c "aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"
-"""
+    ProxyCommand sh -c \"AWS_SSM_IDLE_TIMEOUT=600 aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'\"\n"""
 
     # Check if entry already exists
     host_line = f"Host {engine_name} {SSH_MANAGED_COMMENT}"
@@ -592,13 +591,16 @@ def engine_status(
     engines = response.json().get("engines", [])
     engine = resolve_engine(name_or_id, engines)
 
-    # Get attached studios info
-    response = make_api_request(
-        "GET", f"/engines/{engine['instance_id']}/attached-studios"
-    )
-    attached_studios = []
-    if response.status_code == 200:
-        attached_studios = response.json().get("studios", [])
+    # Get detailed engine status including idle detector info
+    response = make_api_request("GET", f"/engines/{engine['instance_id']}")
+    if response.status_code != 200:
+        console.print("[red]❌ Failed to fetch engine details[/red]")
+        raise typer.Exit(1)
+    
+    engine_details = response.json()
+    engine = engine_details.get("engine", engine)  # Use detailed info if available
+    idle_detector = engine_details.get("idle_detector", {})
+    attached_studios = engine_details.get("attached_studios", [])
 
     # Calculate costs
     launch_time = parse_launch_time(engine["launch_time"])
@@ -614,7 +616,6 @@ def engine_status(
         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)",
@@ -642,6 +643,37 @@ def engine_status(
                 status_lines.append(f"  • Idle Detector: {health.get('idle_detector_timer', 'unknown')}")
         except Exception:
             pass
+    
+    # Idle detector status (from new API endpoint)
+    if idle_detector.get("available"):
+        status_lines.append("")
+        status_lines.append("[bold]Idle Detector:[/bold]")
+        
+        # Overall status
+        if idle_detector["status"] == "active":
+            status_lines.append("  [green]✓ Engine ACTIVE[/green]")
+        else:
+            idle_seconds = idle_detector.get("idle_seconds", 0)
+            idle_threshold = idle_detector.get("idle_threshold", 1800)
+            idle_minutes = idle_seconds // 60
+            threshold_minutes = idle_threshold // 60
+            status_lines.append(f"  [yellow]⏱ Engine IDLE ({idle_minutes}/{threshold_minutes} minutes)[/yellow]")
+        
+        # Coffee lock
+        if idle_detector.get("coffee_lock"):
+            status_lines.append(f"  • [cyan]☕ Caffeinated for another {idle_detector['coffee_lock']}[/cyan]")
+        
+        # SSH sessions
+        ssh_sessions = idle_detector.get("ssh_sessions", [])
+        if ssh_sessions:
+            status_lines.append(f"  • [blue]SSH Sessions ({len(ssh_sessions)}):[/blue]")
+            for session in ssh_sessions:
+                status_lines.append(f"    - {session['tty']} (pid {session['pid']}, idle {session['idle_time']}) from {session['from_ip']}")
+        
+        # IDE connections
+        ide_conn = idle_detector.get("ide_connections")
+        if ide_conn:
+            status_lines.append(f"  • [magenta]🖥  IDE connected ({ide_conn['connection_count']} connections)[/magenta]")
 
     if attached_studios:
         status_lines.append("")
@@ -911,7 +943,7 @@ def config_ssh(
                     f"Host {engine['name']} {SSH_MANAGED_COMMENT}",
                     f"    HostName {engine['instance_id']}",
                     f"    User {ssh_user}",
-                    f"    ProxyCommand sh -c \"aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'\"",
+                    f"    ProxyCommand sh -c \"AWS_SSM_IDLE_TIMEOUT=600 aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'\"",
                 ]
             )
 
@@ -978,7 +1010,7 @@ def coffee(
     else:
         console.print(f"Pouring coffee for [cyan]{engine['name']}[/cyan] for {duration}…")
 
-    # Use SSM to run the engine keep-alive command
+    # Use SSM to run the engine coffee command
     ssm = boto3.client("ssm", region_name="us-east-1")
     try:
         response = ssm.send_command(
@@ -1219,7 +1251,7 @@ def create_ami(
     console.print(f"AMI Name: [cyan]{ami_name}[/cyan]")
     console.print(f"Description: {ami_description}")
     console.print(
-        "\n[yellow]⚠️  Important: This will reboot the engine to ensure a clean snapshot.[/yellow]"
+        "\n[yellow]⚠️  Important: This will detach all studios and reboot the engine to ensure a clean snapshot.[/yellow]"
     )
 
     if not Confirm.ask("\nContinue with AMI creation?"):
@@ -1280,15 +1312,42 @@ def create_ami(
                 )
                 console.print(f"  Excluding volume at {device_name}")
 
+        # --- Check & detach attached studios --------------------------------------------------
+        # If any user studios are still attached we must detach them before the instance reboots
+        # for snapshot consistency; otherwise Studio-Manager metadata becomes stale.
+
+        attached_resp = make_api_request("GET", f"/engines/{engine['instance_id']}/studios")
+        attached_studios = attached_resp.json().get("studios", []) if attached_resp.status_code == 200 else []
+
+        if attached_studios:
+            console.print(f"Detaching {len(attached_studios)} studio(s) from this engine…")
+            for s in attached_studios:
+                console.print(f"  • {s['user']} ({s['studio_id']})")
+
+            for s in attached_studios:
+                resp = make_api_request("POST", f"/studios/{s['studio_id']}/detach")
+                if resp.status_code != 200:
+                    console.print(f"[red]❌ Failed to detach {s['studio_id']} – aborting.[/red]")
+                    return
+
+            # Wait briefly for volumes to become available (max 2 min)
+            # (time is already imported at module level)
+            ec2_wait = boto3.client("ec2", region_name="us-east-1")
+            vol_ids = [s['studio_id'] for s in attached_studios]
+            console.print("Waiting for volumes to detach…")
+            waiter = ec2_wait.get_waiter("volume_available")
+            try:
+                waiter.wait(VolumeIds=vol_ids, WaiterConfig={"Delay": 5, "MaxAttempts": 24})
+            except Exception:
+                console.print("[yellow]Proceeding even though some volumes may still be detaching.[/yellow]")
+
         # Create the AMI
         with Progress(
             SpinnerColumn(),
             TextColumn("[progress.description]{task.description}"),
             transient=True,
         ) as progress:
-            progress.add_task(
-                "Creating AMI (this will take several minutes)...", total=None
-            )
+            progress.add_task("Creating AMI (this will take several minutes)...", total=None)
 
             create_params = {
                 "InstanceId": engine["instance_id"],
@@ -1559,14 +1618,15 @@ def attach_studio(
                 console.print("[red]❌ Failed to start engine[/red]")
                 raise typer.Exit(1)
             console.print("[green]✓ Engine started[/green]")
-            console.print("Waiting for engine to be ready...")
-            import time
-
-            time.sleep(10)
+            # No further waiting here – attachment attempts below handle retry logic while the
+            # engine finishes booting.
         else:
             raise typer.Exit(1)
 
-    # Get SSH key
+    # Track whether this command just started the engine (affects retry length)
+    engine_started_now = False
+
+    # Retrieve SSH public key (required for authorised_keys provisioning)
     try:
         public_key = get_ssh_public_key()
     except FileNotFoundError as e:
@@ -1577,34 +1637,89 @@ def attach_studio(
 
     with Progress(
         SpinnerColumn(),
+        TimeElapsedColumn(),
         TextColumn("[progress.description]{task.description}"),
         transient=True,
-    ) as progress:
-        task = progress.add_task("Attaching studio...", total=100)
+    ) as prog:
+        task = prog.add_task("Attaching studio (engine is still booting)…", total=None)
+        ATTEMPT_LIMIT = 40 if engine_started_now else 6  # shorter retries for already-running engines
+        RETRY_DELAY = 10
+        for attempt in range(ATTEMPT_LIMIT):
+            success, error_msg = _attempt_studio_attach(studio, engine, target_user, public_key)
+
+            if success:
+                break  # success!
+
+            # Update spinner description with attempt number
+            if attempt % 3 == 0:  # update every 3rd attempt to avoid spam
+                prog.update(task, description=f"Attaching studio (engine is still booting)… {attempt+1}/{ATTEMPT_LIMIT}")
+
+            if error_msg:
+                console.print(f"[red]❌ Failed to attach studio: {error_msg}[/red]")
+                return
 
-        response = make_api_request(
-            "POST",
-            f"/studios/{studio['studio_id']}/attach",
-            json_data={
-                "vm_id": engine["instance_id"],
-                "user": target_user,  # Use target_user instead of username
-                "public_key": public_key,
-            },
-        )
+            time.sleep(RETRY_DELAY)
+
+        else:
+            console.print("[yellow]Engine is still starting up – please retry in a minute.[/yellow]")
+            return
+
+    # Successful attach path
+    console.print(f"[green]✓ Studio attached successfully![/green]")
 
-        progress.update(task, completed=100)
+    # 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 _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,
+        },
+    )
 
     if response.status_code == 200:
-        console.print(f"[green]✓ Studio attached successfully![/green]")
+        return True, None
 
-        # 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]")
+    # --- determine if we should retry ---
+    recoverable = False
+    if response.status_code in (409, 503):
+        recoverable = True
     else:
+        err_msg = response.json().get("error", "").lower()
+        RECOVERABLE_PATTERNS = [
+            "not ready",
+            "still starting",
+            "initializing",
+            "failed to mount",
+            "device busy",
+            "not available",
+            "pending",  # VM state pending
+        ]
+        FATAL_PATTERNS = [
+            "in-use",
+            "already attached",
+            "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 and show message
         error = response.json().get("error", "Unknown error")
-        console.print(f"[red]❌ Failed to attach studio: {error}[/red]")
+        return False, error
+
+    # otherwise wait and retry
+    return False, None
 
 
 @studio_app.command("detach")

{dayhoff_tools-1.3.25.dist-info → dayhoff_tools-1.4.10.dist-info}/METADATA

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

{dayhoff_tools-1.3.25.dist-info → dayhoff_tools-1.4.10.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=XVg03q0MG4DpzMxep3VRbEWZqDU_LHRD3j_xKui4Ync,78722
+dayhoff_tools/cli/engine_commands.py,sha256=ul6hdl3oMLXeKqwG5p9bjZbWJisU_l9siObg1Ls_fKs,84109
 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.25.dist-info/METADATA,sha256=-M4-C9Y0ICSbkfUiaIGTw_Voh9QnrIcYNGmklHYyZUU,2825
-dayhoff_tools-1.3.25.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-dayhoff_tools-1.3.25.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.3.25.dist-info/RECORD,,
+dayhoff_tools-1.4.10.dist-info/METADATA,sha256=Fenvb3Ijq1lKIoXpOaqu2Sb3tS82qQSgHN5iLFd1c9E,2825
+dayhoff_tools-1.4.10.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+dayhoff_tools-1.4.10.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.4.10.dist-info/RECORD,,