File size: 4,689 Bytes
31dee3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """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
|