91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import signal
|
||
|
|
import subprocess
|
||
|
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
_stop = False
|
||
|
|
|
||
|
|
|
||
|
|
def _handle_stop(signum: int, frame: object) -> None:
|
||
|
|
global _stop
|
||
|
|
_stop = True
|
||
|
|
|
||
|
|
|
||
|
|
def _run(command: list[str]) -> int:
|
||
|
|
completed = subprocess.run(command, check=False)
|
||
|
|
return int(completed.returncode)
|
||
|
|
|
||
|
|
|
||
|
|
def _rsync(host: str, port: int, remote: str, local: Path, *, include_metrics_only: bool = False) -> int:
|
||
|
|
local.mkdir(parents=True, exist_ok=True)
|
||
|
|
command = [
|
||
|
|
'rsync',
|
||
|
|
'-az',
|
||
|
|
'--delete',
|
||
|
|
'--exclude',
|
||
|
|
'*.pt',
|
||
|
|
'--exclude',
|
||
|
|
'.wandb/',
|
||
|
|
]
|
||
|
|
if include_metrics_only:
|
||
|
|
command.extend([
|
||
|
|
'--include',
|
||
|
|
'*/',
|
||
|
|
'--include',
|
||
|
|
'*.json',
|
||
|
|
'--include',
|
||
|
|
'*.jsonl',
|
||
|
|
'--include',
|
||
|
|
'*.toml',
|
||
|
|
'--include',
|
||
|
|
'*.txt',
|
||
|
|
'--exclude',
|
||
|
|
'*',
|
||
|
|
])
|
||
|
|
command.extend([
|
||
|
|
'-e',
|
||
|
|
f'ssh -p {port} -o StrictHostKeyChecking=no',
|
||
|
|
f'root@{host}:{remote}',
|
||
|
|
str(local),
|
||
|
|
])
|
||
|
|
return _run(command)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser(description='Continuously rsync lightweight sweep artifacts from one node.')
|
||
|
|
parser.add_argument('--host', required=True)
|
||
|
|
parser.add_argument('--port', type=int, required=True)
|
||
|
|
parser.add_argument('--remote-root', default='/root/sky_workdir')
|
||
|
|
parser.add_argument('--local-root', default='artifacts/remote_runs/live_node0')
|
||
|
|
parser.add_argument('--interval-seconds', type=float, default=60.0)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
signal.signal(signal.SIGINT, _handle_stop)
|
||
|
|
signal.signal(signal.SIGTERM, _handle_stop)
|
||
|
|
local_root = Path(args.local_root)
|
||
|
|
remote_root = args.remote_root.rstrip('/')
|
||
|
|
interval = max(5.0, args.interval_seconds)
|
||
|
|
cycle = 0
|
||
|
|
while not _stop:
|
||
|
|
started = time.time()
|
||
|
|
cycle += 1
|
||
|
|
rcodes = [
|
||
|
|
_rsync(args.host, args.port, f'{remote_root}/artifacts/aggressive_oom_sweep/', local_root / 'aggressive_oom_sweep/'),
|
||
|
|
_rsync(args.host, args.port, f'{remote_root}/artifacts/current_run/', local_root / 'current_run/'),
|
||
|
|
_rsync(args.host, args.port, f'{remote_root}/artifacts/runs/', local_root / 'runs/', include_metrics_only=True),
|
||
|
|
]
|
||
|
|
print(f'collector_cycle={cycle} rcodes={rcodes} elapsed_seconds={time.time() - started:.1f}', flush=True)
|
||
|
|
deadline = time.time() + interval
|
||
|
|
while not _stop and time.time() < deadline:
|
||
|
|
time.sleep(min(1.0, deadline - time.time()))
|
||
|
|
print('collector_stopped=true', flush=True)
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
raise SystemExit(main())
|