diffusers-workflow 0.4.0__py3-none-any.whl

dw/server/mcp_mount.py

@@ -0,0 +1,95 @@
+"""Mount the MCP tool surface inside the HTTP server.
+
+`dw_mcp` is a pure HTTP client of dw.serve; mounting it here does not
+change that - the tools reach the REST API over the server's own bind
+address with the same token the API requires, a few milliseconds per
+call. This is the only module under dw/ that imports the mcp SDK, and it
+does so lazily: the package is an optional extra.
+"""
+
+import ipaddress
+
+
+def client_base_url(host, port):
+    """The URL the mounted tools use to reach this same server.
+
+    A loopback or wildcard bind is reached at 127.0.0.1; any other bind
+    address has to be used verbatim, since uvicorn is then listening on
+    that address alone and a loopback connection would be refused - which
+    is exactly the `--host 100.x.y.z` Tailscale setup docs/REMOTE.md
+    recommends. An IPv6 literal is bracketed for the URL's authority.
+    """
+    # imported here rather than at module scope: dw.server.app imports this
+    # module (lazily, inside create_app), so an import back at import time
+    # would be a cycle
+    from .app import LOOPBACK_HOSTS, WILDCARD_HOSTS
+
+    host = (host or "").lower()
+    if host in WILDCARD_HOSTS | LOOPBACK_HOSTS:
+        host = "127.0.0.1"
+    try:
+        if ipaddress.ip_address(host).version == 6:
+            host = f"[{host}]"
+    except ValueError:
+        # a hostname, not an IP literal - nothing to bracket
+        pass
+    return f"http://{host}:{port}"
+
+
+class _SingleRouteApp:
+    """Send every path the parent routed here to the sub-app's one route.
+
+    `dw.server.app` routes `/mcp` and anything under it to this wrapper,
+    so the SDK's Starlette app - which has a single route at `/` - sees
+    `/` whatever spelling the client used. Without it a bare `POST /mcp`
+    either 404s or takes a redirect to `/mcp/`, and `http://box:8765/mcp`
+    is the URL the docs tell people to configure.
+    """
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope.get("type") in ("http", "websocket"):
+            scope = dict(scope, path="/", raw_path=b"/")
+        await self.app(scope, receive, send)
+
+
+def build_mcp_app(*, host, port, token):
+    """The ASGI app to serve at /mcp, the MCPServer behind it, and the
+    HTTP client the tools use to reach this same server (client_base_url).
+
+    The SDK's Starlette app carries its own lifespan (the session
+    manager), which the parent app does not run for a sub-app it routes
+    to - `create_app`'s lifespan enters `server.session_manager.run()`
+    itself.
+    """
+    try:
+        from mcp.server.transport_security import TransportSecuritySettings
+    except ImportError:
+        raise SystemExit(
+            "--mcp needs the mcp extra: pip install 'diffusers-workflow[mcp]'"
+        )
+
+    from dw_mcp.client import DwClient
+    from dw_mcp.server import build_server
+
+    client = DwClient(base_url=client_base_url(host, port), token=token)
+    # The tools now run *on* the server rather than on the agent's machine,
+    # so a tool that writes a local file is writing on the GPU box. It has to
+    # know that to confine where it writes (#113)
+    client.mounted = True
+    server = build_server(client)
+    asgi = server.streamable_http_app(
+        # the SDK app routes at "/"; the parent routes /mcp here and
+        # _SingleRouteApp rewrites the path to "/" on the way in
+        streamable_http_path="/",
+        # one JSON reply per request, no sessions to strand on a restart
+        json_response=True,
+        stateless_http=True,
+        # dw.server.app's own Origin/Host middleware runs first and owns this
+        transport_security=TransportSecuritySettings(
+            enable_dns_rebinding_protection=False
+        ),
+    )
+    return _SingleRouteApp(asgi), server, client

dw/server/netinfo.py

