dayhoff-tools 1.3.5__py3-none-any.whl → 1.3.7__py3-none-any.whl

dayhoff_tools/cli/engine_commands.py

@@ -163,7 +163,7 @@ def format_status(state: str, ready: Optional[bool]) -> str:
         if ready is True:
             return "[green]Running ✓[/green]"
         elif ready is False:
-            return "[yellow]Running ⚠ (Initializing...)[/yellow]"
+            return "[yellow]Running ⚠ (Bootstrapping...)[/yellow]"
         else:
             return "[green]Running[/green]"
     elif state.lower() == "stopped":
@@ -332,17 +332,15 @@ def launch_engine(
     else:
         error = response.json().get("error", "Unknown error")
         console.print(f"[red]❌ Failed to launch engine: {error}[/red]")
-        raise typer.Exit(1)
 
 
 @engine_app.command("list")
 def list_engines(
     user: Optional[str] = typer.Option(None, "--user", "-u", help="Filter by user"),
-    all_users: bool = typer.Option(False, "--all", "-a", help="Show all users' engines"),
     running_only: bool = typer.Option(False, "--running", help="Show only running engines"),
     stopped_only: bool = typer.Option(False, "--stopped", help="Show only stopped engines"),
 ):
-    """List all engines."""
+    """List engines (shows all engines by default)."""
     current_user = check_aws_sso()
 
     params = {}
@@ -618,14 +616,18 @@ def ssh_engine(
 @engine_app.command("config-ssh")
 def config_ssh(
     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"),
 ):
-    """Update SSH config with all available engines."""
+    """Update SSH config with available engines."""
     username = check_aws_sso()
 
     if clean:
         console.print("Removing all managed SSH entries...")
     else:
-        console.print("Updating SSH config with all running engines...")
+        if all_engines:
+            console.print("Updating SSH config with all running engines...")
+        else:
+            console.print(f"Updating SSH config with running engines for [cyan]{username}[/cyan] and [cyan]shared[/cyan]...")
 
     # Get all engines
     response = make_api_request("GET", "/engines")
@@ -635,6 +637,14 @@ def config_ssh(
 
     engines = response.json().get("engines", [])
     running_engines = [e for e in engines if e["state"].lower() == "running"]
+    
+    # Filter engines based on options
+    if not all_engines:
+        # Show only current user's engines and shared engines
+        running_engines = [
+            e for e in running_engines 
+            if e["user"] == username or e["user"] == "shared"
+        ]
 
     # Read existing config
     config_path = Path.home() / ".ssh" / "config"
@@ -685,7 +695,8 @@ def config_ssh(
             f"[green]✓ Updated SSH config with {len(running_engines)} engines[/green]"
         )
         for engine in running_engines:
-            console.print(f"  • {engine['name']} → {engine['instance_id']}")
+            user_display = f"[dim]({engine['user']})[/dim]" if engine['user'] != username else ""
+            console.print(f"  • {engine['name']} → {engine['instance_id']} {user_display}")
 
 
 @engine_app.command("keep-awake")
@@ -936,15 +947,43 @@ def create_ami(
 
             response = ec2.create_image(**create_params)
 
-        ami_id = response["ImageId"]
-        console.print(f"[green]✓ AMI creation initiated![/green]")
-        console.print(f"AMI ID: [cyan]{ami_id}[/cyan]")
-        console.print("\n[dim]The AMI creation process will continue in the background.[/dim]")
-        console.print("[dim]You can monitor progress in the EC2 Console under 'AMIs'.[/dim]")
-        console.print(
-            f"\nOnce complete, run [cyan]terraform apply[/cyan] in "
-            f"terraform/environments/dev to use the new AMI."
-        )
+            ami_id = response["ImageId"]
+            console.print(f"[green]✓ AMI creation initiated![/green]")
+            console.print(f"AMI ID: [cyan]{ami_id}[/cyan]")
+            
+            # Restore the source engine to a normal state
+            console.print("Restoring source engine state...")
+            restore_response = ssm.send_command(
+                InstanceIds=[engine["instance_id"]],
+                DocumentName="AWS-RunShellScript",
+                Parameters={
+                    "commands": [
+                        "sudo touch /opt/dayhoff/first_boot_complete.sentinel",
+                        "sudo systemctl restart engine-idle-detector.timer",
+                    ],
+                    "executionTimeout": ["60"],
+                },
+            )
+
+            # Quick wait to see if it failed immediately
+            time.sleep(1)
+            restore_command_id = restore_response["Command"]["CommandId"]
+            result = ssm.get_command_invocation(
+                CommandId=restore_command_id,
+                InstanceId=engine["instance_id"],
+            )
+            if result["Status"] not in ["Pending", "InProgress", "Success"]:
+                 console.print("[yellow]⚠️  Warning: Failed to restore source engine state.[/yellow]")
+            else:
+                console.print("[green]✓ Source engine restored to normal operation.[/green]")
+
+
+            console.print("\n[dim]The AMI creation process will continue in the background.[/dim]")
+            console.print("[dim]You can monitor progress in the EC2 Console under 'AMIs'.[/dim]")
+            console.print(
+                f"\nOnce complete, run [cyan]terraform apply[/cyan] in "
+                f"terraform/environments/dev to use the new AMI."
+            )
 
     except ClientError as e:
         console.print(f"[red]❌ Failed to create AMI: {e}[/red]")
@@ -1020,7 +1059,7 @@ def studio_status():
     # Format status with colors
     status = studio['status']
     if status == "in-use":
-        status_display = "[magenta]attached[/magenta]"
+        status_display = "[bright_blue]attached[/bright_blue]"
     elif status in ["attaching", "detaching"]:
         status_display = f"[yellow]{status}[/yellow]"
     else:
@@ -1258,7 +1297,7 @@ def list_studios(
         for studio in studios:
             # Change status display
             if studio["status"] == "in-use":
-                status_display = "[magenta]attached[/magenta]"
+                status_display = "[bright_blue]attached[/bright_blue]"
             elif studio["status"] in ["attaching", "detaching"]:
                 status_display = "[yellow]" + studio["status"] + "[/yellow]"
             else:

{dayhoff_tools-1.3.5.dist-info → dayhoff_tools-1.3.7.dist-info}/METADATA

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

{dayhoff_tools-1.3.5.dist-info → dayhoff_tools-1.3.7.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=BU5mi4Q_UTsg2CQLkn2XQaw0anUKwIp_GhMa_Zjlcns,48604
+dayhoff_tools/cli/engine_commands.py,sha256=FBRNwuJ3Q7leRZrUCHn_4FHTA06QU-bqE3GrG87QLt8,50442
 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.5.dist-info/METADATA,sha256=LmjC7aGDDP-Pj-LII3xRjc_G8CELn7KJEozkjMtGeTU,2842
-dayhoff_tools-1.3.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-dayhoff_tools-1.3.5.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.3.5.dist-info/RECORD,,
+dayhoff_tools-1.3.7.dist-info/METADATA,sha256=UzFsOWGLYUgKx_soBtZVUgEhc_4sNoK8DmD6ZWh0_78,2842
+dayhoff_tools-1.3.7.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+dayhoff_tools-1.3.7.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.3.7.dist-info/RECORD,,