Spaces:
Running
Running
| """HF私有临时快照生命周期；创建/回收必须持同目录公共提交锁。""" | |
| from contextlib import contextmanager | |
| import json | |
| import logging | |
| import os | |
| from pathlib import Path | |
| import re | |
| import stat | |
| import uuid | |
| logger=logging.getLogger('ComfyUI-Ranking.DB') | |
| class SnapshotStore: | |
| MANIFEST='snapshot-owner.json' | |
| NAME=re.compile(r'\.hf-sync-[0-9a-f]{32}\Z') | |
| PAYLOAD=re.compile(r'(0|[1-9][0-9]*)\.json\Z') | |
| def __init__(self,directory): | |
| self.root=Path(os.path.abspath(directory)) | |
| def _check(path,directory=False): | |
| info=os.lstat(path) | |
| if (getattr(info,'st_file_attributes',0)&1024 | |
| or not (stat.S_ISDIR(info.st_mode) if directory else stat.S_ISREG(info.st_mode)) | |
| or not info.st_ino): | |
| raise ValueError('invalid snapshot path') | |
| return (info.st_dev,info.st_ino) | |
| def _owner(self,path): | |
| if path.parent!=self.root or not self.NAME.fullmatch(path.name): | |
| raise ValueError('unknown snapshot directory') | |
| self._check(self.root,directory=True) | |
| directory_id=self._check(path,directory=True) | |
| marker=path/self.MANIFEST | |
| marker_id=self._check(marker) | |
| descriptor=os.open(marker,os.O_RDONLY|getattr(os,'O_NOFOLLOW',0) | |
| |getattr(os,'O_NONBLOCK',0)|getattr(os,'O_BINARY',0)) | |
| with os.fdopen(descriptor,'rb') as stream: | |
| info=os.fstat(stream.fileno()) | |
| if not stat.S_ISREG(info.st_mode) or (info.st_dev,info.st_ino)!=marker_id: | |
| raise ValueError('snapshot owner identity changed') | |
| data=stream.read(4097) | |
| if len(data)>4096:raise ValueError('snapshot owner too large') | |
| owner=json.loads(data) | |
| if (not isinstance(owner,dict) or set(owner)!={'schema','kind','directory'} | |
| or type(owner['schema']) is not int or owner['schema']!=1 | |
| or owner['kind']!='ranking-hf-snapshot' or owner['directory']!=path.name): | |
| raise ValueError('invalid snapshot owner') | |
| if self._check(path,directory=True)!=directory_id or self._check(marker)!=marker_id: | |
| raise ValueError('snapshot identity changed') | |
| return directory_id | |
| def remove(self,path): | |
| """只删除验证归属且内容全为普通快照文件的本层目录，不递归删除。""" | |
| path=Path(path) | |
| directory_id=self._owner(path) | |
| entries=[] | |
| for entry in path.iterdir(): | |
| if len(entries)>=4096:raise ValueError('too many snapshot files') | |
| if entry.name!=self.MANIFEST and not self.PAYLOAD.fullmatch(entry.name): | |
| raise ValueError('unknown snapshot content') | |
| self._check(entry) | |
| entries.append(entry) | |
| # 先完整校验，再删除数据，归属记录最后删除，失败可在后续接手。 | |
| for entry in sorted(entries,key=lambda entry:entry.name==self.MANIFEST): | |
| self._check(self.root,directory=True) | |
| if self._check(path,directory=True)!=directory_id: | |
| raise ValueError('snapshot directory changed') | |
| entry.unlink() | |
| path.rmdir() | |
| def cleanup_stale(self): | |
| self._check(self.root,directory=True) | |
| removed=retained=0 | |
| for path in self.root.glob('.hf-sync-*'): | |
| try:self.remove(path);removed+=1 | |
| except FileNotFoundError: | |
| if path.exists() or path.is_symlink():retained+=1 | |
| except (OSError,ValueError):retained+=1 | |
| return removed,retained | |
| def temporary(self): | |
| self._check(self.root,directory=True) | |
| path=self.root/('.hf-sync-'+uuid.uuid4().hex) | |
| path.mkdir(mode=0o700) | |
| try: | |
| with (path/self.MANIFEST).open('xb') as stream: | |
| data=json.dumps({'schema':1,'kind':'ranking-hf-snapshot','directory':path.name}).encode('utf-8') | |
| stream.write(data);stream.flush();os.fsync(stream.fileno()) | |
| yield str(path) | |
| finally: | |
| try:self.remove(path) | |
| except FileNotFoundError: | |
| if path.exists() or path.is_symlink(): | |
| logger.warning('HF snapshot cleanup unavailable; retained for recovery') | |
| except (OSError,ValueError): | |
| logger.warning('HF snapshot cleanup unavailable; retained for recovery') | |