@@ -0,0 +1,124 @@
+"""Which addresses this machine can be reached at.
+
+The Server page answers one question - "what URL do I give the other
+machine?" - and the only part of the answer the server itself knows is the
+list of addresses its interfaces carry. URL composition stays in the UI;
+this module only enumerates.
+
+Two enumeration methods, in order: psutil when it happens to be installed
+(it is a transitive dependency here, not a declared one, so it is used
+opportunistically and never required), and the stdlib otherwise. Neither
+raises: an address list is a convenience, and a machine whose interfaces
+cannot be read still has a working server, so failure yields [].
+"""
+
+import socket
+import ipaddress
+
+__all__ = ["local_addresses"]
+
+
+def _usable(address):
+    """A candidate address, normalized, or None if it is not worth showing.
+
+    Loopback is dropped (it names this machine only - the page's whole
+    point is the other machine) and so is link-local (169.254.x / fe80::,
+    which needs a scope id to be usable at all). Docker and veth addresses
+    survive on purpose: on a box that runs containers they are real
+    routes, and guessing which of a machine's networks the user meant is
+    not the server's job.
+    """
+    if not address:
+        return None
+    # getaddrinfo and psutil both hand back scoped IPv6 ('fe80::1%eth0')
+    address = address.split("%")[0].strip()
+    try:
+        parsed = ipaddress.ip_address(address)
+    except ValueError:
+        return None
+    if parsed.is_loopback or parsed.is_link_local or parsed.is_unspecified:
+        return None
+    return str(parsed), "IPv4" if parsed.version == 4 else "IPv6"
+
+
+def _psutil_addresses():
+    """(address, family, interface) triples from psutil, which is the only
+    method that can name the interface an address belongs to."""
+    import psutil
+
+    found = []
+    for interface, entries in psutil.net_if_addrs().items():
+        for entry in entries:
+            if entry.family not in (socket.AF_INET, socket.AF_INET6):
+                continue
+            usable = _usable(entry.address)
+            if usable is not None:
+                found.append((usable[0], usable[1], interface))
+    return found
+
+
+def _outbound_address():
+    """The address a packet to the outside world would leave from.
+
+    A UDP socket's connect() only sets the peer for later sends; it does
+    no handshake and puts nothing on the wire. It exists because
+    getaddrinfo(gethostname()) is silent on machines whose hostname
+    resolves to 127.0.1.1 - a Debian default - and that is exactly the
+    machine a remote-access page is being read on.
+    """
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        sock.connect(("192.0.2.1", 9))  # TEST-NET-1, routed nowhere
+        return sock.getsockname()[0]
+    finally:
+        sock.close()
+
+
+def _stdlib_addresses():
+    """The psutil-free fallback: whatever this host's name resolves to,
+    plus the primary outbound address. No interface names are available
+    this way, so they come back None."""
+    found = []
+    try:
+        infos = socket.getaddrinfo(socket.gethostname(), None)
+    except OSError:
+        infos = []
+    for family, _type, _proto, _canonical, sockaddr in infos:
+        if family not in (socket.AF_INET, socket.AF_INET6):
+            continue
+        usable = _usable(sockaddr[0])
+        if usable is not None:
+            found.append((usable[0], usable[1], None))
+    try:
+        usable = _usable(_outbound_address())
+    except OSError:
+        usable = None
+    if usable is not None:
+        found.append((usable[0], usable[1], None))
+    return found
+
+
+def local_addresses():
+    """This machine's non-loopback addresses, IPv4 first, as
+    [{"address", "family", "interface"}] - interface is None when the
+    method that found it cannot say. Never raises; [] on failure."""
+    found = []
+    try:
+        found = _psutil_addresses()
+    except Exception:
+        found = []
+    if not found:
+        try:
+            found = _stdlib_addresses()
+        except Exception:
+            found = []
+    seen = set()
+    unique = []
+    for address, family, interface in found:
+        if address in seen:
+            continue
+        seen.add(address)
+        unique.append({"address": address, "family": family, "interface": interface})
+    # IPv4 before IPv6, then by address so the list is stable between calls
+    unique.sort(key=lambda entry: (entry["family"] != "IPv4", entry["address"]))
+    return unique

