dayhoff-tools 1.3.6__tar.gz → 1.3.8__tar.gz

{dayhoff_tools-1.3.6 → dayhoff_tools-1.3.8}/PKG-INFO

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

{dayhoff_tools-1.3.6 → dayhoff_tools-1.3.8}/dayhoff_tools/cli/engine_commands.py

@@ -27,10 +27,10 @@ console = Console()
 # Cost information
 HOURLY_COSTS = {
     "cpu": 0.50,  # r6i.2xlarge
-    "cpumax": 1.00,  # r7i.8xlarge
-    "t4": 1.00,  # g4dn.2xlarge
-    "a10g": 2.00,  # g5.2xlarge
-    "a100": 5.00,  # p4d.24xlarge
+    "cpumax": 2.02,  # r7i.8xlarge
+    "t4": 0.75,  # g4dn.2xlarge
+    "a10g": 1.50,  # g5.2xlarge
+    "a100": 21.96,  # p4d.24xlarge
 }
 
 # SSH config management
@@ -90,7 +90,10 @@ def get_api_url() -> str:
 
 
 def make_api_request(
-    method: str, endpoint: str, json_data: Optional[Dict] = None, params: Optional[Dict] = None
+    method: str,
+    endpoint: str,
+    json_data: Optional[Dict] = None,
+    params: Optional[Dict] = None,
 ) -> requests.Response:
     """Make an API request with error handling."""
     api_url = get_api_url()
@@ -134,14 +137,14 @@ def parse_launch_time(launch_time_str: str) -> datetime:
         "%Y-%m-%dT%H:%M:%S+00:00",  # Explicit UTC offset
         "%Y-%m-%d %H:%M:%S",
     ]
-    
+
     # First try parsing with fromisoformat for better timezone handling
     try:
         # Handle the ISO format properly
-        return datetime.fromisoformat(launch_time_str.replace('Z', '+00:00'))
+        return datetime.fromisoformat(launch_time_str.replace("Z", "+00:00"))
     except (ValueError, AttributeError):
         pass
-    
+
     # Fallback to manual format parsing
     for fmt in formats:
         try:
@@ -152,7 +155,7 @@ def parse_launch_time(launch_time_str: str) -> datetime:
             return parsed
         except ValueError:
             continue
-    
+
     # Fallback: assume it's recent
     return datetime.now(timezone.utc)
 
@@ -163,7 +166,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":
@@ -213,7 +216,9 @@ def resolve_engine(name_or_id: str, engines: List[Dict]) -> Dict:
         while True:
             try:
                 choice = IntPrompt.ask(
-                    "Select engine", default=1, choices=[str(i) for i in range(1, len(matches) + 1)]
+                    "Select engine",
+                    default=1,
+                    choices=[str(i) for i in range(1, len(matches) + 1)],
                 )
                 return matches[choice - 1]
             except (ValueError, IndexError):
