app/backend/fluksio/flow/schemas.py

"""The persisted shape of a flow, shared by the store, the API and the editor.

A flow is structure plus code: this module is the structure. Node logic for
``python`` nodes lives beside it as a plain ``.py`` file.
"""

from __future__ import annotations

import re
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from fluksio.flow.messages import MessageSpec

NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$")

#: Megabytes, gigabytes, or a bare number already in megabytes.
_SIZE = re.compile(r"^\s*(\d+)\s*([mg]b?)?\s*$", re.IGNORECASE)
#: Seconds, minutes, hours or days — one unit, so there is one thing to read.
_SPAN = re.compile(r"^\s*(\d+)\s*([smhd])?\s*$", re.IGNORECASE)
_SPAN_SECONDS = {"s": 1, "m": 60, "h": 3600, "d": 86400}


def _validate_name(value: str) -> str:
    if not NAME_PATTERN.match(value):
        raise ValueError(
            "Use lowercase letters, digits and underscores, starting with a letter"
        )
    return value


def _megabytes(value: Any) -> Any:
    """``2G`` and ``512M`` and ``512`` all mean a number of megabytes."""
    if not isinstance(value, str):
        return value
    match = _SIZE.match(value)
    if not match:
        raise ValueError(f"'{value}' is not a size — write 512M, 2G, or 512")
    amount, unit = int(match.group(1)), (match.group(2) or "m").lower()
    return amount * 1024 if unit.startswith("g") else amount


def _seconds(value: Any) -> Any:
    """``2h``, ``30m``, ``90s`` and ``90`` all mean a number of seconds."""
    if not isinstance(value, str):
        return value
    match = _SPAN.match(value)
    if not match:
        raise ValueError(f"'{value}' is not a duration — write 90s, 30m, 2h, or 90")
    return int(match.group(1)) * _SPAN_SECONDS[(match.group(2) or "s").lower()]


class Resources(BaseModel):
    """What one execution of a node needs to have to itself.

    Declaring nothing is the default and means what it always did: the node
    runs on the shared worker pool and nothing is accounted for it. That is
    right for the kind of node most flows are made of — a poll, a threshold, a
    message on its way somewhere.

    It is wrong for the other kind. A numerical library sizes its thread pool
    to every core it can see, so a handful of them at once oversubscribe the
    machine badly enough to starve the engine's own event loop, and a GPU
    library that preallocates most of the card deadlocks when a second one
    arrives. Both are a node saying how much of the machine it takes, which is
    what this is.
    """

    model_config = ConfigDict(extra="forbid")

    cpus: int = Field(
        default=1,
        ge=1,
        description=(
            "Cores held for the whole execution. Also what the thread-pool "
            "variables are set to, so a library sizing itself to the machine "
            "sizes itself to this instead."
        ),
    )
    gpus: int = Field(
        default=0,
        ge=0,
        description=(
            "Whole devices held for the whole execution, named to the node "
            "through CUDA_VISIBLE_DEVICES. Nothing else is given them while it "
            "runs, which is what keeps two preallocating processes apart."
        ),
    )
    ram: int | None = Field(
        default=None,
        ge=1,
        description=(
            "Megabytes held for the whole execution; accepts '512M' or '2G'. "
            "Counted against machines that said how much they have, and "
            "ignored by those that did not — which is a machine with nothing "
            "to say about memory, not one with none."
        ),
    )
    flavor: str | None = Field(
        default=None,
        description=(
            "A stored size by name, standing in for cpus, gpus and ram. Read "
            "again every time the node is built, so editing the flavor edits "
            "what the next run gets."
        ),
    )
    duration_s: int | None = Field(
        default=None,
        ge=1,
        description=(
            "How long one execution is expected to take; accepts '30m' or "
            "'2h'. A statement about the node for whoever is planning around "
            "it, not a limit — the limit is `timeout`."
        ),
    )
    env: dict[str, str] = Field(
        default_factory=dict,
        description=(
            "Extra environment for the worker this node runs in, applied over "
            "what the allocation derives. Where a library's own tuning goes — "
            "XLA_FLAGS, XLA_PYTHON_CLIENT_MEM_FRACTION — since those are "
            "composed strings the engine must not invent."
        ),
    )

    _parse_ram = field_validator("ram", mode="before")(_megabytes)
    _parse_duration = field_validator("duration_s", mode="before")(_seconds)

    @model_validator(mode="after")
    def _a_flavor_says_it_all(self) -> Resources:
        """A flavor and a number for the same thing is two answers.

        Compared by value rather than by what was set, because an editor that
        writes the whole object back sends the defaults with it — and a node
        that says `flavor` and `cpus: 1` has not actually asked for anything
        the flavor does not already cover.
        """
        if self.flavor and (self.cpus != 1 or self.gpus != 0 or self.ram is not None):
            raise ValueError(
                f"flavor '{self.flavor}' already says how much — "
                "drop cpus, gpus and ram, or drop the flavor"
            )
        return self


