28 lines
2.1 KiB
Python
28 lines
2.1 KiB
Python
import pathlib, hashlib, json, shutil, datetime
|
|
root=pathlib.Path('/home/ps/huobao-gpu-selection')
|
|
out=root/('idle-baseline-fix-'+datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))
|
|
out.mkdir(mode=0o700)
|
|
old=""" if any(d.get('torch_vram_total', 0) > 1024**3 for d in stats.json().get('devices', [])):
|
|
# Release only idle model caches; queue contents and results stay intact."""
|
|
new=""" devices = stats.json().get('devices')
|
|
if not isinstance(devices, list) or not devices or any(
|
|
not isinstance(d, dict) or type(d.get('torch_vram_total')) not in (int, float)
|
|
or not 0 <= d['torch_vram_total'] < float('inf') for d in devices
|
|
):
|
|
return None
|
|
# An idle cudaMallocAsync worker retains ~2.2 GiB of persistent tensors.
|
|
# Allow at most 4 GiB within the new worker's existing 8 GiB reserve;
|
|
# active/queued legacy jobs still block above, and large caches drain below.
|
|
if sum(d['torch_vram_total'] for d in devices) > 4 * 1024**3:
|
|
# Release only idle model caches; queue contents and results stay intact."""
|
|
manifest={}
|
|
for short,name in [('general','canvas-h3-prompt-normalizer-healthy-20261026'),('vdn','huobao-vdn-normalizer')]:
|
|
p=root/(name+'.candidate.py');source=p.read_text();assert source.count(old)==1
|
|
before=out/(short+'.before.py');shutil.copy2(p,before)
|
|
candidate=out/(short+'.after.py');candidate.write_text(source.replace(old,new));candidate.chmod(0o600)
|
|
compile(candidate.read_text(),str(candidate),'exec')
|
|
manifest[short]={'name':name,'path':str(p),'candidate':str(candidate),'before':hashlib.sha256(p.read_bytes()).hexdigest(),'after':hashlib.sha256(candidate.read_bytes()).hexdigest()}
|
|
(out/'manifest.json').write_text(json.dumps(manifest,indent=2))
|
|
(root/'idle-baseline-fix-current.json').write_text(json.dumps({'directory':str(out)}))
|
|
print(json.dumps({'prepared':str(out),'hashes':{k:{x:v[x] for x in ('before','after')} for k,v in manifest.items()}}))
|