dayhoff-tools 1.3.2__py3-none-any.whl → 1.3.4__py3-none-any.whl

dayhoff_tools/cli/engine_commands.py

@@ -3,7 +3,7 @@
 import json
 import subprocess
 import sys
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
@@ -134,11 +134,15 @@ def parse_launch_time(launch_time_str: str) -> datetime:
     ]
     for fmt in formats:
         try:
-            return datetime.strptime(launch_time_str, fmt)
+            parsed = datetime.strptime(launch_time_str, fmt)
+            # If the format includes 'Z', it's UTC
+            if fmt.endswith('Z'):
+                parsed = parsed.replace(tzinfo=timezone.utc)
+            return parsed
         except ValueError:
             continue
     # Fallback: assume it's recent
-    return datetime.utcnow()
+    return datetime.now(timezone.utc)
 
 
 def format_status(state: str, ready: Optional[bool]) -> str:
@@ -330,9 +334,7 @@ def list_engines(
     current_user = check_aws_sso()
 
     params = {}
-    if not all_users and not user:
-        params["user"] = current_user
-    elif user:
+    if user:
         params["user"] = user
 
     response = make_api_request("GET", "/engines", params=params)
@@ -365,7 +367,7 @@ def list_engines(
         total_cost = 0.0
         for engine in engines:
             launch_time = parse_launch_time(engine["launch_time"])
-            uptime = datetime.utcnow() - launch_time
+            uptime = datetime.now(timezone.utc) - launch_time
             hourly_cost = HOURLY_COSTS.get(engine["engine_type"], 0)
 
             if engine["state"].lower() == "running":
@@ -420,7 +422,7 @@ def engine_status(
 
     # Calculate costs
     launch_time = parse_launch_time(engine["launch_time"])
-    uptime = datetime.utcnow() - launch_time
+    uptime = datetime.now(timezone.utc) - launch_time
     hourly_cost = HOURLY_COSTS.get(engine["engine_type"], 0)
     total_cost = hourly_cost * (uptime.total_seconds() / 3600)
 
@@ -552,7 +554,7 @@ def terminate_engine(
 
     # Calculate cost
     launch_time = parse_launch_time(engine["launch_time"])
-    uptime = datetime.utcnow() - launch_time
+    uptime = datetime.now(timezone.utc) - launch_time
     hourly_cost = HOURLY_COSTS.get(engine["engine_type"], 0)
     total_cost = hourly_cost * (uptime.total_seconds() / 3600)
 
@@ -947,7 +949,7 @@ def get_user_studio(username: str) -> Optional[Dict]:
         return None
 
     studios = response.json().get("studios", [])
-    user_studios = [s for s in studios if s["UserID"] == username]
+    user_studios = [s for s in studios if s["user"] == username]
 
     return user_studios[0] if user_studios else None
 
@@ -962,7 +964,7 @@ 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['StudioID']}[/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]...")
@@ -1003,23 +1005,32 @@ def studio_status():
         return
 
     # Create status panel
+    # Format status with colors
+    status = studio['status']
+    if status == "in-use":
+        status_display = "[magenta]attached[/magenta]"
+    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['StudioID']}",
-        f"[bold]User:[/bold]         {studio['UserID']}",
-        f"[bold]Status:[/bold]       {studio['Status']}",
-        f"[bold]Size:[/bold]         {studio['SizeGB']}GB",
-        f"[bold]Created:[/bold]      {studio['CreationDate']}",
+        f"[bold]Studio ID:[/bold]    {studio['studio_id']}",
+        f"[bold]User:[/bold]         {studio['user']}",
+        f"[bold]Status:[/bold]       {status_display}",
+        f"[bold]Size:[/bold]         {studio['size_gb']}GB",
+        f"[bold]Created:[/bold]      {studio['creation_date']}",
     ]
 
-    if studio.get("AttachedVMID"):
-        status_lines.append(f"[bold]Attached to:[/bold]  {studio['AttachedVMID']}")
+    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["AttachedVMID"]),
+                (e for e in engines if e["instance_id"] == studio["attached_vm_id"]),
                 None
             )
             if attached_engine:
