| """Restartable single-GPU prepare → validate → train → probe workflow.""" |
| from pathlib import Path |
| import sys |
| sys.path.insert(0,str(Path(__file__).resolve().parent)) |
| from multitask_data import * |
| import subprocess |
|
|
|
|
| def alive(pid,needle): |
| try:return needle in Path(f'/proc/{pid}/cmdline').read_text() |
| except OSError:return False |
|
|
|
|
| def run_script(name,*args): |
| process=subprocess.Popen([sys.executable,'-u',str(Path(__file__).with_name(name)),*args],env={**os.environ,'TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD':'1'}) |
| write(OUT/'active_process.json',dict(pid=process.pid,script=name,args=list(args),time=time.time())) |
| if name=='multitask_data.py':write(OUT/'prepare_pid.json',dict(pid=process.pid)) |
| code=process.wait() |
| if code:raise RuntimeError(f'{name} exited with code {code}; see driver.log and status.json') |
|
|
|
|
| def wait_existing_stage(script): |
| candidates=[] |
| for name in ['active_process.json','prepare_pid.json','status.json']: |
| try:candidates.append(json.loads((OUT/name).read_text()).get('pid')) |
| except (OSError,ValueError):pass |
| for pid in set(candidates)-{None,os.getpid()}: |
| if alive(pid,script): |
| print('WAITING_FOR_EXISTING_STAGE',script,pid,flush=True) |
| while alive(pid,script):time.sleep(20) |
|
|
|
|
| def publish_protocol(): |
| from multitask_model import CONFIG |
| api=HfApi();api.create_repo(MODEL_REPO,repo_type='model',exist_ok=True) |
| sources=[Path(__file__).with_name(n) for n in ['multitask_data.py','multitask_model.py','multitask_train.py','multitask_evaluate.py','multitask_driver.py','eval_common.py','train_cosmo.py','bigsoldb_evaluate.py','openadmet_evaluate.py']] |
| write(OUT/'source_hashes.json',{p.name:sha(p) for p in sources}) |
| write(OUT/'config.json',CONFIG) |
| readme=OUT/'README.md' |
| if not readme.exists():readme.write_text('''--- |
| tags: |
| - chemistry |
| - multitask |
| --- |
| # REM3DI PCQMep2 COSMOep10 CombiSolv-QM continuation |
| |
| Training is in progress. The exact encoder/preprocessors/COSMO head from the preceding BigSolDB2.0 evaluation initialize this run; only the omitted denoising decoder is restored from PCQM epoch2. Three tasks train jointly: original embedding denoising, COSMO sigma/A/V reconstruction, and CombiSolv-QM solvation free energy. Monroe's uncertainty weights and covariance decorrelation are used. See config.json, initialization.json, validation files, source/ and the pinned prepared dataset reference. |
| |
| Five full joint epochs are planned. Each epoch directory includes all model heads, learned task weights, optimizer, scheduler, RNG, label transforms, architecture and SHA-256 manifest. Each completed checkpoint is backed up and its large-file hashes verified. Every epoch will receive a matched BigSolDB2.0 Kumo Large probe; a final report will replace this progress card after all requested results pass validation. |
| ''') |
| names=['README.md','initialization.json','config.json','source_hashes.json','loss_implementation_validation.json','smoke_validation.json','prepared_data_hub.json','preparation_complete.json'] |
| ops=[CommitOperationAdd(path_in_repo=n,path_or_fileobj=str(OUT/n)) for n in names if (OUT/n).exists()] |
| ops += [CommitOperationAdd(path_in_repo='source/'+p.name,path_or_fileobj=str(p)) for p in sources] |
| c=api.create_commit(repo_id=MODEL_REPO,operations=ops,commit_message='Pin exact evaluated starting checkpoint, objective equivalence tests, source and joint training protocol') |
| write(OUT/'protocol_hub.json',dict(repository=MODEL_REPO,revision=c.oid)) |
|
|
|
|
| def main(): |
| OUT.mkdir(exist_ok=True) |
| lock=(OUT/'driver.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) |
| if (OUT/'complete.json').exists():print('MULTITASK_ALREADY_COMPLETE');return |
| if not (OUT/'preparation_complete.json').exists(): |
| wait_existing_stage('multitask_data.py') |
| if not (OUT/'preparation_complete.json').exists():run_script('multitask_data.py') |
| if not (OUT/'smoke_validation.json').exists(): |
| wait_existing_stage('multitask_train.py') |
| if not (OUT/'smoke_validation.json').exists():run_script('multitask_train.py','--smoke') |
| publish_protocol() |
| if not (OUT/'training_complete.json').exists(): |
| wait_existing_stage('multitask_train.py') |
| if not (OUT/'training_complete.json').exists():run_script('multitask_train.py') |
| wait_existing_stage('multitask_evaluate.py') |
| if not (OUT/'complete.json').exists():run_script('multitask_evaluate.py') |
|
|
|
|
| if __name__=='__main__': |
| try:main() |
| except Exception: |
| error=traceback.format_exc();print(error,flush=True) |
| write(OUT/'driver_failure.json',dict(error=error,time=time.time(),pid=os.getpid()));raise |
|
|
