64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
def _secret(name: str) -> str:
|
||
|
|
value = os.environ.get(name)
|
||
|
|
if value:
|
||
|
|
return value
|
||
|
|
path = Path('.env') / name
|
||
|
|
if path.is_file():
|
||
|
|
value = path.read_text().strip()
|
||
|
|
if value:
|
||
|
|
return value
|
||
|
|
raise RuntimeError(f'{name} is required')
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: list[str]) -> int:
|
||
|
|
if len(argv) != 2:
|
||
|
|
print('usage: run_quick_probe_on_cluster.py CLUSTER', file=sys.stderr)
|
||
|
|
return 2
|
||
|
|
cluster = argv[1]
|
||
|
|
env = os.environ.copy()
|
||
|
|
env.pop('PYTHONPATH', None)
|
||
|
|
env['HF_TOKEN'] = _secret('HF_TOKEN')
|
||
|
|
env['WANDB_API_KEY'] = _secret('WANDB_API_KEY')
|
||
|
|
remote_cmd = (
|
||
|
|
'uv run --no-dev python scripts/aggressive_oom_node_wrapper.py '
|
||
|
|
'--jobs artifacts/aggressive_oom_sweep_quick_probe/jobs.jsonl '
|
||
|
|
'--node quick_probe_node '
|
||
|
|
'--artifact-dir artifacts/current_run_quick_probe '
|
||
|
|
'--utilization artifacts/aggressive_oom_sweep_quick_probe/utilization.jsonl '
|
||
|
|
'--max-jobs 1 '
|
||
|
|
'--sample-interval-seconds 5 '
|
||
|
|
'--require-success'
|
||
|
|
)
|
||
|
|
command = [
|
||
|
|
'uv',
|
||
|
|
'run',
|
||
|
|
'sky',
|
||
|
|
'exec',
|
||
|
|
'--workdir',
|
||
|
|
'.',
|
||
|
|
'--gpus',
|
||
|
|
'RTX4090:1',
|
||
|
|
'--secret',
|
||
|
|
'HF_TOKEN',
|
||
|
|
'--secret',
|
||
|
|
'WANDB_API_KEY',
|
||
|
|
cluster,
|
||
|
|
remote_cmd,
|
||
|
|
]
|
||
|
|
print(f'quick_probe_cluster={cluster}', flush=True)
|
||
|
|
print('quick_probe_command=sky exec --workdir . --gpus RTX4090:1 --secret HF_TOKEN --secret WANDB_API_KEY <cluster> <wrapper>', flush=True)
|
||
|
|
return subprocess.run(command, env=env, check=False).returncode
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
raise SystemExit(main(sys.argv))
|