stroblmeandClaude Opus 5
81649dbfca
A module defining nodes could not be run directly: the decorator refused `__main__` while the module body was still executing, so a `if __name__ == "__main__"` self-check beside the nodes was impossible and the checks had to live in a separate pytest file. The refusal now fires where the generated body is written — document() and shims(), both, since sync writes the document first — and the decorator hands the function back as it always did. Syncing a __main__-defined node is still refused, with the same words. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019Hra4ndWMCLU5F3KjUuVAc
816 lines
30 KiB
Python
816 lines
30 KiB
Python
"""Declare flows and nodes in the repository you already have.
|
|
|
|
A research repository keeps its functions where they are; the decorators here
|
|
say which of them are nodes and which nodes make up a flow. ``fluksio sync``
|
|
uploads the flow document and, per node, a generated import shim — so the flow
|
|
store still holds a complete, runnable, git-versioned definition, and the code
|
|
it imports stays yours.
|
|
|
|
Nothing here imports the engine. A ``from fluksio import node`` in a training
|
|
script must not drag FastAPI and SQLAlchemy in behind it, and the same names
|
|
have to resolve inside a worker, where ``fluksio`` is the reporter module.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import inspect
|
|
import json
|
|
import re
|
|
import textwrap
|
|
import warnings
|
|
from collections.abc import Callable, Sequence
|
|
from pathlib import Path
|
|
from typing import Any, TypeVar
|
|
|
|
from fluksio.sdk.imports import digest_of, reached
|
|
|
|
__all__ = [
|
|
"FLOWS",
|
|
"Flow",
|
|
"Port",
|
|
"SyncError",
|
|
"build",
|
|
"node",
|
|
"shims",
|
|
"use",
|
|
]
|
|
|
|
F = TypeVar("F", bound=Callable[..., Any])
|
|
|
|
#: Mirrors :class:`fluksio.flow.messages.DType`. Mirrored rather than imported:
|
|
#: importing it would pull the engine into a research process.
|
|
DTYPES = frozenset(
|
|
{
|
|
"float",
|
|
"int",
|
|
"str",
|
|
"bool",
|
|
"json",
|
|
"series",
|
|
"record",
|
|
"list",
|
|
"artifact",
|
|
"image",
|
|
"audio",
|
|
"video",
|
|
}
|
|
)
|
|
#: What a list may hold, mirroring ``messages._ITEM_TYPES``.
|
|
ITEM_TYPES = frozenset({"float", "int", "str", "bool", "json", "record"})
|
|
#: Settings the engine reads itself, mirroring ``nodes.base.RESERVED_SETTINGS``.
|
|
RESERVED_SETTINGS = frozenset({"synchronous"})
|
|
#: Mirrors ``fluksio.flow.schemas.NAME_PATTERN``.
|
|
NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$")
|
|
#: First line of every generated node body. Its absence is how ``sync`` knows a
|
|
#: node's source was edited on the canvas.
|
|
MARKER = "# generated by fluksio sync"
|
|
|
|
|
|
class SyncError(Exception):
|
|
"""A declaration the engine would refuse, caught where it was written."""
|
|
|
|
|
|
#: Every flow declared in the modules imported so far. ``fluksio sync`` reads it.
|
|
FLOWS: dict[str, Flow] = {}
|
|
|
|
|
|
def _name(value: str, what: str) -> str:
|
|
if not NAME_PATTERN.match(value):
|
|
raise SyncError(
|
|
f"{what} '{value}': use lowercase letters, digits and underscores, "
|
|
"starting with a letter"
|
|
)
|
|
return value
|
|
|
|
|
|
class Port:
|
|
"""One port of a node, and the message it binds to.
|
|
|
|
A thin mirror of :class:`fluksio.flow.messages.MessageSpec`. ``name`` is
|
|
the wiring — two nodes are connected because one provides the name another
|
|
requires — and ``port`` is what the function calls its parameter, which
|
|
defaults to the name.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
dtype: str | None = None,
|
|
*,
|
|
port: str = "",
|
|
stream: bool = False,
|
|
initial: Any = None,
|
|
item: str | None = None,
|
|
interval: float = 0.0,
|
|
trigger: bool = True,
|
|
) -> None:
|
|
#: Whether the caller said the type or left it to be resolved from
|
|
#: whoever provides the message. A message has one type, declared where
|
|
#: it is produced; requiring it does not mean declaring it again.
|
|
self.explicit = dtype is not None
|
|
dtype = dtype or "float"
|
|
if dtype not in DTYPES:
|
|
raise SyncError(f"port '{name}': unknown type '{dtype}'")
|
|
if item is not None and item not in ITEM_TYPES:
|
|
raise SyncError(f"port '{name}': a list cannot hold '{item}' items")
|
|
if interval < 0:
|
|
raise SyncError(f"port '{name}': interval cannot be negative")
|
|
self.name = name
|
|
self.dtype = dtype
|
|
self.port = port or name.rsplit(".", 1)[-1]
|
|
self.stream = stream
|
|
self.initial = initial
|
|
self.item = item
|
|
self.interval = interval
|
|
self.trigger = trigger
|
|
if not self.port.isidentifier():
|
|
raise SyncError(
|
|
f"port '{name}': '{self.port}' is not a Python identifier, so no "
|
|
"parameter can carry it — pass port='...'"
|
|
)
|
|
|
|
def spec(self) -> dict[str, Any]:
|
|
"""The ``MessageSpec`` this port serialises to."""
|
|
out: dict[str, Any] = {
|
|
"name": self.name,
|
|
"port": self.port,
|
|
"dtype": self.dtype,
|
|
"interval": self.interval,
|
|
"trigger": self.trigger,
|
|
"stream": self.stream,
|
|
}
|
|
if self.item is not None:
|
|
out["item"] = self.item
|
|
return out
|
|
|
|
def replace(self, name: str = "") -> Port:
|
|
"""A copy of this port, optionally bound to a different message."""
|
|
copy = Port(
|
|
name or self.name,
|
|
self.dtype,
|
|
port=self.port,
|
|
stream=self.stream,
|
|
initial=self.initial,
|
|
item=self.item,
|
|
interval=self.interval,
|
|
trigger=self.trigger,
|
|
)
|
|
copy.explicit = self.explicit
|
|
return copy
|
|
|
|
def __repr__(self) -> str:
|
|
return f"Port({self.name!r}, {self.dtype!r})"
|
|
|
|
|
|
def _ports(value: Sequence[str | Port] | str | Port) -> tuple[list[Port], bool]:
|
|
"""Normalise a ports declaration; the flag says it was a single port.
|
|
|
|
A single ``Port`` rather than a list of one is the opt-in to a bare return:
|
|
the shim wraps the value in the dict the engine expects. It has to be
|
|
explicit, because an artifact reference is itself a dict and a dict whose
|
|
keys match no port is dropped in silence.
|
|
"""
|
|
if isinstance(value, (str, Port)):
|
|
return [value if isinstance(value, Port) else Port(value)], True
|
|
return [item if isinstance(item, Port) else Port(item) for item in value], False
|
|
|
|
|
|
class NodeSpec:
|
|
"""What ``@node`` recorded about a function."""
|
|
|
|
def __init__(
|
|
self,
|
|
fn: Callable[..., Any],
|
|
*,
|
|
id: str,
|
|
requires: list[Port],
|
|
provides: list[Port],
|
|
single: bool,
|
|
settings: dict[str, Any],
|
|
title: str,
|
|
timeout: float | None,
|
|
device: str | None,
|
|
device_policy: str,
|
|
cache: bool = True,
|
|
resources: dict[str, Any] | None = None,
|
|
) -> None:
|
|
self.fn = fn
|
|
self.id = id
|
|
self.requires = requires
|
|
self.provides = provides
|
|
self.single = single
|
|
self.settings = settings
|
|
self.title = title
|
|
self.timeout = timeout
|
|
self.device = device
|
|
self.device_policy = device_policy
|
|
self.cache = cache
|
|
self.resources = resources
|
|
|
|
def rebind(
|
|
self, *, id: str = "", wire: dict[str, str] | None = None, **settings: Any
|
|
) -> NodeSpec:
|
|
"""A copy of this node wired and configured for one flow."""
|
|
wire = wire or {}
|
|
known = {port.port for port in self.requires} | {p.port for p in self.provides}
|
|
unknown = sorted(set(wire) - known)
|
|
if unknown:
|
|
raise SyncError(
|
|
f"node '{self.id}': wire={{'{unknown[0]}': ...}} names no port of it "
|
|
f"(it has {', '.join(sorted(known)) or 'none'})"
|
|
)
|
|
|
|
# Copied rather than shared: the same decorated function appears in
|
|
# several flows, and each resolves its own port types.
|
|
def rewire(ports: list[Port]) -> list[Port]:
|
|
return [port.replace(wire.get(port.port, "")) for port in ports]
|
|
|
|
return NodeSpec(
|
|
self.fn,
|
|
id=_name(id, "node id") if id else self.id,
|
|
requires=rewire(self.requires),
|
|
provides=rewire(self.provides),
|
|
single=self.single,
|
|
settings={**self.settings, **settings},
|
|
title=self.title,
|
|
timeout=self.timeout,
|
|
device=self.device,
|
|
device_policy=self.device_policy,
|
|
cache=self.cache,
|
|
resources=self.resources,
|
|
)
|
|
|
|
|
|
class NodeUse:
|
|
"""One use of a decorated function in one flow — see :func:`use`."""
|
|
|
|
def __init__(
|
|
self,
|
|
fn: Callable[..., Any],
|
|
*,
|
|
id: str = "",
|
|
wire: dict[str, str] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
) -> None:
|
|
self.fn = fn
|
|
self.id = id
|
|
self.wire = wire
|
|
self.settings = settings or {}
|
|
|
|
|
|
def node(
|
|
*,
|
|
requires: Sequence[str | Port] = (),
|
|
provides: Sequence[str | Port] | str | Port = (),
|
|
id: str = "",
|
|
settings: dict[str, Any] | None = None,
|
|
title: str = "",
|
|
timeout: float | None = None,
|
|
device: str | None = None,
|
|
device_policy: str = "require",
|
|
cache: bool = True,
|
|
resources: dict[str, Any] | str | None = None,
|
|
) -> Callable[[F], F]:
|
|
"""Mark a function as a node, declaring its ports.
|
|
|
|
The function is returned untouched: it is still callable, testable and
|
|
importable as what it was. A bare string in ``requires`` or ``provides`` is
|
|
shorthand for a ``float`` port of that name.
|
|
|
|
Parameters with a default that are not ports become the node's settings, so
|
|
the canvas can tune them without editing code. ``device`` picks the worker
|
|
the code runs on — ``"gpu"``, say — and ``device_policy="prefer"`` runs it
|
|
locally when no such worker is attached rather than waiting for one.
|
|
|
|
A batch run skips this node when an earlier one already ran the same source
|
|
with the same settings and the same input values, and restores what it
|
|
returned. ``cache=False`` says not to: the answer can change on its own.
|
|
|
|
``timeout`` is seconds of silence — a yield or an emit resets it — after
|
|
which the node is stopped. Left out it inherits the engine's default, which
|
|
is no limit; set one where silence means stuck rather than working.
|
|
|
|
``resources`` says how much of the machine one execution takes, for a node
|
|
heavy enough that the answer is not "a share of it"::
|
|
|
|
@node(..., resources={"cpus": 4})
|
|
@node(..., resources={"gpus": 1,
|
|
"env": {"XLA_PYTHON_CLIENT_MEM_FRACTION": "0.9"}})
|
|
|
|
The engine holds that much while the node runs, so the rest of them wait
|
|
their turn rather than oversubscribing the box, and the worker is started
|
|
with thread limits matching what it was given — which is the only way a
|
|
library that sizes itself to every visible core is told otherwise. ``env``
|
|
is for the tuning the engine must not invent, such as ``XLA_FLAGS``.
|
|
Declaring nothing is the default and changes nothing.
|
|
|
|
A string is a stored flavor — ``resources="gpu-small"`` — which is the same
|
|
thing as ``{"flavor": "gpu-small"}`` and stands in for cpus, gpus and ram.
|
|
Sizes accept ``"2G"``; ``duration_s`` accepts ``"30m"``.
|
|
"""
|
|
if device_policy not in ("require", "prefer"):
|
|
raise SyncError("device_policy is 'require' or 'prefer'")
|
|
if isinstance(resources, str):
|
|
resources = {"flavor": resources}
|
|
if resources is not None:
|
|
known = {"cpus", "gpus", "ram", "flavor", "duration_s", "env"}
|
|
unknown = sorted(set(resources) - known)
|
|
if unknown:
|
|
raise SyncError(
|
|
f"resources={{'{unknown[0]}': ...}} is not one of "
|
|
f"{', '.join(sorted(known))}"
|
|
)
|
|
if timeout is not None and timeout < 0:
|
|
raise SyncError(
|
|
f"node timeout must be 0 or more (0 disables it), got {timeout}"
|
|
)
|
|
|
|
def decorate(fn: F) -> F:
|
|
required, _ = _ports(requires)
|
|
provided, single = _ports(provides)
|
|
spec = NodeSpec(
|
|
fn,
|
|
id=_name(id or fn.__name__, "node id"),
|
|
requires=required,
|
|
provides=provided,
|
|
single=single,
|
|
settings=dict(settings or {}),
|
|
title=title,
|
|
timeout=timeout,
|
|
device=device,
|
|
device_policy=device_policy,
|
|
cache=cache,
|
|
resources=dict(resources) if resources is not None else None,
|
|
)
|
|
_check_signature(spec)
|
|
fn.__fluksio__ = spec # type: ignore[attr-defined]
|
|
return fn
|
|
|
|
return decorate
|
|
|
|
|
|
def use(
|
|
fn: Callable[..., Any],
|
|
*,
|
|
id: str = "",
|
|
wire: dict[str, str] | None = None,
|
|
**settings: Any,
|
|
) -> NodeUse:
|
|
"""The same node, wired or configured differently in this flow.
|
|
|
|
``wire={"dataset": "augmented"}`` binds the ``dataset`` port to another
|
|
message; keyword arguments override settings. Give ``id`` when one function
|
|
appears twice in one flow.
|
|
"""
|
|
return NodeUse(fn, id=id, wire=wire, settings=settings)
|
|
|
|
|
|
def _spec_of(entry: Callable[..., Any] | NodeUse) -> NodeSpec:
|
|
fn = entry.fn if isinstance(entry, NodeUse) else entry
|
|
spec = getattr(fn, "__fluksio__", None)
|
|
if not isinstance(spec, NodeSpec):
|
|
raise SyncError(
|
|
f"{getattr(fn, '__name__', fn)!r} is not a node — decorate it with @node()"
|
|
)
|
|
if isinstance(entry, NodeUse):
|
|
return spec.rebind(id=entry.id, wire=entry.wire, **entry.settings)
|
|
return spec.rebind()
|
|
|
|
|
|
def _check_signature(spec: NodeSpec) -> None:
|
|
"""Whether the declaration matches the function it is on.
|
|
|
|
Raised where the decorator is rather than where the flow is: this is a
|
|
disagreement between a function and its own ports, and the engine would
|
|
otherwise only find it when the node is first called.
|
|
"""
|
|
fn, where = spec.fn, f"node '{spec.id}'"
|
|
if inspect.iscoroutinefunction(fn):
|
|
raise SyncError(f"{where}: async functions cannot be nodes")
|
|
generator = inspect.isgeneratorfunction(fn)
|
|
if generator and spec.single:
|
|
raise SyncError(
|
|
f"{where}: {fn.__name__}() yields, so it says which port each value is "
|
|
f"on — pass provides=[Port(...)] rather than a single port"
|
|
)
|
|
|
|
signature = inspect.signature(fn)
|
|
parameters = signature.parameters
|
|
var_keyword = any(
|
|
p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()
|
|
)
|
|
ports = {port.port for port in spec.requires}
|
|
|
|
for port in ports:
|
|
if port not in parameters and not var_keyword:
|
|
raise SyncError(
|
|
f"{where}: port '{port}' has no parameter to arrive in — "
|
|
f"{fn.__name__}({', '.join(parameters)}) takes no such argument"
|
|
)
|
|
for name, parameter in parameters.items():
|
|
if parameter.kind is inspect.Parameter.VAR_POSITIONAL:
|
|
raise SyncError(f"{where}: *{name} cannot be filled from ports")
|
|
if parameter.kind is inspect.Parameter.VAR_KEYWORD or name in ports:
|
|
continue
|
|
if parameter.default is inspect.Parameter.empty:
|
|
raise SyncError(
|
|
f"{where}: parameter '{name}' is neither a port nor a setting — "
|
|
f"add Port('{name}', ...) to requires=, or give it a default"
|
|
)
|
|
|
|
_check_emissions(spec)
|
|
|
|
|
|
def _check_emissions(spec: NodeSpec) -> None:
|
|
"""Whether what the function yields is what it says it provides.
|
|
|
|
The engine refuses an undeclared key at the first yield, which is right but
|
|
late: a sweep can be an hour in. Only literal keys are read here — a yield
|
|
of a name the code computes is left to the engine — so this catches the
|
|
mistake it is for, which is a typo.
|
|
"""
|
|
provided = {port.port for port in spec.provides}
|
|
try:
|
|
tree = ast.parse(textwrap.dedent(inspect.getsource(spec.fn)))
|
|
except (OSError, TypeError, SyntaxError):
|
|
# No source to read: a function built at runtime, or one from a
|
|
# console. Not a reason to refuse it.
|
|
return
|
|
|
|
for key in _emitted_keys(tree):
|
|
if key not in provided:
|
|
declared = ", ".join(sorted(provided)) or "none"
|
|
raise SyncError(
|
|
f"node '{spec.id}': {spec.fn.__name__}() emits '{key}', which no "
|
|
f"port declares. Its ports are: {declared}."
|
|
)
|
|
|
|
|
|
def _emitted_keys(tree: ast.AST) -> list[str]:
|
|
"""Literal port names a function yields or emits, its own only.
|
|
|
|
A nested function is somebody else's ports — a helper generator yielding
|
|
pairs is not this node emitting them — so the walk stops at one.
|
|
"""
|
|
keys: list[str] = []
|
|
# Each branch carries whether it is already inside the node's own function,
|
|
# so the first def is the node and any def below it is a helper.
|
|
stack: list[tuple[ast.AST, bool]] = [(tree, False)]
|
|
while stack:
|
|
node, inside = stack.pop()
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
|
|
if inside:
|
|
continue
|
|
inside = True
|
|
if isinstance(node, ast.Yield) and isinstance(node.value, ast.Dict):
|
|
keys += [
|
|
key.value
|
|
for key in node.value.keys
|
|
if isinstance(key, ast.Constant) and isinstance(key.value, str)
|
|
]
|
|
elif isinstance(node, ast.Call) and _is_emit(node.func):
|
|
keys += [kw.arg for kw in node.keywords if kw.arg]
|
|
stack += [(child, inside) for child in ast.iter_child_nodes(node)]
|
|
return keys
|
|
|
|
|
|
def _is_emit(func: ast.expr) -> bool:
|
|
"""``emit(...)`` or ``fluksio.emit(...)``, however it was imported."""
|
|
if isinstance(func, ast.Name):
|
|
return func.id == "emit"
|
|
return isinstance(func, ast.Attribute) and func.attr == "emit"
|
|
|
|
|
|
def _check(spec: NodeSpec, mode: str) -> dict[str, Any]:
|
|
"""The node as stored, with the settings and wiring of one flow.
|
|
|
|
Mirrors ``FlowController._build_node`` and ``flow.runs.batch_issues``. Kept
|
|
apart from the signature check because ``use(fn, epochs=3)`` adds settings
|
|
per flow, so this is the first point at which they are all known.
|
|
"""
|
|
fn, where = spec.fn, f"node '{spec.id}'"
|
|
parameters = inspect.signature(fn).parameters
|
|
var_keyword = any(
|
|
p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()
|
|
)
|
|
ports = {port.port for port in spec.requires}
|
|
|
|
params: dict[str, Any] = {}
|
|
for name, parameter in parameters.items():
|
|
if name in ports or parameter.default is inspect.Parameter.empty:
|
|
continue
|
|
if name in RESERVED_SETTINGS:
|
|
continue
|
|
if _serialisable(parameter.default):
|
|
params[name] = parameter.default
|
|
for name, value in spec.settings.items():
|
|
if name in ports:
|
|
raise SyncError(
|
|
f"{where}: '{name}' is both a port and a setting — a node reads its "
|
|
"settings beside its ports, so the names have to differ"
|
|
)
|
|
if name in RESERVED_SETTINGS:
|
|
raise SyncError(f"{where}: '{name}' is a setting the engine reads itself")
|
|
if name not in parameters and not var_keyword:
|
|
raise SyncError(
|
|
f"{where}: setting '{name}' is not a parameter of {fn.__name__}()"
|
|
)
|
|
if not _serialisable(value):
|
|
raise SyncError(
|
|
f"{where}: setting '{name}' is not JSON — settings are stored with "
|
|
"the flow, so they have to survive a round trip"
|
|
)
|
|
params[name] = value
|
|
|
|
if mode == "batch":
|
|
for port in spec.requires + spec.provides:
|
|
if port.interval > 0 and not port.stream:
|
|
raise SyncError(
|
|
f"{where}: port '{port.port}' has interval={port.interval} but a "
|
|
"batch run delivers once — drop it, or declare stream=True"
|
|
)
|
|
|
|
return {
|
|
"id": spec.id,
|
|
"type": "python",
|
|
"title": spec.title,
|
|
"params": params,
|
|
"requires": [port.spec() for port in spec.requires],
|
|
"provides": [port.spec() for port in spec.provides],
|
|
"timeout": spec.timeout,
|
|
"device": spec.device,
|
|
"device_policy": spec.device_policy,
|
|
"cache": spec.cache,
|
|
"resources": spec.resources,
|
|
}
|
|
|
|
|
|
def _refuse_main(spec: NodeSpec) -> None:
|
|
"""A node the generated body would have no way to import.
|
|
|
|
Refused where the body is written rather than where the decorator is: a
|
|
module defining nodes is then still runnable as a script, which is what a
|
|
`__main__` self-check beside them needs.
|
|
"""
|
|
if spec.fn.__module__ == "__main__":
|
|
raise SyncError(
|
|
f"node '{spec.id}': {spec.fn.__name__}() is defined in a script run "
|
|
"directly, so the generated node could not import it — put it in an "
|
|
"importable module"
|
|
)
|
|
|
|
|
|
def _code_of(spec: NodeSpec) -> dict[str, Any]:
|
|
"""What this node's function calls into, which its shim does not say.
|
|
|
|
Read here rather than when the flow is declared, so a second `sync()` in
|
|
one process — a notebook — sees the edit that happened between them. This
|
|
side is the only one that can work it out at all: it has imported the
|
|
code, and the engine never does.
|
|
"""
|
|
_refuse_main(spec)
|
|
files = reached(spec.fn)
|
|
return {"code_files": files, "code_digest": digest_of(files)}
|
|
|
|
|
|
def _serialisable(value: Any) -> bool:
|
|
try:
|
|
json.dumps(value)
|
|
except (TypeError, ValueError):
|
|
return False
|
|
return True
|
|
|
|
|
|
class Flow:
|
|
"""A flow: which nodes are in it, what it takes and what it reports."""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
*,
|
|
nodes: Sequence[Callable[..., Any] | NodeUse],
|
|
mode: str = "batch",
|
|
inputs: Sequence[Port] = (),
|
|
outputs: Sequence[str] = (),
|
|
title: str = "",
|
|
) -> None:
|
|
"""Declare a flow over decorated functions, and register it for sync.
|
|
|
|
``nodes`` is the membership: any decorated function, from anywhere in
|
|
the repository. Wiring is not membership — nodes connect because one
|
|
provides a message another requires, never because one imported the
|
|
other. Wrap an entry in :func:`use` to rebind or reconfigure it here
|
|
without touching the other flows using it.
|
|
"""
|
|
if mode not in ("batch", "live"):
|
|
raise SyncError(f"flow '{name}': mode is 'batch' or 'live'")
|
|
self.name = _name(name, "flow name")
|
|
self.title = title
|
|
self.mode = mode
|
|
self.inputs = list(inputs)
|
|
self.outputs = list(outputs)
|
|
self.nodes = [_spec_of(entry) for entry in nodes]
|
|
self._resolve_types()
|
|
self._validate()
|
|
self._nodes = [_check(spec, mode) for spec in self.nodes]
|
|
self._register()
|
|
|
|
def _register(self) -> None:
|
|
"""Put this flow where `fluksio sync` will find it."""
|
|
existing = FLOWS.get(self.name)
|
|
if existing is not None and existing is not self:
|
|
warnings.warn(
|
|
f"two flows are called '{self.name}'; the later one wins, and only "
|
|
"it will be synced",
|
|
stacklevel=3,
|
|
)
|
|
FLOWS[self.name] = self
|
|
|
|
def _resolve_types(self) -> None:
|
|
"""Give every required port the type of the message it reads.
|
|
|
|
A message has one type, declared once by whoever produces it. Requiring
|
|
it is not a second declaration, so `requires=["dataset"]` takes the
|
|
artifact type from the node providing `dataset` rather than defaulting
|
|
to a float the engine would then refuse at run time. Spelling the type
|
|
out on both sides is still allowed — and a disagreement is an error,
|
|
because one of the two is wrong about what it is handling.
|
|
"""
|
|
declared: dict[str, Port] = {}
|
|
for port in self.inputs:
|
|
declared.setdefault(port.name, port)
|
|
for spec in self.nodes:
|
|
for port in spec.provides:
|
|
declared.setdefault(port.name, port)
|
|
for spec in self.nodes:
|
|
for port in spec.requires:
|
|
source = declared.get(port.name)
|
|
if source is None:
|
|
continue
|
|
if not port.explicit:
|
|
port.dtype = source.dtype
|
|
port.item = source.item
|
|
elif port.dtype != source.dtype:
|
|
raise SyncError(
|
|
f"flow '{self.name}': node '{spec.id}' reads '{port.name}' as "
|
|
f"{port.dtype}, but it is provided as {source.dtype}"
|
|
)
|
|
|
|
def _validate(self) -> None:
|
|
seen: set[str] = set()
|
|
for spec in self.nodes:
|
|
if spec.id in seen:
|
|
raise SyncError(
|
|
f"flow '{self.name}': two nodes are called '{spec.id}' — "
|
|
"pass use(fn, id='...') to tell them apart"
|
|
)
|
|
seen.add(spec.id)
|
|
|
|
produced: dict[str, list[str]] = {}
|
|
for spec in self.nodes:
|
|
for port in spec.provides:
|
|
produced.setdefault(port.name, []).append(spec.id)
|
|
for name, writers in produced.items():
|
|
if len(writers) > 1:
|
|
warnings.warn(
|
|
f"flow '{self.name}': {', '.join(writers)} all publish "
|
|
f"'{name}', so a reader gets whichever ran last",
|
|
stacklevel=4,
|
|
)
|
|
|
|
available = set(produced) | {port.name for port in self.inputs}
|
|
for spec in self.nodes:
|
|
for port in spec.requires:
|
|
if "." in port.name or port.name in available:
|
|
continue
|
|
raise SyncError(
|
|
f"flow '{self.name}': node '{spec.id}' requires '{port.name}', "
|
|
"which no node in it provides and no input declares"
|
|
)
|
|
for output in self.outputs:
|
|
if output not in available:
|
|
raise SyncError(
|
|
f"flow '{self.name}': output '{output}' is not produced by any "
|
|
"of its nodes"
|
|
)
|
|
|
|
def document(self, origin: dict[str, Any] | None = None) -> dict[str, Any]:
|
|
"""The flow document to store, without a version — sync sets that."""
|
|
doc: dict[str, Any] = {
|
|
"name": self.name,
|
|
"title": self.title,
|
|
"mode": self.mode,
|
|
"nodes": [
|
|
node | _code_of(spec)
|
|
for spec, node in zip(self.nodes, self._nodes, strict=True)
|
|
],
|
|
"inputs": [
|
|
{"spec": port.spec(), "initial": port.initial} for port in self.inputs
|
|
],
|
|
"outputs": self.outputs,
|
|
}
|
|
if origin is not None:
|
|
doc["origin"] = origin
|
|
return doc
|
|
|
|
def shims(self) -> dict[str, str]:
|
|
"""One generated node body per node, keyed by node id."""
|
|
return {spec.id: _shim(spec) for spec in self.nodes}
|
|
|
|
def submit(
|
|
self, *, seed: int | None = None, client: Any = None, **params: Any
|
|
) -> Any:
|
|
"""Start a run of this flow and hand back a handle to it."""
|
|
from fluksio.sdk.client import Client
|
|
|
|
return (client or Client()).submit(self.name, params, seed=seed)
|
|
|
|
def runs(self, limit: int = 20, client: Any = None) -> list[dict[str, Any]]:
|
|
"""This flow's runs, newest first."""
|
|
from fluksio.sdk.client import Client
|
|
|
|
rows: list[dict[str, Any]] = (client or Client()).runs(
|
|
flow=self.name, limit=limit
|
|
)
|
|
return rows
|
|
|
|
def __repr__(self) -> str:
|
|
return f"Flow({self.name!r}, nodes={[n.id for n in self.nodes]})"
|
|
|
|
|
|
def build(target: Flow, origin: dict[str, Any] | None = None) -> dict[str, Any]:
|
|
"""The flow document for a flow — checked, without a version."""
|
|
return target.document(origin)
|
|
|
|
|
|
def shims(target: Flow) -> dict[str, str]:
|
|
"""The generated node bodies for a flow."""
|
|
return target.shims()
|
|
|
|
|
|
def import_root(fn: Callable[..., Any]) -> str:
|
|
"""The directory that has to be on the path for ``fn`` to be importable.
|
|
|
|
Its module's dotted name says how deep in a package it sits, so counting
|
|
the segments back up from the file lands on the directory the import
|
|
resolves against — which is not always the repository root, and is what a
|
|
generated shim actually needs.
|
|
"""
|
|
source = inspect.getsourcefile(fn)
|
|
if not source:
|
|
return ""
|
|
path = Path(source).resolve()
|
|
for _ in fn.__module__.split("."):
|
|
path = path.parent
|
|
return str(path)
|
|
|
|
|
|
def _shim(spec: NodeSpec) -> str:
|
|
"""The node body the store keeps: an import of the real function.
|
|
|
|
The store therefore still holds a complete, runnable definition — the body
|
|
simply happens to be generated, which is why it says so and says where the
|
|
real thing is.
|
|
"""
|
|
_refuse_main(spec)
|
|
fn = spec.fn
|
|
where = inspect.getsourcefile(fn) or fn.__module__
|
|
repo = import_root(fn)
|
|
ports = [port.port for port in spec.requires]
|
|
signature = ", ".join(ports + ["**settings"])
|
|
arguments = ", ".join([f"{port}={port}" for port in ports] + ["**settings"])
|
|
call = f"{fn.__name__}({arguments})"
|
|
|
|
lines = [
|
|
f"{MARKER} from {where} — edit that file instead",
|
|
"import sys",
|
|
"",
|
|
]
|
|
if repo:
|
|
# ponytail: the repository is on the path because the shim puts it
|
|
# there. An editable install (`-e /repo` in the Modules manifest, or a
|
|
# VCS requirement on a remote worker) makes these three lines a no-op.
|
|
lines += [
|
|
f"_REPO = {repo!r}",
|
|
"if _REPO not in sys.path:",
|
|
" sys.path.insert(0, _REPO)",
|
|
"",
|
|
]
|
|
lines += [
|
|
f"from {fn.__module__} import {fn.__name__}",
|
|
"",
|
|
"",
|
|
f"def process({signature}):",
|
|
]
|
|
if inspect.isgeneratorfunction(fn):
|
|
# `return (yield from ...)`, not a bare `yield from`: the delegating
|
|
# form is what carries the generator's own return value out, and that
|
|
# return value is everything the node produces at the end.
|
|
lines.append(f" return (yield from {call})")
|
|
elif spec.single:
|
|
lines.append(f" return {{{spec.provides[0].port!r}: {call}}}")
|
|
else:
|
|
lines.append(f" return {call}")
|
|
return "\n".join(lines) + "\n"
|