fullFold 0.1.0__py3-none-any.whl

fullFold/__init__.py

@@ -0,0 +1,3 @@
+"""Wrapper-only multi-GPU scheduling layer around AlphaFold 3."""
+
+__version__ = '0.1.0'

fullFold/__main__.py

@@ -0,0 +1,4 @@
+from fullFold.cli import main
+
+if __name__ == '__main__':
+    raise SystemExit(main())

fullFold/banner.py

@@ -0,0 +1,20 @@
+"""Startup banner for the fullFold CLI. Printed to stderr."""
+
+from __future__ import annotations
+
+import sys
+
+BANNER = """\
+  __       _ _ _____     _     _ 
+ / _|_   _| | |  ___|__ | | __| |
+| |_| | | | | | |_ / _ \\| |/ _` |
+|  _| |_| | | |  _| (_) | | (_| |
+|_|  \\__,_|_|_|_|  \\___/|_|\\__,_|
+
+fullFold: unlocking the full speed of AF3!
+Please cite AlphaFold3 and fullFold (see our repository).
+"""
+
+
+def print_banner(file=None) -> None:
+    print(BANNER, file=file or sys.stderr, end='')

fullFold/benchmark.py

@@ -0,0 +1,232 @@
+"""GPU discovery and 1024-token compile/throughput probe."""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import warnings
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+from fullFold.config import Config, host_id, worker_environ, write_atomic
+from fullFold.scheduling import compilation_modifier, inference_modifier
+
+T_REF_S = 59.423  # A100 tokamax inference-only seconds at bucket 1024
+
+
+@dataclass(frozen=True)
+class Gpu:
+    slot: int
+    physical_id: str
+    uuid: str
+    pci_bus_id: str
+    device_kind: str
+    memory_bytes: int
+
+
+@dataclass(frozen=True)
+class Bench:
+    t_ms: tuple[float, float, float, float]
+    s_ms: float
+    r_ms: float
+    contaminated: bool
+    host: str
+    gpu_key: str
+
+
+def gpu_multiplier(bench: Bench) -> float:
+    """m_gpu = S_1024 / T_REF from the probe's inference-only time."""
+    s_s = bench.s_ms / 1000.0
+    if T_REF_S <= 0 or s_s <= 0:
+        return 1.0
+    return s_s / T_REF_S
+
+
+def predict_inference_s(m_gpu: float, bucket: int) -> float:
+    """T_predicted = m_gpu * inference_modifier(bucket) * T_REF, in seconds."""
+    return m_gpu * inference_modifier(bucket) * T_REF_S
+
+
+def compile_overhead_ms(bench: Bench, shape: int) -> float:
+    """Probe R at 1024 scaled by the CSV compile modifier (last-row above 5216)."""
+    return bench.r_ms * compilation_modifier(shape)
+
+
+def summarise_timings(t_ms: tuple[float, ...] | list[float]) -> tuple[float, float, bool]:
+    """Tokamax pays compile on seeds 1 and 2; seeds 3–4 are steady-state inference.
+
+    ``contaminated`` means S is unusable (t3 vs t4 disagree). Negative R is a
+    warm cache or inverted compile, not a reason to ignore throughput.
+    """
+    t = tuple(float(x) for x in t_ms)
+    s = (t[2] + t[3]) / 2.0
+    r = t[0] + t[1] - 2.0 * s
+    contaminated = s > 0 and abs(t[2] - t[3]) / s > 0.25
+    return s, r, contaminated
+
+
+def nvidia_smi() -> list[dict]:
+    cmd = [
+        'nvidia-smi',
+        '--query-gpu=index,uuid,pci.bus_id,name,memory.total',
+        '--format=csv,noheader,nounits',
+    ]
+    try:
+        out = subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL)
+    except (OSError, subprocess.CalledProcessError):
+        return []
+    rows = []
+    for line in out.splitlines():
+        parts = [p.strip() for p in line.split(',')]
+        if len(parts) < 5:
+            continue
+        try:
+            mem = int(float(parts[4])) * 1024 * 1024
+        except ValueError:
+            mem = 0
+        rows.append({
+            'index': parts[0], 'uuid': parts[1], 'pci': parts[2],
+            'name': parts[3], 'memory_bytes': mem,
+        })
+    rows.sort(key=lambda r: r['pci'])
+    return rows
+
+
+def discover_gpus(cfg: Config, rows: list[dict] | None = None) -> list[Gpu]:
+    rows = list(rows if rows is not None else nvidia_smi())
+    env = os.environ.get('CUDA_VISIBLE_DEVICES')
+    wanted = list(cfg.gpus) if cfg.gpus else (
+        [x.strip() for x in env.split(',') if x.strip()]
+        if env not in (None, '') else []
+    )
+    if wanted:
+        selected = []
+        for w in wanted:
+            hit = next((r for r in rows if r['index'] == w or r['uuid'] == w), None)
+            if hit is None and rows:
+                raise ValueError(f'GPU {w!r} is not in nvidia-smi output')
+            if hit:
+                selected.append(hit)
+        rows = selected
+    gpus = [
+        Gpu(slot=i, physical_id=r['index'], uuid=r['uuid'],
+            pci_bus_id=r['pci'], device_kind=r['name'],
+            memory_bytes=int(r['memory_bytes']))
+        for i, r in enumerate(rows)
+    ]
+    if not gpus:
+        raise RuntimeError('no GPUs discovered (nvidia-smi empty and no --gpus)')
+    return gpus
+
+
+def gpu_key(gpu: Gpu, jax_version: str = '', af3_version: str = '') -> str:
+    return f'{gpu.pci_bus_id}|{gpu.device_kind}|{jax_version}|{af3_version}'
+
+
+def cache_path(cfg: Config, gpu: Gpu, jax_version: str = '', af3_version: str = '') -> Path:
+    return cfg.cache_dir / f'{host_id()}__{gpu_key(gpu, jax_version, af3_version)}.json'
+
+
+def load_bench(path: Path) -> Bench | None:
+    if not path.is_file():
+        return None
+    d = json.loads(path.read_text())
+    t = tuple(d['t_ms'])
+    s, r, cont = summarise_timings(t)
+    return Bench(t_ms=t, s_ms=s, r_ms=r, contaminated=cont,
+                 host=d.get('host', ''), gpu_key=d.get('gpu_key', ''))
+
+
+def save_bench(path: Path, bench: Bench) -> None:
+    write_atomic(path, json.dumps(asdict(bench), indent=2))
+
+
+def build_probe_json(n_tokens: int = 1024, rng_seed: int = 42) -> dict:
+    import random
+    aa = list('ACDEFGHIKLMNPQRSTVWY')
+    rng = random.Random(rng_seed)
+    seq = ''.join(rng.choices(aa, k=n_tokens))
+    msa = [f'>query\n{seq}\n']
+    for i in range(31):
+        msa.append(f'>seq_{i+1}\n' + ''.join(rng.choices(aa, k=n_tokens)) + '\n')
+    return {
+        'name': f'fullFold_bench_{n_tokens}',
+        'modelSeeds': [0, 1, 2, 3],
+        'sequences': [{
+            'protein': {
+                'id': 'A', 'sequence': seq,
+                'unpairedMsa': ''.join(msa), 'pairedMsa': '', 'templates': [],
+            }
+        }],
+        'dialect': 'alphafold3', 'version': 3,
+    }
+
+
+def to_scheduler_gpu(bench: Bench) -> tuple[int, int]:
+    """(compile_us, infer_us_1024) for scheduling.py.
+
+    ``m_gpu = S_1024 / T_REF`` from the probe; inference at bucket b is
+    ``m_gpu * f(b/1024) * T_REF``.
+    """
+    compile_us = max(int(round(compile_overhead_ms(bench, 1024) * 1000)), 0)
+    infer_us_1024 = max(1, int(round(gpu_multiplier(bench) * T_REF_S * 1e6)))
+    return compile_us, infer_us_1024
+
+
+def measure(gpu: Gpu, cfg: Config) -> Bench:
+    """Spawn a per-GPU probe subprocess so the parent never initialises JAX."""
+    import sys
+    env = worker_environ(
+        cfg, gpu.physical_id, pci_bus_id=gpu.pci_bus_id, device_kind=gpu.device_kind)
+    cmd = [
+        sys.executable, '-m', 'fullFold.worker', '--probe',
+        '--gpu', str(gpu.physical_id),
+        '--model-dir', str(cfg.model_dir),
+        '--output-dir', str(cfg.output_dir),
+        '--flash-attention', cfg.flash_attention,
+    ]
+    r = subprocess.run(cmd, env=env, capture_output=True, text=True)
+    if r.returncode != 0:
+        detail = (r.stderr or r.stdout or '(no output)').strip()
+        raise RuntimeError(
+            f'probe on GPU {gpu.physical_id} exited {r.returncode}:\n{detail}'
+        )
+    t = json.loads(r.stdout.strip().splitlines()[-1])
+    s, r_ms, cont = summarise_timings(t)
+    if cont:
+        warnings.warn(f'contaminated benchmark on GPU {gpu.physical_id}: t={t}')
+    b = Bench(t_ms=tuple(t), s_ms=s, r_ms=r_ms, contaminated=cont,
+              host=host_id(), gpu_key=gpu_key(gpu))
+    save_bench(cache_path(cfg, gpu), b)
+    return b
+
+
+def get_or_measure(
+    gpus: list[Gpu], cfg: Config, measure_fn=None,
+) -> list[tuple[Gpu, Bench]]:
+    measure_fn = measure_fn or measure
+    out: list[tuple[Gpu, Bench]] = []
+    missing: list[Gpu] = []
+    for g in gpus:
+        p = cache_path(cfg, g)
+        b = None if cfg.force_benchmark else load_bench(p)
+        if b is None:
+            missing.append(g)
+        else:
+            out.append((g, b))
+    if missing:
+        n = len(missing)
+        gpu_word = 'GPU' if n == 1 else 'GPUs'
+        print(
+            f'A few-minute benchmark of the available GPUs is running '
+            f'({n} {gpu_word}).',
+            file=sys.stderr,
+        )
+        from concurrent.futures import ThreadPoolExecutor
+        with ThreadPoolExecutor(max_workers=n) as ex:
+            benches = list(ex.map(lambda g: measure_fn(g, cfg), missing))
+        out.extend(zip(missing, benches))
+        out.sort(key=lambda x: x[0].slot)
+    return out