dw/server/observed_cost.py

@@ -0,0 +1,379 @@
+"""What this server's own history says a workflow costs.
+
+`cost` is a maintainer's claim: runs they measured, on cards they name,
+written into the workflow. `dw/workflow_schema.json` says of it *"Never
+derived; absent means unknown"*, and that sentence is load-bearing - so this
+module does not touch it. It reports a sibling, `observed`, which is only ever
+this box's own finished jobs (#93, split out of #91).
+
+The report that asked for it: `list_workflows(shape="shot",
+traits="identity-referenced")` answered `cost: null` for seven of eight
+entries, on a box that had run `templates/minimax/reference-to-video` five
+times at one size - 511, 464.9, 478, ~460, ~467 seconds. The consumer is told
+to quote a price before spending GPU minutes and the catalog answered
+"unknown" for the template their series actually used.
+
+Four rules, each of which is a way the naive median would lie:
+
+- **Comparability.** A 141-frame run does not inform a 124-frame estimate.
+  A workflow declares `cost_drivers` - the variables that move its cost - and
+  runs are bucketed by those values; the bucket reported is the one a caller
+  gets by *default*, so it stays comparable to the curated figure. A workflow
+  declaring no drivers falls back to default-arguments-only runs, which is
+  thin but never wrong.
+- **Cold is not warm.** The same `text-to-image` run is 13.6 s with the model
+  on disk and 6.3 s with it resident. Averaging those describes neither, so
+  `cold_minutes` and `warm_minutes` are separate, each with its own run
+  count, and the one comparable to curated `cost` (wall clock *including*
+  model load) is the one named `cold`.
+- **A cached run is not a run.** A seeded workflow whose every step hit the
+  step cache finishes in seconds and wrote nothing. Counting those would
+  collapse the figure toward zero for exactly the templates that get re-run
+  most.
+- **Legibility over stability.** History is prunable, so the number moves.
+  `runs` and `since` are what make that readable rather than surprising.
+
+Everything comes off the job row - one query, no per-run JSON reads from the
+run directories, so the figures survive a pruned run directory.
+"""
+
+import json
+import logging
+import time
+
+logger = logging.getLogger("dw")
+
+COST_DRIVERS_KEY = "cost_drivers"
+# A phase event as `json.dumps` wrote it. Matched in SQL so a run's 200
+# persisted events are never parsed just to ask whether the weights loaded.
+LOADING_MARKER = '"phase": "loading"'
+# Persisted events are trimmed to the newest 200 (`MAX_PERSISTED_EVENTS`), so
+# a long run's `loading` phase can be gone from the record. Such a run counts
+# toward `runs` and toward neither side of the split - claiming it was warm
+# because the evidence was trimmed is the one answer that would be wrong.
+EVENT_CAP = 200
+
+
+def declared_drivers(definition):
+    """The variables the author says move this workflow's cost.
+
+    A name the workflow declares no variable for is dropped rather than
+    bucketed on: its effective value is None for every run, so it would put
+    every run in one bucket while looking like it had partitioned them.
+    `tests/test_observed_cost.py` sweeps the catalog for one, since a driver
+    that quietly checks nothing is the failure this whole module exists to
+    avoid.
+    """
+    drivers = definition.get(COST_DRIVERS_KEY)
+    if not isinstance(drivers, list):
+        return []
+    variables = definition.get("variables") or {}
+    return [name for name in drivers if isinstance(name, str) and name in variables]
+
+
+def _bucket_key(definition, arguments):
+    """What makes two runs of this workflow comparable, or None when they
+    are not comparable at all.
+
+    With declared drivers: the effective value of each, which is what the run
+    passed or else the stored default. Without: only a run that overrode
+    nothing that matters is comparable, and since nothing is declared to
+    matter, that means a run that overrode nothing at all.
+    """
+    variables = definition.get("variables") or {}
+    drivers = declared_drivers(definition)
+    if drivers:
+        return tuple(
+            (name, _comparable(arguments.get(name, variables.get(name))))
+            for name in sorted(drivers)
+        )
+    if any(key in variables for key in arguments):
+        return None
+    return ()
+
+
+def _comparable(value):
+    """A driver value as something hashable and stable across JSON round
+    trips - 124 and "124" are one bucket, since the engine coerces.
+
+    A *list* driver buckets on its length, not its contents: a `shots` list
+    costs what its number of entries costs, and two four-shot runs of one
+    template are comparable however different their prompts. Bucketing on
+    contents would give every run its own bucket and so a permanent
+    `runs: 1`, which is the same as reporting nothing.
+    """
+    if isinstance(value, bool) or value is None:
+        return value
+    if isinstance(value, (int, float)):
+        return float(value)
+    if isinstance(value, list):
+        return ("entries", len(value))
+    if isinstance(value, str):
+        try:
+            return float(value)
+        except ValueError:
+            return value
+    return json.dumps(value, sort_keys=True, default=str)
+
+
+def _is_task_only(definition):
+    """Whether every declared step is a plain task - no `pipeline`,
+    `pipeline_reference` or composed `workflow` step anywhere.
+
+    Such a workflow never logs a `loading` phase, so every run of it lands
+    in `had_load: False` and the cold/warm split below would bucket it all
+    as warm forever - silently withholding `cold_minutes`, the only figure
+    `_observed` (dw/plan.py) will quote as `basis: "observed"` (#439). A
+    task-only run has no model-load component to separate from its wall
+    clock, so its whole duration is what a curated `cost` figure would be
+    measured against, and is reported as `cold` rather than `warm`.
+
+    False when `steps` is missing or empty - an inline or otherwise
+    step-less definition is not evidence either way, so the ordinary
+    `had_load` split stays in force rather than guessing.
+    """
+    steps = definition.get("steps")
+    if not isinstance(steps, list) or not steps:
+        return False
+    for step in steps:
+        if not isinstance(step, dict):
+            return False
+        if "pipeline" in step or "pipeline_reference" in step or "workflow" in step:
+            return False
+    return True
+
+
+def _every_step_was_reused(manifest_text):
+    """Whether this run generated nothing - every manifest entry a step-cache
+    hit republishing an earlier run's files."""
+    try:
+        manifest = json.loads(manifest_text or "[]")
+    except (TypeError, ValueError):
+        return False
+    if not isinstance(manifest, list) or not manifest:
+        return False
+    return all(
+        isinstance(entry, dict) and entry.get("reused") is True for entry in manifest
+    )
+
+
+def _median(values):
+    ordered = sorted(values)
+    middle = len(ordered) // 2
+    if len(ordered) % 2:
+        return ordered[middle]
+    return (ordered[middle - 1] + ordered[middle]) / 2
+
+
+def _minutes_block(prefix, durations):
+    """One side of the cold/warm split, or {} when it has no runs."""
+    if not durations:
+        return {}
+    return {
+        f"{prefix}_minutes": round(_median(durations) / 60, 2),
+        f"{prefix}_runs": len(durations),
+        f"{prefix}_range_minutes": [
+            round(min(durations) / 60, 2),
+            round(max(durations) / 60, 2),
+        ],
+    }
+
+
+def observed_for(definition, rows, device=None, device_name=None, arguments=None):
+    """This server's history for one workflow, as the `observed` block, or
+    None when it has nothing comparable to report.
+
+    `rows` are that workflow's finished runs as `history_rows` yields them.
+    The device is the server's current one: the history holds no device
+    column, so a figure is about whatever accelerator this box has now - true
+    for every box that has not had its card swapped, and `runs`/`since` are
+    what let a reader notice if it has.
+
+    `arguments` are the run being asked about, and they choose the bucket:
+    the listing asks with none, so its figure is the one the *defaults* give
+    and stays comparable to a curated `cost`, while a plan asks with the
+    caller's own values and gets the figure for the run it is quoting, or
+    None when this box has never run that shape (#154).
+    """
+    wanted = _bucket_key(definition, arguments or {})
+    if wanted is None:
+        # No declared drivers and the caller overrode something: nothing here
+        # is comparable to the run being asked about
+        return None
+    task_only = _is_task_only(definition)
+    cold, warm, unclassified, earliest = [], [], 0, None
+    for row in rows:
+        if _every_step_was_reused(row["manifest"]):
+            continue
+        try:
+            row_arguments = json.loads(row["arguments"] or "{}")
+        except (TypeError, ValueError):
+            continue
+        if not isinstance(row_arguments, dict):
+            continue
+        if _bucket_key(definition, row_arguments) != wanted:
+            continue
+        duration = row["duration"]
+        if duration is None or duration <= 0:
+            continue
+        if task_only or row["had_load"]:
+            cold.append(duration)
+        elif row["events_at_cap"]:
+            unclassified += 1
+        else:
+            warm.append(duration)
+        finished = row["finished_at"]
+        if finished and (earliest is None or finished < earliest):
+            earliest = finished
+
+    runs = len(cold) + len(warm) + unclassified
+    if not runs:
+        return None
+    variables = definition.get("variables") or {}
+    drivers = declared_drivers(definition)
+    observed = {
+        "device": device,
+        "name": device_name,
+        "runs": runs,
+        "comparable": "drivers" if drivers else "default-arguments",
+        **_minutes_block("cold", cold),
+        **_minutes_block("warm", warm),
+    }
+    if drivers:
+        # A list driver is reported as the length it buckets on, so the
+        # figure says what it is a figure *for* without carrying a whole
+        # default shot list into the listing
+        # The *effective* driver values - the caller's where they supplied
+        # one, else the stored default - so the block says which run it is a
+        # figure for rather than always describing the defaults
+        effective = {
+            name: (arguments or {}).get(name, variables.get(name))
+            for name in sorted(drivers)
+        }
+        observed["drivers"] = {
+            name: len(value) if isinstance(value, list) else value
+            for name, value in effective.items()
+        }
+    if unclassified:
+        # Says the numbers do not add up, and why, rather than letting a
+        # reader assume a trimmed run was a warm one
+        observed["unclassified_runs"] = unclassified
+    if earliest:
+        observed["since"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(earliest))
+    return observed
+
+
+class ObservedCosts:
+    """The `observed` block per workflow name, recomputed when the job table
+    has moved.
+
+    One query for the whole listing rather than one per workflow, cached
+    against the jobs table's own high-water mark - the same reason
+    `workflow_details` caches against a file's mtime, except that a job
+    landing changes every figure and no file changes at all.
+    """
+
+    def __init__(self, history):
+        self.history = history
+        self._mark = None
+        self._rows = {}
+        self._device = None
+
+    def device(self):
+        """The accelerator every figure here is about, as (type, card name).
+
+        Computed once: it cannot change inside a process, and the listing is
+        the agent's hot path - asking torch for the card's name on every call
+        would be a torch round trip per catalog read. The history holds no
+        device column, so this is the server's current accelerator rather
+        than the one each run used; a box whose card was swapped mid-history
+        is the one case that misreports, and `runs`/`since` are what let a
+        reader notice.
+        """
+        if self._device is None:
+            from .. import device_memory_stats, get_device, get_device_type
+
+            kind, card = None, None
+            try:
+                kind = get_device_type(get_device())
+            except Exception:
+                logger.debug("observed cost: could not read the device type")
+            try:
+                # The card's marketing name, which only CUDA reports. Its
+                # absence is not the device's absence, so the two are asked
+                # for separately - one try around both let an ImportError on
+                # the second answer `device: null` for a box plainly running
+                # on CUDA, which is the silent-null shape this whole field
+                # exists to replace
+                card = device_memory_stats().get("device_name")
+            except Exception:
+                logger.debug("observed cost: could not read the device name")
+            self._device = (kind, card)
+        return self._device
+
+    def refresh(self):
+        """Bring the cached rows up to the table's watermark, once.
+
+        Returns False when there is no history to read or the read failed -
+        the caller then has no rows and knows why. A listing calls this
+        once and then asks `rows_for(name, fresh=False)` per entry, so the
+        COUNT(*) under the history lock happens once per request rather
+        than once per workflow."""
+        if self.history is None:
+            return False
+        try:
+            mark = self.history.watermark()
+        except Exception:
+            logger.debug("observed cost: could not read the job watermark")
+            return False
+        if mark != self._mark:
+            try:
+                self._rows = self.history.finished_runs()
+            except Exception:
+                logger.debug("observed cost: could not read job history")
+                self._rows = {}
+                return False
+            self._mark = mark
+        return True
+
+    def rows_for(self, name, *, fresh=True, workspace=None):
+        """That workflow's comparable-candidate runs, refreshing the cache
+        when the table has moved (unless the caller already did).
+
+        `workspace` is the identity a workspace-writable workflow's own copy
+        is scoped to (#274): two workspaces that each save a workflow called
+        the same thing are different workflows, so passing a workspace name
+        answers only for that one's own history and no other. Omitting it -
+        what a shared, read-only catalog source (a template or example) is
+        asked with - keeps #154's behavior: every workspace's runs of that
+        catalog entry are one pool, unioned across whichever workspaces
+        happen to have run it.
+        """
+        if fresh and not self.refresh():
+            return []
+        if workspace is not None:
+            return self._rows.get((workspace, name), [])
+        rows = []
+        for (row_workspace, row_name), workspace_rows in self._rows.items():
+            if row_name == name:
+                rows.extend(workspace_rows)
+        return rows
+
+    def observed(self, name, definition, arguments=None, *, fresh=True, workspace=None):
+        rows = self.rows_for(name, fresh=fresh, workspace=workspace)
+        if not rows:
+            return None
+        device, device_name = self.device()
+        try:
+            return observed_for(
+                definition,
+                rows,
+                device=device,
+                device_name=device_name,
+                arguments=arguments,
+            )
+        except Exception:
+            # A figure is a nicety; a listing that 500s because one row had a
+            # shape nobody expected is not
+            logger.debug("observed cost: could not aggregate %s", name, exc_info=True)
+            return None

