59 lines
2.3 KiB
Python
59 lines
2.3 KiB
Python
"""Supervise the one-model job; no automatic download of subsequent models."""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--root', type=pathlib.Path, required=True)
|
|
args = p.parse_args()
|
|
root = args.root.resolve()
|
|
root.mkdir(exist_ok=True, parents=True)
|
|
status_file = root / 'status.json'
|
|
start = time.time()
|
|
env = dict(os.environ, OPENBLAS_NUM_THREADS='1', OMP_NUM_THREADS='1')
|
|
|
|
def status(phase, **extra):
|
|
status_file.write_text(json.dumps({'phase': phase, 'updated_unix': time.time(),
|
|
'supervisor_started_unix': start, **extra}, indent=2))
|
|
print(phase, flush=True)
|
|
|
|
def command(module, log, timeout, extra=()):
|
|
with (root / log).open('w') as stream:
|
|
subprocess.run([sys.executable, '-B', '-m', module, '--root', str(root), *extra],
|
|
check=True, timeout=timeout, env=env, stdout=stream, stderr=subprocess.STDOUT)
|
|
|
|
try:
|
|
status('waiting_for_download')
|
|
download_wait_start = time.monotonic()
|
|
while True:
|
|
manifest_path = root / 'manifest.json'
|
|
if manifest_path.exists():
|
|
manifest = json.loads(manifest_path.read_text())
|
|
if manifest.get('completed_unix') and all((root / 'weights' / f['name']).exists() for f in manifest['files']):
|
|
break
|
|
if time.monotonic() - download_wait_start > 7200:
|
|
raise TimeoutError('Download wait exceeded 2 hours; no analysis started')
|
|
time.sleep(10)
|
|
status('analyzing', analysis_hard_limit_seconds=7200)
|
|
command('pilot.analyze', 'analysis.log', 7200, ('--seconds', '6600'))
|
|
status('plotting')
|
|
command('pilot.plot', 'plot.log', 300)
|
|
status('validating')
|
|
command('pilot.validate', 'validate.log', 120)
|
|
with (root / 'environment.txt').open('w') as f:
|
|
subprocess.run([sys.executable, '-B', '-m', 'pip', 'freeze'], check=True, stdout=f, env=env)
|
|
status('results_ready_weights_retained_until_archive_copy')
|
|
except Exception as e:
|
|
status('failed', error_type=type(e).__name__, error=str(e))
|
|
raise
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|