fullFold/cli.py

@@ -0,0 +1,103 @@
+"""argparse -> Config -> one function. No scheduling logic here."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from fullFold.banner import print_banner
+from fullFold.config import DEFAULT_BUCKETS, Config, config_from_dict
+
+
+def _buckets(s: str) -> tuple[int, ...]:
+    return tuple(int(x) for x in s.split(',')) if s else DEFAULT_BUCKETS
+
+
+def _add_io(p: argparse.ArgumentParser) -> None:
+    p.add_argument('--input-dir', type=Path, required=True)
+    p.add_argument('--output-dir', type=Path, required=True)
+
+
+def _add_sched(p: argparse.ArgumentParser) -> None:
+    p.add_argument('--model-dir', type=Path, default=Config.model_dir)
+    p.add_argument('--gpus', default='')
+    p.add_argument('--no-prefetch', dest='prefetch', action='store_const',
+                   const=0, default=1)
+    p.add_argument('--no-background-extract', dest='background_extract',
+                   action='store_false')
+    p.add_argument('--policy', choices=('contiguous', 'roundrobin'), default='contiguous')
+    p.add_argument('--bucket-mode', choices=('ladder', 'free'), default='free')
+    p.add_argument('--buckets', type=_buckets, default=DEFAULT_BUCKETS)
+    p.add_argument('--bucket-margin', type=float, default=0.05)
+    p.add_argument('--stale-lock-seconds', type=int, default=900)
+    p.add_argument('--retry-failed', action='store_true')
+    p.add_argument('--exact-split-threshold', type=int, default=512)
+    p.add_argument('--force-benchmark', action='store_true')
+    p.add_argument('--bench-seed', type=int, default=42)
+    p.add_argument('--cache-dir', type=Path, default=Config.cache_dir)
+    p.add_argument('--exact-tokens', action='store_true')
+    p.add_argument('--dry-run', action='store_true')
+    p.add_argument('--jax-compilation-cache-dir', type=Path, default=None)
+    p.add_argument('--xla-mem-fraction', type=float, default=0.97)
+    p.add_argument('--save-embeddings', action='store_true')
+    p.add_argument('--save-distogram', action='store_true')
+    p.add_argument('--compress-large-output-files', action='store_true')
+    p.add_argument('--num-recycles', type=int, default=10)
+    p.add_argument('--num-diffusion-samples', type=int, default=5)
+    p.add_argument('--flash-attention', default='triton')
+
+
+def _cfg(ns: argparse.Namespace) -> Config:
+    kw = {n: getattr(ns, n) for n in Config.__dataclass_fields__ if hasattr(ns, n)}
+    g = kw.get('gpus', ())
+    if isinstance(g, str):
+        kw['gpus'] = tuple(x for x in g.split(',') if x)
+    return config_from_dict(kw)
+
+
+def main(argv: list[str] | None = None) -> int:
+    print_banner()
+    parser = argparse.ArgumentParser(prog='fullFold')
+    sub = parser.add_subparsers(dest='cmd', required=True)
+
+    for name in ('scan', 'plan', 'run'):
+        p = sub.add_parser(name)
+        _add_io(p)
+        _add_sched(p)
+        if name == 'run':
+            p.add_argument('-v', '--verbose', action='store_true',
+                           help='Show per-GPU progress bars during the run')
+
+    p = sub.add_parser('benchmark')
+    p.add_argument('--output-dir', type=Path, default=Path('.'))
+    p.add_argument('--input-dir', type=Path, default=Path('.'))
+    _add_sched(p)
+
+    p = sub.add_parser('template')
+    p.add_argument('--template', type=Path, required=True)
+    p.add_argument('--records', type=Path, required=True)
+    p.add_argument('--output-dir', type=Path, required=True)
+    p.add_argument('--type', dest='record_type', default='protein')
+
+    args = parser.parse_args(argv)
+    cfg = _cfg(args)
+    if args.cmd == 'scan':
+        from fullFold.engine import cmd_scan
+        return cmd_scan(cfg)
+    if args.cmd == 'benchmark':
+        from fullFold.engine import cmd_benchmark
+        return cmd_benchmark(cfg)
+    if args.cmd == 'plan':
+        from fullFold.engine import cmd_plan
+        return cmd_plan(cfg)
+    if args.cmd == 'run':
+        from fullFold.engine import cmd_run
+        return cmd_run(cfg)
+    if args.cmd == 'template':
+        from fullFold.templates import cmd_template
+        return cmd_template(cfg)
+    return 2
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())

fullFold/config.py

@@ -0,0 +1,148 @@
+"""Single frozen Config plus hashing helpers."""
+
+from __future__ import annotations
+
+import dataclasses
+import hashlib
+import json
+import os
+import socket
+from pathlib import Path
+
+DEFAULT_BUCKETS = (
+    128, 256, 384, 512, 768, 1024, 1280, 1536, 2048, 2560, 3072, 3584, 4096,
+    4608, 5120,
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class Config:
+    input_dir: Path = Path('.')
+    output_dir: Path = Path('.')
+    model_dir: Path = Path('~/models').expanduser()
+    gpus: tuple[str, ...] = ()
+    dry_run: bool = False
+    verbose: bool = False
+    prefetch: int = 1
+    background_extract: bool = True
+    policy: str = 'contiguous'  # contiguous | roundrobin
+    bucket_mode: str = 'free'  # free | ladder
+    buckets: tuple[int, ...] = DEFAULT_BUCKETS
+    bucket_margin: float = 0.05
+    stale_lock_seconds: int = 900
+    retry_failed: bool = False
+    exact_split_threshold: int = 512
+    force_benchmark: bool = False
+    bench_seed: int = 42
+    cache_dir: Path = Path('~/.cache/fullFold/bench').expanduser()
+    jax_compilation_cache_dir: Path | None = None
+    xla_mem_fraction: float = 0.97
+    xla_preallocate: bool = True
+    exact_tokens: bool = False
+    template: Path | None = None
+    records: Path | None = None
+    record_type: str = 'protein'
+    save_embeddings: bool = False
+    save_distogram: bool = False
+    compress_large_output_files: bool = False
+    save_terms_of_use: bool = True
+    num_recycles: int = 10
+    num_diffusion_samples: int = 5
+    flash_attention: str = 'triton'
+
+
+def host_id() -> str:
+    p = Path('/etc/machine-id')
+    if p.is_file():
+        return p.read_text().strip()
+    return hashlib.sha256(socket.gethostname().encode()).hexdigest()[:16]
+
+
+def sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def sha256_file(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open('rb') as f:
+        for chunk in iter(lambda: f.read(1 << 16), b''):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def sanitise(name: str) -> str:
+    allowed = set('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-.')
+    return ''.join(c for c in name.replace(' ', '_') if c in allowed)
+
+
+def config_hash(cfg: Config) -> str:
+    return sha256_bytes(json.dumps(config_to_dict(cfg), sort_keys=True).encode())[:16]
+
+
+def config_to_dict(cfg: Config) -> dict:
+    d = dataclasses.asdict(cfg)
+    for k, v in d.items():
+        if isinstance(v, Path):
+            d[k] = str(v)
+        elif isinstance(v, tuple):
+            d[k] = list(v)
+    return d
+
+
+def config_from_dict(d: dict) -> Config:
+    fields = Config.__dataclass_fields__
+    kw = {}
+    for k, v in d.items():
+        if k not in fields:
+            continue
+        hint = str(fields[k].type)
+        if v is None:
+            kw[k] = None
+        elif 'Path' in hint and not isinstance(v, Path):
+            kw[k] = Path(v)
+        elif 'tuple' in hint and not isinstance(v, tuple):
+            kw[k] = tuple(v)
+        else:
+            kw[k] = v
+    return Config(**kw)
+
+
+def write_atomic(path: Path, text: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_name(path.name + '.tmp')
+    tmp.write_text(text)
+    os.replace(tmp, path)
+
+
+def jax_cache_dir(
+    cfg: Config, *, physical_id: str, pci_bus_id: str = '', device_kind: str = '',
+) -> Path:
+    """Persistent XLA cache root for this host + GPU name. Always enabled.
+
+    PCI / physical_id are accepted for call-site compatibility but are not
+    part of the path: identical nvidia-smi names share a directory. JAX's
+    file keys (HLO, jaxlib, GPU-name topology) still miss across models.
+    """
+    del physical_id, pci_bus_id
+    root = Path(cfg.jax_compilation_cache_dir) if cfg.jax_compilation_cache_dir else (
+        cfg.cache_dir / 'jax')
+    kind_part = sanitise(device_kind) if device_kind else 'gpu'
+    return root / f'{host_id()}__{kind_part}'
+
+
+def worker_environ(
+    cfg: Config, gpu_physical_id: str, *,
+    pci_bus_id: str = '', device_kind: str = '',
+) -> dict[str, str]:
+    """Env for a one-GPU worker/probe subprocess (parent never initialises JAX)."""
+    env = os.environ.copy()
+    env['CUDA_VISIBLE_DEVICES'] = str(gpu_physical_id)
+    env['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    env['XLA_CLIENT_MEM_FRACTION'] = str(cfg.xla_mem_fraction)
+    env['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'true' if cfg.xla_preallocate else 'false'
+    cache = jax_cache_dir(
+        cfg, physical_id=gpu_physical_id, pci_bus_id=pci_bus_id,
+        device_kind=device_kind)
+    cache.mkdir(parents=True, exist_ok=True)
+    env['AF3SCHED_JAX_CACHE'] = str(cache)
+    return env