84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
import shlex
|
||
|
|
import subprocess
|
||
|
|
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() -> int:
|
||
|
|
parser = argparse.ArgumentParser(description='Run an aggressive sweep node directly over SSH, outside Sky/Ray job workers.')
|
||
|
|
parser.add_argument('--host', required=True)
|
||
|
|
parser.add_argument('--port', required=True)
|
||
|
|
parser.add_argument('--node', required=True)
|
||
|
|
parser.add_argument('--jobs', required=True)
|
||
|
|
parser.add_argument('--artifact-dir', default='artifacts/current_run')
|
||
|
|
parser.add_argument('--utilization', required=True)
|
||
|
|
parser.add_argument('--remote-root', default='/root/sky_workdir')
|
||
|
|
parser.add_argument('--max-jobs', type=int)
|
||
|
|
parser.add_argument('--sample-interval-seconds', type=float, default=30.0)
|
||
|
|
parser.add_argument('--stale-after-seconds', type=float, default=21600.0)
|
||
|
|
parser.add_argument('--max-attempts', type=int, default=2)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
hf_token = _secret('HF_TOKEN')
|
||
|
|
wandb_api_key = _secret('WANDB_API_KEY')
|
||
|
|
parts = [
|
||
|
|
'uv run --no-dev python scripts/aggressive_oom_node_wrapper.py',
|
||
|
|
f'--jobs {shlex.quote(args.jobs)}',
|
||
|
|
f'--node {shlex.quote(args.node)}',
|
||
|
|
f'--artifact-dir {shlex.quote(args.artifact_dir)}',
|
||
|
|
f'--utilization {shlex.quote(args.utilization)}',
|
||
|
|
f'--sample-interval-seconds {args.sample_interval_seconds:g}',
|
||
|
|
f'--stale-after-seconds {args.stale_after_seconds:g}',
|
||
|
|
f'--max-attempts {args.max_attempts}',
|
||
|
|
]
|
||
|
|
if args.max_jobs is not None:
|
||
|
|
parts.append(f'--max-jobs {args.max_jobs}')
|
||
|
|
remote_cmd = ' '.join(parts)
|
||
|
|
remote_script = '\n'.join(
|
||
|
|
[
|
||
|
|
'set -euo pipefail',
|
||
|
|
f'cd {shlex.quote(args.remote_root)}',
|
||
|
|
f'export HF_TOKEN={shlex.quote(hf_token)}',
|
||
|
|
f'export WANDB_API_KEY={shlex.quote(wandb_api_key)}',
|
||
|
|
'export PATH=\"$HOME/.local/bin:/root/.local/bin:$PATH\"',
|
||
|
|
'export PYTHONUNBUFFERED=1',
|
||
|
|
remote_cmd,
|
||
|
|
'',
|
||
|
|
]
|
||
|
|
)
|
||
|
|
command = [
|
||
|
|
'ssh',
|
||
|
|
'-o',
|
||
|
|
'StrictHostKeyChecking=no',
|
||
|
|
'-p',
|
||
|
|
str(args.port),
|
||
|
|
f'root@{args.host}',
|
||
|
|
'bash',
|
||
|
|
'-s',
|
||
|
|
]
|
||
|
|
print(f'sweep_direct_host={args.host}', flush=True)
|
||
|
|
print(f'sweep_direct_port={args.port}', flush=True)
|
||
|
|
print(f'sweep_direct_node={args.node}', flush=True)
|
||
|
|
print('sweep_direct_command=<redacted secrets> uv run --no-dev python scripts/aggressive_oom_node_wrapper.py', flush=True)
|
||
|
|
return subprocess.run(command, input=remote_script, text=True, check=False).returncode
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
raise SystemExit(main())
|