diffusers-workflow 0.4.0__py3-none-any.whl

dw/for_each.py

@@ -0,0 +1,409 @@
+"""Expand a step's 'for_each' list into one ordinary step per entry.
+
+A template that generates one shot per entry of a list is otherwise written
+the long way - 'shot_1', 'shot_2', ... each a near-copy of the one before -
+and a six-shot episode is a different file from a five-shot one. This pass
+runs on the definition after variable substitution and before the run id is
+computed (Workflow.run) and before the reference check (validation_errors),
+and produces a definition with no 'for_each' in it: every member is an
+ordinary step, so the step loop, the step cache, the manifest and the
+reference checker never learn a new reference kind.
+
+Inside a member:
+  - 'item:'       is the whole entry; 'item:field' one field of an object
+                  entry, spliced in whole whatever its type
+  - a reference to another for_each group over the SAME list resolves to
+    the member with the same key ('slice' inside 'shot@open' -> 'slice@open')
+Outside (or inside, for any other group):
+  - 'gather:shot' is the list of every member's result, as explicit
+    'previous_result:shot@<key>' strings; inside a list it splices
+  - 'previous_result:shot' naming a group is an error that says to gather
+
+Members are named '<group>@<key>' - the entry's own 'name' when it carries
+one, else its index - and '@' is reserved in every step name. Names rather
+than indexes because the step cache (dw/step_cache.py) keys on the step
+name: inserting a shot in the middle of a list must not shift every later
+member onto a different entry's cache line.
+"""
+
+import copy
+
+from .arguments import FROM_PREVIOUS_RESULT_KEY, PREVIOUS_RESULT_PREFIX
+from .security import InvalidInputError, validate_variable_name
+from .step_cache import reference_resolves_to
+from .variables import argument_errors, set_variables
+
+FOR_EACH_KEY = "for_each"
+ITEM_PREFIX = "item:"
+GATHER_PREFIX = "gather:"
+MEMBER_SEPARATOR = "@"
+# Each entry is a full generation. Stated against the step cache's bound
+# (DEFAULT_MAX_ENTRIES = 128): a run whose expanded steps exceed the cache
+# evicts its own earlier members, so this is kept well under it
+MAX_FOR_EACH_ENTRIES = 32
+# release_pipeline / release_models would drop the model after the first
+# member and reload it for the second, so they are carried onto the last one
+_LAST_MEMBER_ONLY = ("release_pipeline", "release_models")
+
+
+class ForEachError(ValueError):
+    """A for_each step that cannot be expanded, with the JSON path at fault."""
+
+    def __init__(self, path, message):
+        super().__init__(message)
+        self.path = path
+
+
+def member_name(group, key):
+    return f"{group}{MEMBER_SEPARATOR}{key}"
+
+
+def expand_for_each(definition, source_indices=None):
+    """The definition with every 'for_each' step replaced by its members.
+
+    Returns a new structure; `definition` is left as it was passed in.
+    Raises ForEachError for anything that cannot be expanded.
+
+    `source_indices`, when a list is passed, has the index of the step each
+    expanded step was written as appended to it, so a later check can report
+    an error at a path in the file the author wrote rather than at an
+    expanded index that exists nowhere. A parallel list rather than a key on
+    the step: step_data is what the step cache keys on and what the schema
+    validates, and neither may learn a new field.
+    """
+    steps = definition.get("steps") if isinstance(definition, dict) else None
+    if not isinstance(steps, list):
+        return definition
+
+    # group name -> {"keys": [...], "entries": [...]} for every group
+    # expanded so far, in step order, so a reference can only reach an
+    # earlier group - the same rule previous_result: has always had
+    groups = {}
+    expanded = []
+    for index, step in enumerate(steps):
+        if not isinstance(step, dict):
+            expanded.append(copy.deepcopy(step))
+            _record(source_indices, index)
+            continue
+        path = ("steps", index)
+        name = step.get("name")
+        if isinstance(name, str) and MEMBER_SEPARATOR in name:
+            raise ForEachError(
+                render_path(path + ("name",)),
+                f"Step name '{name}' contains '{MEMBER_SEPARATOR}', which is "
+                f"reserved for the members of a for_each step",
+            )
+        if FOR_EACH_KEY not in step:
+            expanded.append(_rewrite(step, path, groups, member=None))
+            _record(source_indices, index)
+            continue
+
+        entries = step[FOR_EACH_KEY]
+        keys = _entry_keys(entries, path + (FOR_EACH_KEY,))
+        template = {k: v for k, v in step.items() if k != FOR_EACH_KEY}
+        last = len(keys) - 1
+        for position, (key, entry) in enumerate(zip(keys, entries)):
+            member = {
+                "group": name,
+                "key": key,
+                "entry": entry,
+                "entries": entries,
+                "index": position,
+            }
+            expanded_step = _rewrite(template, path, groups, member)
+            expanded_step["name"] = member_name(name, key)
+            if position != last:
+                for flag in _LAST_MEMBER_ONLY:
+                    expanded_step.pop(flag, None)
+            expanded.append(expanded_step)
+            _record(source_indices, index)
+        groups[name] = {"keys": keys, "entries": entries}
+
+    result = {k: v for k, v in definition.items() if k != "steps"}
+    result["steps"] = expanded
+    return result
+
+
+def _record(source_indices, index):
+    if source_indices is not None:
+        source_indices.append(index)
+
+
+def _entry_keys(entries, path):
+    """The key of every entry - its 'name' when it is an object carrying
+    one, else its index - validated and unique."""
+    if not isinstance(entries, list):
+        if isinstance(entries, str) and entries.startswith("variable:"):
+            hint = f" - '{entries}' was not substituted; is the variable declared?"
+        else:
+            hint = ""
+        raise ForEachError(
+            render_path(path),
+            f"for_each must be a list, got {type(entries).__name__}{hint}",
+        )
+    if not entries:
+        raise ForEachError(
+            render_path(path),
+            "for_each over an empty list would run no steps - a workflow that "
+            "generates nothing is never what was asked for",
+        )
+    if len(entries) > MAX_FOR_EACH_ENTRIES:
+        raise ForEachError(
+            render_path(path),
+            f"for_each has {len(entries)} entries; the limit is {MAX_FOR_EACH_ENTRIES}",
+        )
+    keys = []
+    for index, entry in enumerate(entries):
+        key = str(index)
+        if isinstance(entry, dict) and "name" in entry:
+            key = entry["name"]
+            key_path = render_path(path + (index, "name"))
+            if not isinstance(key, str):
+                raise ForEachError(key_path, "An entry's name must be a string")
+            try:
+                validate_variable_name(key)
+            except InvalidInputError as e:
+                raise ForEachError(key_path, f"Invalid entry name '{key}': {e}") from e
+            if key in keys:
+                raise ForEachError(key_path, f"Duplicate entry name '{key}'")
+        keys.append(key)
+    return keys
+
+
+def _rewrite(value, path, groups, member):
+    """Rebuild `value` with item:, gather: and group references resolved.
+
+    `member` is None outside a for_each step; inside one it carries the
+    group, key and entry of the member being built.
+    """
+    if isinstance(value, dict):
+        rebuilt = {}
+        for key, item in value.items():
+            if key == FROM_PREVIOUS_RESULT_KEY and isinstance(item, str):
+                rebuilt[key] = _rewrite_reference(item, path + (key,), groups, member)
+            else:
+                rebuilt[key] = _rewrite(item, path + (key,), groups, member)
+        return rebuilt
+    if isinstance(value, list):
+        rebuilt = []
+        for index, item in enumerate(value):
+            if isinstance(item, str) and item.startswith(GATHER_PREFIX):
+                # A gather inside a list splices into it
+                rebuilt.extend(_gather(item, path + (index,), groups))
+            else:
+                rebuilt.append(_rewrite(item, path + (index,), groups, member))
+        return rebuilt
+    if isinstance(value, str):
+        if value.startswith(GATHER_PREFIX):
+            return _gather(value, path, groups)
+        if value.startswith(ITEM_PREFIX):
+            return _item(value, path, member)
+        if value.startswith(PREVIOUS_RESULT_PREFIX):
+            reference = value[len(PREVIOUS_RESULT_PREFIX) :]
+            return PREVIOUS_RESULT_PREFIX + _rewrite_reference(
+                reference, path, groups, member
+            )
+        return value
+    # A leaf is only copied where the copy is needed: inside a member, where
+    # the same template value is about to appear in every one of them.
+    # Outside, the leaf is handed back as it is - this pass runs on every run
+    # of every workflow, after realize_args has turned 'asset:' arguments
+    # into loaded images and decoded frame lists, and copying all of that
+    # would multiply the media a run holds. The 'input untouched' contract
+    # still holds because nothing here ever mutates a leaf
+    return _copy_leaf(value) if member is not None else value
+
+
+def _copy_leaf(value):
+    """A copy of a leaf, or the leaf itself when it cannot be copied.
+
+    An open handle or a live model object reaching a member is not a reason
+    to fail a run - the step cache makes the same choice for a realized
+    argument it cannot deep-copy (dw/workflow.py).
+    """
+    try:
+        return copy.deepcopy(value)
+    except Exception:
+        return value
+
+
+def _item(value, path, member):
+    if member is None:
+        raise ForEachError(
+            render_path(path), f"'{value}' is only meaningful inside a for_each step"
+        )
+    field = value[len(ITEM_PREFIX) :]
+    entry = member["entry"]
+    if field == "":
+        return _copy_leaf(entry)
+    if not isinstance(entry, dict):
+        raise ForEachError(
+            render_path(path),
+            f"'{value}' asks for a field of entry '{member['key']}' of "
+            f"for_each step '{member['group']}', which is not an object",
+        )
+    if field not in entry:
+        raise ForEachError(
+            render_path(path),
+            f"'{value}' names no field of entry '{member['key']}' of for_each "
+            f"step '{member['group']}'; it has: {sorted(entry)}",
+        )
+    return _copy_leaf(entry[field])
+
+
+def _gather(value, path, groups):
+    group = value[len(GATHER_PREFIX) :]
+    if group not in groups:
+        raise ForEachError(
+            render_path(path),
+            f"'{value}' names no earlier for_each step. "
+            f"for_each steps available here: {sorted(groups)}",
+        )
+    return [
+        PREVIOUS_RESULT_PREFIX + member_name(group, key)
+        for key in groups[group]["keys"]
+    ]
+
+
+def _rewrite_reference(reference, path, groups, member):
+    """A previous_result reference (without its prefix) as the expanded
+    definition spells it: unchanged unless it names a for_each group."""
+    if reference.startswith("variable:"):
+        return reference
+    group = next((g for g in groups if reference_resolves_to(reference, g)), None)
+    if group is None:
+        if member is not None and reference_resolves_to(reference, member["group"]):
+            raise ForEachError(
+                render_path(path),
+                f"'{reference}' names its own for_each step '{member['group']}'",
+            )
+        return reference
+    if member is not None and member["entries"] == groups[group]["entries"]:
+        # Same list: shot@open reads slice@open
+        return member_name(group, member["key"]) + reference[len(group) :]
+    where = (
+        f"from inside for_each step '{member['group']}', which runs over a different list"
+        if member is not None
+        else "from outside a for_each step"
+    )
+    raise ForEachError(
+        render_path(path),
+        f"'{reference}' names the for_each step '{group}' {where}. Use "
+        f"'{GATHER_PREFIX}{group}' for every member's result, or a reference "
+        f"from a for_each step over the same list for the same-keyed member",
+    )
+
+
+def list_fields(definition):
+    """What an entry of each list-driven variable has to carry, read off
+    the definition: for every step whose for_each is 'variable:<name>',
+    the fields its 'item:<field>' references name.
+
+    Returns {<variable>: {"fields": [...] or None, "steps": [...]}} -
+    fields sorted with 'name' first, or None when a step splices the
+    whole entry with a bare 'item:' (the entries are values, not
+    objects). A literal for_each list is not an argument and is skipped.
+    Reads the raw definition, no substitution, so the catalog and the
+    validator derive the same answer from the file as written.
+    """
+    steps = definition.get("steps") if isinstance(definition, dict) else None
+    if not isinstance(steps, list):
+        return {}
+    found = {}
+    for step in steps:
+        if not isinstance(step, dict):
+            continue
+        target = step.get(FOR_EACH_KEY)
+        if not (isinstance(target, str) and target.startswith("variable:")):
+            continue
+        variable = target.removeprefix("variable:")
+        entry = found.setdefault(variable, {"fields": set(), "steps": []})
+        entry["steps"].append(step.get("name"))
+        for value in _strings(step):
+            if not value.startswith(ITEM_PREFIX):
+                continue
+            field = value[len(ITEM_PREFIX) :]
+            if field == "":
+                entry["fields"] = None
+            elif entry["fields"] is not None:
+                entry["fields"].add(field)
+    return {
+        variable: {
+            "fields": (
+                None
+                if entry["fields"] is None
+                else ["name"] + sorted(entry["fields"] - {"name"})
+            ),
+            "steps": entry["steps"],
+        }
+        for variable, entry in found.items()
+    }
+
+
+def entry_field_warnings(definition, arguments=None):
+    """Every entry key of a list-driven variable that no step reads.
+
+    A caller who writes 'num_frame' for 'num_frames' gets the template's
+    value for the field they meant to set, in silence; this names the
+    key, at the entry it sits in, with the fields the list takes. A
+    warning rather than an error: an entry may carry a note on purpose.
+    Good `arguments` are folded in first, and a list the caller supplied
+    is reported under 'arguments.', where they wrote it.
+    """
+    if not isinstance(definition, dict):
+        return []
+    variables = definition.get("variables")
+    if not isinstance(variables, dict):
+        return []
+    variables = copy.deepcopy(variables)
+    supplied = set()
+    if arguments and not argument_errors(definition, arguments):
+        set_variables(arguments, variables)
+        supplied = set(arguments)
+    warnings = []
+    for variable, spec in list_fields(definition).items():
+        fields = spec["fields"]
+        entries = variables.get(variable)
+        if fields is None or not isinstance(entries, list):
+            continue
+        where = "arguments" if variable in supplied else "variables"
+        for index, entry in enumerate(entries):
+            if not isinstance(entry, dict):
+                continue
+            unknown = sorted(set(entry) - set(fields))
+            if not unknown:
+                continue
+            label = repr(entry["name"]) if isinstance(entry.get("name"), str) else index
+            warnings.append(
+                f"{where}.{variable}[{index}]: entry {label} carries "
+                f"{', '.join(repr(k) for k in unknown)}, which no step reads; "
+                f"entries of '{variable}' take: {', '.join(fields)}"
+            )
+    return warnings
+
+
+def _strings(value):
+    """Every string anywhere inside a JSON value, except the for_each key
+    itself."""
+    if isinstance(value, str):
+        yield value
+    elif isinstance(value, list):
+        for item in value:
+            yield from _strings(item)
+    elif isinstance(value, dict):
+        for key, item in value.items():
+            if key != FOR_EACH_KEY:
+                yield from _strings(item)
+
+
+def render_path(path):
+    """'steps[3].task.arguments.videos[1]' - the same shape schema errors use."""
+    rendered = ""
+    for part in path:
+        if isinstance(part, int):
+            rendered += f"[{part}]"
+        elif rendered:
+            rendered += f".{part}"
+        else:
+            rendered = str(part)
+    return rendered