@@ -307,7 +312,9 @@ def launch_engine(
         raise typer.Exit(1)
 
     cost = HOURLY_COSTS.get(engine_type, 0)
-    console.print(f"Launching [cyan]{name}[/cyan] ({engine_type}) for ${cost:.2f}/hour...")
+    console.print(
+        f"Launching [cyan]{name}[/cyan] ({engine_type}) for ${cost:.2f}/hour..."
+    )
 
     with Progress(
         SpinnerColumn(),
@@ -332,14 +339,17 @@ 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"),
-    running_only: bool = typer.Option(False, "--running", help="Show only running engines"),
-    stopped_only: bool = typer.Option(False, "--stopped", help="Show only stopped 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 engines (shows all engines by default)."""
     current_user = check_aws_sso()
@@ -426,7 +436,9 @@ def engine_status(
     engine = resolve_engine(name_or_id, engines)
 
     # Get attached studios info
-    response = make_api_request("GET", f"/engines/{engine['instance_id']}/attached-studios")
+    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", [])
@@ -469,7 +481,9 @@ def engine_status(
 @engine_app.command("stop")
 def stop_engine(
     name_or_id: str = typer.Argument(help="Engine name or instance ID"),
-    force: bool = typer.Option(False, "--force", "-f", help="Force stop and detach all studios"),
+    force: bool = typer.Option(
+        False, "--force", "-f", help="Force stop and detach all studios"
+    ),
 ):
     """Stop an engine."""
     check_aws_sso()
@@ -569,7 +583,9 @@ def terminate_engine(
     hourly_cost = HOURLY_COSTS.get(engine["engine_type"], 0)
     total_cost = hourly_cost * (uptime.total_seconds() / 3600)
 
-    console.print(f"\n[yellow]⚠️  This will permanently terminate engine '{engine['name']}'[/yellow]")
+    console.print(
+        f"\n[yellow]⚠️  This will permanently terminate engine '{engine['name']}'[/yellow]"
+    )
     console.print(f"Total cost for this session: ${total_cost:.2f}")
 
     if not Confirm.ask("\nAre you sure you want to terminate this engine?"):
@@ -617,7 +633,9 @@ 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"),
+    all_engines: bool = typer.Option(
+        False, "--all", "-a", help="Include all engines from all users"
+    ),
 ):
     """Update SSH config with available engines."""
     username = check_aws_sso()
@@ -628,7 +646,9 @@ def config_ssh(
         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]...")
+            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")
@@ -638,13 +658,12 @@ 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"
+            e for e in running_engines if e["user"] == username or e["user"] == "shared"
         ]
 
     # Read existing config
@@ -681,7 +700,7 @@ def config_ssh(
                     f"Host {engine['name']} {SSH_MANAGED_COMMENT}",
                     f"    HostName {engine['instance_id']}",
                     f"    User {username}",
-                    f'    ProxyCommand sh -c "aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters \'portNumber=%p\'"',
+                    f"    ProxyCommand sh -c \"aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'\"",
                 ]
             )
 
@@ -696,8 +715,12 @@ def config_ssh(
             f"[green]✓ Updated SSH config with {len(running_engines)} engines[/green]"
         )
         for engine in running_engines:
-            user_display = f"[dim]({engine['user']})[/dim]" if engine['user'] != username else ""
-            console.print(f"  • {engine['name']} → {engine['instance_id']} {user_display}")
+            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")
@@ -710,6 +733,7 @@ def keep_awake(
 
     # Parse duration
     import re
+
     match = re.match(r"(?:(\d+)h)?(?:(\d+)m)?", duration)
     if not match or (not match.group(1) and not match.group(2)):
         console.print(f"[red]❌ Invalid duration format: {duration}[/red]")
@@ -753,6 +777,7 @@ def keep_awake(
 
         # Wait for command to complete
         import time
+
         for _ in range(10):
             time.sleep(1)
             result = ssm.get_command_invocation(
@@ -771,7 +796,9 @@ def keep_awake(
                 "[dim]Use keep-awake for nohup operations or other background tasks.[/dim]"
             )
         else:
-            console.print(f"[red]❌ Failed to set keep-awake: {result.get('StatusDetails', 'Unknown error')}[/red]")
+            console.print(
+                f"[red]❌ Failed to set keep-awake: {result.get('StatusDetails', 'Unknown error')}[/red]"
+            )
 
     except ClientError as e:
         console.print(f"[red]❌ Failed to set keep-awake: {e}[/red]")
@@ -811,6 +838,7 @@ def cancel_keep_awake(
 
         # Wait for command to complete
         import time
+
         for _ in range(10):
             time.sleep(1)
             result = ssm.get_command_invocation(
@@ -821,9 +849,13 @@ def cancel_keep_awake(
                 break
 
         if result["Status"] == "Success":
-            console.print("[green]✓ Keep-awake cancelled, auto-shutdown re-enabled[/green]")
+            console.print(
+                "[green]✓ Keep-awake cancelled, auto-shutdown re-enabled[/green]"
+            )
         else:
-            console.print(f"[red]❌ Failed to cancel keep-awake: {result.get('StatusDetails', 'Unknown error')}[/red]")
+            console.print(
+                f"[red]❌ Failed to cancel keep-awake: {result.get('StatusDetails', 'Unknown error')}[/red]"
+            )
 
     except ClientError as e:
         console.print(f"[red]❌ Failed to cancel keep-awake: {e}[/red]")
@@ -831,7 +863,9 @@ def cancel_keep_awake(
 
 @engine_app.command("create-ami")
 def create_ami(
-    name_or_id: str = typer.Argument(help="Engine name or instance ID to create AMI from"),
+    name_or_id: str = typer.Argument(
+        help="Engine name or instance ID to create AMI from"
+    ),
 ):
     """Create a Golden AMI from an engine."""
     check_aws_sso()
@@ -860,7 +894,9 @@ 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]")
+    console.print(
+        "\n[yellow]⚠️  Important: This will reboot the engine to ensure a clean snapshot.[/yellow]"
+    )
 
     if not Confirm.ask("\nContinue with AMI creation?"):
         console.print("AMI creation cancelled.")
@@ -873,7 +909,7 @@ def create_ami(
         # First, we need to clean up the sentinel file via SSM
         console.print("Cleaning up bootstrap sentinel file...")
         ssm = boto3.client("ssm", region_name="us-east-1")
-        
+
         cleanup_response = ssm.send_command(
             InstanceIds=[engine["instance_id"]],
             DocumentName="AWS-RunShellScript",
@@ -889,6 +925,7 @@ def create_ami(
 
         # Wait for cleanup to complete
         import time
+
         command_id = cleanup_response["Command"]["CommandId"]
         for _ in range(10):
             time.sleep(1)
@@ -900,21 +937,25 @@ def create_ami(
                 break
 
         if result["Status"] != "Success":
-            console.print("[yellow]⚠️  Warning: Cleanup command may have failed[/yellow]")
+            console.print(
+                "[yellow]⚠️  Warning: Cleanup command may have failed[/yellow]"
+            )
 
         # 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": ""})
+                block_device_mappings.append(
+                    {"DeviceName": device_name, "NoDevice": ""}
+                )
                 console.print(f"  Excluding volume at {device_name}")
 
         # Create the AMI
@@ -923,7 +964,9 @@ def create_ami(
             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"],
@@ -948,15 +991,50 @@ 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]")
@@ -988,7 +1066,9 @@ def create_studio(
     # Check if user already has a studio
     existing = get_user_studio(username)
     if existing:
-        console.print(f"[yellow]You already have a studio: {existing['studio_id']}[/yellow]")
+        console.print(
+            f"[yellow]You already have a studio: {existing['studio_id']}[/yellow]"
+        )
         return
 
     console.print(f"Creating {size_gb}GB studio for user [cyan]{username}[/cyan]...")
@@ -1030,14 +1110,14 @@ def studio_status():
 
     # Create status panel
     # Format status with colors
-    status = studio['status']
+    status = studio["status"]
     if status == "in-use":
         status_display = "[bright_blue]attached[/bright_blue]"
     elif status in ["attaching", "detaching"]:
         status_display = f"[yellow]{status}[/yellow]"
     else:
         status_display = f"[green]{status}[/green]"
-    
+
     status_lines = [
         f"[bold]Studio ID:[/bold]    {studio['studio_id']}",
         f"[bold]User:[/bold]         {studio['user']}",
@@ -1048,17 +1128,19 @@ def studio_status():
 
     if studio.get("attached_vm_id"):
         status_lines.append(f"[bold]Attached to:[/bold]  {studio['attached_vm_id']}")
-        
+
         # Try to get engine details
         response = make_api_request("GET", "/engines")
         if response.status_code == 200:
             engines = response.json().get("engines", [])
             attached_engine = next(
                 (e for e in engines if e["instance_id"] == studio["attached_vm_id"]),
-                None
+                None,
             )
             if attached_engine:
-                status_lines.append(f"[bold]Engine Name:[/bold]  {attached_engine['name']}")
+                status_lines.append(
+                    f"[bold]Engine Name:[/bold]  {attached_engine['name']}"
+                )
 
     panel = Panel(
         "\n".join(status_lines),
@@ -1118,14 +1200,19 @@ def attach_studio(
 
     if engine["state"].lower() != "running":
         console.print(f"[yellow]⚠️  Engine is {engine['state']}[/yellow]")
-        if engine["state"].lower() == "stopped" and Confirm.ask("Start the engine first?"):
-            response = make_api_request("POST", f"/engines/{engine['instance_id']}/start")
+        if engine["state"].lower() == "stopped" and Confirm.ask(
+            "Start the engine first?"
+        ):
+            response = make_api_request(
+                "POST", f"/engines/{engine['instance_id']}/start"
+            )
             if response.status_code != 200:
                 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)
         else:
             raise typer.Exit(1)
@@ -1206,7 +1293,9 @@ def delete_studio():
         console.print("[yellow]You don't have a studio to delete.[/yellow]")
         return
 
-    console.print("[red]⚠️  WARNING: This will permanently delete your studio and all data![/red]")
+    console.print(
+        "[red]⚠️  WARNING: This will permanently delete your studio and all data![/red]"
+    )
     console.print(f"Studio ID: {studio['studio_id']}")
     console.print(f"Size: {studio['size_gb']}GB")
 
@@ -1219,9 +1308,7 @@ def delete_studio():
         console.print("Deletion cancelled.")
         return
 
-    typed_confirm = Prompt.ask(
-        'Type "DELETE" to confirm permanent deletion'
-    )
+    typed_confirm = Prompt.ask('Type "DELETE" to confirm permanent deletion')
     if typed_confirm != "DELETE":
         console.print("Deletion cancelled.")
         return
@@ -1237,7 +1324,9 @@ def delete_studio():
 
 @studio_app.command("list")
 def list_studios(
-    all_users: bool = typer.Option(False, "--all", "-a", help="Show all users' studios"),
+    all_users: bool = typer.Option(
+        False, "--all", "-a", help="Show all users' studios"
+    ),
 ):
     """List studios."""
     username = check_aws_sso()
@@ -1275,24 +1364,24 @@ def list_studios(
                 status_display = "[yellow]" + studio["status"] + "[/yellow]"
             else:
                 status_display = "[green]available[/green]"
-            
+
             # Format attached engine info
             attached_to = "-"
             if studio.get("attached_vm_id"):
                 vm_id = studio["attached_vm_id"]
                 engine_name = engines.get(vm_id, "unknown")
                 attached_to = f"{engine_name} ({vm_id})"
-            
+
             # Format creation date (remove microseconds and timezone info)
             created = studio["creation_date"]
             try:
                 # Parse and reformat to just show date and time
-                if 'T' in created:
-                    created_dt = datetime.fromisoformat(created.replace('Z', '+00:00'))
+                if "T" in created:
+                    created_dt = datetime.fromisoformat(created.replace("Z", "+00:00"))
                     created = created_dt.strftime("%Y-%m-%d %H:%M")
             except:
                 # If parsing fails, just truncate
-                created = created.split('T')[0] if 'T' in created else created[:16]
+                created = created.split("T")[0] if "T" in created else created[:16]
 
             table.add_row(
                 studio["studio_id"],

{dayhoff_tools-1.3.6 → dayhoff_tools-1.3.8}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "dayhoff-tools"
-version = "1.3.6"
+version = "1.3.8"
 description = "Common tools for all the repos at Dayhoff Labs"
 authors = [
     {name = "Daniel Martin-Alarcon", email = "dma@dayhofflabs.com"}