class NodeDef(BaseModel):
    """A node as stored: identity, configuration and ports.

    Deliberately no canvas position. The editor lays a flow out itself, so
    where a node sits is a fact about the drawing rather than about the flow —
    and a graph nobody can arrange is one worth keeping small.
    """

    id: str
    type: str = "python"
    title: str = ""
    #: This node's settings: constants of its function, stored with the flow.
    #: A function node reads them as keyword arguments beside its ports, so a
    #: setting cannot share a name with one.
    params: dict[str, Any] = Field(default_factory=dict)
    requires: list[MessageSpec] = Field(default_factory=list)
    provides: list[MessageSpec] = Field(default_factory=list)
    #: Name of a shared source in the library, instead of this node's own file.
    #: Editing it edits the copy every flow using it runs.
    source_ref: str | None = None
    timeout: float | None = Field(
        default=None,
        ge=0,
        description=(
            "Seconds this node's code may be silent before it is stopped. A "
            "yield or an emit resets the clock, and the first call's imports "
            "are not charged to it. 0 disables the limit: the node runs until "
            "it finishes, and only a dead worker fails the call. Empty "
            "inherits the engine default."
        ),
    )
    device: str | None = Field(
        default=None,
        description=(
            "Label of the worker this node's code must run on, such as 'gpu'. "
            "Empty means the engine's own workers. A run needing a label no "
            "attached worker carries waits rather than failing."
        ),
    )
    device_policy: Literal["require", "prefer"] = Field(
        default="require",
        description=(
            "What to do when no worker carries `device`: wait for one, or run "
            "locally anyway."
        ),
    )
    cache: bool = Field(
        default=True,
        description=(
            "Whether a batch run may reuse an earlier execution of this node "
            "with the same source, settings and inputs. Turn it off for a "
            "function whose answer can change on its own."
        ),
    )
    resources: Resources | None = Field(
        default=None,
        description=(
            "What one execution of this node holds while it runs. Absent — the "
            "default — means it is not accounted for and shares the engine's "
            "workers, which is right for everything that is not compute-heavy."
        ),
    )
    code_files: dict[str, str] = Field(
        default_factory=dict,
        description=(
            "The project modules this node's function reaches, dotted name to "
            "the file on the machine that synced it. Written by `fluksio "
            "sync`, which is the only side that can import the code and see "
            "what it imports. Empty for a node drawn on the canvas."
        ),
    )
    code_digest: str = Field(
        default="",
        description=(
            "What those modules hashed to at sync. The engine hashes them "
            "again when a run is claimed, and falls back to this when it "
            "cannot see the files — a worker on another machine."
        ),
    )

    @field_validator("id")
    @classmethod
    def _check_id(cls, value: str) -> str:
        return _validate_name(value)


class FlowOrigin(BaseModel):
    """Where a flow was declared, when that was somewhere other than here.

    A flow drawn on the canvas has no origin: the store is where it lives. One
    stamped with this was declared with the decorators in somebody's own
    repository and put here by ``fluksio sync``, so the node bodies below it
    are generated imports and the code they run is versioned twice — once here
    and once there. Its presence is what makes a flow code-defined.

    Deliberately no timestamp. The store commits every change it is given, so
    when a flow was last synced is a fact its own history already holds — and
    one that would otherwise change on every sync, making an unchanged upload
    look like a new version of the flow.
    """

    kind: Literal["python"] = "python"
    #: The repository root on the machine that ran ``sync``.
    repo: str = ""
    #: Its commit, and whether the tree had uncommitted changes at the time —
    #: a run stamped with a dirty commit names code that was never stored.
    commit: str = ""
    dirty: bool = False


class FlowInput(BaseModel):
    """A message the flow starts with rather than computes."""

    spec: MessageSpec
    initial: Any | None = None


class FlowDef(BaseModel):
    """One atomic flow."""

    name: str
    title: str = ""
    nodes: list[NodeDef] = Field(default_factory=list)
    inputs: list[FlowInput] = Field(default_factory=list)
    version: int = 1
    mode: Literal["live", "batch"] = Field(
        default="live",
        description=(
            "A live flow reacts to what arrives: its subscriptions, schedules "
            "and webhooks run until it is stopped. A batch flow only runs when "
            "a run asks it to, from its inputs to its outputs, and is never "
            "activated."
        ),
    )
    outputs: list[str] = Field(
        default_factory=list,
        description=(
            "Messages a batch run reports as its result, unqualified. Empty "
            "means every message the flow ends up holding."
        ),
    )
    origin: FlowOrigin | None = Field(
        default=None,
        description=(
            "Set when the flow was declared in code elsewhere and uploaded by "
            "`fluksio sync`. Absent for a flow drawn on the canvas."
        ),
    )

    @field_validator("name")
    @classmethod
    def _check_name(cls, value: str) -> str:
        return _validate_name(value)


