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.
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/PKG-INFO +2 -2
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/engine_commands.py +419 -58
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/pyproject.toml +2 -2
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/README.md +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/__init__.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/chemistry/standardizer.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/chemistry/utils.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/__init__.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/cloud_commands.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/main.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/swarm_commands.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/cli/utility_commands.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/base.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/deploy_aws.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/deploy_gcp.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/deploy_utils.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/job_runner.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/processors.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/deployment/swarm.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/embedders.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/fasta.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/file_ops.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/h5.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/gcp.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/gtdb.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/kegg.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/mmseqs.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/structure.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/intake/uniprot.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/logs.py +0 -0
- {dayhoff_tools-1.3.9 → dayhoff_tools-1.3.11}/dayhoff_tools/sqlite.py +0 -0
- {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.
|
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)
|
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
|
-
|
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
|
-
|
1108
|
-
|
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(
|
1391
|
+
studio = get_user_studio(target_user)
|
1162
1392
|
if not studio:
|
1163
|
-
|
1164
|
-
|
1165
|
-
|
1166
|
-
|
1167
|
-
|
1168
|
-
|
1169
|
-
|
1170
|
-
|
1171
|
-
|
1172
|
-
|
1173
|
-
|
1174
|
-
|
1175
|
-
|
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
|
-
|
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"],
|
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"
|
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(
|
1512
|
+
studio = get_user_studio(target_user)
|
1267
1513
|
if not studio:
|
1268
|
-
|
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
|
-
|
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
|
-
|
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
|
-
|
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
|
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
|
-
|
1376
|
-
|
1377
|
-
|
1378
|
-
|
1379
|
-
|
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
|
-
|
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
|
-
|
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
|
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.
|
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
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|