dw/host_memory.py

@@ -0,0 +1,258 @@
+"""What the worker process and its machine are using in host RAM.
+
+`device_memory_stats` answers for the accelerator, which on this project is
+often the least informative half of the question: the templates here keep
+weights in host memory by design (`offload: "sequential"`, `group_offload`),
+so a card can sit near-empty through a generation whose weights are very
+much resident somewhere. This module is the other half.
+
+Three methods, in order, none of them required: psutil when it happens to be
+installed (a transitive dependency here, not a declared one), Linux's
+`/proc` otherwise, and `resource.getrusage` for the process's peak, which is
+available on every POSIX platform. Nothing raises - a reading that cannot be
+taken is reported as None rather than failing the call that asked for it.
+"""
+
+import logging
+import os
+
+logger = logging.getLogger("dw")
+
+__all__ = [
+    "host_memory_stats",
+    "host_memory_fields",
+    "trim_host_memory",
+    "release_host_caches",
+    "pinned_host_memory_fields",
+]
+
+_MB = 1024.0 * 1024.0
+
+
+def host_memory_stats():
+    """Snapshot of host memory for this process and the machine it is on.
+
+    Returns:
+        dict with keys, any of which may be None when the platform cannot
+        answer:
+            rss_mb (float or None): resident set size of *this* process -
+                in the worker, the weights it is holding
+            peak_rss_mb (float or None): the high-water mark of the above,
+                which is what says whether a run that has finished released
+                what it took
+            total_mb (float or None): the machine's physical memory
+            available_mb (float or None): what can be handed out without
+                swapping - MemAvailable on Linux, not free memory, since
+                page cache is reclaimable
+    """
+    stats = {
+        "rss_mb": None,
+        "peak_rss_mb": _peak_rss_mb(),
+        "total_mb": None,
+        "available_mb": None,
+    }
+
+    for method in (_psutil_stats, _proc_stats):
+        try:
+            reading = method()
+        except Exception as e:  # a memory reading is never worth an exception
+            logger.debug(f"Host memory via {method.__name__} failed: {e}")
+            continue
+        for key, value in reading.items():
+            if stats.get(key) is None and value is not None:
+                stats[key] = value
+        if all(stats[key] is not None for key in ("rss_mb", "total_mb")):
+            break
+
+    return _hold_the_high_water_mark(stats)
+
+
+def _hold_the_high_water_mark(stats):
+    """Keep peak_rss_mb >= rss_mb, which is what a high-water mark means.
+
+    The two readings come from different places - getrusage's ru_maxrss,
+    quantized to whole pages and taken first, against psutil's rss taken a
+    moment later - so a process that has never peaked meaningfully above its
+    current size reports them within a megabyte of each other in either
+    order. `peak - rss` is the whole point of the pair (what a run took and
+    did not give back), and a small negative there reads as "these fields
+    are not comparable" rather than "nothing leaked" (#83).
+    """
+    peak, rss = stats["peak_rss_mb"], stats["rss_mb"]
+    if peak is not None and rss is not None and peak < rss:
+        stats["peak_rss_mb"] = rss
+    return stats
+
+
+def _psutil_stats():
+    import psutil
+
+    virtual = psutil.virtual_memory()
+    return {
+        "rss_mb": psutil.Process().memory_info().rss / _MB,
+        "total_mb": virtual.total / _MB,
+        "available_mb": virtual.available / _MB,
+    }
+
+
+def _proc_stats():
+    """Linux without psutil: /proc/self/statm for this process, /proc/meminfo
+    for the machine. Both absent elsewhere, which is what the None default
+    covers."""
+    reading = {"rss_mb": None, "total_mb": None, "available_mb": None}
+
+    try:
+        with open("/proc/self/statm", "r") as f:
+            pages = int(f.read().split()[1])
+        reading["rss_mb"] = pages * os.sysconf("SC_PAGE_SIZE") / _MB
+    except (OSError, ValueError, IndexError, AttributeError):
+        pass
+
+    wanted = {"MemTotal:": "total_mb", "MemAvailable:": "available_mb"}
+    try:
+        with open("/proc/meminfo", "r") as f:
+            for line in f:
+                parts = line.split()
+                if len(parts) >= 2 and parts[0] in wanted:
+                    # meminfo is in kB
+                    reading[wanted[parts[0]]] = int(parts[1]) / 1024.0
+    except (OSError, ValueError):
+        pass
+
+    return reading
+
+
+def _peak_rss_mb():
+    """getrusage's high-water mark, in kB on Linux and bytes on macOS - the
+    one place in this module where the unit depends on the platform."""
+    try:
+        import resource
+        import sys
+
+        peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    except Exception:
+        return None
+    if not peak:
+        return None
+    return peak / _MB if sys.platform == "darwin" else peak / 1024.0
+
+
+# The memory payload's own names for the above, beside its gpu_* keys
+FIELD_NAMES = {
+    "rss_mb": "host_memory_rss_mb",
+    "peak_rss_mb": "host_memory_peak_rss_mb",
+    "total_mb": "host_memory_total_mb",
+    "available_mb": "host_memory_available_mb",
+}
+
+
+def host_memory_fields():
+    """host_memory_stats under the names the memory payload reports, with
+    the readings this platform cannot take left out rather than sent as
+    null - a key that is absent says "not measurable here", where a null
+    would read as "measured, and nothing"."""
+    stats = host_memory_stats()
+    return {
+        FIELD_NAMES[key]: value for key, value in stats.items() if value is not None
+    }
+
+
+def trim_host_memory():
+    """Hand memory the process has already freed back to the operating
+    system, and report how much that was in MB (0.0 when the platform has no
+    way to ask).
+
+    Dropping the last reference to a model frees it inside the process, not
+    back to the kernel: glibc keeps the arenas the weights were read into and
+    hands them out again to *this* process. That is normally invisible and
+    correct - until a template needs 96% of host RAM to run at all, at which
+    point the several GB the previous job's arenas are sitting on is the
+    difference between a run and a SIGKILL five minutes in (#98). The
+    templates here load tens of GB of weights through host memory, so the
+    arenas in question are large ones and the fragmentation that keeps
+    `malloc_trim` from returning them is the exception rather than the rule.
+
+    Linux/glibc only: `malloc_trim` is a GNU extension. Everywhere else -
+    macOS included - this is a no-op that reports 0.0, because there is
+    nothing to ask and a fabricated number would be worse than none.
+    """
+    import ctypes
+    import ctypes.util
+    import sys
+
+    if not sys.platform.startswith("linux"):
+        return 0.0
+    before = host_memory_stats()["rss_mb"]
+    try:
+        libc = ctypes.CDLL(ctypes.util.find_library("c") or "libc.so.6", use_errno=True)
+        libc.malloc_trim(ctypes.c_size_t(0))
+    except (OSError, AttributeError) as e:
+        # musl and friends have no malloc_trim; not having one is not an error
+        logger.debug(f"malloc_trim unavailable: {e}")
+        return 0.0
+    after = host_memory_stats()["rss_mb"]
+    if before is None or after is None:
+        return 0.0
+    return max(0.0, before - after)
+
+
+def pinned_host_memory_fields():
+    """What CUDA's pinned-host allocator is holding, in MB, or {} where
+    there is none to report.
+
+    Group offloading with `use_stream` stages a component's weights through
+    *pinned* host memory, which torch caches per process exactly as it
+    caches device memory: freeing the tensors returns the blocks to that
+    cache, not to the OS, so they stay in this process's RSS and count
+    against the next job's host budget. It is invisible in every figure the
+    memory payload carried before - `rss_mb` includes it without saying so
+    and `gpu_memory_*` does not see it at all (#98).
+    """
+    try:
+        import torch
+
+        stats = torch.cuda.host_memory_stats()
+    except Exception:
+        return {}
+    fields = {}
+    for key, name in (
+        ("allocated_bytes.all.current", "host_pinned_allocated_mb"),
+        ("reserved_bytes.all.current", "host_pinned_reserved_mb"),
+    ):
+        value = stats.get(key)
+        if value is not None:
+            fields[name] = value / _MB
+    return fields
+
+
+def release_host_caches():
+    """Give back host memory this process is holding but no longer using,
+    and report what came back in MB.
+
+    Two caches, neither of which `gc.collect()` touches:
+
+    - torch's pinned-host allocator, where group offloading's staging
+      buffers live. `_host_emptyCache` frees the blocks nothing is using;
+      blocks a still-loaded pipeline is staging through are in use and are
+      not touched, so this is safe to call with models resident.
+    - glibc's heap arenas, via `malloc_trim`. Freeing a large block inside
+      the process does not hand its pages back to the kernel.
+
+    Together they are why a worker that has released every model still sat
+    on 14.5 GB, which is the difference between the next job running and
+    being OOM-killed five minutes in on a template that needs 96% of host
+    RAM (#98).
+    """
+    before = host_memory_stats()["rss_mb"]
+    try:
+        import torch
+
+        if hasattr(torch._C, "_host_emptyCache"):
+            torch._C._host_emptyCache()
+    except Exception as e:  # a cleanup is never worth failing the run for
+        logger.debug(f"Could not empty the pinned host cache: {e}")
+    trim_host_memory()
+    after = host_memory_stats()["rss_mb"]
+    if before is None or after is None:
+        return 0.0
+    return max(0.0, before - after)