diffusers-workflow 0.4.0__py3-none-any.whl

dw/host_memory_projection.py

@@ -0,0 +1,230 @@
+"""Whether a list-driven or composed run is projected to exceed host RAM.
+
+`dw/host_memory.py` reads what a *running* worker is holding; this module is
+the pre-flight half, for the failure mode `device_memory_stats`/VRAM warnings
+never covered - a `for_each` that keeps its pipeline resident across a
+12-entry list, or a per-clip step whose footprint scales with the list, can
+SIGKILL on host RAM with the accelerator nowhere near full (#243).
+
+Scope is v1, by the repo owner's own sign-off on the issue: observed, not
+curated (no author-declared memory figure - `host_memory_job_peak_rss_mb` is
+a worker-reported field, not a schema key); warn, not refuse (host RAM headroom
+is a property of *this machine*, not something a caller chose, so it never
+blocks a run); no cross-machine normalization; and a cold start - no history
+for this workflow at all - means no check, the same rule `observed_cost.py`
+uses for GPU minutes.
+"""
+
+import json
+import logging
+
+from .for_each import FOR_EACH_KEY
+from .plan import _has_seedable_step
+
+logger = logging.getLogger("dw")
+
+# The fraction of physical RAM a projected peak may reach before this module
+# says anything - matches nothing else in the codebase; there is no existing
+# "safe headroom" constant to share, because nothing else projects against
+# the machine's total RAM rather than a curated figure.
+CEILING_FRACTION = 0.9
+
+_RELEASE_FLAGS = ("release_pipeline", "release_models")
+
+
+def releases_between_iterations(definition):
+    """Whether every `for_each` step in this workflow drops its pipeline (or
+    its models) between entries, rather than holding it resident for the
+    whole list.
+
+    `release_pipeline`/`release_models` survive on a for_each step's *last*
+    member only (`dw/for_each.py`), so their presence on the template step is
+    what says the author asked for the resident-vs-released shape here - the
+    same field, read for a different question. A workflow with no `for_each`
+    step at all answers True: there is no list to hold anything resident
+    across, so the "keeps everything resident" projection would not describe
+    it.
+    """
+    for_each_steps = [
+        step
+        for step in definition.get("steps") or []
+        if isinstance(step, dict) and FOR_EACH_KEY in step
+    ]
+    if not for_each_steps:
+        return True
+    return all(
+        any(step.get(flag) for flag in _RELEASE_FLAGS) for step in for_each_steps
+    )
+
+
+def _requested_count(list_entries):
+    """The largest list length the caller's own arguments drive, or None
+    when this run has no list at all - such a run is never the failure mode
+    this module exists for."""
+    if not list_entries:
+        return None
+    counts = [value for value in list_entries.values() if isinstance(value, int)]
+    return max(counts) if counts else None
+
+
+def _row_count(row, list_entries, variables):
+    """The list length a historical row ran with, read from its own stored
+    arguments rather than the current request's - a row that ran a shorter
+    or longer list is still comparable, once divided out, for the resident
+    projection's per-entry figure.
+
+    A row's stored `arguments` are only the caller's *overrides*
+    (`job.spec["arguments"]`), never merged with the workflow's declared
+    defaults - a run of the catalog default (the common case for this
+    module's own regression case) is recorded as `{}`. Falling back to
+    `variables[name]` the same way `observed_cost._bucket_key` does is what
+    lets such a row resolve to its actual list length rather than to no
+    length at all, which used to drop every default-arguments row out of the
+    per-entry figure and left the resident shape silent regardless of
+    history (#264).
+
+    Only the variable names the *current* request's `list_entries` names are
+    read back, on the assumption that a workflow's list-driving variables
+    are stable across its history - the same assumption `observed_cost.py`
+    makes bucketing by declared `cost_drivers`.
+    """
+    try:
+        arguments = json.loads(row.get("arguments") or "{}")
+    except (TypeError, ValueError):
+        arguments = {}
+    if not isinstance(arguments, dict):
+        arguments = {}
+    counts = []
+    for name in list_entries:
+        value = arguments.get(name, variables.get(name))
+        if isinstance(value, list):
+            counts.append(len(value))
+    return max(counts) if counts else None
+
+
+def _peak_by_count(rows, list_entries, variables):
+    """This workflow's history grouped by list length, each count's peak
+    reduced to its median - the input a fixed+marginal fit reads."""
+    by_count = {}
+    for row in rows:
+        peak = row.get("host_memory_job_peak_rss_mb")
+        count = _row_count(row, list_entries, variables)
+        if isinstance(peak, (int, float)) and count:
+            by_count.setdefault(count, []).append(peak)
+    medians = {}
+    for count, peaks in by_count.items():
+        peaks.sort()
+        medians[count] = peaks[len(peaks) // 2]
+    return medians
+
+
+def _fit_peak_model(medians):
+    """A run's peak as fixed (one model load) plus marginal (per extra
+    entry), fit from this workflow's own history rather than assumed.
+
+    A single model load dominates a resident run's peak - dividing the
+    whole observed figure by the entry count and multiplying back out
+    (the old per-entry model) counted that load once per entry instead of
+    once per run, projecting a 12-entry batch at roughly five times its
+    measured peak (#348). With history at two or more distinct list
+    lengths, the base and the marginal are fit through the smallest and
+    largest observed counts - the same "two extremes" a least-squares fit
+    would reduce to with only two points, and cheaper than one with more.
+    With history at only one list length, there is nothing to fit a slope
+    from: the measured peak is projected flat, which is the conservative
+    reading of "the measured behavior" the run has actually shown, rather
+    than guessing a growth rate with a single data point.
+
+    Returns `(base_mb, slope_mb_per_entry, low_count, high_count)`.
+    """
+    counts = sorted(medians)
+    if len(counts) == 1:
+        only = counts[0]
+        return medians[only], 0.0, only, only
+    low, high = counts[0], counts[-1]
+    slope = (medians[high] - medians[low]) / (high - low)
+    slope = max(slope, 0.0)
+    base = medians[low] - slope * low
+    return base, slope, low, high
+
+
+def _already_survived(rows, list_entries, requested, projected_mb, variables):
+    """Whether a run at least as large as this request already finished on
+    this box at or above the projected peak.
+
+    A projection built from this workflow's own history can land above the
+    ceiling for a shape that is itself part of that history - the run being
+    projected already happened and succeeded (#254). `rows` only ever holds
+    successful runs (`JobHistory.finished_runs()`), so a matching row is a
+    completed data point, not a guess: warning about a peak already reached
+    without incident tells a caller nothing they don't already know from ten
+    minutes ago.
+    """
+    for row in rows:
+        peak = row.get("host_memory_job_peak_rss_mb")
+        if not isinstance(peak, (int, float)) or peak < projected_mb:
+            continue
+        count = _row_count(row, list_entries, variables)
+        if count is not None and count >= requested:
+            return True
+    return False
+
+
+def host_memory_warnings(definition, list_entries, rows, ceiling_mb):
+    """Warnings for a projected host-memory peak this machine cannot hold,
+    or [] when there is nothing to project from or nothing to warn about.
+
+    `rows` are this workflow's finished runs as `JobHistory.finished_runs()`
+    groups them - the same history `observed_cost.py` reads, extended with
+    `host_memory_job_peak_rss_mb` per row (#243, repointed from the
+    process-lifetime `host_memory_peak_rss_mb` by #272 - a small job run
+    right after a heavy one no longer inherits the heavy job's peak).
+    `ceiling_mb` is this box's own RAM, scaled by `CEILING_FRACTION`; the
+    caller reads that once per request rather than this module importing
+    `host_memory` for a per-validate syscall.
+    """
+    requested = _requested_count(list_entries)
+    if requested is None or not rows or not ceiling_mb:
+        return []
+    if not _has_seedable_step(definition):
+        # A task-only workflow loads no model to accumulate across entries -
+        # the failure mode this module projects for doesn't apply to it
+        return []
+    variables = definition.get("variables") or {}
+    peaks = [row["host_memory_job_peak_rss_mb"] for row in rows]
+    peaks = [value for value in peaks if isinstance(value, (int, float))]
+    if not peaks:
+        # Cold start: history exists for this workflow, but no run of it
+        # ever reported a host-memory reading - nothing to project from
+        return []
+    if releases_between_iterations(definition):
+        projected_mb = max(peaks)
+        shape = f"~{round(projected_mb)} MB, the largest single iteration observed"
+    else:
+        medians = _peak_by_count(rows, list_entries, variables)
+        if not medians:
+            return []
+        base, slope, low_count, high_count = _fit_peak_model(medians)
+        projected_mb = base + slope * requested
+        if low_count == high_count:
+            shape = (
+                f"~{round(projected_mb)} MB at {requested} entries, based on "
+                f"this workflow's own {low_count}-entry runs with no larger "
+                "history yet to project growth from"
+            )
+        else:
+            shape = (
+                f"~{round(projected_mb)} MB at {requested} entries, "
+                f"extrapolated from runs of {low_count} and {high_count} entries"
+            )
+    if projected_mb <= ceiling_mb:
+        return []
+    if _already_survived(rows, list_entries, requested, projected_mb, variables):
+        return []
+    return [
+        f"Projected host memory for this run ({shape}) exceeds this "
+        f"machine's usable RAM (~{round(ceiling_mb)} MB) - based on "
+        f"{len(rows)} run(s) of this workflow's own history on this "
+        "machine, not a curated figure. The run is not blocked, but it "
+        "may be killed by the OS partway through."
+    ]

dw/hub_cache.py

@@ -0,0 +1,432 @@
+"""Inventory and deletion for the Hugging Face hub cache.
+
+Every from_pretrained download lands in the hub cache and stays there until
+something deletes it - a few video models is a few hundred GB. This wraps
+huggingface_hub's own cache scanner so the server can show what is on disk
+and free it, without inventing any path handling of its own: deletion goes
+through scan_cache_dir's delete_revisions strategy, which only ever removes
+revisions it found inside the cache directory.
+"""
+
+import json
+import logging
+import os
+import shutil
+import threading
+import time
+import uuid
+
+from huggingface_hub import constants, scan_cache_dir
+from huggingface_hub.utils import CacheNotFound
+
+from .security import SecurityError, validate_path
+
+try:
+    # Xet-backed downloads aggregate into two bars built from our tracker
+    # class: reconstruction (file bytes written) and transfer (network
+    # bytes). Only reconstruction matches the file-size total the manager
+    # reports against - counting both would double the progress - and the
+    # transfer bar is recognizable by the bar_format it is created with
+    from huggingface_hub.utils._xet_progress_reporting import XET_TRANSFER_BAR_FORMAT
+except ImportError:  # older huggingface_hub - no xet aggregation
+    XET_TRANSFER_BAR_FORMAT = None
+
+logger = logging.getLogger("dw")
+
+
+def _resolved_cache_dir(cache_dir):
+    return str(cache_dir) if cache_dir else constants.HF_HUB_CACHE
+
+
+def scan_models(cache_dir=None):
+    """The cache's contents as plain data: repos sorted largest-first,
+    with per-revision detail, plus disk totals for the volume it lives on."""
+    resolved = _resolved_cache_dir(cache_dir)
+    try:
+        scan = scan_cache_dir(resolved)
+    except CacheNotFound:
+        # No cache directory yet - nothing downloaded is a state, not an error
+        return {
+            "cache_dir": resolved,
+            "size_on_disk": 0,
+            "repos": [],
+            "warnings": [],
+            "disk_free": None,
+            "disk_total": None,
+        }
+
+    repos = []
+    for repo in scan.repos:
+        revisions = sorted(
+            repo.revisions, key=lambda rev: rev.last_modified or 0, reverse=True
+        )
+        repos.append(
+            {
+                "repo_id": repo.repo_id,
+                "repo_type": repo.repo_type,
+                "size_on_disk": repo.size_on_disk,
+                "nb_files": repo.nb_files,
+                "last_accessed": repo.last_accessed,
+                "last_modified": repo.last_modified,
+                "revisions": [
+                    {
+                        "commit_hash": revision.commit_hash,
+                        "size_on_disk": revision.size_on_disk,
+                        "refs": sorted(revision.refs),
+                        "last_modified": revision.last_modified,
+                    }
+                    for revision in revisions
+                ],
+            }
+        )
+    repos.sort(key=lambda entry: entry["size_on_disk"], reverse=True)
+
+    usage = shutil.disk_usage(resolved)
+    return {
+        "cache_dir": resolved,
+        "size_on_disk": scan.size_on_disk,
+        "repos": repos,
+        "warnings": [str(warning) for warning in scan.warnings],
+        "disk_free": usage.free,
+        "disk_total": usage.total,
+    }
+
+
+_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".msgpack", ".onnx", ".pt")
+
+
+def _repo_folder_name(repo_id, repo_type="model"):
+    return f"{repo_type}s--{repo_id.replace('/', '--')}"
+
+
+def _snapshot_dir(repo_dir):
+    """The snapshot directory `refs/main` currently points at, or None -
+    a repo dir with no readable ref or no matching snapshot folder is
+    itself a sign the pull never finished."""
+    try:
+        with open(os.path.join(repo_dir, "refs", "main")) as f:
+            commit = f.read().strip()
+    except OSError:
+        return None
+    snapshot = os.path.join(repo_dir, "snapshots", commit)
+    return snapshot if os.path.isdir(snapshot) else None
+
+
+def _component_incomplete(folder, variant):
+    """Whether a diffusers pipeline component's folder is missing files a
+    load would need. A component with no weight files at all (a scheduler
+    or tokenizer config) is complete once its folder exists; one with
+    weight files is checked against the requested `variant` only when that
+    variant is actually in play, since an unvarianted component (most
+    schedulers, safety checkers) never carries a tagged file."""
+    if not os.path.isdir(folder):
+        return True
+    try:
+        entries = os.listdir(folder)
+    except OSError:
+        return True
+    if not entries:
+        return True
+    if not variant:
+        return False
+    weight_files = [e for e in entries if e.endswith(_WEIGHT_SUFFIXES)]
+    if not weight_files:
+        return False
+    return not any(f".{variant}." in name for name in weight_files)
+
+
+def repo_download_incomplete(repo_id, cache_dir=None, variant=None):
+    """Whether repo_id, though listed by `scan_models`, is left over from an
+    interrupted pull rather than fully fetched (#382): a cancelled
+    `snapshot_download` leaves the revision folder in place - the repo still
+    shows up in `scan_cache_dir` - but with an in-progress blob or with whole
+    components never fetched.
+
+    Two checks, in order of how cheap they are to rule out: any
+    `*.incomplete` blob (huggingface_hub's own marker for a file mid-transfer)
+    anywhere in the repo's `blobs/` makes the repo incomplete outright. Past
+    that, a diffusers pipeline repo (one with a `model_index.json` in its
+    current snapshot) is checked component by component, for the requested
+    `variant`; a plain checkpoint/LoRA repo (no `model_index.json`) has
+    nothing further to check beyond its blobs.  This is not a full hub
+    file-list diff - it does not fetch the repo's file list from the hub, so
+    a file that was never attempted at all on an otherwise-untouched repo
+    (nothing downloaded, no blobs, no snapshot) is caught by the "no snapshot"
+    case, not enumerated.
+    """
+    resolved = _resolved_cache_dir(cache_dir)
+    try:
+        # _is_repo_id has already held repo_id to Hub's one-segment shape;
+        # this keeps the cache folder inside the cache on its own terms
+        repo_dir = validate_path(
+            os.path.join(resolved, _repo_folder_name(repo_id)), resolved
+        )
+    except SecurityError:
+        return True
+    if not os.path.isdir(repo_dir):
+        return True
+
+    blobs_dir = os.path.join(repo_dir, "blobs")
+    if os.path.isdir(blobs_dir):
+        try:
+            if any(name.endswith(".incomplete") for name in os.listdir(blobs_dir)):
+                return True
+        except OSError:
+            return True
+
+    snapshot = _snapshot_dir(repo_dir)
+    if snapshot is None:
+        return True
+
+    index_path = os.path.join(snapshot, "model_index.json")
+    if not os.path.isfile(index_path):
+        # Not a diffusers pipeline repo - blob completeness is all there is
+        return False
+    try:
+        with open(index_path) as f:
+            index = json.load(f)
+    except (OSError, ValueError):
+        return False
+
+    for key, value in index.items():
+        if key.startswith("_") or not isinstance(value, list):
+            continue
+        # A component is a folder beside model_index.json. The file comes
+        # from whoever published the repo, so an absolute or '..' key would
+        # point this listdir anywhere and make validate_workflow's
+        # downloads_required a yes/no oracle for directories on the box
+        if os.path.basename(key) != key or key in (".", ".."):
+            continue
+        if _component_incomplete(os.path.join(snapshot, key), variant):
+            return True
+    return False
+
+
+def delete_model(repo_id, cache_dir=None):
+    """Delete every cached revision of repo_id. Returns the bytes freed.
+
+    Raises ValueError when the repo is not in the cache - the caller typed
+    or raced something, and nothing was deleted.
+    """
+    scan = scan_cache_dir(_resolved_cache_dir(cache_dir))
+    repo = next((r for r in scan.repos if r.repo_id == repo_id), None)
+    if repo is None:
+        raise ValueError(f"'{repo_id}' is not in the hub cache")
+
+    strategy = scan.delete_revisions(
+        *[revision.commit_hash for revision in repo.revisions]
+    )
+    freed = strategy.expected_freed_size
+    strategy.execute()
+    return freed
+
+
+# ------------------------------------------------------------------ downloads
+
+
+class DownloadCancelled(Exception):
+    """Raised inside a download's progress callback to abort it."""
+
+
+class DownloadManager:
+    """Background snapshot downloads into the hub cache, with progress.
+
+    One thread per download; progress is fed by a tqdm-compatible tracker
+    that snapshot_download instantiates per file, so the counters aggregate
+    across the file pool. Cancellation raises out of the next progress tick;
+    huggingface_hub's partial files remain resumable, so a cancelled or
+    failed download picks up where it stopped when retried.
+    """
+
+    KEEP_FINISHED = 20
+
+    def __init__(self, download_fn=None, info_fn=None):
+        # Injectable for tests - the defaults reach the network
+        if download_fn is None or info_fn is None:
+            from huggingface_hub import HfApi, snapshot_download
+
+            download_fn = download_fn or snapshot_download
+            info_fn = info_fn or (
+                lambda repo_id: HfApi().repo_info(repo_id, files_metadata=True)
+            )
+        self._download_fn = download_fn
+        self._info_fn = info_fn
+        self._lock = threading.Lock()
+        self._downloads = {}
+
+    def start(self, repo_id):
+        """Begin downloading repo_id; returns the download's status dict.
+        Raises ValueError for an invalid repo id or one already in flight."""
+        from huggingface_hub.utils import HFValidationError, validate_repo_id
+
+        try:
+            validate_repo_id(repo_id)
+        except HFValidationError as e:
+            raise ValueError(str(e))
+
+        cancel_event = threading.Event()
+        with self._lock:
+            for entry in self._downloads.values():
+                if entry["repo_id"] == repo_id and entry["status"] == "downloading":
+                    raise ValueError(f"'{repo_id}' is already downloading")
+            entry = {
+                "id": uuid.uuid4().hex[:12],
+                "repo_id": repo_id,
+                "status": "downloading",
+                "downloaded": 0,
+                "total": None,
+                "error": None,
+                "started_at": time.time(),
+                "finished_at": None,
+                # Set before the entry is published: cancel() reaches for this
+                # the moment the id is visible, and assigning it after the
+                # lock left a window where cancelling raised KeyError
+                "_cancel": cancel_event,
+            }
+            self._downloads[entry["id"]] = entry
+            self._prune()
+
+        thread = threading.Thread(
+            target=self._run, args=(entry, cancel_event), daemon=True
+        )
+        thread.start()
+        return self.status(entry["id"])
+
+    def _run(self, entry, cancel_event):
+        try:
+            try:
+                info = self._info_fn(entry["repo_id"])
+                total = sum(sibling.size or 0 for sibling in (info.siblings or []))
+                with self._lock:
+                    entry["total"] = total or None
+            except Exception as e:
+                # Size is cosmetic; the download itself decides success
+                logger.debug(f"No size metadata for {entry['repo_id']}: {e}")
+
+            self._download_fn(
+                entry["repo_id"], tqdm_class=_tracker_class(self, entry, cancel_event)
+            )
+            self._finish(entry, "completed")
+        except DownloadCancelled:
+            self._finish(entry, "cancelled")
+        except Exception as e:
+            self._finish(entry, "failed", str(e))
+
+    def _finish(self, entry, status, error=None):
+        with self._lock:
+            entry["status"] = status
+            entry["error"] = error
+            entry["finished_at"] = time.time()
+
+    def _add_progress(self, entry, n):
+        with self._lock:
+            entry["downloaded"] += n
+
+    def cancel(self, download_id):
+        """Request cancellation; returns the status dict or None if unknown.
+        Takes effect at the download's next progress tick."""
+        with self._lock:
+            entry = self._downloads.get(download_id)
+            if entry is None:
+                return None
+            if entry["status"] == "downloading":
+                entry["_cancel"].set()
+        return self.status(download_id)
+
+    def status(self, download_id):
+        with self._lock:
+            entry = self._downloads.get(download_id)
+            return _public(entry) if entry else None
+
+    def status_list(self):
+        """Every tracked download, newest first."""
+        with self._lock:
+            entries = sorted(
+                self._downloads.values(),
+                key=lambda e: e["started_at"],
+                reverse=True,
+            )
+            return [_public(entry) for entry in entries]
+
+    def is_active(self):
+        with self._lock:
+            return any(
+                entry["status"] == "downloading" for entry in self._downloads.values()
+            )
+
+    def _prune(self):
+        # Called with the lock held: drop the oldest finished entries
+        finished = sorted(
+            (e for e in self._downloads.values() if e["status"] != "downloading"),
+            key=lambda e: e["started_at"],
+        )
+        excess = len(finished) - self.KEEP_FINISHED
+        for entry in finished[: max(0, excess)]:
+            del self._downloads[entry["id"]]
+
+
+def _public(entry):
+    return {key: value for key, value in entry.items() if not key.startswith("_")}
+
+
+def _tracker_class(manager, entry, cancel_event):
+    """A tqdm stand-in snapshot_download instantiates per file; every update
+    feeds the shared counters and honours cancellation."""
+
+    class Tracker:
+        def __init__(self, *args, **kwargs):
+            self.n = 0
+            self.total = kwargs.get("total")
+            # The xet transfer bar still ticks (cancellation works through
+            # it) but its network bytes stay out of the shared counters
+            self._counted = (
+                XET_TRANSFER_BAR_FORMAT is None
+                or kwargs.get("bar_format") != XET_TRANSFER_BAR_FORMAT
+            )
+
+        def update(self, n=1):
+            if cancel_event.is_set():
+                raise DownloadCancelled()
+            if n:
+                self.n += n
+                if self._counted:
+                    manager._add_progress(entry, n)
+            return True
+
+        def close(self):
+            pass
+
+        def refresh(self):
+            pass
+
+        def set_description(self, *args, **kwargs):
+            pass
+
+        def set_postfix(self, *args, **kwargs):
+            pass
+
+        def set_postfix_str(self, *args, **kwargs):
+            pass
+
+        @property
+        def format_dict(self):
+            # What tqdm exposes for rate rendering; huggingface_hub reads
+            # 'rate' from it when composing the xet speed postfix
+            return {"n": self.n, "total": self.total, "elapsed": 0, "rate": None}
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            self.close()
+            return False
+
+        @staticmethod
+        def get_lock():
+            return threading.RLock()
+
+        @staticmethod
+        def set_lock(lock):
+            pass
+
+    return Tracker