dayhoff-tools 1.6.1__py3-none-any.whl → 1.6.2__py3-none-any.whl

dayhoff_tools/cli/engine_commands.py

@@ -1357,282 +1357,134 @@ def create_ami(
         help="Engine name or instance ID to create AMI from"
     ),
 ):
-    """Create a Golden AMI from an engine."""
+    """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
-    response = make_api_request("GET", "/engines")
+    # 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)
 
-    console.print(f"Creating AMI from engine [cyan]{engine['name']}[/cyan]...")
+    # --- 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)
 
-    # Get engine type from the engine
-    engine_type = engine["engine_type"]
+    # 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)
 
-    # Generate AMI name
-    date_str = datetime.now().strftime("%Y%m%d")
-    ami_name = f"prewarmed-engine-{engine_type}-{date_str}"
-    ami_description = (
-        f"Amazon Linux 2023 with NVIDIA drivers, Docker, and pre-pulled dev container "
-        f"image for {engine_type} engines"
+    # 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"AMI Name: [cyan]{ami_name}[/cyan]")
-    console.print(f"Description: {ami_description}")
+    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[yellow]⚠️  Important: This will detach all studios and reboot the engine to ensure a clean snapshot.[/yellow]"
+        "\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("\nContinue with AMI creation?"):
-        console.print("AMI creation cancelled.")
-        return
+    if not Confirm.ask("Continue with AMI creation?"):
+        raise typer.Exit()
 
-    # Create AMI using EC2 client
+    # 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:
-        # First, we need to clean up the instance before snapshotting
+        # Clean up instance state before snapshotting
         console.print("Cleaning up instance for AMI creation...")
-        ssm = boto3.client("ssm", region_name="us-east-1")
-
-        # Clean up instance state, stopping SSM agent last to allow proper status reporting
         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/*",
-            # Stop SSM agent last with a delay to allow status reporting
-            "sleep 2 && sudo systemctl stop amazon-ssm-agent &",
+            "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"],
-            },
+            Parameters={"commands": cleanup_commands, "executionTimeout": ["120"]},
         )
 
-        # Wait for cleanup to complete
-        command_id = cleanup_response["Command"]["CommandId"]
-        for _ in range(10):
-            time.sleep(1)
-            result = ssm.get_command_invocation(
-                CommandId=command_id,
-                InstanceId=engine["instance_id"],
-            )
-            if result["Status"] in ["Success", "Failed"]:
-                break
-
-        # Note: InProgress status is expected when SSM agent stops itself
-        if result["Status"] not in ["Success", "InProgress"]:
-            console.print(
-                f"[yellow]⚠️  Warning: Cleanup command status: {result['Status']}[/yellow]"
-            )
-        elif result["Status"] == "InProgress":
-            console.print(
-                "[dim]ℹ️  Cleanup command still in progress (expected when SSM agent stops itself)[/dim]"
-            )
-
-        # Get instance details to find volumes to exclude
-        instances = ec2.describe_instances(InstanceIds=[engine["instance_id"]])
-        instance = instances["Reservations"][0]["Instances"][0]
-
-        root_device = instance.get("RootDeviceName", "/dev/xvda")
-        block_mappings = instance.get("BlockDeviceMappings", [])
-
-        # Build exclusion list for non-root volumes
-        block_device_mappings = []
-        for mapping in block_mappings:
-            device_name = mapping.get("DeviceName", "")
-            if device_name != root_device:
-                block_device_mappings.append(
-                    {"DeviceName": device_name, "NoDevice": ""}
-                )
-                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']})")
-
-            detach_failed = []
-            for s in attached_studios:
-                resp = make_api_request("POST", f"/studios/{s['studio_id']}/detach")
-                if resp.status_code != 200:
-                    # Check if actually detached despite error
-                    time.sleep(2)
-                    check_resp = make_api_request("GET", f"/studios/{s['studio_id']}")
-                    if check_resp.status_code == 200:
-                        studio_data = check_resp.json()
-                        if studio_data.get("status") == "available":
-                            console.print(
-                                f"  [yellow]⚠ {s['studio_id']} reported error but is detached[/yellow]"
-                            )
-                            continue  # It's actually detached, continue
-                    
-                    console.print(
-                        f"  [red]❌ Failed to detach {s['studio_id']}[/red]"
-                    )
-                    detach_failed.append(s['studio_id'])
-
-            if detach_failed:
-                console.print(
-                    f"[red]Failed to detach {len(detach_failed)} studio(s). Aborting AMI creation.[/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…")
-            
-            # Check volume states directly instead of using waiter
-            for attempt in range(24):  # Max 2 minutes
-                time.sleep(5)
-                volumes = ec2_wait.describe_volumes(VolumeIds=vol_ids)["Volumes"]
-                all_available = all(v["State"] == "available" for v in volumes)
-                if all_available:
-                    console.print("[green]✓ All studios detached[/green]")
-                    break
-            else:
-                console.print(
-                    "[yellow]⚠ Some volumes may still be detaching, but proceeding...[/yellow]"
-                )
-
+        # 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:
-            progress.add_task(
-                "Creating AMI (this will take several minutes)...", total=None
-            )
-
-            create_params = {
-                "InstanceId": engine["instance_id"],
-                "Name": ami_name,
-                "Description": ami_description,
-                "NoReboot": False,  # Important: reboot for clean snapshot
-                "TagSpecifications": [
+            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_type},
+                            {"Key": "EngineType", "Value": engine['engine_type']},
                             {"Key": "Name", "Value": ami_name},
                         ],
                     }
                 ],
-            }
-
-            if block_device_mappings:
-                create_params["BlockDeviceMappings"] = block_device_mappings
-
-            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]")
-
-            # Restore the source engine to a normal state
-            console.print("Restoring source engine state...")
-            
-            # Wait for instance to come back after reboot (AMI creation reboots by default)
-            console.print("[dim]Waiting for engine to reboot after snapshot...[/dim]")
-            ec2_waiter = ec2.get_waiter('instance_status_ok')
-            try:
-                ec2_waiter.wait(
-                    InstanceIds=[engine["instance_id"]],
-                    WaiterConfig={'Delay': 10, 'MaxAttempts': 30}  # Wait up to 5 minutes
-                )
-            except Exception as e:
-                console.print(f"[yellow]⚠️  Warning: Engine may still be rebooting: {e}[/yellow]")
-            
-            # Now restore the sentinel and restart services
-            restore_response = ssm.send_command(
-                InstanceIds=[engine["instance_id"]],
-                DocumentName="AWS-RunShellScript",
-                Parameters={
-                    "commands": [
-                        # Ensure the directories exist
-                        "sudo mkdir -p /opt/dayhoff /opt/dayhoff/state",
-                        # Recreate the sentinel file
-                        "sudo touch /opt/dayhoff/first_boot_complete.sentinel",
-                        # Mark bootstrap as finished
-                        "echo 'finished' | sudo tee /opt/dayhoff/state/engine-init.stage > /dev/null",
-                        # Restart idle detector if it exists
-                        "sudo systemctl restart engine-idle-detector.timer 2>/dev/null || true",
-                        # Ensure SSM agent is running
-                        "sudo systemctl start amazon-ssm-agent 2>/dev/null || true",
-                    ],
-                    "executionTimeout": ["60"],
-                },
             )
 
-            # Wait for restore command to complete
-            restore_command_id = restore_response["Command"]["CommandId"]
-            for _ in range(10):
-                time.sleep(2)
-                result = ssm.get_command_invocation(
-                    CommandId=restore_command_id,
-                    InstanceId=engine["instance_id"],
-                )
-                if result["Status"] in ["Success", "Failed"]:
-                    break
-            
-            if result["Status"] == "Success":
-                console.print(
-                    "[green]✓ Source engine restored to normal operation.[/green]"
-                )
-            else:
-                console.print(
-                    "[yellow]⚠️  Warning: Engine state restoration incomplete. You may need to run:[/yellow]"
-                )
-                console.print(
-                    f"[dim]  dh engine repair {engine['name']}[/dim]"
-                )
+            ami_id = response["ImageId"]
+            progress.update(task, completed=True, description=f"[green]✓ AMI creation initiated![/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."
-            )
+        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)
 
 
+
 # ==================== STUDIO COMMANDS ====================
 
 

{dayhoff_tools-1.6.1.dist-info → dayhoff_tools-1.6.2.dist-info}/METADATA

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

{dayhoff_tools-1.6.1.dist-info → dayhoff_tools-1.6.2.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=iZdl1eyJqwGVVm0vukY8sbnMv_k7CaesdnGBjE3cF0c,98283
+dayhoff_tools/cli/engine_commands.py,sha256=Pqh-x2dycEuAE-Ts762KWMdXFZXhClbMmSWTEjsb45o,92065
 dayhoff_tools/cli/main.py,sha256=tRN7WCBHg6uyNp6rA54pKTCoVmBntta2i0Yas3bUpZ4,4853
 dayhoff_tools/cli/swarm_commands.py,sha256=5EyKj8yietvT5lfoz8Zx0iQvVaNgc3SJX1z2zQR6o6M,5614
 dayhoff_tools/cli/utility_commands.py,sha256=FRZTPrjsG_qmIIqoNxd1Q1vVkS_5w8aY33IrVYVNCLg,18131
@@ -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=heaYc64qplgN3_1WVPFmqj53goStioWwY5NqlWc4c0s,24453
-dayhoff_tools-1.6.1.dist-info/METADATA,sha256=HL_pHrcjy_JbaTdP8O0mAP0xUzrEDp7VWpcFfuWwzbM,2914
-dayhoff_tools-1.6.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-dayhoff_tools-1.6.1.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.6.1.dist-info/RECORD,,
+dayhoff_tools-1.6.2.dist-info/METADATA,sha256=wIKq0EVcXH3E0fz12fWus4_LYESTopdNP-5aT2rjkVk,2914
+dayhoff_tools-1.6.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+dayhoff_tools-1.6.2.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.6.2.dist-info/RECORD,,