dayhoff-tools 1.3.9__tar.gz → 1.3.11__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (32) hide show
  1. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/PKG-INFO +2 -2
  2. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/engine_commands.py +419 -58
  3. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/pyproject.toml +2 -2
  4. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/README.md +0 -0
  5. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/__init__.py +0 -0
  6. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/chemistry/standardizer.py +0 -0
  7. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/chemistry/utils.py +0 -0
  8. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/__init__.py +0 -0
  9. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/cloud_commands.py +0 -0
  10. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/main.py +0 -0
  11. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/swarm_commands.py +0 -0
  12. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/utility_commands.py +0 -0
  13. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/base.py +0 -0
  14. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/deploy_aws.py +0 -0
  15. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/deploy_gcp.py +0 -0
  16. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/deploy_utils.py +0 -0
  17. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/job_runner.py +0 -0
  18. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/processors.py +0 -0
  19. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/swarm.py +0 -0
  20. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/embedders.py +0 -0
  21. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/fasta.py +0 -0
  22. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/file_ops.py +0 -0
  23. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/h5.py +0 -0
  24. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/gcp.py +0 -0
  25. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/gtdb.py +0 -0
  26. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/kegg.py +0 -0
  27. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/mmseqs.py +0 -0
  28. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/structure.py +0 -0
  29. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/uniprot.py +0 -0
  30. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/logs.py +0 -0
  31. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/sqlite.py +0 -0
  32. {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/warehouse.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: dayhoff-tools
3
- Version: 1.3.9
3
+ Version: 1.3.11
4
4
  Summary: Common tools for all the repos at Dayhoff Labs
5
5
  Author: Daniel Martin-Alarcon
6
6
  Author-email: dma@dayhofflabs.com
@@ -14,7 +14,7 @@ Provides-Extra: embedders
14
14
  Provides-Extra: full
15
15
  Requires-Dist: biopython (>=1.84) ; extra == "full"
16
16
  Requires-Dist: biopython (>=1.85) ; extra == "embedders"
17
- Requires-Dist: boto3 (>=1.36.8) ; extra == "full"
17
+ Requires-Dist: boto3 (>=1.36.8)
18
18
  Requires-Dist: docker (>=7.1.0) ; extra == "full"
19
19
  Requires-Dist: fair-esm (>=2.0.0) ; extra == "embedders"
20
20
  Requires-Dist: fair-esm (>=2.0.0) ; extra == "full"
@@ -3,6 +3,7 @@
3
3
  import json
4
4
  import subprocess
5
5
  import sys
6
+ import time
6
7
  from datetime import datetime, timedelta, timezone
7
8
  from pathlib import Path
8
9
  from typing import Dict, List, Optional, Tuple
@@ -127,6 +128,95 @@ def format_duration(duration: timedelta) -> str:
127
128
  return f"{minutes}m"
128
129
 
129
130
 
131
+ def get_disk_usage_via_ssm(instance_id: str) -> Optional[str]:
132
+ """Get disk usage for an engine via SSM.
133
+
134
+ Returns:
135
+ String like "17/50 GB" or None if failed
136
+ """
137
+ try:
138
+ ssm = boto3.client("ssm", region_name="us-east-1")
139
+
140
+ # Run df command to get disk usage
141
+ response = ssm.send_command(
142
+ InstanceIds=[instance_id],
143
+ DocumentName="AWS-RunShellScript",
144
+ Parameters={
145
+ "commands": [
146
+ # Get root filesystem usage in GB
147
+ "df -BG / | tail -1 | awk '{gsub(/G/, \"\", $2); gsub(/G/, \"\", $3); print $3 \"/\" $2 \" GB\"}'"
148
+ ],
149
+ "executionTimeout": ["10"],
150
+ },
151
+ )
152
+
153
+ command_id = response["Command"]["CommandId"]
154
+
155
+ # Wait for command to complete (with timeout)
156
+ for _ in range(5): # 5 second timeout
157
+ time.sleep(1)
158
+ result = ssm.get_command_invocation(
159
+ CommandId=command_id,
160
+ InstanceId=instance_id,
161
+ )
162
+ if result["Status"] in ["Success", "Failed"]:
163
+ break
164
+
165
+ if result["Status"] == "Success":
166
+ output = result["StandardOutputContent"].strip()
167
+ return output if output else None
168
+
169
+ return None
170
+
171
+ except Exception as e:
172
+ # logger.debug(f"Failed to get disk usage for {instance_id}: {e}") # Original code had this line commented out
173
+ return None
174
+
175
+
176
+ def get_studio_disk_usage_via_ssm(instance_id: str, username: str) -> Optional[str]:
177
+ """Get disk usage for a studio via SSM.
178
+
179
+ Returns:
180
+ String like "333/500 GB" or None if failed
181
+ """
182
+ try:
183
+ ssm = boto3.client("ssm", region_name="us-east-1")
184
+
185
+ # Run df command to get studio disk usage
186
+ response = ssm.send_command(
187
+ InstanceIds=[instance_id],
188
+ DocumentName="AWS-RunShellScript",
189
+ Parameters={
190
+ "commands": [
191
+ # Get studio filesystem usage in GB
192
+ f"df -BG /studios/{username} 2>/dev/null | tail -1 | awk '{{gsub(/G/, \"\", $2); gsub(/G/, \"\", $3); print $3 \"/\" $2 \" GB\"}}'"
193
+ ],
194
+ "executionTimeout": ["10"],
195
+ },
196
+ )
197
+
198
+ command_id = response["Command"]["CommandId"]
199
+
200
+ # Wait for command to complete (with timeout)
201
+ for _ in range(5): # 5 second timeout
202
+ time.sleep(1)
203
+ result = ssm.get_command_invocation(
204
+ CommandId=command_id,
205
+ InstanceId=instance_id,
206
+ )
207
+ if result["Status"] in ["Success", "Failed"]:
208
+ break
209
+
210
+ if result["Status"] == "Success":
211
+ output = result["StandardOutputContent"].strip()
212
+ return output if output else None
213
+
214
+ return None
215
+
216
+ except Exception:
217
+ return None
218
+
219
+
130
220
  def parse_launch_time(launch_time_str: str) -> datetime:
131
221
  """Parse launch time from API response."""
132
222
  # Try different datetime formats
@@ -381,6 +471,7 @@ def list_engines(
381
471
  table.add_column("Type")
382
472
  table.add_column("User")
383
473
  table.add_column("Status")
474
+ table.add_column("Disk Usage")
384
475
  table.add_column("Uptime/Since")
385
476
  table.add_column("$/hour", justify="right")
386
477
  table.add_column("Cost Today", justify="right", style="yellow")
@@ -395,9 +486,12 @@ def list_engines(
395
486
  daily_cost = hourly_cost * min(uptime.total_seconds() / 3600, 24)
396
487
  total_cost += daily_cost
397
488
  time_str = format_duration(uptime)
489
+ # Get disk usage for running engines
490
+ disk_usage = get_disk_usage_via_ssm(engine["instance_id"]) or "-"
398
491
  else:
399
492
  daily_cost = 0
400
493
  time_str = launch_time.strftime("%Y-%m-%d %H:%M")
494
+ disk_usage = "-"
401
495
 
402
496
  table.add_row(
403
497
  engine["name"],
@@ -405,6 +499,7 @@ def list_engines(
405
499
  engine["engine_type"],
406
500
  engine["user"],
407
501
  format_status(engine["state"], engine.get("ready")),
502
+ disk_usage,
408
503
  time_str,
409
504
  f"${hourly_cost:.2f}",
410
505
  f"${daily_cost:.2f}" if daily_cost > 0 else "-",
@@ -776,8 +871,6 @@ def keep_awake(
776
871
  command_id = response["Command"]["CommandId"]
777
872
 
778
873
  # Wait for command to complete
779
- import time
780
-
781
874
  for _ in range(10):
782
875
  time.sleep(1)
783
876
  result = ssm.get_command_invocation(
@@ -837,8 +930,6 @@ def cancel_keep_awake(
837
930
  command_id = response["Command"]["CommandId"]
838
931
 
839
932
  # Wait for command to complete
840
- import time
841
-
842
933
  for _ in range(10):
843
934
  time.sleep(1)
844
935
  result = ssm.get_command_invocation(
@@ -861,6 +952,124 @@ def cancel_keep_awake(
861
952
  console.print(f"[red]❌ Failed to cancel keep-awake: {e}[/red]")
862
953
 
863
954
 
955
+ @engine_app.command("resize")
956
+ def resize_engine(
957
+ name_or_id: str = typer.Argument(help="Engine name or instance ID"),
958
+ size: int = typer.Option(..., "--size", "-s", help="New size in GB"),
959
+ online: bool = typer.Option(
960
+ False, "--online", help="Resize while running (requires manual filesystem expansion)"
961
+ ),
962
+ ):
963
+ """Resize an engine's boot disk."""
964
+ check_aws_sso()
965
+
966
+ # Get all engines to resolve name
967
+ response = make_api_request("GET", "/engines")
968
+ if response.status_code != 200:
969
+ console.print("[red]❌ Failed to fetch engines[/red]")
970
+ raise typer.Exit(1)
971
+
972
+ engines = response.json().get("engines", [])
973
+ engine = resolve_engine(name_or_id, engines)
974
+
975
+ # Get current volume info to validate size
976
+ ec2 = boto3.client("ec2", region_name="us-east-1")
977
+
978
+ try:
979
+ # Get instance details to find root volume
980
+ instance_info = ec2.describe_instances(InstanceIds=[engine["instance_id"]])
981
+ instance = instance_info["Reservations"][0]["Instances"][0]
982
+
983
+ # Find root volume
984
+ root_device = instance.get("RootDeviceName", "/dev/xvda")
985
+ root_volume_id = None
986
+
987
+ for bdm in instance.get("BlockDeviceMappings", []):
988
+ if bdm["DeviceName"] == root_device:
989
+ root_volume_id = bdm["Ebs"]["VolumeId"]
990
+ break
991
+
992
+ if not root_volume_id:
993
+ console.print("[red]❌ Could not find root volume[/red]")
994
+ raise typer.Exit(1)
995
+
996
+ # Get current volume size
997
+ volumes = ec2.describe_volumes(VolumeIds=[root_volume_id])
998
+ current_size = volumes["Volumes"][0]["Size"]
999
+
1000
+ if size <= current_size:
1001
+ console.print(f"[red]❌ New size ({size}GB) must be larger than current size ({current_size}GB)[/red]")
1002
+ raise typer.Exit(1)
1003
+
1004
+ console.print(f"[yellow]Resizing engine boot disk from {current_size}GB to {size}GB[/yellow]")
1005
+
1006
+ # Check if we need to stop the instance
1007
+ if not online and engine["state"].lower() == "running":
1008
+ console.print("Stopping engine for offline resize...")
1009
+ stop_response = make_api_request(
1010
+ "POST", f"/engines/{engine['instance_id']}/stop", json_data={"detach_studios": False}
1011
+ )
1012
+ if stop_response.status_code != 200:
1013
+ console.print("[red]❌ Failed to stop engine[/red]")
1014
+ raise typer.Exit(1)
1015
+
1016
+ # Wait for instance to stop
1017
+ console.print("Waiting for engine to stop...")
1018
+ waiter = ec2.get_waiter("instance_stopped")
1019
+ waiter.wait(InstanceIds=[engine["instance_id"]])
1020
+ console.print("[green]✓ Engine stopped[/green]")
1021
+
1022
+ # Call the resize API
1023
+ console.print("Resizing volume...")
1024
+ resize_response = make_api_request(
1025
+ "POST",
1026
+ f"/engines/{engine['instance_id']}/resize",
1027
+ json_data={"size": size}
1028
+ )
1029
+
1030
+ if resize_response.status_code != 200:
1031
+ error = resize_response.json().get("error", "Unknown error")
1032
+ console.print(f"[red]❌ Failed to resize engine: {error}[/red]")
1033
+ raise typer.Exit(1)
1034
+
1035
+ # Wait for volume modification to complete
1036
+ console.print("Waiting for volume modification to complete...")
1037
+ while True:
1038
+ mod_state = ec2.describe_volumes_modifications(VolumeIds=[root_volume_id])
1039
+ if not mod_state["VolumesModifications"]:
1040
+ break # Modification complete
1041
+ state = mod_state["VolumesModifications"][0]["ModificationState"]
1042
+ if state == "completed":
1043
+ break
1044
+ elif state == "failed":
1045
+ console.print("[red]❌ Volume modification failed[/red]")
1046
+ raise typer.Exit(1)
1047
+ time.sleep(5)
1048
+
1049
+ console.print("[green]✓ Volume resized successfully[/green]")
1050
+
1051
+ # If offline resize, start the instance back up
1052
+ if not online and engine["state"].lower() == "running":
1053
+ console.print("Starting engine back up...")
1054
+ start_response = make_api_request("POST", f"/engines/{engine['instance_id']}/start")
1055
+ if start_response.status_code != 200:
1056
+ console.print("[yellow]⚠️ Failed to restart engine automatically[/yellow]")
1057
+ console.print(f"Please start it manually: [cyan]dh engine start {engine['name']}[/cyan]")
1058
+ else:
1059
+ console.print("[green]✓ Engine started[/green]")
1060
+ console.print("The filesystem will be automatically expanded on boot.")
1061
+
1062
+ elif online and engine["state"].lower() == "running":
1063
+ console.print("\n[yellow]⚠️ Online resize complete. You must now expand the filesystem:[/yellow]")
1064
+ console.print(f"1. SSH into the engine: [cyan]ssh {engine['name']}[/cyan]")
1065
+ console.print("2. Run: [cyan]sudo growpart /dev/nvme0n1 1[/cyan]")
1066
+ console.print("3. Run: [cyan]sudo xfs_growfs /[/cyan]")
1067
+
1068
+ except ClientError as e:
1069
+ console.print(f"[red]❌ Failed to resize engine: {e}[/red]")
1070
+ raise typer.Exit(1)
1071
+
1072
+
864
1073
  @engine_app.command("create-ami")
865
1074
  def create_ami(
866
1075
  name_or_id: str = typer.Argument(
@@ -924,8 +1133,6 @@ def create_ami(
924
1133
  )
925
1134
 
926
1135
  # Wait for cleanup to complete
927
- import time
928
-
929
1136
  command_id = cleanup_response["Command"]["CommandId"]
930
1137
  for _ in range(10):
931
1138
  time.sleep(1)
@@ -1098,14 +1305,26 @@ def create_studio(
1098
1305
 
1099
1306
 
1100
1307
  @studio_app.command("status")
1101
- def studio_status():
1308
+ def studio_status(
1309
+ user: Optional[str] = typer.Option(None, "--user", "-u", help="Check status for a different user (admin only)"),
1310
+ ):
1102
1311
  """Show status of your studio."""
1103
1312
  username = check_aws_sso()
1104
-
1105
- studio = get_user_studio(username)
1313
+
1314
+ # Use specified user if provided, otherwise use current user
1315
+ target_user = user if user else username
1316
+
1317
+ # Add warning when checking another user's studio
1318
+ if target_user != username:
1319
+ console.print(f"[yellow]⚠️ Checking studio status for user: {target_user}[/yellow]")
1320
+
1321
+ studio = get_user_studio(target_user)
1106
1322
  if not studio:
1107
- console.print("[yellow]You don't have a studio yet.[/yellow]")
1108
- console.print("Create one with: [cyan]dh studio create[/cyan]")
1323
+ if target_user == username:
1324
+ console.print("[yellow]You don't have a studio yet.[/yellow]")
1325
+ console.print("Create one with: [cyan]dh studio create[/cyan]")
1326
+ else:
1327
+ console.print(f"[yellow]User {target_user} doesn't have a studio.[/yellow]")
1109
1328
  return
1110
1329
 
1111
1330
  # Create status panel
@@ -1153,28 +1372,43 @@ def studio_status():
1153
1372
  @studio_app.command("attach")
1154
1373
  def attach_studio(
1155
1374
  engine_name_or_id: str = typer.Argument(help="Engine name or instance ID"),
1375
+ user: Optional[str] = typer.Option(None, "--user", "-u", help="Attach a different user's studio (admin only)"),
1156
1376
  ):
1157
1377
  """Attach your studio to an engine."""
1158
1378
  username = check_aws_sso()
1379
+
1380
+ # Use specified user if provided, otherwise use current user
1381
+ target_user = user if user else username
1382
+
1383
+ # Add confirmation when attaching another user's studio
1384
+ if target_user != username:
1385
+ console.print(f"[yellow]⚠️ Managing studio for user: {target_user}[/yellow]")
1386
+ if not Confirm.ask(f"Are you sure you want to attach {target_user}'s studio?"):
1387
+ console.print("Operation cancelled.")
1388
+ return
1159
1389
 
1160
1390
  # Get user's studio
1161
- studio = get_user_studio(username)
1391
+ studio = get_user_studio(target_user)
1162
1392
  if not studio:
1163
- console.print("[yellow]You don't have a studio yet.[/yellow]")
1164
- if Confirm.ask("Would you like to create one now?"):
1165
- size = IntPrompt.ask("Studio size (GB)", default=50)
1166
- response = make_api_request(
1167
- "POST",
1168
- "/studios",
1169
- json_data={"user": username, "size_gb": size},
1170
- )
1171
- if response.status_code != 201:
1172
- console.print("[red]❌ Failed to create studio[/red]")
1173
- raise typer.Exit(1)
1174
- studio = response.json()
1175
- studio["studio_id"] = studio["studio_id"] # Normalize key
1393
+ if target_user == username:
1394
+ console.print("[yellow]You don't have a studio yet.[/yellow]")
1395
+ if Confirm.ask("Would you like to create one now?"):
1396
+ size = IntPrompt.ask("Studio size (GB)", default=50)
1397
+ response = make_api_request(
1398
+ "POST",
1399
+ "/studios",
1400
+ json_data={"user": username, "size_gb": size},
1401
+ )
1402
+ if response.status_code != 201:
1403
+ console.print("[red]❌ Failed to create studio[/red]")
1404
+ raise typer.Exit(1)
1405
+ studio = response.json()
1406
+ studio["studio_id"] = studio["studio_id"] # Normalize key
1407
+ else:
1408
+ raise typer.Exit(0)
1176
1409
  else:
1177
- raise typer.Exit(0)
1410
+ console.print(f"[red]❌ User {target_user} doesn't have a studio.[/red]")
1411
+ raise typer.Exit(1)
1178
1412
 
1179
1413
  # Check if already attached
1180
1414
  if studio.get("status") == "in-use":
@@ -1238,7 +1472,7 @@ def attach_studio(
1238
1472
  f"/studios/{studio['studio_id']}/attach",
1239
1473
  json_data={
1240
1474
  "vm_id": engine["instance_id"],
1241
- "user": username,
1475
+ "user": target_user, # Use target_user instead of username
1242
1476
  "public_key": public_key,
1243
1477
  },
1244
1478
  )
@@ -1248,28 +1482,46 @@ def attach_studio(
1248
1482
  if response.status_code == 200:
1249
1483
  console.print(f"[green]✓ Studio attached successfully![/green]")
1250
1484
 
1251
- # Update SSH config
1252
- update_ssh_config_entry(engine["name"], engine["instance_id"], username)
1485
+ # Update SSH config - use target_user for the connection
1486
+ update_ssh_config_entry(engine["name"], engine["instance_id"], target_user)
1253
1487
  console.print(f"[green]✓ SSH config updated[/green]")
1254
1488
  console.print(f"\nConnect with: [cyan]ssh {engine['name']}[/cyan]")
1255
- console.print(f"Your files are at: [cyan]/studios/{username}[/cyan]")
1489
+ console.print(f"Files are at: [cyan]/studios/{target_user}[/cyan]")
1256
1490
  else:
1257
1491
  error = response.json().get("error", "Unknown error")
1258
1492
  console.print(f"[red]❌ Failed to attach studio: {error}[/red]")
1259
1493
 
1260
1494
 
1261
1495
  @studio_app.command("detach")
1262
- def detach_studio():
1496
+ def detach_studio(
1497
+ user: Optional[str] = typer.Option(None, "--user", "-u", help="Detach a different user's studio (admin only)"),
1498
+ ):
1263
1499
  """Detach your studio from its current engine."""
1264
1500
  username = check_aws_sso()
1501
+
1502
+ # Use specified user if provided, otherwise use current user
1503
+ target_user = user if user else username
1504
+
1505
+ # Add confirmation when detaching another user's studio
1506
+ if target_user != username:
1507
+ console.print(f"[yellow]⚠️ Managing studio for user: {target_user}[/yellow]")
1508
+ if not Confirm.ask(f"Are you sure you want to detach {target_user}'s studio?"):
1509
+ console.print("Operation cancelled.")
1510
+ return
1265
1511
 
1266
- studio = get_user_studio(username)
1512
+ studio = get_user_studio(target_user)
1267
1513
  if not studio:
1268
- console.print("[yellow]You don't have a studio.[/yellow]")
1514
+ if target_user == username:
1515
+ console.print("[yellow]You don't have a studio.[/yellow]")
1516
+ else:
1517
+ console.print(f"[yellow]User {target_user} doesn't have a studio.[/yellow]")
1269
1518
  return
1270
1519
 
1271
1520
  if studio.get("status") != "in-use":
1272
- console.print("[yellow]Your studio is not attached to any engine.[/yellow]")
1521
+ if target_user == username:
1522
+ console.print("[yellow]Your studio is not attached to any engine.[/yellow]")
1523
+ else:
1524
+ console.print(f"[yellow]{target_user}'s studio is not attached to any engine.[/yellow]")
1273
1525
  return
1274
1526
 
1275
1527
  console.print(f"Detaching studio from {studio.get('attached_vm_id')}...")
@@ -1284,23 +1536,36 @@ def detach_studio():
1284
1536
 
1285
1537
 
1286
1538
  @studio_app.command("delete")
1287
- def delete_studio():
1539
+ def delete_studio(
1540
+ user: Optional[str] = typer.Option(None, "--user", "-u", help="Delete a different user's studio (admin only)"),
1541
+ ):
1288
1542
  """Delete your studio permanently."""
1289
1543
  username = check_aws_sso()
1290
-
1291
- studio = get_user_studio(username)
1544
+
1545
+ # Use specified user if provided, otherwise use current user
1546
+ target_user = user if user else username
1547
+
1548
+ # Extra warning when deleting another user's studio
1549
+ if target_user != username:
1550
+ console.print(f"[red]⚠️ ADMIN ACTION: Deleting studio for user: {target_user}[/red]")
1551
+
1552
+ studio = get_user_studio(target_user)
1292
1553
  if not studio:
1293
- console.print("[yellow]You don't have a studio to delete.[/yellow]")
1554
+ if target_user == username:
1555
+ console.print("[yellow]You don't have a studio to delete.[/yellow]")
1556
+ else:
1557
+ console.print(f"[yellow]User {target_user} doesn't have a studio to delete.[/yellow]")
1294
1558
  return
1295
1559
 
1296
1560
  console.print(
1297
- "[red]⚠️ WARNING: This will permanently delete your studio and all data![/red]"
1561
+ "[red]⚠️ WARNING: This will permanently delete the studio and all data![/red]"
1298
1562
  )
1299
1563
  console.print(f"Studio ID: {studio['studio_id']}")
1564
+ console.print(f"User: {target_user}")
1300
1565
  console.print(f"Size: {studio['size_gb']}GB")
1301
1566
 
1302
1567
  # Multiple confirmations
1303
- if not Confirm.ask("\nAre you sure you want to delete your studio?"):
1568
+ if not Confirm.ask(f"\nAre you sure you want to delete {target_user}'s studio?" if target_user != username else "\nAre you sure you want to delete your studio?"):
1304
1569
  console.print("Deletion cancelled.")
1305
1570
  return
1306
1571
 
@@ -1353,8 +1618,8 @@ def list_studios(
1353
1618
  table.add_column("User")
1354
1619
  table.add_column("Status")
1355
1620
  table.add_column("Size", justify="right")
1621
+ table.add_column("Disk Usage", justify="right")
1356
1622
  table.add_column("Attached To")
1357
- table.add_column("Created")
1358
1623
 
1359
1624
  for studio in studios:
1360
1625
  # Change status display
@@ -1367,29 +1632,25 @@ def list_studios(
1367
1632
 
1368
1633
  # Format attached engine info
1369
1634
  attached_to = "-"
1635
+ disk_usage = "?/?"
1370
1636
  if studio.get("attached_vm_id"):
1371
1637
  vm_id = studio["attached_vm_id"]
1372
1638
  engine_name = engines.get(vm_id, "unknown")
1373
1639
  attached_to = f"{engine_name} ({vm_id})"
1374
-
1375
- # Format creation date (remove microseconds and timezone info)
1376
- created = studio["creation_date"]
1377
- try:
1378
- # Parse and reformat to just show date and time
1379
- if "T" in created:
1380
- created_dt = datetime.fromisoformat(created.replace("Z", "+00:00"))
1381
- created = created_dt.strftime("%Y-%m-%d %H:%M")
1382
- except:
1383
- # If parsing fails, just truncate
1384
- created = created.split("T")[0] if "T" in created else created[:16]
1640
+
1641
+ # Try to get disk usage if attached
1642
+ if studio["status"] == "in-use":
1643
+ usage = get_studio_disk_usage_via_ssm(vm_id, studio["user"])
1644
+ if usage:
1645
+ disk_usage = usage
1385
1646
 
1386
1647
  table.add_row(
1387
1648
  studio["studio_id"],
1388
1649
  studio["user"],
1389
1650
  status_display,
1390
1651
  f"{studio['size_gb']}GB",
1652
+ disk_usage,
1391
1653
  attached_to,
1392
- created,
1393
1654
  )
1394
1655
 
1395
1656
  console.print(table)
@@ -1399,16 +1660,28 @@ def list_studios(
1399
1660
 
1400
1661
 
1401
1662
  @studio_app.command("reset")
1402
- def reset_studio():
1663
+ def reset_studio(
1664
+ user: Optional[str] = typer.Option(None, "--user", "-u", help="Reset a different user's studio"),
1665
+ ):
1403
1666
  """Reset a stuck studio (admin operation)."""
1404
1667
  username = check_aws_sso()
1405
-
1406
- studio = get_user_studio(username)
1668
+
1669
+ # Use specified user if provided, otherwise use current user
1670
+ target_user = user if user else username
1671
+
1672
+ # Add warning when resetting another user's studio
1673
+ if target_user != username:
1674
+ console.print(f"[yellow]⚠️ Resetting studio for user: {target_user}[/yellow]")
1675
+
1676
+ studio = get_user_studio(target_user)
1407
1677
  if not studio:
1408
- console.print("[yellow]You don't have a studio.[/yellow]")
1678
+ if target_user == username:
1679
+ console.print("[yellow]You don't have a studio.[/yellow]")
1680
+ else:
1681
+ console.print(f"[yellow]User {target_user} doesn't have a studio.[/yellow]")
1409
1682
  return
1410
1683
 
1411
- console.print(f"[yellow]⚠️ This will force-reset your studio state[/yellow]")
1684
+ console.print(f"[yellow]⚠️ This will force-reset the studio state[/yellow]")
1412
1685
  console.print(f"Current status: {studio['status']}")
1413
1686
  if studio.get("attached_vm_id"):
1414
1687
  console.print(f"Listed as attached to: {studio['attached_vm_id']}")
@@ -1461,3 +1734,91 @@ def reset_studio():
1461
1734
 
1462
1735
  except ClientError as e:
1463
1736
  console.print(f"[red]❌ Failed to reset studio: {e}[/red]")
1737
+
1738
+
1739
+ @studio_app.command("resize")
1740
+ def resize_studio(
1741
+ size: int = typer.Option(..., "--size", "-s", help="New size in GB"),
1742
+ user: Optional[str] = typer.Option(None, "--user", "-u", help="Resize a different user's studio (admin only)"),
1743
+ ):
1744
+ """Resize your studio volume (requires detachment)."""
1745
+ username = check_aws_sso()
1746
+
1747
+ # Use specified user if provided, otherwise use current user
1748
+ target_user = user if user else username
1749
+
1750
+ # Add warning when resizing another user's studio
1751
+ if target_user != username:
1752
+ console.print(f"[yellow]⚠️ Resizing studio for user: {target_user}[/yellow]")
1753
+
1754
+ studio = get_user_studio(target_user)
1755
+ if not studio:
1756
+ if target_user == username:
1757
+ console.print("[yellow]You don't have a studio yet.[/yellow]")
1758
+ else:
1759
+ console.print(f"[yellow]User {target_user} doesn't have a studio.[/yellow]")
1760
+ return
1761
+
1762
+ current_size = studio["size_gb"]
1763
+
1764
+ if size <= current_size:
1765
+ console.print(f"[red]❌ New size ({size}GB) must be larger than current size ({current_size}GB)[/red]")
1766
+ raise typer.Exit(1)
1767
+
1768
+ # Check if studio is attached
1769
+ if studio["status"] == "in-use":
1770
+ console.print("[yellow]⚠️ Studio must be detached before resizing[/yellow]")
1771
+ console.print(f"Currently attached to: {studio.get('attached_vm_id')}")
1772
+
1773
+ if not Confirm.ask("\nDetach studio and proceed with resize?"):
1774
+ console.print("Resize cancelled.")
1775
+ return
1776
+
1777
+ # Detach the studio
1778
+ console.print("Detaching studio...")
1779
+ response = make_api_request("POST", f"/studios/{studio['studio_id']}/detach")
1780
+ if response.status_code != 200:
1781
+ console.print("[red]❌ Failed to detach studio[/red]")
1782
+ raise typer.Exit(1)
1783
+
1784
+ console.print("[green]✓ Studio detached[/green]")
1785
+
1786
+ # Wait a moment for detachment to complete
1787
+ time.sleep(5)
1788
+
1789
+ console.print(f"[yellow]Resizing studio from {current_size}GB to {size}GB[/yellow]")
1790
+
1791
+ # Call the resize API
1792
+ resize_response = make_api_request(
1793
+ "POST",
1794
+ f"/studios/{studio['studio_id']}/resize",
1795
+ json_data={"size": size}
1796
+ )
1797
+
1798
+ if resize_response.status_code != 200:
1799
+ error = resize_response.json().get("error", "Unknown error")
1800
+ console.print(f"[red]❌ Failed to resize studio: {error}[/red]")
1801
+ raise typer.Exit(1)
1802
+
1803
+ # Wait for volume modification to complete
1804
+ ec2 = boto3.client("ec2", region_name="us-east-1")
1805
+ console.print("Waiting for volume modification to complete...")
1806
+ while True:
1807
+ try:
1808
+ mod_state = ec2.describe_volumes_modifications(VolumeIds=[studio["studio_id"]])
1809
+ if not mod_state["VolumesModifications"]:
1810
+ break # Modification complete
1811
+ state = mod_state["VolumesModifications"][0]["ModificationState"]
1812
+ if state == "completed":
1813
+ break
1814
+ elif state == "failed":
1815
+ console.print("[red]❌ Volume modification failed[/red]")
1816
+ raise typer.Exit(1)
1817
+ time.sleep(5)
1818
+ except ClientError:
1819
+ # Modification might be complete
1820
+ break
1821
+
1822
+ console.print(f"[green]✓ Studio resized successfully to {size}GB![/green]")
1823
+ console.print("\n[dim]The filesystem will be automatically expanded when you next attach the studio.[/dim]")
1824
+ console.print(f"To attach: [cyan]dh studio attach <engine-name>[/cyan]")
@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
5
5
 
6
6
  [project]
7
7
  name = "dayhoff-tools"
8
- version = "1.3.9"
8
+ version = "1.3.11"
9
9
  description = "Common tools for all the repos at Dayhoff Labs"
10
10
  authors = [
11
11
  {name = "Daniel Martin-Alarcon", email = "dma@dayhofflabs.com"}
@@ -20,13 +20,13 @@ dependencies = [
20
20
  "toml>=0.10",
21
21
  "questionary>=2.0.1",
22
22
  "tzdata>=2025.2",
23
+ "boto3>=1.36.8",
23
24
  ]
24
25
  requires-python = ">=3.10,<4.0"
25
26
 
26
27
  [project.optional-dependencies]
27
28
  full = [
28
29
  "biopython>=1.84",
29
- "boto3>=1.36.8",
30
30
  "docker>=7.1.0",
31
31
  "fair-esm>=2.0.0",
32
32
  "h5py>=3.11.0",
File without changes