hfdask 0.1.0__py3-none-any.whl

hfdask/__init__.py

@@ -0,0 +1,30 @@
+"""Run Dask workloads on Hugging Face Jobs.
+
+`JobSpec` describes a workload and `submit` launches it on one paid HF Job.
+Multi-machine and persistent clusters use `hfdask.cluster`; the CLI ships a
+locked Git project and runs an ordinary Python script from a YAML definition.
+
+## Module reference
+
+- [Jobs](hfdask/jobs.html) — workload specifications and single-Job lifecycle
+- [Clusters](hfdask/cluster.html) — worker groups, identities, launch, and recovery
+- [Clients](hfdask/client.html) — persistent Dask connections over the encrypted mesh
+- [Routing](hfdask/routing.html) — categorical worker affinity and resource reservations
+- [Configuration](hfdask/config.html) — strict input schemas and validation contracts
+- [CLI](hfdask/cli.html) — source packaging, staging, and script submission
+- [Runner](hfdask/runner.html) — in-Job workload execution
+- [Hardware](hfdask/hardware.html) — detected worker profiles and resources
+- [Network](hfdask/network.html) — authenticated transport and loopback proxies
+
+## Lifecycle and safety
+
+Submission reserves paid capacity. Keep recovery manifests until termination
+is verified; a local timeout or client disconnect does not release remote Jobs.
+HF credentials stay with the submitter. Only run trusted workloads and grant
+mounts the minimum access they need.
+"""
+
+from .cluster import WorkerGroup
+from .jobs import Job, JobFailed, JobSpec, submit
+
+__all__ = ["Job", "JobFailed", "JobSpec", "WorkerGroup", "submit"]

hfdask/bootstrap.py

@@ -0,0 +1,86 @@
+"""Stage and launch a locked project using only the image's standard library."""
+
+import hashlib
+import io
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+from pathlib import Path, PurePosixPath
+
+SOURCE = Path("/tmp/hfdask-source/project.tar.gz")
+ROOT = Path("/tmp/hfdask-project")
+MAX_ARCHIVE_BYTES = 8 * 1024 * 1024
+MAX_SOURCE_BYTES = 8 * 1024 * 1024
+MAX_FILES = 2000
+
+
+def require_verified_payload(payload: bytes, checksum: str) -> None:
+    if len(payload) > MAX_ARCHIVE_BYTES:
+        raise SystemExit("hfdask compressed source exceeds 8 MiB")
+    if hashlib.sha256(payload).hexdigest() != checksum:
+        raise SystemExit("hfdask source SHA256 checksum mismatch")
+
+
+def require_extraction_budget(file_count: int, total: int) -> None:
+    if file_count > MAX_FILES or total > MAX_SOURCE_BYTES:
+        raise SystemExit("hfdask source archive exceeds extraction limits")
+
+
+def require_safe_archive_member(
+    member: tarfile.TarInfo, names: set[PurePosixPath]
+) -> PurePosixPath:
+    path = PurePosixPath(member.name)
+    if (
+        not member.isfile()
+        or path.is_absolute()
+        or ".." in path.parts
+        or not path.parts
+        or member.size < 0
+        or path in names
+    ):
+        raise SystemExit("Unsafe hfdask source archive")
+    return path
+
+
+def main() -> None:
+    groups = json.loads(sys.argv[1])
+    with SOURCE.open("rb") as source:
+        payload = source.read(MAX_ARCHIVE_BYTES + 1)
+    require_verified_payload(payload, sys.argv[2])
+    print("hfdask: staging project source", flush=True)
+    ROOT.mkdir(mode=0o700, parents=True, exist_ok=False)
+    with tarfile.open(fileobj=io.BytesIO(payload), mode="r:gz") as archive:
+        members: list[tarfile.TarInfo] = []
+        total = 0
+        names: set[PurePosixPath] = set()
+        for member in archive:
+            total += member.size
+            require_extraction_budget(len(members) + 1, total)
+            path = require_safe_archive_member(member, names)
+            names.add(path)
+            members.append(member)
+        for member in members:
+            target = ROOT.joinpath(*PurePosixPath(member.name).parts)
+            target.parent.mkdir(parents=True, exist_ok=True)
+            extracted = archive.extractfile(member)
+            if extracted is None:
+                raise SystemExit("Unsafe hfdask source archive")
+            with extracted, target.open("xb") as output:
+                shutil.copyfileobj(extracted, output)
+            target.chmod(member.mode & 0o777)
+    os.chdir(ROOT)
+    uv = shutil.which("uv")
+    if uv is None:
+        raise SystemExit("environment.image must contain uv and Python")
+    print("hfdask: syncing locked environment", flush=True)
+    group_args = [arg for group in groups for arg in ("--group", group)]
+    subprocess.run([uv, "sync", "--locked", "--no-dev", *group_args], check=True)
+    print("hfdask: starting runner", flush=True)
+    os.execv(uv, [uv, "run", "--no-sync", *sys.argv[3:]])
+
+
+if __name__ == "__main__":
+    main()

