| """Where a run's files come from: run-name resolution and the `Source` the ingest reads. |
| |
| A Source has `list(prefix, recursive=True) -> [(path, size, key, mtime)]` (key changes whenever the content does: |
| `xet_hash:size` on a bucket, the manifest diff the ingest keeps in its `files` table) and |
| `fetch(path, dest)`. `BucketSource` reads an HF bucket prefix, `LocalDirSource` a directory (tests, |
| offline copies). A non-recursive listing also returns directories, as ('<path>/', None, None, None). |
| """ |
| import os, re, shutil |
|
|
| from huggingface_hub import HfApi |
|
|
| RUN_RE = re.compile(r'^[A-Za-z0-9][\w.-]*/[A-Za-z0-9][\w.-]*/(runs/)?[A-Za-z0-9][\w.-]*$') |
| NAME_RE = re.compile(r'^[A-Za-z0-9][\w.-]*$') |
| CTL_RE = re.compile(r'^[A-Za-z0-9][\w.-]*/[A-Za-z0-9][\w.-]*(/[A-Za-z0-9][\w.-]*)+$') |
|
|
|
|
| def _strip(x): |
| x = (x or '').strip().strip('/') |
| return x[len('hf://buckets/'):] if x.startswith('hf://buckets/') else x |
|
|
|
|
| def resolve_run(arg, default_bucket): |
| """'ns/bucket/prefix[@controller=ns/bucket/path]', 'hf://buckets/...' or a bare run name |
| -> (bucket_id, prefix, controller 'ns/bucket/path' or None) or None.""" |
| arg, _, ctl = (arg or '').partition('@controller=') |
| arg, ctl = _strip(arg), _strip(ctl) or None |
| if '..' in arg or (ctl and ('..' in ctl or not CTL_RE.match(ctl))): |
| return None |
| if RUN_RE.match(arg): |
| ns, b, prefix = arg.split('/', 2) |
| return f'{ns}/{b}', prefix, ctl |
| if NAME_RE.match(arg) and default_bucket: |
| return default_bucket, arg, ctl |
| return None |
|
|
|
|
| def discover(source, known=()): |
| """Run prefixes 'runs/<name>' of a bucket root Source: subdirectories of runs/ holding a run.json |
| (non-recursive listings; prefixes in `known` are not listed again).""" |
| out = [] |
| for p, size, *_ in source.list('runs/', recursive=False): |
| name = p.rstrip('/') |
| if size is None and NAME_RE.match(name[5:]) and (name in known or any( |
| q == p + 'run.json' for q, *_ in source.list(p, recursive=False))): |
| out.append(name) |
| return out |
|
|
|
|
| def default_controller(bucket_id, prefix): |
| """flock's jobs mode syncs the controller (state.db, events/) to <bucket>/controller/<run>.""" |
| return f'{bucket_id}/controller/{prefix}' |
|
|
|
|
| class BucketSource: |
| def __init__(self, bucket_id, prefix, token=None): |
| self.api, self.bucket, self.prefix = HfApi(token=token), bucket_id, prefix.strip('/') |
| self._seen = {} |
|
|
| def list(self, prefix='', recursive=True): |
| out, skip = [], len(self.prefix) + 1 if self.prefix else 0 |
| for f in self.api.list_bucket_tree(self.bucket, prefix=f'{self.prefix}/{prefix}' if self.prefix else prefix, recursive=recursive): |
| p = f.path[skip:] |
| if f.type != 'file': |
| if not recursive: |
| out.append((p + '/', None, None, None)) |
| continue |
| self._seen[p] = f |
| out.append((p, f.size, f'{f.xet_hash}:{f.size}', (f.mtime or f.uploaded_at).timestamp())) |
| return out |
|
|
| def sub(self, prefix): |
| return BucketSource(self.bucket, '/'.join(x for x in (self.prefix, prefix) if x), self.api.token) |
|
|
| def fetch(self, path, dest): |
| self.api.download_bucket_files(self.bucket, [(self._seen.get(path) or f'{self.prefix}/{path}', dest)], |
| raise_on_missing_files=True) |
|
|
|
|
| class LocalDirSource: |
| def __init__(self, root): |
| self.root = root |
|
|
| def list(self, prefix='', recursive=True): |
| out = [] |
| if not recursive: |
| d = os.path.join(self.root, prefix) |
| for n in sorted(os.listdir(d)) if os.path.isdir(d) else []: |
| st = os.stat(os.path.join(d, n)) |
| out.append((prefix + n + '/', None, None, None) if os.path.isdir(os.path.join(d, n)) else |
| (prefix + n, st.st_size, f'{st.st_size}:{st.st_mtime_ns}', st.st_mtime)) |
| return out |
| for d, _, names in os.walk(os.path.join(self.root, prefix)): |
| for n in names: |
| p = os.path.join(d, n) |
| st = os.stat(p) |
| out.append((os.path.relpath(p, self.root).replace(os.sep, '/'), st.st_size, |
| f'{st.st_size}:{st.st_mtime_ns}', st.st_mtime)) |
| return out |
|
|
| def sub(self, prefix): |
| return type(self)(os.path.join(self.root, prefix)) |
|
|
| def fetch(self, path, dest): |
| shutil.copyfile(os.path.join(self.root, path), dest) |
|
|
