| """One-request stdlib SSH agent. No server, listening socket, or credentials.""" |
| import base64, contextlib, datetime, hashlib, json, os, re, shlex, subprocess, sys, tempfile, urllib.request, uuid |
| from pathlib import Path |
| REPO = 'will-taggart/AFFECTOR-ARC' |
|
|
| def atomic(path, obj): |
| path = Path(path); path.parent.mkdir(parents=True, exist_ok=True) |
| fd, name = tempfile.mkstemp(dir=path.parent) |
| try: |
| with os.fdopen(fd, 'w') as f: |
| json.dump(obj, f, ensure_ascii=False, allow_nan=False); f.flush(); os.fsync(f.fileno()) |
| os.replace(name, path) |
| finally: |
| if os.path.exists(name): os.unlink(name) |
|
|
| def command(args, **kw): |
| p = subprocess.run(args, capture_output=True, text=True, **kw) |
| if p.returncode: raise RuntimeError((p.stderr or p.stdout or 'Command failed').strip()) |
| return p.stdout.strip() |
|
|
| def token(value, label): |
| if not isinstance(value, str) or not re.fullmatch(r'[A-Za-z0-9_.,:+@/-]+', value) or value.startswith('-'): |
| raise ValueError('Invalid '+label) |
| return value |
|
|
| def rootpath(value): |
| if not isinstance(value, str) or not value.strip(): raise ValueError('Choose a remote workspace') |
| p = Path(value).expanduser() |
| if not p.is_absolute() or p == Path('/') or len(p.parts)<3: raise ValueError('Use an absolute path or ~/folder for the workspace') |
| return p.resolve() |
|
|
| def jobdir(root, identity): |
| uuid.UUID(identity); p=root/'jobs'/identity |
| if not p.is_dir(): raise ValueError('Job record not found') |
| return p |
|
|
| def fetch(url): |
| with urllib.request.urlopen(urllib.request.Request(url,headers={'User-Agent':'AFFECTOR-ARC-Desktop/1.0'}), timeout=120) as r: return r.read() |
|
|
| def current(root): |
| p=root/'current' |
| if not p.exists(): raise ValueError('Install/update AFFECTOR-ARC first') |
| release=p.resolve() |
| if release.parent != (root/'releases').resolve(): raise ValueError('Unexpected release path') |
| return release |
|
|
| def envpath(root,release): |
| digest=hashlib.sha256((release/'setup.sh').read_bytes()).hexdigest()[:16] |
| return root/'environments'/digest |
|
|
| def scheduler(q, root, kind, payload=None): |
| release=current(root); env=envpath(root,release) |
| if kind!='setup' and not (env/'.affector-ready').exists(): raise ValueError('Run GPU setup successfully before submitting a transformation') |
| identity=str(uuid.uuid4()); d=root/'jobs'/identity; d.mkdir(parents=True) |
| account=token(q['account'],'allocation'); partition=token(q['partition'],'partition') |
| wall=q.get('walltime','01:00:00') |
| if not re.fullmatch(r'(?:\d+-)?\d{1,3}:\d{2}:\d{2}',wall): raise ValueError('Use HH:MM:SS or D-HH:MM:SS for wall time') |
| cpus=int(q.get('cpus',8)); memory=int(q.get('memory',32)) |
| if not 1<=cpus<=128 or not 1<=memory<=1024: raise ValueError('Invalid CPU or memory request') |
| cache=rootpath(q.get('cache',str(root/'model-cache'))) |
| lines=['#!/bin/bash','set -euo pipefail','umask 077','module load Miniconda3',f'export HF_HOME={shlex.quote(str(cache))}',f'cd {shlex.quote(str(release))}'] |
| if kind=='setup': |
| |
| lines += [f'ENV={shlex.quote(str(env))}', 'mkdir -p "$(dirname "$ENV")"', 'exec 9>"$ENV.lock"', 'flock 9', 'if [ ! -f "$ENV/.affector-ready" ]; then', |
| 'if [ ! -x "$ENV/bin/python" ]; then conda create -y --override-channels -c conda-forge -p "$ENV" python=3.12 pip; fi', |
| '"$ENV/bin/python" -m pip install torch --index-url https://download.pytorch.org/whl/cu126', |
| '"$ENV/bin/python" -m pip install "transformers==5.22.2" accelerate huggingface_hub safetensors', |
| '"$ENV/bin/python" -m unittest discover -s tests', |
| '"$ENV/bin/python" -m pip freeze > "$ENV/environment-resolved.txt"', |
| '"$ENV/bin/python" -c \'import torch; assert torch.cuda.is_available(), "CUDA unavailable"; print(torch.__version__, torch.cuda.get_device_name(0))\'', |
| 'touch "$ENV/.affector-ready"', 'fi', |
| '"$ENV/bin/python" -m unittest discover -s tests', |
| '"$ENV/bin/python" -c \'import torch; assert torch.cuda.is_available(), "CUDA unavailable"; print(torch.__version__, torch.cuda.get_device_name(0))\''] |
| else: |
| text=q.get('text','') |
| if not isinstance(text,str) or not text.strip(): raise ValueError('Input text is empty') |
| if len(text.encode())>20_000_000: raise ValueError('Input exceeds 20 MB') |
| (d/'input.txt').write_text(text,encoding='utf-8') |
| geometry=q.get('geometry') |
| if not isinstance(geometry,dict): raise ValueError('Missing shape') |
| |
| import importlib.util |
| spec=importlib.util.spec_from_file_location('arc_geometry',release/'engine/geometry.py'); m=importlib.util.module_from_spec(spec); spec.loader.exec_module(m) |
| geometry=m.import_atlas(geometry,geometry.get('name','Imported shape')); m.validate(geometry) |
| atomic(d/'shape.json',geometry) |
| args=[str(env/'bin/python'),'-u','run_arc.py','--mode',kind,'--text-file',str(d/'input.txt'),'--geometry',str(d/'shape.json'),'--output',str(d/'report.json')] |
| style=q.get('style','Restrained') |
| if style not in ['Minimal','Restrained','Balanced','Creative']: raise ValueError('Invalid rewrite style') |
| args += ['--style',style] |
| if kind=='rewrite': |
| for key,lo,hi in [('temperature',0,2),('top_p',0.00001,1),('top_k',0,100000),('max_tokens',32,2048),('time_limit',1,3600),('seed',0,2147483647)]: |
| val=q.get(key,{'temperature':0,'top_p':.95,'top_k':40,'max_tokens':512,'time_limit':120,'seed':0}[key]) |
| if type(val) not in (float,int) or not lo<=val<=hi: raise ValueError('Invalid '+key) |
| if key in ['top_k','max_tokens','seed'] and int(val)!=val: raise ValueError(key+' must be an integer') |
| args += ['--'+key.replace('_','-'),str(val)] |
| lines += [shlex.join(args)] |
| script='\n'.join(lines)+'\n'; (d/'job.sh').write_text(script) |
| jid=command(['sbatch','--parsable','--account='+account,'--partition='+partition,'--nodes=1','--ntasks=1','--cpus-per-task='+str(cpus),'--gres=gpu:1','--mem='+str(memory)+'G','--time='+wall,'--job-name=affector-'+kind,'--output='+str(d/'job.log'),'--error='+str(d/'job.log'),str(d/'job.sh')]).split(';')[0] |
| if not jid.isdigit(): raise RuntimeError('Unexpected Slurm job ID: '+jid) |
| record={'id':identity,'job_id':jid,'kind':kind,'created_at':datetime.datetime.now(datetime.timezone.utc).isoformat(),'release':release.name,'directory':str(d),'account':account,'partition':partition} |
| atomic(d/'job.json',record); return record |
|
|
|
|
| def response_text(text): |
| if not isinstance(text,str):return '' |
| text=re.sub(r'<think>[\s\S]*?</think>\s*','',text,flags=re.I) |
| if re.search(r'<think>',text,re.I):return '' |
| label=r'(?:transformed(?:\s+(?:text|passage|document))?|rewritten(?:\s+(?:text|passage|document))?|rewrite|revised(?:\s+(?:text|passage|document))?|final(?:\s+(?:answer|text))?|attempt|output)' |
| pattern=r'(?im)^\s*(?:\#{1,6}\s*)?(?:here is (?:the |a )?)?(?:\*\*)?'+label+r'(?:\*\*)?\s*:\s*(?:\*\*)?[^\S\r\n]*' |
| matches=list(re.finditer(pattern,text)) |
| if matches: |
| text=text[matches[-1].end():] |
| suffix=re.search(r'(?im)^\s*(?:\#{1,6}\s*)?(?:\*\*)?(?:original(?:\s+(?:text|target|passage))?|source(?:\s+text)?|prompt|target|instructions?|explanation|notes?)(?:\*\*)?\s*:',text) |
| if suffix:text=text[:suffix.start()] |
| elif re.search(r'(?im)^\s*(?:\#{1,6}\s*)?(?:\*\*)?(?:original(?:\s+(?:text|target|passage))?|source\s+text|prompt|target)(?:\*\*)?\s*:',text):return '' |
| text=text.strip() |
| if text.startswith('```') and text.endswith('```'): |
| lines=text.splitlines();text='\n'.join(lines[1:-1]).strip() |
| if re.match(r"(?is)^(?:the prompt (?:asks|requests|requires|instructs)|this is (?:a |an )?prohibited attempt|i (?:cannot|can't|am unable to) (?:rewrite|comply with|fulfill this request))\b",text):return '' |
| return text |
|
|
| def transformed_text(report): |
| records=report.get('records') |
| if not isinstance(records,dict):return response_text(report.get('text','')) |
| spans=report.get('spans',[]);source=report.get('input','');parts=[];previous=None |
| for index in sorted(int(k) for k in records if str(k).isdigit()): |
| record=records[str(index)];text=record.get('text','') |
| if record.get('original_retained') or not isinstance(text,str) or not text.strip():continue |
| if previous is not None: |
| gap='\n\n' |
| if index==previous+1 and index<len(spans) and len(spans[previous])==2 and len(spans[index])==2: |
| a,b=spans[previous][1],spans[index][0] |
| if 0<=a<=b<=len(source) and not source[a:b].strip():gap=source[a:b] |
| parts.append(gap) |
| clean=response_text(text) |
| if not clean: |
| if previous is not None:parts.pop() |
| continue |
| parts.append(clean);previous=index |
| return ''.join(parts) |
|
|
| def dispatch(q): |
| op=q['op']; root=rootpath(q['workspace']) |
| if op=='probe': |
| return {'host':command(['hostname']),'home':str(Path.home()),'workspace':str(root),'slurm':command(['sbatch','--version']),'installed':(root/'current').exists(),'revision':current(root).name if (root/'current').exists() else None,'ready':(envpath(root,current(root))/'.affector-ready').exists() if (root/'current').exists() else False} |
| root.mkdir(parents=True,exist_ok=True) |
| if op=='update': |
| meta=json.loads(fetch('https://huggingface.co/api/models/'+REPO));rev=meta['sha'] |
| if not re.fullmatch('[0-9a-f]{40}',rev):raise RuntimeError('Invalid HF revision') |
| release=root/'releases'/rev |
| if not (release/'.snapshot-complete').exists(): |
| tree=json.loads(fetch(f'https://huggingface.co/api/models/{REPO}/tree/{rev}?recursive=true&limit=1000')) |
| selected=[x for x in tree if x['type']=='file' and (x['path'].startswith(('engine/','tests/')) or x['path'] in ['run_arc.py','document_run.py','setup.sh','smoke.sbatch','README.md','README.txt','ATLAS_LICENSE','ENGINE_MANIFEST.json'])] |
| if not selected or not {'run_arc.py','setup.sh'}.issubset({x['path'] for x in selected}):raise RuntimeError('Incomplete HF snapshot') |
| release.mkdir(parents=True,exist_ok=True) |
| for entry in selected: |
| rel=Path(entry['path']) |
| if rel.is_absolute() or '..' in rel.parts:raise RuntimeError('Unsafe snapshot path') |
| dest=release/rel;dest.parent.mkdir(parents=True,exist_ok=True) |
| data=fetch(f'https://huggingface.co/{REPO}/resolve/{rev}/{rel.as_posix()}') |
| if data.startswith(b'version https://git-lfs.github.com/spec/'):raise RuntimeError('Unresolved LFS pointer') |
| dest.write_bytes(data) |
| (release/'.snapshot-complete').write_text(rev) |
| temp=root/('current-'+str(uuid.uuid4()));temp.symlink_to(release,target_is_directory=True);os.replace(temp,root/'current') |
| return {'revision':rev,'path':str(release),'ready':(envpath(root,release)/'.affector-ready').exists()} |
| if op in ['setup','rewrite','smoke']:return scheduler(q,root,op) |
| if op=='list': |
| result=[] |
| for p in (root/'jobs').glob('*/job.json'): |
| with contextlib.suppress(Exception):result.append(json.loads(p.read_text())) |
| return {'jobs':sorted(result,key=lambda x:x['created_at'],reverse=True)} |
| if op in ['status','cancel','download']: |
| d=jobdir(root,q['id']);record=json.loads((d/'job.json').read_text());jid=record['job_id'] |
| if not jid.isdigit():raise ValueError('Invalid job ID') |
| if op=='cancel':command(['scancel',jid]);return {'message':'Cancellation requested','job_id':jid} |
| if op=='download': |
| files={} |
| for name in ['input.txt','shape.json','report.json','report.txt','job.log','job.json','job.sh']: |
| p=d/name |
| if p.exists():files[name]=base64.b64encode(p.read_bytes()).decode() |
| return {'files':files} |
| try: queued=command(['squeue','-h','-j',jid,'-o','%T|%R']) |
| except RuntimeError: queued='' |
| state=queued.split('|',1)[0] if queued else '' |
| if not state: |
| try:state=command(['sacct','-n','-X','-j',jid,'--format=State','--parsable2']).splitlines()[0].split('|')[0].strip() |
| except Exception:state='UNKNOWN' |
| log=(d/'job.log').read_text(errors='replace')[-100000:] if (d/'job.log').exists() else '' |
| data={'job':record,'state':state,'queue':queued,'log':log} |
| p=d/'report.json' |
| if p.exists(): |
| try: |
| report=json.loads(p.read_text()) |
| |
| data['text']=report.get('text','') |
| data['transformed_text']=transformed_text(report) |
| preview=dict(report) |
| if isinstance(preview.get('geometry'),dict): |
| preview['geometry']=dict(preview['geometry'],basis='Full calibration weights are in the downloaded report.') |
| data['report']=preview |
| except (ValueError,OSError):pass |
| textfile=d/'report.txt' |
| if not data.get('text') and textfile.exists(): |
| data['text']=textfile.read_text(errors='replace') |
| return data |
| raise ValueError('Unknown operation') |
|
|
| if __name__=='__main__': |
| try: |
| q=json.load(sys.stdin);print(json.dumps({'ok':True,'data':dispatch(q)},ensure_ascii=False,allow_nan=False)) |
| except Exception as e: |
| print(json.dumps({'ok':False,'error':str(e)}));sys.exit(1) |
|
|
