"""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