@@ -1055,19 +1066,19 @@ def attach_studio(
                 console.print("[red]❌ Failed to create studio[/red]")
                 raise typer.Exit(1)
             studio = response.json()
-            studio["StudioID"] = studio["studio_id"]  # Normalize key
+            studio["studio_id"] = studio["studio_id"]  # Normalize key
         else:
             raise typer.Exit(0)
 
     # Check if already attached
-    if studio.get("Status") == "in-use":
+    if studio.get("status") == "in-use":
         console.print(
-            f"[yellow]Studio is already attached to {studio.get('AttachedVMID')}[/yellow]"
+            f"[yellow]Studio is already attached to {studio.get('attached_vm_id')}[/yellow]"
         )
         if not Confirm.ask("Detach and reattach to new engine?"):
             return
         # Detach first
-        response = make_api_request("POST", f"/studios/{studio['StudioID']}/detach")
+        response = make_api_request("POST", f"/studios/{studio['studio_id']}/detach")
         if response.status_code != 200:
             console.print("[red]❌ Failed to detach studio[/red]")
             raise typer.Exit(1)
@@ -1113,7 +1124,7 @@ def attach_studio(
 
         response = make_api_request(
             "POST",
-            f"/studios/{studio['StudioID']}/attach",
+            f"/studios/{studio['studio_id']}/attach",
             json_data={
                 "vm_id": engine["instance_id"],
                 "user": username,
@@ -1146,13 +1157,13 @@ def detach_studio():
         console.print("[yellow]You don't have a studio.[/yellow]")
         return
 
-    if studio.get("Status") != "in-use":
+    if studio.get("status") != "in-use":
         console.print("[yellow]Your studio is not attached to any engine.[/yellow]")
         return
 
-    console.print(f"Detaching studio from {studio.get('AttachedVMID')}...")
+    console.print(f"Detaching studio from {studio.get('attached_vm_id')}...")
 
-    response = make_api_request("POST", f"/studios/{studio['StudioID']}/detach")
+    response = make_api_request("POST", f"/studios/{studio['studio_id']}/detach")
 
     if response.status_code == 200:
         console.print(f"[green]✓ Studio detached successfully![/green]")
@@ -1172,8 +1183,8 @@ def delete_studio():
         return
 
     console.print("[red]⚠️  WARNING: This will permanently delete your studio and all data![/red]")
-    console.print(f"Studio ID: {studio['StudioID']}")
-    console.print(f"Size: {studio['SizeGB']}GB")
+    console.print(f"Studio ID: {studio['studio_id']}")
+    console.print(f"Size: {studio['size_gb']}GB")
 
     # Multiple confirmations
     if not Confirm.ask("\nAre you sure you want to delete your studio?"):
@@ -1191,7 +1202,7 @@ def delete_studio():
         console.print("Deletion cancelled.")
         return
 
-    response = make_api_request("DELETE", f"/studios/{studio['StudioID']}")
+    response = make_api_request("DELETE", f"/studios/{studio['studio_id']}")
 
     if response.status_code == 200:
         console.print(f"[green]✓ Studio deleted successfully![/green]")
@@ -1212,13 +1223,17 @@ def list_studios(
     if response.status_code == 200:
         studios = response.json().get("studios", [])
 
-        if not all_users:
-            studios = [s for s in studios if s["UserID"] == username]
-
         if not studios:
             console.print("No studios found.")
             return
 
+        # Get all engines to map instance IDs to names
+        engines_response = make_api_request("GET", "/engines")
+        engines = {}
+        if engines_response.status_code == 200:
+            for engine in engines_response.json().get("engines", []):
+                engines[engine["instance_id"]] = engine["name"]
+
         # Create table
         table = Table(title="Studios", box=box.ROUNDED)
         table.add_column("Studio ID", style="cyan")
@@ -1229,14 +1244,39 @@ def list_studios(
         table.add_column("Created")
 
         for studio in studios:
-            status_color = "green" if studio["Status"] == "available" else "yellow"
+            # Change status display
+            if studio["status"] == "in-use":
+                status_display = "[magenta]attached[/magenta]"
+            elif studio["status"] in ["attaching", "detaching"]:
+                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'))
+                    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]
+
             table.add_row(
-                studio["StudioID"],
-                studio["UserID"],
-                f"[{status_color}]{studio['Status']}[/{status_color}]",
-                f"{studio['SizeGB']}GB",
-                studio.get("AttachedVMID", "-"),
-                studio["CreationDate"],
+                studio["studio_id"],
+                studio["user"],
+                status_display,
+                f"{studio['size_gb']}GB",
+                attached_to,
+                created,
             )
 
         console.print(table)
@@ -1256,9 +1296,9 @@ def reset_studio():
         return
 
     console.print(f"[yellow]⚠️  This will force-reset your studio state[/yellow]")
-    console.print(f"Current status: {studio['Status']}")
-    if studio.get("AttachedVMID"):
-        console.print(f"Listed as attached to: {studio['AttachedVMID']}")
+    console.print(f"Current status: {studio['status']}")
+    if studio.get("attached_vm_id"):
+        console.print(f"Listed as attached to: {studio['attached_vm_id']}")
 
     if not Confirm.ask("\nReset studio state?"):
         console.print("Reset cancelled.")
@@ -1273,7 +1313,7 @@ def reset_studio():
     try:
         # Check if volume is actually attached
         ec2 = boto3.client("ec2", region_name="us-east-1")
-        volumes = ec2.describe_volumes(VolumeIds=[studio["StudioID"]])
+        volumes = ec2.describe_volumes(VolumeIds=[studio["studio_id"]])
 
         if volumes["Volumes"]:
             volume = volumes["Volumes"][0]
@@ -1284,19 +1324,19 @@ def reset_studio():
                 )
                 if Confirm.ask("Force-detach the volume?"):
                     ec2.detach_volume(
-                        VolumeId=studio["StudioID"],
+                        VolumeId=studio["studio_id"],
                         InstanceId=attachments[0]["InstanceId"],
                         Force=True,
                     )
                     console.print("Waiting for volume to detach...")
                     waiter = ec2.get_waiter("volume_available")
-                    waiter.wait(VolumeIds=[studio["StudioID"]])
+                    waiter.wait(VolumeIds=[studio["studio_id"]])
 
         # Reset in DynamoDB
         table.update_item(
-            Key={"StudioID": studio["StudioID"]},
-            UpdateExpression="SET #status = :status, AttachedVMID = :vm_id, AttachedDevice = :device",
-            ExpressionAttributeNames={"#status": "Status"},
+            Key={"studio_id": studio["studio_id"]},
+            UpdateExpression="SET #status = :status, attached_vm_id = :vm_id, attached_device = :device",
+            ExpressionAttributeNames={"#status": "status"},
             ExpressionAttributeValues={
                 ":status": "available",
                 ":vm_id": None,

{dayhoff_tools-1.3.2.dist-info → dayhoff_tools-1.3.4.dist-info}/METADATA

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

{dayhoff_tools-1.3.2.dist-info → dayhoff_tools-1.3.4.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=zHyisrYz0kgm_EXFIRqSR10jPPXrBmhMpZnN296x86U,46408
+dayhoff_tools/cli/engine_commands.py,sha256=w0wzs4aQUdzeyAyhHH9lHxpJEG8acrbguQ4Rd0u5bA4,48179
 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.2.dist-info/METADATA,sha256=i0AQJh4nEsStInWpyTviptfaOltjscM-MObda7o0vlI,2842
-dayhoff_tools-1.3.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-dayhoff_tools-1.3.2.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.3.2.dist-info/RECORD,,
+dayhoff_tools-1.3.4.dist-info/METADATA,sha256=YZw82MtQpm4aWNzOfEl62Iov6ZRXtw6LfZ43eI9psKE,2842
+dayhoff_tools-1.3.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+dayhoff_tools-1.3.4.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.3.4.dist-info/RECORD,,