class NodeSource(BaseModel):
    """The Python source of a node."""

    code: str
    #: True when nothing is stored and `code` is the new-node template. An
    #: editor opens on it either way; a client deciding whether somebody wrote
    #: that code needs to know it was nobody. Read-only — set on the way out.
    missing: bool = False


Health = Literal["ok", "degraded", "down"]


class NodeStatusPublic(BaseModel):
    """Whether a node loaded, and how its connection is doing."""

    id: str
    status: str = "active"
    error: str | None = None
    health: Health = "ok"
    health_detail: str | None = None
    #: The node's last runtime failure, kept after it runs again: a failure
    #: that fired an alert should leave a trace of what it was.
    last_error: str = ""
    last_error_ts: float | None = None


class MessageValue(BaseModel):
    """The last payload seen on a message."""

    value: Any = None
    ts: float | None = None


class HistoryPoint(BaseModel):
    """One numeric value a message carried, and when."""

    ts: float
    value: float


class MessageHistory(BaseModel):
    """A message's recent numeric values, oldest first.

    Only numbers are recorded, so ``numeric`` tells the panel whether an empty
    series means "nothing plottable here" or "nothing has arrived yet".
    """

    message: str
    numeric: bool = False
    points: list[HistoryPoint] = Field(default_factory=list)


class FlowSummary(BaseModel):
    name: str
    title: str = ""
    node_count: int = 0
    error_count: int = 0
    has_draft: bool = False
    enabled: bool = True
    paused: bool = False
    # Its background tasks kept crashing, so the engine stopped restarting them.
    quarantined: bool = False
    #: Of the working copy, so publishing from a list needs no second read.
    version: int = 1


class FlowsPublic(BaseModel):
    data: list[FlowSummary]
    count: int


class LibraryNode(BaseModel):
    """A node source shared across flows, and who is using it."""

    name: str
    used_by: list[str] = Field(default_factory=list)


class FlowStatePublic(BaseModel):
    values: dict[str, MessageValue] = Field(default_factory=dict)
    nodes: list[NodeStatusPublic] = Field(default_factory=list)


class ModulePackage(BaseModel):
    """One package installed in the venv node code runs on."""

    name: str
    version: str


class ModulesInfo(BaseModel):
    """The venv node code imports from, and the manifest that describes it."""

    python_version: str = ""
    venv_path: str = ""
    requirements: str = ""
    packages: list[ModulePackage] = Field(default_factory=list)
    #: Whether what is installed matches the manifest.
    applied: bool = False
    #: True when node code runs on the venv Fluksio was installed into rather
    #: than one the engine built. That venv belongs to whoever made it, so the
    #: manifest does not describe it and nothing here installs into it.
    adopted: bool = False


class ApplyRequest(BaseModel):
    """A pip manifest, one requirement per line."""

    requirements: str = ""


class ApplyResult(BaseModel):
    ok: bool
    output: str = ""


class BrainNode(BaseModel):
    """One neuron: a thing the engine talks to, or a node that only computes.

    Nodes of the same type pointing at the same outside thing — one broker
    topic, one URL, one bucket — are a single entry here, whichever flows they
    sit in. ``members`` are the ``flow.node_id`` names behind it, which is also
    what the live events are keyed by.
    """

    id: str
    label: str
    kind: str
    members: list[str] = Field(default_factory=list)
    flows: list[str] = Field(default_factory=list)
    issue: str | None = Field(
        default=None,
        description=(
            "Why this neuron cannot run, if validation found something. A "
            "failure the engine hits while running arrives over the socket "
            "instead; this is the part that is already true before anything "
            "fires, and so has to travel with the graph."
        ),
    )


class BrainEdge(BaseModel):
    """Messages carrying values from one neuron to another."""

    source: str
    target: str
    messages: list[str] = Field(default_factory=list)


class BrainGraph(BaseModel):
    """Every published flow at once, merged on what its nodes talk to."""

    nodes: list[BrainNode] = Field(default_factory=list)
    edges: list[BrainEdge] = Field(default_factory=list)


class NodeTypeInfo(BaseModel):
    """A node type the editor can offer, with its parameter schema."""

    type: str
    title: str
    description: str
    params_schema: dict[str, Any] = Field(default_factory=dict)
    has_source: bool = False
    #: Whether this type takes settings beyond the ones its schema declares.
    #: A function node's settings are its author's to name, and reach `process`
    #: as keyword arguments beside its ports.
    free_params: bool = False
    #: The package a connector came from; empty for the built-in types.
    plugin: str | None = None