hfdask/cli.py

@@ -0,0 +1,348 @@
+"""YAML-driven batch submission of a locked, Git-selected source project."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import io
+import json
+import os
+import secrets
+import stat
+import subprocess
+import sys
+import tarfile
+import tempfile
+import tomllib
+from dataclasses import replace
+from functools import partial
+from importlib import import_module
+from importlib.resources import files
+from pathlib import Path, PurePosixPath
+from typing import Any
+from uuid import uuid4
+
+from huggingface_hub import HfApi, Volume
+
+from .cluster import Cluster, Identity, LaunchError, submit_cluster
+from .config import MOUNT_SOURCE, ClusterConfig, PackageConfig, PackagePlan, ProjectConfig
+from .jobs import JobSpec
+
+# Bound both upload size and extracted working-tree size.
+MAX_ARCHIVE_BYTES = 8 * 1024 * 1024
+MAX_SOURCE_BYTES = 8 * 1024 * 1024
+MAX_FILES = 2000
+
+
+def _secret_path(path: PurePosixPath) -> bool:
+    return any(
+        part
+        in {
+            ".git",
+            ".venv",
+            ".ssh",
+            ".aws",
+            ".gnupg",
+            ".hfdask",
+            "mise.local.toml",
+            ".mise.local.toml",
+            ".netrc",
+            ".pypirc",
+            "credentials",
+            "credentials.json",
+            "id_rsa",
+            "id_ed25519",
+            "secrets",
+            ".secrets",
+        }
+        or part.startswith((".env", "secrets."))
+        or part.endswith((".env", ".pem", ".key", ".p12", ".pfx"))
+        for part in path.parts
+    )
+
+
+def require_safe_project_file(root: Path, name: str) -> Path | None:
+    relative = PurePosixPath(name)
+    if relative.is_absolute() or ".." in relative.parts:
+        raise ValueError(f"Unsafe project path: {name}")
+    if _secret_path(relative):
+        return None
+    candidate = root
+    for part in relative.parts:
+        candidate = candidate / part
+        if candidate.is_symlink():
+            raise ValueError(f"Project symlinks are not supported: {name}; exclude with .gitignore")
+    if not candidate.exists():  # Tracked files may be deleted in the working tree.
+        return None
+    if not stat.S_ISREG(candidate.stat().st_mode):
+        raise ValueError(f"Project entry is not a regular file: {name}")
+    return candidate
+
+
+def require_source_budget(total: int, file_count: int) -> None:
+    if total > MAX_SOURCE_BYTES or file_count > MAX_FILES:
+        raise ValueError(
+            "Project source exceeds 8 MiB/2000 files; move data to bucket mounts "
+            "and exclude generated files with .gitignore"
+        )
+
+
+def require_project_inputs(files: dict[str, bytes], script: str) -> None:
+    for required in ("pyproject.toml", "uv.lock", script):
+        if required not in files:
+            raise ValueError(
+                f"Required project file {required} is missing or excluded; "
+                "run uv lock and check .gitignore"
+            )
+
+
+def require_upload_budget(payload: bytes) -> bytes:
+    if len(payload) > MAX_ARCHIVE_BYTES:
+        raise ValueError(
+            "Compressed source exceeds 8 MiB upload limit; move data to bucket "
+            "mounts and exclude generated files with .gitignore"
+        )
+    return payload
+
+
+def package_project(root: Path, script: str, groups: list[str]) -> bytes:
+    """Snapshot working-tree files, not HEAD; never follow project symlinks."""
+    inputs = PackageConfig(script=script, groups=groups)
+    script = inputs.script
+    try:
+        selected = subprocess.run(
+            ["git", "ls-files", "--cached", "--others", "--exclude-standard", "-z", "--", "."],
+            cwd=root,
+            check=True,
+            capture_output=True,
+        ).stdout.split(b"\0")
+        # --exclude-standard only filters untracked files; also exclude ignored tracked files.
+        ignored = set(
+            subprocess.run(
+                ["git", "ls-files", "--cached", "--ignored", "--exclude-standard", "-z", "--", "."],
+                cwd=root,
+                check=True,
+                capture_output=True,
+            ).stdout.split(b"\0")
+        )
+    except (OSError, subprocess.CalledProcessError) as error:
+        raise ValueError("Run from a Git project containing pyproject.toml and uv.lock") from error
+    files: dict[str, bytes] = {}
+    modes: dict[str, int] = {}
+    total = 0
+    for raw in sorted(set(selected) - ignored - {b""}):
+        name = os.fsdecode(raw)
+        candidate = require_safe_project_file(root, name)
+        if candidate is None:
+            continue
+        metadata = candidate.stat()
+        total += metadata.st_size
+        require_source_budget(total, len(files) + 1)
+        files[name] = candidate.read_bytes()
+        modes[name] = 0o755 if metadata.st_mode & 0o111 else 0o644
+    require_project_inputs(files, script)
+    metadata = tomllib.loads(files["pyproject.toml"].decode())
+    PackagePlan.model_validate(
+        {
+            "script": script,
+            "groups": inputs.groups,
+            "project": ProjectConfig.model_validate(metadata.get("project", {})),
+            "dependency-groups": metadata.get("dependency-groups", {}),
+        }
+    )
+    buffer = io.BytesIO()
+    with tarfile.open(fileobj=buffer, mode="w:gz") as archive:
+        for name, data in files.items():
+            member = tarfile.TarInfo(name)
+            member.size = len(data)
+            member.mode = modes[name]
+            archive.addfile(member, io.BytesIO(data))
+    return require_upload_budget(buffer.getvalue())
+
+
+def require_private_source_bucket(api: HfApi, bucket_id: str) -> None:
+    if api.bucket_info(bucket_id).private is not True:
+        raise ValueError(
+            f"Source bucket {bucket_id} must be verified private; "
+            "refusing upload without changing visibility"
+        )
+
+
+def prepare_spec(spec: JobSpec, root: Path, script: str, groups: list[str], api: HfApi) -> JobSpec:
+    """Validate and stage source once; retained artifacts incur storage until deleted."""
+    payload = package_project(root, script, groups)
+    bucket_id = f"{spec.namespace}/jobs-artifacts"
+    folder = f"hfdask-source/{uuid4()}"
+    remote_path = f"{folder}/project.tar.gz"
+    api.create_bucket(bucket_id, private=True, exist_ok=True)
+    require_private_source_bucket(api, bucket_id)
+    bootstrap = files("hfdask").joinpath("bootstrap.py").read_bytes()
+    api.batch_bucket_files(
+        bucket_id,
+        add=[
+            (payload, remote_path),
+            (bootstrap, f"{folder}/bootstrap.py"),
+        ],
+    )
+    print(
+        f"Source artifact: hf://buckets/{bucket_id}/{remote_path} "
+        "(retained after run; storage charges may apply until manually deleted)",
+        file=sys.stderr,
+    )
+    return replace(
+        spec,
+        bootstrap=(
+            "python3",
+            "/tmp/hfdask-source/bootstrap.py",
+            json.dumps(groups),
+            hashlib.sha256(payload).hexdigest(),
+        ),
+        volumes=[
+            *spec.volumes,
+            Volume(
+                type="bucket",
+                source=bucket_id,
+                path=folder,
+                mount_path="/tmp/hfdask-source",
+                read_only=True,
+            ),
+        ],
+    )
+
+
+def load_cluster(
+    path: Path,
+    root: Path,
+    script: str,
+) -> tuple[JobSpec, dict[str, Any], float, list[str]]:
+    # Load lazily so non-CLI library users do not need PyYAML at import time.
+    yaml = import_module("yaml")
+
+    config = ClusterConfig.model_validate(yaml.safe_load(path.read_text()))
+    scheduler_worker = config.coordinator.worker
+    volumes: list[Volume] = []
+    for mount in config.mounts:
+        location = MOUNT_SOURCE.fullmatch(mount.source)
+        assert location is not None  # Validated by MountConfig; conversion only below.
+        kind, namespace, name, subfolder = location.groups()
+        volumes.append(
+            Volume(
+                type=kind[:-1],
+                source=f"{namespace}/{name}",
+                revision=mount.revision,
+                path=subfolder or "",
+                mount_path=mount.target,
+                read_only=mount.read_only,
+            )
+        )
+    spec = JobSpec(
+        namespace=config.namespace,
+        image=config.environment.image,
+        entrypoint="hfdask.runner:run_script",
+        kwargs={"script": script},
+        # YAML counts remote workers; JobSpec includes the colocated worker.
+        workers=config.workers.count + int(scheduler_worker),
+        flavor=config.workers.flavor,
+        timeout=config.timeout,
+        volumes=volumes,
+        env={
+            "PYTHONPATH": "/tmp/hfdask-project:"
+            + str(PurePosixPath("/tmp/hfdask-project") / PurePosixPath(script).parent),
+            "DASK_DISTRIBUTED__WORKER__DAEMON": "False",
+            "DASK_DISTRIBUTED__WORKER__MULTIPROCESSING_METHOD": "spawn",
+            "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
+        },
+    )
+    options: dict[str, Any] = {
+        "public_relays": True,
+        "scheduler_worker": scheduler_worker,
+        "scheduler_flavor": config.coordinator.flavor,
+    }
+    return spec, options, config.timeout_seconds, config.environment.groups
+
+
+def save_manifest(path: Path, cluster: Cluster) -> None:
+    """Atomic replacement; manifests contain only public recovery handles."""
+    with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as output:
+        temporary = Path(output.name)
+        try:
+            json.dump(cluster.manifest(), output, indent=2)
+            output.write("\n")
+            output.flush()
+            os.fsync(output.fileno())
+        except BaseException:
+            temporary.unlink(missing_ok=True)
+            raise
+    try:
+        temporary.replace(path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="hfdask")
+    commands = parser.add_subparsers(dest="command", required=True)
+    run = commands.add_parser("run", help="Run a project-relative Dask script on a YAML cluster")
+    run.add_argument("--cluster", type=Path, required=True)
+    run.add_argument(
+        "--manifest", type=Path, help="Public recovery manifest (default: .hfdask/run-*.json)"
+    )
+    run.add_argument("script")
+    args = parser.parse_args(argv)
+    cluster: Cluster | None = None
+    manifest: Path = (
+        args.manifest
+        if args.manifest is not None
+        else Path(".hfdask") / f"run-{secrets.token_hex(8)}.json"
+    )
+    try:
+        spec, options, timeout, groups = load_cluster(args.cluster, Path.cwd(), args.script)
+        nodes = spec.workers + 1 - int(options["scheduler_worker"])
+        identities = [Identity(secrets.token_bytes(32)) for _ in range(nodes)]
+        # Validate transport identities before reserving the manifest or submitting.
+        for identity in identities:
+            identity.public_id()
+
+        manifest.parent.mkdir(parents=True, exist_ok=True)
+        with manifest.open("x") as output:
+            output.write("{}\n")
+        print(f"Recovery manifest: {manifest}", file=sys.stderr)
+        api = HfApi()
+        spec = prepare_spec(spec, Path.cwd(), args.script, groups, api)
+        try:
+            cluster = submit_cluster(
+                spec, identities, api=api, **options, on_submitted=partial(save_manifest, manifest)
+            )
+        except LaunchError as error:
+            cluster = error.cluster
+            raise
+        cluster.wait(timeout=timeout)
+        return 0
+    # CLI boundary: report failures and release known jobs before returning an exit code.
+    except (Exception, KeyboardInterrupt) as error:  # noqa: BLE001
+        print(f"hfdask: {error or 'interrupted'}", file=sys.stderr)
+        if isinstance(error, LaunchError) and error.__cause__ is not None:
+            print(f"Submission cause: {error.__cause__}", file=sys.stderr)
+        if cluster is not None:
+            try:
+                save_manifest(manifest, cluster)
+            # A manifest failure must not prevent cleanup; print recovery handles instead.
+            except Exception as save_error:  # noqa: BLE001
+                print(
+                    f"Manifest write failed: {save_error}; known jobs: "
+                    f"{json.dumps(cluster.manifest())}",
+                    file=sys.stderr,
+                )
+            try:
+                cluster.close()
+            # Report any cleanup failure without claiming that capacity was released.
+            except Exception as close_error:  # noqa: BLE001
+                print(f"Cleanup unverified: {close_error}; retain {manifest}", file=sys.stderr)
+        interrupted = isinstance(error, KeyboardInterrupt) or (
+            isinstance(error, LaunchError) and isinstance(error.__cause__, KeyboardInterrupt)
+        )
+        return 130 if interrupted else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

hfdask/client.py

@@ -0,0 +1,135 @@
+"""External synchronous Dask clients over the cluster's encrypted mesh."""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any
+
+from .cluster import Identity
+from .config import ConnectConfig, ConnectionConfig
+
+_connection_lock = threading.Lock()
+
+
+async def _serve(
+    config: ConnectionConfig,
+    identity: Identity,
+    ready: concurrent.futures.Future[tuple[int, ...]],
+    stop: threading.Event,
+    timeout: float,
+) -> None:
+    import iroh
+
+    from .hardware import service_owners
+    from .network import ALPN, Mesh, fetch_worker_counts
+
+    # FFI annotates BaseEventLoop but uses the standard AbstractEventLoop interface.
+    iroh.iroh_ffi.uniffi_set_event_loop(asyncio.get_running_loop())  # ty: ignore[invalid-argument-type]
+    mode = (
+        iroh.RelayMode.custom_from_urls(config.relays)
+        if config.relays
+        else iroh.RelayMode.default_mode()
+    )
+    endpoint = await iroh.Endpoint.bind(
+        iroh.EndpointOptions(
+            preset=iroh.preset_n0(), secret_key=identity.secret, alpns=[ALPN], relay_mode=mode
+        )
+    )
+    try:
+        await asyncio.wait_for(endpoint.online(), timeout=timeout)
+        peers = [
+            iroh.EndpointAddr(iroh.EndpointId.from_bytes(bytes.fromhex(peer)), None, [])
+            for peer in config.peers
+        ]
+        nodes = config.job_nodes
+        workers_per_node = await fetch_worker_counts(endpoint, peers[0], timeout)
+        if len(workers_per_node) != nodes:
+            raise ConnectionError("Live worker topology does not match the manifest")
+        async with Mesh(
+            endpoint,
+            peers,
+            nodes,
+            services=service_owners(workers_per_node),
+            workers_per_node=workers_per_node,
+            max_connections=max(256, 4 * sum(workers_per_node)),
+        ):
+            ready.set_result(workers_per_node)
+            while not stop.is_set():
+                await asyncio.sleep(0.1)
+    finally:
+        await endpoint.close()
+
+
+@contextmanager
+def connect(
+    manifest: dict[str, Any], identity: Identity, *, timeout: float = 1200
+) -> Iterator[Any]:
+    """Connect to a persistent cluster without taking ownership of its HF Jobs.
+
+    Args:
+        manifest: Public `hfdask.cluster.Cluster.manifest` from `boot_cluster`.
+        identity: Private client identity supplied at boot, not a Job node identity.
+        timeout: Positive timeout in seconds used for mesh startup, Dask connection,
+            and worker readiness. These phases do not share one total deadline.
+
+    Yields:
+        A synchronous Dask client after the expected worker topology is ready.
+        It is not installed as Dask's default client.
+
+    Raises:
+        ValueError: If the manifest, identity, or timeout fails validation.
+        TimeoutError: If mesh startup or worker readiness exceeds its deadline.
+        RuntimeError: If this process already has a mesh client or cleanup cannot finish.
+
+    One mesh connection per host: fixed loopback ports start at 21000. Keep the
+    context open while using futures. Exiting closes the client and local mesh,
+    never the HF Jobs; explicitly call `hfdask.cluster.Cluster.close` to release them.
+    """
+    from distributed import Client
+
+    from .runner import wait_topology
+
+    connection = ConnectionConfig.model_validate(manifest.get("connection"))
+    ConnectConfig(timeout=timeout, connection=connection, public_id=identity.public_id())
+
+    if not _connection_lock.acquire(blocking=False):
+        raise RuntimeError("A mesh client is already connected in this process")
+    ready: concurrent.futures.Future[tuple[int, ...]] = concurrent.futures.Future()
+    stopped = threading.Event()
+    loop = asyncio.new_event_loop()
+    task: asyncio.Task[None] | None = None
+
+    def serve() -> None:
+        nonlocal task
+        asyncio.set_event_loop(loop)
+        task = loop.create_task(_serve(connection, identity, ready, stopped, timeout))
+        try:
+            loop.run_until_complete(task)
+        except BaseException as error:  # noqa: BLE001 - propagate startup failures across threads.
+            if not ready.done():
+                ready.set_exception(error)
+        finally:
+            loop.run_until_complete(loop.shutdown_asyncgens())
+            loop.close()
+
+    thread = threading.Thread(target=serve, name="hfdask-client-mesh", daemon=True)
+    try:
+        thread.start()
+        workers_per_node = ready.result(timeout=timeout + 5)
+        with Client("tcp://127.0.0.1:21000", timeout=timeout, set_as_default=False) as client:
+            wait_topology(client, workers_per_node, timeout)
+            yield client
+    finally:
+        stopped.set()
+        if thread.ident is not None:
+            thread.join(timeout=5)
+            if thread.is_alive() and not loop.is_closed() and task is not None:
+                loop.call_soon_threadsafe(task.cancel)
+                thread.join(timeout=5)
+        if thread.is_alive():
+            raise RuntimeError("Mesh cleanup did not finish; restart this client process")
+        _connection_lock.release()