dw/server/sysinfo.py

@@ -0,0 +1,71 @@
+"""Runtime environment details - Python, torch, CUDA/driver and other
+installed package versions - for diagnosing environment mismatches between
+boxes (#222). Read-only, gathered fresh on every call rather than cached,
+since an install can change underneath a long-running server (see
+`update_diffusers`).
+"""
+
+import platform
+import subprocess
+from importlib.metadata import PackageNotFoundError, version
+
+# Packages worth reporting beyond torch (which gets its own field): the
+# ones most likely to cause a version-mismatch failure or be silently
+# absent on a given box.
+PACKAGES = (
+    "diffusers",
+    "transformers",
+    "accelerate",
+    "bitsandbytes",
+    "peft",
+    "safetensors",
+    "sentencepiece",
+)
+
+
+def _package_version(name):
+    try:
+        return version(name)
+    except PackageNotFoundError:
+        return None
+
+
+def _nvidia_driver_version():
+    """The NVIDIA driver version, via nvidia-smi when it's on PATH. Distinct
+    from torch's own `cuda_version`, which is the CUDA toolkit torch was
+    built against, not the driver actually installed on this box."""
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+    if result.returncode != 0:
+        return None
+    lines = result.stdout.strip().splitlines()
+    return lines[0].strip() if lines else None
+
+
+def runtime_info():
+    """Python, torch, CUDA/driver and other installed package versions -
+    the detail `device`/`version` alone doesn't answer, like "is
+    bitsandbytes even installed here" or "which CUDA build is torch"."""
+    try:
+        import torch
+
+        torch_version = torch.__version__
+        cuda_version = torch.version.cuda
+    except ImportError:
+        torch_version = None
+        cuda_version = None
+
+    return {
+        "python_version": platform.python_version(),
+        "torch_version": torch_version,
+        "cuda_version": cuda_version,
+        "driver_version": _nvidia_driver_version(),
+        "packages": {name: _package_version(name) for name in PACKAGES},
+    }