diff --git a/DashAI/alembic/versions/82fb7a6b8ac2_merge_pipeline_tracking_and_cross_.py b/DashAI/alembic/versions/82fb7a6b8ac2_merge_pipeline_tracking_and_cross_.py new file mode 100644 index 000000000..cdb31300c --- /dev/null +++ b/DashAI/alembic/versions/82fb7a6b8ac2_merge_pipeline_tracking_and_cross_.py @@ -0,0 +1,37 @@ +"""Merge the pipeline tracking and cross-validation heads. + +Two lines of work reached the database at the same time and neither knew about +the other, so the revision graph ended with two heads and ``upgrade head`` +refused to pick one: + +``e7b4d1a9c206`` added the pipeline run tracking tables (``pipeline_run``, +``pipeline_node_run``, ``pipeline_node_artifact``); ``b7e4d2a19c63`` is the tip +of the branch that brought cross-validation and the evaluation reports, which +added ``model_session.evaluation_strategy``, ``run.nested``, the fold columns on +``metric`` and the ``report`` table. + +There is nothing to do here. The two sides touch disjoint tables and disjoint +columns, so joining the graph is the whole of the change: this revision exists +to give the two heads a single successor. + +Revision ID: 82fb7a6b8ac2 +Revises: b7e4d2a19c63, e7b4d1a9c206 +Create Date: 2026-09-10 13:03:45.322840 + +""" + +from typing import Sequence, Union + +# revision identifiers, used by Alembic. +revision: str = "82fb7a6b8ac2" +down_revision: Union[str, None] = ("b7e4d2a19c63", "e7b4d1a9c206") +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Join the two revision lines. Neither side needs anything applied.""" + + +def downgrade() -> None: + """Split them again, which is likewise nothing to undo.""" diff --git a/DashAI/alembic/versions/e7b4d1a9c206_add_pipeline_run_tracking.py b/DashAI/alembic/versions/e7b4d1a9c206_add_pipeline_run_tracking.py new file mode 100644 index 000000000..fdbeb079e --- /dev/null +++ b/DashAI/alembic/versions/e7b4d1a9c206_add_pipeline_run_tracking.py @@ -0,0 +1,122 @@ +"""Add pipeline run tracking tables. + +Separates the definition of a graph from its executions. ``Pipeline`` keeps the +definition, which changes as the user edits it; ``pipeline_run`` freezes the +steps and edges of one execution, ``pipeline_node_run`` tracks each node of it, +and ``pipeline_node_artifact`` holds what a node emitted, keyed by a key from +the unit's PROVIDES rather than by a column per node type. + +The three JSON result columns on ``pipeline`` (exploration, train, prediction) +are deliberately left in place: the pipelines endpoints and the front's results +view still read them, so removing them belongs with rewriting those. + +Revision ID: e7b4d1a9c206 +Revises: d5b3c8f2a041 +Create Date: 2026-08-20 00:00:00.000000 +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +revision: str = "e7b4d1a9c206" +down_revision: Union[str, None] = "d5b3c8f2a041" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "pipeline_run", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("pipeline_id", sa.Integer(), nullable=False), + sa.Column("steps", sa.JSON(), nullable=True), + sa.Column("edges", sa.JSON(), nullable=True), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("last_modified", sa.DateTime(), nullable=False), + sa.Column("delivery_time", sa.DateTime(), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "status", + sa.Enum( + "NOT_STARTED", + "DELIVERED", + "STARTED", + "FINISHED", + "ERROR", + name="pipelinerunstatus", + ), + nullable=False, + ), + sa.Column("error_message", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["pipeline_id"], + ["pipeline.id"], + name=op.f("fk_pipeline_run_pipeline_id_pipeline"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_pipeline_run")), + ) + + op.create_table( + "pipeline_node_run", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("pipeline_run_id", sa.Integer(), nullable=False), + sa.Column("node_id", sa.String(), nullable=False), + sa.Column("block_id", sa.String(), nullable=False), + sa.Column("node_type", sa.String(), nullable=False), + sa.Column("config", sa.JSON(), nullable=True), + sa.Column("input", sa.JSON(), nullable=True), + sa.Column("output", sa.JSON(), nullable=True), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("last_modified", sa.DateTime(), nullable=False), + sa.Column("delivery_time", sa.DateTime(), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "status", + sa.Enum( + "NOT_STARTED", + "DELIVERED", + "STARTED", + "FINISHED", + "ERROR", + "CANCELLED", + name="noderunstatus", + ), + nullable=False, + ), + sa.Column("error_message", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["pipeline_run_id"], + ["pipeline_run.id"], + name=op.f("fk_pipeline_node_run_pipeline_run_id_pipeline_run"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_pipeline_node_run")), + ) + + op.create_table( + "pipeline_node_artifact", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("node_run_id", sa.Integer(), nullable=False), + sa.Column("key", sa.String(), nullable=False), + sa.Column("value", sa.JSON(), nullable=True), + sa.Column("created", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["node_run_id"], + ["pipeline_node_run.id"], + name=op.f("fk_pipeline_node_artifact_node_run_id_pipeline_node_run"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_pipeline_node_artifact")), + ) + + +def downgrade() -> None: + op.drop_table("pipeline_node_artifact") + op.drop_table("pipeline_node_run") + op.drop_table("pipeline_run") diff --git a/DashAI/back/api/api_v1/endpoints/predict.py b/DashAI/back/api/api_v1/endpoints/predict.py index da3778af3..3f1cf58f6 100644 --- a/DashAI/back/api/api_v1/endpoints/predict.py +++ b/DashAI/back/api/api_v1/endpoints/predict.py @@ -330,7 +330,6 @@ async def delete_prediction( @inject async def preview_manual_prediction( request: Request, - component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), ): """Run a synchronous manual prediction and return results without persisting. @@ -413,7 +412,6 @@ async def preview_manual_prediction( run_manual_prediction, run_id=run_id_int, manual_input_data=rows_data, - component_registry=component_registry, session_factory=session_factory, ) return {"columns": columns, "rows": rows} diff --git a/DashAI/back/core/enums/status.py b/DashAI/back/core/enums/status.py index 20ec6f22e..0f31d99f1 100644 --- a/DashAI/back/core/enums/status.py +++ b/DashAI/back/core/enums/status.py @@ -67,3 +67,24 @@ class DatafileStatus(Enum): DOWNLOADING = "downloading" READY = "ready" ERROR = "error" + + +class PipelineRunStatus(Enum): + NOT_STARTED = 0 + DELIVERED = 1 + STARTED = 2 + FINISHED = 3 + ERROR = 4 + + +class NodeRunStatus(Enum): + NOT_STARTED = 0 + DELIVERED = 1 + STARTED = 2 + FINISHED = 3 + ERROR = 4 + # A node that never ran because an earlier one failed. Distinct from + # NOT_STARTED, which is a node still waiting its turn: without the + # distinction a run that died halfway is indistinguishable from one still + # in flight. + CANCELLED = 5 diff --git a/DashAI/back/dag/__init__.py b/DashAI/back/dag/__init__.py new file mode 100644 index 000000000..f144f1d12 --- /dev/null +++ b/DashAI/back/dag/__init__.py @@ -0,0 +1,28 @@ +"""A sequential DAG engine over the atomic units. + +Nothing here imports from ``DashAI/back/pipeline/``. That subsystem does not +run — its nodes implement two of ``BaseJob``'s four abstract methods, so they +cannot even be instantiated — and it is replaced rather than repaired. + +The engine is deliberately sequential. Its predecessor was concurrent, and the +concurrency was not buying what it cost: the most expensive node in a pipeline +was already serialised behind an exclusive lock, every context and database +write was wrapped in another, and SQLite answered the remaining concurrent +writes with "database is locked" often enough to need retries with exponential +backoff. Units open a database session each, which is safe in sequence and was +the source of that contention in parallel. +""" + +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node, connect, sinks +from DashAI.back.dag.validate import resolve_unit_class, validate + +__all__ = [ + "Edge", + "Graph", + "GraphError", + "Node", + "connect", + "resolve_unit_class", + "sinks", + "validate", +] diff --git a/DashAI/back/dag/engine.py b/DashAI/back/dag/engine.py new file mode 100644 index 000000000..2d1a3a221 --- /dev/null +++ b/DashAI/back/dag/engine.py @@ -0,0 +1,184 @@ +"""Sequential execution of a validated graph. + +The engine gives every node an ``ExecutionContext`` of its own, preloaded with +exactly the keys that node declares in ``REQUIRES``, and the renaming lives on +the edge. There is no shared context: merging the outputs of two predecessors +into one dictionary loses a key whenever both publish the same one, which is +how its predecessor lost values silently and why it grew a parallel list of raw +branch dictionaries to work around itself. + +Nothing here touches the database. Tracking goes through a sink, so the engine +can run with no persistence at all — which is how it is tested. +""" + +import logging +from typing import Any, Dict, List, Mapping, Optional, Protocol + +from DashAI.back.dag.graph import Graph, GraphError +from DashAI.back.dag.validate import instantiate, resolve_unit_class, validate +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class TrackingSink(Protocol): + """Where a run reports what it is doing. + + Every method is a notification, not a decision: a sink may persist, log, or + do nothing, and the engine's behaviour does not depend on which. + """ + + def run_started(self, order: List[str]) -> None: + """The graph passed validation and is about to run, in this order.""" + + def node_started(self, node_id: str, payload: Mapping[str, Any]) -> None: + """``payload`` is the serializable half of the node's input context.""" + + def node_finished( + self, node_id: str, artifacts: Mapping[str, Any], payload: Mapping[str, Any] + ) -> None: + """``artifacts`` are the node's serializable outputs, by PROVIDES key.""" + + def node_failed(self, node_id: str, message: str) -> None: ... + + def nodes_cancelled(self, node_ids: List[str]) -> None: + """Nodes that will never run because an earlier one failed.""" + + def run_finished(self) -> None: ... + + def run_failed(self, message: str) -> None: ... + + +class NullSink: + """A sink that records nothing. The engine's default.""" + + def run_started(self, order: List[str]) -> None: + pass + + def node_started(self, node_id: str, payload: Mapping[str, Any]) -> None: + pass + + def node_finished( + self, node_id: str, artifacts: Mapping[str, Any], payload: Mapping[str, Any] + ) -> None: + pass + + def node_failed(self, node_id: str, message: str) -> None: + pass + + def nodes_cancelled(self, node_ids: List[str]) -> None: + pass + + def run_finished(self) -> None: + pass + + def run_failed(self, message: str) -> None: + pass + + +def artifacts_of(unit_class: type, ctx: ExecutionContext) -> Dict[str, Any]: + """The outputs of a node that can be persisted, by ``PROVIDES`` key. + + Only the reference half: the cache holds live datasets, models and tasks, + which are not serializable and are always derivable again. A unit whose + real output is on disk publishes the path as a reference, so the path is + what gets recorded — which is the whole of what a caller needs later. + + ``origin`` is what makes this cheap. Asking ``key in ctx.refs`` answers the + same question but deep-copies every reference the context holds, once per + key. + """ + return { + key: ctx.get(key) for key in unit_class.PROVIDES if ctx.origin(key) == "ref" + } + + +def run( + graph: Graph, sink: Optional[TrackingSink] = None +) -> Dict[str, ExecutionContext]: + """Execute a graph sequentially, returning the contexts still in use. + + Parameters + ---------- + graph : Graph + The graph to run. It is validated first, so a graph that cannot work + fails before any unit does. + sink : Optional[TrackingSink] + Where progress is reported. Defaults to recording nothing. + + Returns + ------- + Dict[str, ExecutionContext] + The context of every node whose values were still needed at the end, + keyed by node id. Contexts nobody reads any more are dropped as the run + proceeds, so intermediate datasets do not pile up for the length of the + run. + + Raises + ------ + GraphError + If the graph does not validate. + """ + sink = sink if sink is not None else NullSink() + + order = validate(graph) + by_id = graph.by_id() + classes = {node.id: resolve_unit_class(node.unit) for node in graph.nodes} + + # One instance per node, built before anything runs: a configuration that + # cannot build its unit is a mistake in the graph, not a failure of a node + # halfway through it. + units = {node.id: instantiate(node) for node in graph.nodes} + + consumers: Dict[str, int] = dict.fromkeys(by_id, 0) + for edge in graph.edges: + consumers[edge.src] += 1 + + sink.run_started(order) + + contexts: Dict[str, ExecutionContext] = {} + for position, node_id in enumerate(order): + ctx = ExecutionContext() + + for edge in graph.edges: + if edge.dst != node_id: + continue + source = contexts[edge.src] + # The half a value lives in decides how it may be moved: put_ref + # validates with json.dumps and deep-copies, put stores the live + # object by reference. A dataset handed to put_ref raises, and a + # reference handed to put loses its copy-on-read guarantee. + if source.origin(edge.src_key) == "ref": + ctx.put_ref(edge.dst_key, source.get(edge.src_key)) + else: + ctx.put(edge.dst_key, source.require(edge.src_key)) + + sink.node_started(node_id, ctx.to_dict()) + try: + units[node_id](ctx) + except Exception as e: + log.exception(e) + sink.node_failed(node_id, str(e)) + # Everything after this point in the order never runs. Reporting + # them keeps a run that died halfway distinguishable from one still + # in flight. + sink.nodes_cancelled(list(order[position + 1 :])) + sink.run_failed(str(e)) + raise + + contexts[node_id] = ctx + sink.node_finished(node_id, artifacts_of(classes[node_id], ctx), ctx.to_dict()) + + # Release what nothing downstream will read again. + for edge in graph.edges: + if edge.dst != node_id: + continue + consumers[edge.src] -= 1 + if consumers[edge.src] == 0 and edge.src in contexts: + contexts.pop(edge.src).clear_cache() + + sink.run_finished() + return contexts + + +__all__ = ["GraphError", "NullSink", "TrackingSink", "artifacts_of", "run"] diff --git a/DashAI/back/dag/expand.py b/DashAI/back/dag/expand.py new file mode 100644 index 000000000..f5f300183 --- /dev/null +++ b/DashAI/back/dag/expand.py @@ -0,0 +1,217 @@ +"""Turning what a canvas holds into the graph the engine runs. + +A canvas cannot show a training as six nodes and fifteen edges, so a block on +it stands for a sequence of units. The engine works on units, so a run is over +the expanded graph, and it is the expanded graph a run freezes: replaying one +never depends on re-running the expansion that produced it. + +The first version has one unit per block, and expansion is close to a rename. +The shape is here from the start because the persistence has to be right before +there is anything to migrate. +""" + +import re +from typing import Any, Dict, List, Mapping, Sequence + +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node +from DashAI.back.dag.validate import resolve_unit_class + +#: Anything outside what can name a directory under RUNS_PATH, plus the escape +#: character itself. ``SaveModelUnit.validate`` refuses the rest, so a prefix +#: built from a node id a user chose has to be brought into that alphabet here. +#: +#: ``_`` is in here on purpose. Replacing unsafe characters with a fixed ``_`` +#: is not injective -- ``save.a`` and ``save_a`` would collapse to the same +#: name, and the second saving node would overwrite the first one's model +#: directory in silence. Escaping to ``_`` instead, with ``_`` itself +#: escaped, is reversible, so two different node ids cannot produce one prefix. +_NEEDS_ESCAPE = re.compile(r"[^A-Za-z0-9-]") + +#: The runtime params this module knows how to answer. A unit declares which +#: ones it takes in its own ``RUNTIME_PARAMS``; these are the two that are about +#: the pipeline run rather than about a job, so a pipeline is what supplies them. +#: Any other runtime param belongs to whoever else runs that unit. +ENGINE_SUPPLIED = ("artifact_prefix", "run_id") + + +def artifact_prefix(pipeline_run_id: int, node_id: str) -> str: + """Name this node's artifacts so nothing else can collide with them. + + The runs directory is shared with every real ``Run``, and a pipeline run id + is a different sequence from a run id: both start at 1, so a pipeline that + used its own id as the name would write over the model of the run with that + id. Not a risk -- a certainty. The ``pipeline-`` prefix is what keeps the + two apart, the run id keeps two executions of the same pipeline apart, and + the node id keeps two saving nodes in one graph apart. + + The node id is escaped rather than cleaned, so the mapping is injective and + two different ids can never produce one prefix. + """ + escaped = _NEEDS_ESCAPE.sub( + lambda match: f"_{ord(match.group()):02x}", str(node_id) + ) + return f"pipeline-{pipeline_run_id}-{escaped}" + + +def expand( + steps: Sequence[Mapping[str, Any]], + edges: Sequence[Mapping[str, Any]], + pipeline_run_id: int, +) -> Graph: + """Expand a canvas into a unit-level graph. + + Parameters + ---------- + steps : Sequence[Mapping[str, Any]] + The blocks, each ``{id, units: [{id, unit, config}], ...}``. + edges : Sequence[Mapping[str, Any]] + The wires between blocks, each ``{source, target}``. + pipeline_run_id : int + The run being expanded, which is what names its artifacts. + + Returns + ------- + Graph + Nodes are units, and every node carries the block it came from. + + Raises + ------ + GraphError + If a block is malformed, or if it holds more than one unit — see the + note below. + + Notes + ----- + A block holding several units is refused rather than guessed at. Wiring + across a boundary where either side is a sequence has more than one + defensible answer, and the rule that bundles keys does not cover every wire + a real graph needs: ``ApplyConverterUnit`` deliberately does not republish + ``dataset_id``, so a node needing both the converted dataset and the id + takes the id on an edge that skips the converter. Choosing a rule for that + without a case to check it against would be inventing semantics. The + persistence already carries N units per block, so adding it later costs no + migration. + """ + blocks = _read_blocks(steps) + + nodes: List[Node] = [] + for block_id, units in blocks.items(): + for unit in units: + nodes.append( + Node( + id=unit["id"], + unit=unit["unit"], + config=_with_engine_config( + unit["unit"], + unit.get("config") or {}, + pipeline_run_id, + unit["id"], + ), + block_id=block_id, + ) + ) + + wires: List[Edge] = [] + for edge in edges: + source, target = edge.get("source"), edge.get("target") + if source not in blocks or target not in blocks: + raise GraphError( + f"An edge connects '{source}' to '{target}', and one of them is " + "not a block in this pipeline." + ) + wires.extend(_between(blocks[source][-1], blocks[target][0])) + + return Graph(nodes, wires) + + +def _read_blocks( + steps: Sequence[Mapping[str, Any]], +) -> Dict[str, List[Mapping[str, Any]]]: + """Validate the shape of the blocks and index them by id.""" + blocks: Dict[str, List[Mapping[str, Any]]] = {} + + for step in steps: + block_id = step.get("id") + if not block_id: + raise GraphError("A block has no id.") + if block_id in blocks: + raise GraphError(f"There is more than one block with the id '{block_id}'.") + + units = step.get("units") + if not units: + raise GraphError( + f"Block '{block_id}' declares no units. Blocks saved by the " + "previous pipeline subsystem name a node type and carry a " + "single config instead, and cannot be run by this engine." + ) + if len(units) > 1: + raise GraphError( + f"Block '{block_id}' holds {len(units)} units. Only one unit " + "per block is supported so far." + ) + for unit in units: + if not unit.get("id") or not unit.get("unit"): + raise GraphError( + f"A unit in block '{block_id}' is missing its id or its unit name." + ) + + blocks[block_id] = list(units) + + return blocks + + +def _with_engine_config( + unit_name: str, + config: Mapping[str, Any], + pipeline_run_id: int, + node_id: str, +) -> Dict[str, Any]: + """Set the configuration the engine owns rather than the user. + + Whatever the stored graph holds for these fields is **discarded**, not + merged. Neither is a value a user could get right: + + * ``artifact_prefix`` is built from the id of the run, and there is no run + when a node is configured -- so any value already there was chosen + without the one thing that decides it. Overriding is also what makes two + nodes sharing a prefix impossible rather than merely unlikely, so the + collision needs no validator to catch it. + * ``run_id`` is null because a pipeline has no ``Run`` row. Forcing it is + what makes the sandbox mechanical instead of conventional: a stored graph + cannot point a training node at a real run and start writing ``Metric`` + rows into it. + + Both are declared in the unit's ``RUNTIME_PARAMS`` rather than in its + schema, so a form never offers them in the first place. This is the other + half of the same statement, on the side a hand-edited row could still reach. + """ + resolved = dict(config) + takes = set(resolve_unit_class(unit_name).RUNTIME_PARAMS) + + if "artifact_prefix" in takes: + resolved["artifact_prefix"] = artifact_prefix(pipeline_run_id, node_id) + + if "run_id" in takes: + resolved["run_id"] = None + + return resolved + + +def _between(source: Mapping[str, Any], target: Mapping[str, Any]) -> List[Edge]: + """Wire the keys the two units agree on: ``PROVIDES`` against ``REQUIRES``. + + One drawn edge stands for the whole set, because a canvas cannot draw one + per key. + """ + provides = set(resolve_unit_class(source["unit"]).PROVIDES) + requires = set(resolve_unit_class(target["unit"]).REQUIRES) + shared = sorted(provides & requires) + + if not shared: + raise GraphError( + f"'{source['id']}' ({source['unit']}) has nothing " + f"'{target['id']}' ({target['unit']}) needs, so the edge between " + "them would carry nothing." + ) + + return [Edge(source["id"], key, target["id"], key) for key in shared] diff --git a/DashAI/back/dag/graph.py b/DashAI/back/dag/graph.py new file mode 100644 index 000000000..7ee2cd75a --- /dev/null +++ b/DashAI/back/dag/graph.py @@ -0,0 +1,180 @@ +"""The shape of a graph: nodes, edges, and the rule that bundles them.""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Mapping, Optional, Sequence, Set, Tuple + + +class GraphError(Exception): + """Raised when a graph is not executable as declared.""" + + +@dataclass(frozen=True) +class Node: + """One unit placed in a graph. + + A node carries the *name* of its unit rather than an instance. Unit + instances hold state of their own — ``FitModelUnit`` memoizes its resolved + optimizer, ``FitConverterUnit`` its converter class — so two nodes sharing + one instance would have the second silently run with the first one's. The + engine builds one instance per node from ``(unit, config)``, which is also + exactly what a persisted graph can hold. + + Parameters + ---------- + id : str + Identifier of this node inside the graph. + unit : str + Class name of the unit, as the component registry knows it. + config : Mapping[str, Any] + The unit's own configuration, validated by its ``SCHEMA``. Constants a + caller chooses live here — not in the context, where nothing upstream + could ever supply them. + block_id : Optional[str] + The visual block this node belongs to, for an editor that groups + several units into one shape on a canvas. Left out, it becomes ``id``: + a node always belongs to some block, so nothing downstream — the + tracking column least of all — has a null to handle, and a node has one + representation rather than two. + """ + + id: str + unit: str + config: Mapping[str, Any] = field(default_factory=dict) + block_id: Optional[str] = None + + def __post_init__(self) -> None: + if self.block_id is None: + object.__setattr__(self, "block_id", self.id) + + +@dataclass(frozen=True) +class Edge: + """One key travelling from a node's output to another's input. + + The pair of names is what makes the fixed strings in ``REQUIRES`` and + ``PROVIDES`` behave as port names: ``src_key`` is drawn from the source's + ``PROVIDES`` and ``dst_key`` from the target's ``REQUIRES``, and the edge + maps one onto the other. That is what lets two ``LoadDatasetUnit`` nodes, + which both write the fixed key ``dataset``, coexist in one graph — the + renaming lives here rather than in the unit. + """ + + src: str + src_key: str + dst: str + dst_key: str + + +@dataclass(frozen=True) +class Graph: + """A set of nodes and the edges between them.""" + + nodes: Sequence[Node] + edges: Sequence[Edge] + + def node(self, node_id: str) -> Node: + for node in self.nodes: + if node.id == node_id: + return node + raise GraphError(f"Unknown node id: {node_id}") + + def by_id(self) -> Dict[str, Node]: + return {node.id: node for node in self.nodes} + + +def connect(src: Node, dst: Node) -> Tuple[Edge, ...]: + """Wire every key two nodes agree on: ``PROVIDES`` against ``REQUIRES``. + + A canvas cannot draw one edge per key — ``FitModelUnit`` alone requires + seven — so one drawn edge between two nodes stands for this whole set. + + The rule does not cover every wire a real graph needs, and that is by + design rather than an omission: ``ApplyConverterUnit`` does not republish + ``dataset_id``, because nothing derived from the object being transformed + may cross the context boundary. A downstream node that needs both the + converted dataset and the id therefore takes the id on an edge that skips + the converter entirely. + """ + from DashAI.back.dag.validate import resolve_unit_class + + provides = set(resolve_unit_class(src.unit).PROVIDES) + requires = set(resolve_unit_class(dst.unit).REQUIRES) + return tuple(Edge(src.id, key, dst.id, key) for key in sorted(provides & requires)) + + +def sinks(graph: Graph) -> Set[str]: + """Node ids nothing consumes. They are what a run exists to produce.""" + consumed = {edge.src for edge in graph.edges} + return {node.id for node in graph.nodes} - consumed + + +def dump(graph: Graph) -> Tuple[list, list]: + """Serialize a graph to the plain data a run freezes. + + Unit level, not block level: a run records the graph it actually executed, + so replaying it never depends on re-running the expansion that produced it. + """ + steps = [ + { + "id": node.id, + "block_id": node.block_id, + "unit": node.unit, + "config": dict(node.config), + } + for node in graph.nodes + ] + edges = [ + { + "src": edge.src, + "src_key": edge.src_key, + "dst": edge.dst, + "dst_key": edge.dst_key, + } + for edge in graph.edges + ] + return steps, edges + + +def load( + steps: Sequence[Mapping[str, Any]], edges: Sequence[Mapping[str, Any]] +) -> Graph: + """Rebuild a graph from what :func:`dump` produced. + + Raises + ------ + GraphError + If a step or an edge is missing a field. Rows written by the previous + subsystem land here: their steps carry a node ``type`` and a single + ``config``, with no unit to resolve, so they are refused with a message + that says so rather than failing somewhere deeper. + """ + try: + nodes = [ + Node( + id=step["id"], + unit=step["unit"], + config=step.get("config") or {}, + block_id=step.get("block_id"), + ) + for step in steps + ] + except KeyError as e: + raise GraphError( + f"A step is missing the field {e}. Steps saved by the previous " + "pipeline subsystem name a node type rather than a unit and cannot " + "be run by this engine." + ) from e + + try: + wires = [ + Edge(edge["src"], edge["src_key"], edge["dst"], edge["dst_key"]) + for edge in edges + ] + except KeyError as e: + raise GraphError( + f"An edge is missing the field {e}. Edges saved by the previous " + "pipeline subsystem connect nodes without naming the keys they " + "carry, so they cannot be run by this engine." + ) from e + + return Graph(nodes, wires) diff --git a/DashAI/back/dag/tracking.py b/DashAI/back/dag/tracking.py new file mode 100644 index 000000000..929079b52 --- /dev/null +++ b/DashAI/back/dag/tracking.py @@ -0,0 +1,155 @@ +"""Persistence of a run: the sink the engine reports to. + +Kept apart from the engine on purpose. The engine decides what happens and in +what order; this decides what is written down about it, and the engine runs +identically with a sink that writes nothing. + +The run row is created by the caller, not here. Naming a run's artifacts needs +its id, and that has to happen before the graph is built rather than after, so +the caller creates the row, expands with its id, and hands the id over. This +writes progress into it and nothing else -- a graph that fails to validate +leaves the row exactly as the caller made it, and saying why is the caller's +job. +""" + +import logging +from typing import Any, Dict, List, Mapping + +from DashAI.back.dag.graph import Graph, dump + +log = logging.getLogger(__name__) + + +class DatabaseSink: + """Writes a run to ``PipelineRun``, ``NodeRun`` and ``NodeArtifact``. + + Each notification opens its own short session and commits, so the canvas + can colour a node the moment it changes rather than at the end of the run. + That is safe here and was not in the concurrent predecessor, which needed a + process-wide lock and retries with backoff to survive SQLite answering + "database is locked". + + Parameters + ---------- + pipeline_run_id : int + The run to write into. Already created by the caller. + graph : Graph + The expanded, unit-level graph. It is frozen into the run, so the run + stays readable after the pipeline it came from is edited. + """ + + def __init__(self, pipeline_run_id: int, graph: Graph) -> None: + self.pipeline_run_id = pipeline_run_id + self._graph = graph + self._nodes = graph.by_id() + #: node id -> NodeRun id, so later notifications need no lookup by name. + self._node_run_ids: Dict[str, int] = {} + + @property + def _session_factory(self): + from kink import di + + return di["session_factory"] + + def run_started(self, order: List[str]) -> None: + """Freeze the graph into the run, and create a row per node. + + Every node gets a row up front, in ``NOT_STARTED``. Creating them on + demand instead would leave a node that never ran with no row at all, + and "no row" cannot distinguish a node still waiting from one whose run + died before reaching it. + """ + from DashAI.back.dependencies.database.models import NodeRun, PipelineRun + + steps, edges = dump(self._graph) + + with self._session_factory() as db: + pipeline_run = db.get(PipelineRun, self.pipeline_run_id) + pipeline_run.steps = steps + pipeline_run.edges = edges + pipeline_run.set_status_as_started() + db.commit() + + node_runs = [ + NodeRun( + pipeline_run_id=self.pipeline_run_id, + node_id=node_id, + block_id=self._nodes[node_id].block_id, + node_type=self._nodes[node_id].unit, + config=dict(self._nodes[node_id].config), + ) + for node_id in order + ] + db.add_all(node_runs) + db.commit() + self._node_run_ids = {row.node_id: row.id for row in node_runs} + + def node_started(self, node_id: str, payload: Mapping[str, Any]) -> None: + from DashAI.back.dependencies.database.models import NodeRun + + with self._session_factory() as db: + node_run = db.get(NodeRun, self._node_run_ids[node_id]) + node_run.set_status_as_started() + node_run.input = dict(payload) + db.commit() + + def node_finished( + self, node_id: str, artifacts: Mapping[str, Any], payload: Mapping[str, Any] + ) -> None: + """Record the node as finished, with one artifact row per output key. + + The key is a key from the unit's own ``PROVIDES``. That is what + replaced a column per node type: adding a kind of node used to mean + adding a column to the pipeline table and a branch to an ``if + node_type`` chain. + """ + from DashAI.back.dependencies.database.models import NodeArtifact, NodeRun + + node_run_id = self._node_run_ids[node_id] + with self._session_factory() as db: + node_run = db.get(NodeRun, node_run_id) + node_run.set_status_as_finished() + node_run.output = dict(payload) + db.add_all( + [ + NodeArtifact(node_run_id=node_run_id, key=key, value=value) + for key, value in artifacts.items() + ] + ) + db.commit() + + def node_failed(self, node_id: str, message: str) -> None: + from DashAI.back.dependencies.database.models import NodeRun + + with self._session_factory() as db: + node_run = db.get(NodeRun, self._node_run_ids[node_id]) + node_run.set_status_as_error(message) + db.commit() + + def nodes_cancelled(self, node_ids: List[str]) -> None: + from DashAI.back.dependencies.database.models import NodeRun + + if not node_ids: + return + + with self._session_factory() as db: + for node_id in node_ids: + node_run = db.get(NodeRun, self._node_run_ids[node_id]) + node_run.set_status_as_cancelled() + db.commit() + + def run_finished(self) -> None: + from DashAI.back.dependencies.database.models import PipelineRun + + with self._session_factory() as db: + pipeline_run = db.get(PipelineRun, self.pipeline_run_id) + pipeline_run.set_status_as_finished() + db.commit() + + def run_failed(self, message: str) -> None: + from DashAI.back.dependencies.database.models import PipelineRun + + with self._session_factory() as db: + pipeline_run = db.get(PipelineRun, self.pipeline_run_id) + pipeline_run.set_status_as_error(message) + db.commit() diff --git a/DashAI/back/dag/validate.py b/DashAI/back/dag/validate.py new file mode 100644 index 000000000..435bf7ae7 --- /dev/null +++ b/DashAI/back/dag/validate.py @@ -0,0 +1,199 @@ +"""Static validation of a graph, and the registry lookup it rests on. + +Everything here runs before any unit does. ``REQUIRES`` and ``PROVIDES`` are +class attributes, so the whole check needs the unit *classes* and never an +instance: a graph that cannot work is rejected without a dataset being read or +a model being built. +""" + +from typing import Dict, List, Mapping + +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node + + +def resolve_unit_class(name: str) -> type: + """Look a unit class up in the component registry by name. + + Parameters + ---------- + name : str + Class name of the unit. + + Returns + ------- + type + The unit class. Not an instance: this is what static validation reads + ``REQUIRES`` and ``PROVIDES`` off, and building an instance would mean + running the unit's own configuration checks too early. + + Raises + ------ + GraphError + If the name is unknown, or names something that is not a unit. + """ + from kink import di + + from DashAI.back.units.base_unit import BaseUnit + + component_registry = di["component_registry"] + + try: + unit_class = component_registry[name]["class"] + except Exception as e: + raise GraphError(f"There is no unit named '{name}'.") from e + + if not (isinstance(unit_class, type) and issubclass(unit_class, BaseUnit)): + raise GraphError( + f"'{name}' is registered but it is not a unit, so it cannot be a " + "node in a pipeline." + ) + + return unit_class + + +def instantiate(node: Node): + """Build this node's own unit instance. + + One instance per node, never shared: a unit keeps state on itself — a + memoized optimizer, a resolved converter class — so two nodes sharing an + instance would have the second silently run with the first one's. + + Raises + ------ + GraphError + If the configuration does not build the unit. + """ + unit_class = resolve_unit_class(node.unit) + try: + return unit_class(**dict(node.config)) + except Exception as e: + raise GraphError( + f"Node '{node.id}' could not be configured as a {node.unit}: {e}" + ) from e + + +def validate(graph: Graph) -> List[str]: + """Check the graph and return a topological execution order. + + Every problem found is reported at once rather than one per run, because a + user fixing a graph on a canvas wants the whole list. + + Two kinds of missing input are caught here: a context key no edge feeds, + and a runtime param nobody supplied. + + There is no escape hatch for an unfed input. Every key in ``REQUIRES`` has + a unit somewhere that publishes it — the audit in + ``tests/back/units/test_unit_contracts.py`` enforces that — so "nothing + supplies this key" means a missing edge and nothing else. That was not true + while ``run_id`` was a context key no unit published: an injected constant + and a wire the user forgot to draw were indistinguishable. + + Raises + ------ + GraphError + If any check fails, listing every problem found. + """ + problems: List[str] = [] + + ids = [node.id for node in graph.nodes] + duplicates = sorted({node_id for node_id in ids if ids.count(node_id) > 1}) + if duplicates: + problems.append(f"Duplicate node ids in the graph: {duplicates}.") + + by_id = graph.by_id() + + # Resolve every unit up front: a name that is not a unit makes every + # contract check below meaningless, so it is reported on its own. + classes: Dict[str, type] = {} + for node in graph.nodes: + try: + classes[node.id] = resolve_unit_class(node.unit) + except GraphError as e: + problems.append(f"Node '{node.id}': {e}") + + if problems and not classes: + raise GraphError("\n".join(problems)) + + inbound: Dict[str, Dict[str, List[Edge]]] = {node_id: {} for node_id in by_id} + + for edge in graph.edges: + if edge.src not in by_id or edge.dst not in by_id: + problems.append(f"Edge {edge} points at a node that is not in the graph.") + continue + if edge.src in classes and edge.src_key not in classes[edge.src].PROVIDES: + problems.append( + f"'{edge.src}' does not provide '{edge.src_key}' " + f"(provides: {list(classes[edge.src].PROVIDES)})." + ) + if edge.dst in classes and edge.dst_key not in classes[edge.dst].REQUIRES: + problems.append( + f"'{edge.dst}' does not require '{edge.dst_key}' " + f"(requires: {list(classes[edge.dst].REQUIRES)})." + ) + inbound[edge.dst].setdefault(edge.dst_key, []).append(edge) + + for node in graph.nodes: + if node.id not in classes: + continue + # A runtime param nobody supplied is a KeyError halfway through the + # run otherwise -- after earlier nodes already wrote to disk. It also + # tells a user something they need to know: a unit whose runtime params + # only a particular job knows how to fill is not usable as a node yet. + missing = [ + param + for param in classes[node.id].RUNTIME_PARAMS + if param not in node.config + ] + if missing: + problems.append( + f"'{node.id}' ({node.unit}) needs {missing} supplied by " + "whatever runs it, and nothing in this pipeline supplies them." + ) + for key in classes[node.id].REQUIRES: + sources = len(inbound[node.id].get(key, [])) + if sources == 0: + problems.append( + f"'{node.id}' requires '{key}' and no edge supplies it." + ) + elif sources > 1: + problems.append( + f"'{node.id}' gets '{key}' from {sources} edges; a port " + "takes one. Merging two values into one input is the " + "ambiguity this rejects rather than resolving silently." + ) + + order = _topological_order(graph, by_id, problems) + + if problems: + raise GraphError("\n".join(problems)) + + return order + + +def _topological_order( + graph: Graph, by_id: Mapping[str, Node], problems: List[str] +) -> List[str]: + """Kahn's algorithm. Nodes left over are a cycle, found before running.""" + pending = dict.fromkeys(by_id, 0) + successors: Dict[str, List[str]] = {node_id: [] for node_id in by_id} + for edge in graph.edges: + if edge.src in by_id and edge.dst in by_id: + pending[edge.dst] += 1 + successors[edge.src].append(edge.dst) + + ready = [node_id for node_id, count in pending.items() if count == 0] + order: List[str] = [] + while ready: + node_id = ready.pop(0) + order.append(node_id) + for successor in successors[node_id]: + pending[successor] -= 1 + if pending[successor] == 0: + ready.append(successor) + + if len(order) != len(by_id): + problems.append( + "The graph has a cycle: " + ", ".join(sorted(set(by_id) - set(order))) + ) + + return order diff --git a/DashAI/back/dependencies/database/models.py b/DashAI/back/dependencies/database/models.py index 1b41d8d5e..c5518bf62 100644 --- a/DashAI/back/dependencies/database/models.py +++ b/DashAI/back/dependencies/database/models.py @@ -28,6 +28,8 @@ DatasetStatus, ExplainerStatus, ExplorerStatus, + NodeRunStatus, + PipelineRunStatus, PluginStatus, PredictionStatus, ReportStatus, @@ -107,10 +109,14 @@ def set_status_as_delivered(self) -> None: def set_status_as_started(self) -> None: """ Update the status of the dataset to started and set created to now. + + Unlike ``Run`` or ``Explorer``, this table has no ``start_time`` / + ``end_time`` columns, so there is nothing to stamp beyond the timestamps + it does declare. Assigning them anyway only set an attribute on the + instance that was dropped at commit and read back as missing. """ self.status = DatasetStatus.STARTED self.created = datetime.now() - self.start_time = datetime.now() def set_status_as_finished(self) -> None: """ @@ -118,7 +124,6 @@ def set_status_as_finished(self) -> None: """ self.status = DatasetStatus.FINISHED self.last_modified = datetime.now() - self.end_time = datetime.now() def set_status_as_error(self) -> None: """ @@ -613,10 +618,171 @@ class Pipeline(Base): name: Mapped[str] = mapped_column(String, nullable=False) steps: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) edges: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + # Dead columns, kept for compatibility. Results live in NodeArtifact now, + # keyed by a PROVIDES key rather than by node type. They are still read by + # the pipelines endpoints (dataexploration/results, filter_models) and by + # the results view in the front, so dropping them is a change to those. exploration: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) train: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) prediction: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + pipeline_runs: Mapped[List["PipelineRun"]] = relationship( + "PipelineRun", cascade="all, delete-orphan", back_populates="pipeline" + ) + + +class PipelineRun(Base): + __tablename__ = "pipeline_run" + """ + Table to store one execution of a pipeline. + + The definition of a graph lives in ``Pipeline`` and changes as the user + edits it; a run freezes the ``steps`` and ``edges`` it actually executed, + so a past execution stays readable after the pipeline it came from was + rewritten. + """ + id: Mapped[int] = mapped_column(primary_key=True) + pipeline_id: Mapped[int] = mapped_column( + ForeignKey("pipeline.id", ondelete="CASCADE"), nullable=False + ) + steps: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + edges: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + last_modified: Mapped[DateTime] = mapped_column( + DateTime, default=datetime.now, onupdate=datetime.now + ) + delivery_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + start_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + end_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + status: Mapped[Enum] = mapped_column( + Enum(PipelineRunStatus), + nullable=False, + default=PipelineRunStatus.NOT_STARTED, + ) + error_message: Mapped[str] = mapped_column(String, nullable=True) + + pipeline: Mapped["Pipeline"] = relationship( + "Pipeline", back_populates="pipeline_runs" + ) + node_runs: Mapped[List["NodeRun"]] = relationship( + "NodeRun", cascade="all, delete-orphan", back_populates="pipeline_run" + ) + + def set_status_as_delivered(self) -> None: + """Update the status of the pipeline run to delivered.""" + self.status = PipelineRunStatus.DELIVERED + self.delivery_time = datetime.now() + + def set_status_as_started(self) -> None: + """Update the status of the pipeline run to started.""" + self.status = PipelineRunStatus.STARTED + self.start_time = datetime.now() + + def set_status_as_finished(self) -> None: + """Update the status of the pipeline run to finished.""" + self.status = PipelineRunStatus.FINISHED + self.end_time = datetime.now() + + def set_status_as_error(self, error_message: Optional[str] = None) -> None: + """Update the status of the pipeline run to error.""" + self.status = PipelineRunStatus.ERROR + self.error_message = error_message + self.end_time = datetime.now() + + +class NodeRun(Base): + __tablename__ = "pipeline_node_run" + """ + Table to store the execution of one node of a pipeline run. + + A node is one unit. ``block_id`` is the visual block the node belongs to: + an editor groups several units into one block on the canvas, so a block + maps to N node runs and its status is an aggregate of theirs. In the first + version every block holds exactly one unit and ``block_id`` equals + ``node_id``, but the column is here from the start so growing to N never + needs a migration. + """ + id: Mapped[int] = mapped_column(primary_key=True) + pipeline_run_id: Mapped[int] = mapped_column( + ForeignKey("pipeline_run.id", ondelete="CASCADE"), nullable=False + ) + node_id: Mapped[str] = mapped_column(String, nullable=False) + block_id: Mapped[str] = mapped_column(String, nullable=False) + #: Class name of the unit this node runs, as the component registry knows it. + node_type: Mapped[str] = mapped_column(String, nullable=False) + config: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + input: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + output: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + last_modified: Mapped[DateTime] = mapped_column( + DateTime, default=datetime.now, onupdate=datetime.now + ) + delivery_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + start_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + end_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + status: Mapped[Enum] = mapped_column( + Enum(NodeRunStatus), + nullable=False, + default=NodeRunStatus.NOT_STARTED, + ) + error_message: Mapped[str] = mapped_column(String, nullable=True) + + pipeline_run: Mapped["PipelineRun"] = relationship( + "PipelineRun", back_populates="node_runs" + ) + artifacts: Mapped[List["NodeArtifact"]] = relationship( + "NodeArtifact", cascade="all, delete-orphan", back_populates="node_run" + ) + + def set_status_as_delivered(self) -> None: + """Update the status of the node run to delivered.""" + self.status = NodeRunStatus.DELIVERED + self.delivery_time = datetime.now() + + def set_status_as_started(self) -> None: + """Update the status of the node run to started.""" + self.status = NodeRunStatus.STARTED + self.start_time = datetime.now() + + def set_status_as_finished(self) -> None: + """Update the status of the node run to finished.""" + self.status = NodeRunStatus.FINISHED + self.end_time = datetime.now() + + def set_status_as_error(self, error_message: Optional[str] = None) -> None: + """Update the status of the node run to error.""" + self.status = NodeRunStatus.ERROR + self.error_message = error_message + self.end_time = datetime.now() + + def set_status_as_cancelled(self) -> None: + """Update the status of the node run to cancelled. + + A node that never ran because an earlier one failed, as opposed to one + still waiting its turn. + """ + self.status = NodeRunStatus.CANCELLED + self.end_time = datetime.now() + + +class NodeArtifact(Base): + __tablename__ = "pipeline_node_artifact" + """ + Table to store one output of a node run. + + ``key`` is a key from the unit's ``PROVIDES``, so what a node emits is + named by its own declared contract rather than by a column per node type. + """ + id: Mapped[int] = mapped_column(primary_key=True) + node_run_id: Mapped[int] = mapped_column( + ForeignKey("pipeline_node_run.id", ondelete="CASCADE"), nullable=False + ) + key: Mapped[str] = mapped_column(String, nullable=False) + value: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + + node_run: Mapped["NodeRun"] = relationship("NodeRun", back_populates="artifacts") + class Converter(Base): __tablename__ = "converter" diff --git a/DashAI/back/evaluation/base_evaluation_strategy.py b/DashAI/back/evaluation/base_evaluation_strategy.py index f100bad2d..6b07240b9 100644 --- a/DashAI/back/evaluation/base_evaluation_strategy.py +++ b/DashAI/back/evaluation/base_evaluation_strategy.py @@ -1,31 +1,56 @@ -import os -import pickle -from abc import ABCMeta, abstractmethod -from typing import Callable, Final, List, Optional +"""What a run records, declared per strategy. -from kink import di +These classes used to run the training too: ``execute`` took the run row and the +database session and did the fitting, the search, the scoring, the aggregation +and the persistence behind one method. That work is the units' now -- see +``FitModelUnit``, ``FitModelOverFoldsUnit`` and their siblings -- and what is +left here is what was underneath it all along: a declaration of how a run is +carved and which partitions it records a score for. + +**They stay registered even though they no longer do anything.** The frontend +reads them in four places, and only one is about metrics: + +- the session wizard lists them so the user can choose one, and + ``ModelSession.evaluation_strategy`` is NOT NULL, so without that listing a + session cannot be created at all; +- it starts on the first one whose ``kind`` is holdout; +- ``kind`` decides the shape of the splits payload and which controls are shown; +- ``scored_splits`` tells the metric charts which partitions exist to plot. + +There is precedent for a class in this codebase that declares and does not +execute: ``BaseSplitter.PARTITIONING`` and ``explainable_partitions`` are read +the same way, by the backend and by the frontend, and nothing calls them to do +work. +""" + +from typing import Final -from DashAI.back.core.artifacts import normalize_artifacts from DashAI.back.core.enums.metrics import SplitEnum -from DashAI.back.dependencies.database.models import Run -from DashAI.back.models.base_model import BaseModel -from DashAI.back.models.model_factory import ModelFactory -from DashAI.back.optimizers.base_optimizer import BaseOptimizer -class BaseEvaluationStrategy(metaclass=ABCMeta): - """Abstract base class defining the interface for model evaluation strategies. +class BaseEvaluationStrategy: + """How a run is carved, and what it records. - Concrete implementations (e.g., CrossValidationEvaluationStrategy, - HoldoutEvaluationStrategy) inherit from this class and provide specific - strategies for model evaluation. + Subclasses declare; none of them execute. """ TYPE: Final[str] = "EvaluationStrategy" + #: Whether this strategy splits the dataset once or into folds. It decides + #: which unit prepares the data, because the two publish different shapes, + #: and which controls the session wizard offers. KIND: str = "holdout" + + #: The partitions a run records a score for. Not the same as which + #: partitions have metrics configured: a forecaster has training metrics + #: and still must not be judged on the dates it was fitted on, because an + #: in-sample fit statistic is not comparable with a forecast. A screen that + #: offers one control per partition reads this instead of assuming all + #: three exist. SCORED_SPLITS: tuple = (SplitEnum.TRAIN, SplitEnum.VALIDATION, SplitEnum.TEST) + #: The partition the kept model was fitted on, which is what decides which + #: partitions of a finished run can still be predicted. FINAL_FIT_PARTITIONS: tuple = ("train",) @classmethod @@ -37,209 +62,9 @@ def get_metadata(cls) -> dict: dict Mapping with ``kind``, which says whether this strategy splits the dataset once or into folds, and ``scored_splits``, the partitions - it writes metrics for. A screen that offers one control per - partition reads the latter instead of assuming all three exist: - a forecasting strategy scores no training partition, so asking it - for train metrics finds nothing. + it writes metrics for. """ return { "kind": cls.KIND, "scored_splits": [split.value for split in cls.SCORED_SPLITS], } - - def __init__( - self, - factory: ModelFactory, - optimizer: BaseOptimizer, - goal_metric, - **kwargs, - ): - """Initialize the evaluation strategy with model and optimization configuration. - - Parameters - ---------- - factory : ModelFactory - Factory owning the model to be trained/evaluated and the - hyperparameters that are eligible for optimization. - optimizer : BaseOptimizer - The hyperparameter optimizer instance. Can be None if no HPO is needed. - goal_metric : dict (obtained from Metric component registry) - The target metric to optimize during hyperparameter search. - **kwargs - Additional keyword arguments passed from subclasses (ignored). - """ - self.factory: ModelFactory = factory - self.model: BaseModel = factory.model - self.run_optimizable_parameters = factory.optimizable_parameters - self.optimizer: BaseOptimizer = optimizer - self.goal_metric = goal_metric - self._progress_reporter: Optional[ - Callable[[Optional[float], Optional[str]], None] - ] = None - - def set_progress_reporter( - self, - progress_reporter: Optional[Callable[[Optional[float], Optional[str]], None]], - ) -> None: - """Register a callback that will receive progress updates.""" - self._progress_reporter = progress_reporter - - def _report_progress( - self, fraction: Optional[float], message: Optional[str] = None - ): - """Emit progress updates when a reporter has been registered.""" - if self._progress_reporter is not None: - self._progress_reporter(fraction, message) - - @abstractmethod - def execute(self, x, y, run: Run, db): - """Execute the evaluation strategy on the provided data. - - This is the main entry point for the evaluation process. Subclasses implement - strategy-specific logic for: - - Model training across folds/splits - - Metric computation and persistence - - HPO execution and result handling - - Parameters - ---------- - x : DatasetDict or list of DastasetDict - Input features. Structure depends on the evaluation strategy: - - For holdout: DatasetDict with train/validation/test splits - - For CV: List of DatasetDicts, one per fold with train/test splits - y : dict or list - Target labels. Same structure as x. - run : Run - Database model representing the current experiment run. - db : Session - SQLAlchemy database session for persisting results. - - Returns - ------- - tuple - (trained_model, plot_paths) where: - - trained_model : BaseModel - The trained model after evaluation - - plot_paths : list[str] - Paths to generated HPO visualization files - """ - raise NotImplementedError("Subclasses must implement this method") - - @abstractmethod - def evaluate(self, model: BaseModel, x, y, metric): - """Evaluate the model on the given data and return the score. - - This method is called during hyperparameter optimization to compute - the objective function value for a given set of hyperparameters. - Different strategies may compute metrics differently (e.g., across CV folds - or on a validation split). - - Parameters - ---------- - model : BaseModel - The model instance to evaluate. - x : DatasetDict or list of DastasetDict - Input features for evaluation (structure depends on strategy). - y : DatasetDict or list of DastasetDict - Target labels for evaluation (structure depends on strategy). - metric : Metric - The metric instance to compute. - - Returns - ------- - float - The computed metric value used as the optimization objective. - """ - raise NotImplementedError("Subclasses must implement this method") - - def _do_hpo(self, model: BaseModel, x, y, run: Run, db): - """Execute hyperparameter optimization using the configured optimizer. - - The optimizer uses the self.evaluate method as the objective function, - allowing each strategy to define its own evaluation logic. - - Parameters - ---------- - model : BaseModel - The model instance to optimize. - x : DatasetDict or list of DatasetDict - Training input features (structure varies by strategy). - y : DatasetDict or list of DatasetDict - Training target labels (structure varies by strategy). - run : Run - Database run instance to update with optimized parameters. - db : Session - SQLAlchemy database session for transactions. - - Returns - ------- - BaseModel - The model with the best hyperparameters found during optimization. - """ - from sqlalchemy.orm.attributes import flag_modified - - # Execute hyperparameter optimization and get best model with parameters - self.optimizer.optimize( - model, - x, - y, - self.run_optimizable_parameters, - self.goal_metric, - strategy=self.evaluate, - ) - model = self.optimizer.get_model() - best_params = self.optimizer.get_best_params() - - # Update the run's parameters with the optimized hyperparameters - old_parameters = run.parameters.copy() - updated_parameters = self.factory.update_parameters(old_parameters, best_params) - - # Persist the updated parameters to the database - run.parameters = updated_parameters - flag_modified(run, "parameters") - db.commit() - - return model - - def _generate_hpo_plots(self, run: Run) -> List[str]: - """Generate and pickle the hyperparameter optimization plots to disk. - - Shared by every evaluation strategy that runs HPO, so the plot - generation logic only needs to be maintained in one place. - - Parameters - ---------- - run : Run - The run the plots belong to (used for the plot filenames). - - Returns - ------- - list[str] - Paths to the pickled plot files, in the order produced by the - optimizer. - """ - config = di["config"] - plot_paths: List[str] = [] - - # Retrieve optimization trial data from the optimizer - trials = self.optimizer.get_trials_values() - - # Generate plot visualizations from the trial data - # Plots typically show parameter importance, optimization history, etc. - plot_filenames, plots = self.optimizer.create_plots( - trials, - run.id, - n_params=len(self.run_optimizable_parameters), - goal_metric=self.goal_metric, - ) - - # Convert plots to serializable format (handles special objects, arrays, etc.) - normalized_plots = normalize_artifacts(plots) - - # Serialize and persist each plot to disk - for filename, plot in zip(plot_filenames, normalized_plots, strict=False): - plot_path = os.path.join(config["RUNS_PATH"], filename) - # Serialize the plot object using pickle and write to disk - with open(plot_path, "wb") as file: - pickle.dump(plot, file) - plot_paths.append(plot_path) - - return plot_paths diff --git a/DashAI/back/evaluation/cv.py b/DashAI/back/evaluation/cv.py index 00ea17788..39e0e8b7e 100644 --- a/DashAI/back/evaluation/cv.py +++ b/DashAI/back/evaluation/cv.py @@ -1,421 +1,31 @@ -from functools import partial +"""Cross-validation: several train and validation pairs, scored per fold.""" -import numpy as np -from kink import di - -from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum -from DashAI.back.dependencies.database.models import Metric, Run from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy -from DashAI.back.splitters.base_splitter import BaseSplitter class FoldEvaluationStrategy(BaseEvaluationStrategy): - """Evaluation strategy implementing k-fold cross-validation with optional - nested CV and HPO. - - This strategy partitions the dataset into k folds and performs k rounds of - training and evaluation. Each fold is split into a train and a validation - partition, so the scores obtained by resampling are validation estimates. - When the session reserved rows, the final model is fitted on everything the - folds could use and scored once against those reserved rows. - - The strategy handles metric aggregation at multiple levels: - - FOLD level: Individual metrics from each fold - - TRIAL level: Metrics during HPO trials - - LAST/LAST_OUTER: Aggregated metrics (mean and std) for simple/nested CV + """Carve the dataset into folds and score each of them. + + Each fold is split into a train and a validation partition, so the scores + obtained by resampling are validation estimates. When the session reserved + rows, the kept model is fitted on everything the folds could use and scored + once against those reserved rows -- the only data no fold and no trial ever + saw. + + A run carved this way is prepared by ``PrepareAndFoldUnit`` and fitted by + ``FitModelOverFoldsUnit``, or by ``FitModelOverNestedFoldsUnit`` when the + run also asks for a search inside each fold. + + The levels a fold run records, which is why the enum has more of them than + a holdout run needs: ``FOLD`` per fold, aggregated to ``LAST`` with a + standard deviation; ``TRIAL`` once per trial of a search, holding the mean + over that trial's folds; and, for a nested run, ``OUTER_FOLD`` aggregated + to ``LAST_OUTER``, kept apart because it answers a different question -- + how the procedure does rather than how this model does. """ KIND: str = "cv" - def execute(self, x, y, run: Run, db): - """Execute k-fold cross-validation with optional nested CV and HPO. - - Trains and evaluates a model using k-fold cross-validation. Optionally performs - hyperparameter optimization and nested CV to prevent overfitting. Aggregates - metrics across folds and returns the trained model. - - Parameters - ---------- - x : list of DatasetDict - List of fold DatasetDict each containing: - {"train": X_train, "validation": X_validation} - The last element is not a fold: it holds every row the folds could - use as {"train": X_pool, "test": X_test}, where the test partition - is empty when the session reserved nothing. - y : list of DatasetDict - List of fold label DatasetDicts with same structure as x. - run : Run - Database run instance containing configuration (nested CV settings, etc.). - db : Session - SQLAlchemy database session for persisting metrics and parameters. - - Returns - ------- - tuple - (trained_model, plot_paths) where: - - trained_model : BaseModel - The trained model - - plot_paths : list[str] - Paths to HPO visualization plot files - """ - plot_paths = [] - model = self.model - - # STEP 1: Hyperparameter Optimization (if enabled) - if self.optimizer and self.run_optimizable_parameters: - # Initialize nested CV if required - if run.nested: - try: - registry = di["component_registry"] - inner_splits = run.nested - splitter_name = inner_splits.get("splitter_name", None) - - # Create inner splitter for nested CV fold generation - self.inner_splitter: BaseSplitter = registry[splitter_name][ - "class" - ](inner_splits) - except Exception as e: - raise ValueError( - f"Error configuring inner splitter for nested CV: {e}" - ) from e - - # Execute nested cross-validation for HPO - self._report_progress(0.1, "Nested cross-validation") - self._nested_cv(run.id, model, x, y, db) - - # Perform hyperparameter optimization and update run.parameters - self._report_progress(0.25, "Hyperparameter optimization") - model = self._do_hpo(model, x, y, run, db) - - # Generate and serialize HPO visualization plots - plot_paths = self._generate_hpo_plots(run) - - total_folds = len(x) - 1 - - # STEP 2: Main k-fold Cross-Validation Loop - # Note: Last fold (index len(x)-1) is reserved for final training, - # not CV evaluation - for i in range(total_folds): - self._report_progress( - 0.4 + ((i + 1) / total_folds) * 0.4, - f"Evaluating fold {i + 1}/{total_folds}", - ) - x_fold = x[i] - y_fold = y[i] - - # Set model's internal references to current fold data - model.x_data = x_fold - model.y_data = y_fold - - # Train model on fold's training partition - model.train(x_fold["train"], y_fold["train"]) - - # Compute and store metrics for this fold - if SplitEnum.TRAIN in self.SCORED_SPLITS: - model.calculate_metrics( - split=SplitEnum.TRAIN, level=LevelEnum.FOLD, fold_index=i - ) - model.calculate_metrics( - split=SplitEnum.VALIDATION, level=LevelEnum.FOLD, fold_index=i - ) - - # STEP 3: Aggregate metrics across all folds - # Compute mean and std of fold metrics and store as LAST level metrics - self._aggregate_fold_metrics( - run_id=run.id, - db=db, - level_to_agg=LevelEnum.FOLD, - level_to_save=LevelEnum.LAST, - ) - - # STEP 4: Final model training on every row the folds could use, which - # excludes the rows reserved when the session asked for them. - self._report_progress(0.85, "Training final model") - model.x_data = x[-1] - model.y_data = y[-1] - model.train(x[-1]["train"], y[-1]["train"]) - - # STEP 5: Score the final model once on the reserved rows. No fold and - # no hyperparameter trial ever saw them, so this is the only estimate - # that stays honest after a model is picked out of the comparison - # table. It is a single measurement, hence no standard deviation. - if len(x[-1]["test"]) > 0: - self._report_progress(0.90, "Evaluating on the reserved rows") - model.calculate_metrics(split=SplitEnum.TEST, level=LevelEnum.LAST) - - return model, plot_paths - - def evaluate(self, model, input_dataset, output_dataset, metric, **kwargs): - """Evaluate model using k-fold cross-validation (used as HPO objective - function). - - This method implements cross-validation evaluation for hyperparameter - optimization. It trains on each fold's train partition, scores the fold's - validation partition and averages the results. When used in nested CV, it - evaluates on inner folds within a specific outer fold. - - Parameters - ---------- - model : BaseModel - The model instance to evaluate (with specific hyperparameters). - input_dataset : list of DatasetDict - List of fold data {"train": X_train, "validation": X_validation}. - output_dataset : list of DatasetDict - List of fold labels {"train": y_train, "validation": y_validation}. - metric : Metric - The metric function to optimize. - **kwargs - Additional arguments including: - - fold_index : int or None - Inner outer fold index in nested CV (None for simple CV) - - Returns - ------- - float - Mean metric value across all k folds (objective value for HPO). - - Note: When fold_index is provided (nested CV inner loop), - intermediate metrics are NOT being saved (only outer loop metrics are saved). - """ - # Extract context: fold_index indicates if we're in nested CV inner loop - # None means simple CV; an integer means nested CV on that outer fold - fold_index = kwargs.get("fold_index") - - # List to collect the goal metric value from each fold - folds_results = [] - - # Dictionaries to accumulate all metrics across folds for averaging - train_results = {} - validation_results = {} - - # Cross-validation loop - # Iterate through k-1 folds (last fold is the complete dataset) - # This loop represents either: - # - Main CV loop (if fold_index is None) - # - Inner CV loop within outer fold i (if fold_index is set) - for i in range(len(input_dataset) - 1): - x_fold = input_dataset[i] - y_fold = output_dataset[i] - - # Set model's internal data references for this fold - model.x_data = x_fold - model.y_data = y_fold - - # Train model on this fold's training partition - model.train(x_fold["train"], y_fold["train"]) - - # Compute metrics on both training and validation sets - train_scores = ( - model.compute_metrics(split=SplitEnum.TRAIN) - if SplitEnum.TRAIN in self.SCORED_SPLITS - else {} - ) - validation_scores = model.compute_metrics(split=SplitEnum.VALIDATION) - - # Collect the goal metric value from this fold - folds_results.append(validation_scores[metric.__name__]) - - # Accumulate all metrics only if NOT in nested CV inner loop - if fold_index is None: - for results, scores in [ - (train_results, train_scores), - (validation_results, validation_scores), - ]: - for metric_name, value in scores.items(): - if metric_name not in results: - results[metric_name] = [] - results[metric_name].append(value) - - # Save intermediate metrics (simple CV only) - if fold_index is None: - # Compute average metrics across all folds - averaged_train_results = { - metric: np.mean(values) for metric, values in train_results.items() - } - averaged_validation_results = { - metric: np.mean(values) for metric, values in validation_results.items() - } - - # Persist averaged metrics as TRIAL level (intermediate HPO result) - if SplitEnum.TRAIN in self.SCORED_SPLITS: - model._save_metrics( - results=averaged_train_results, - split=SplitEnum.TRAIN, - level=LevelEnum.TRIAL, - ) - model._save_metrics( - results=averaged_validation_results, - split=SplitEnum.VALIDATION, - level=LevelEnum.TRIAL, - ) - - # Return the mean of the goal metric across folds - # This is the objective value used by the optimizer - return np.mean(folds_results) - - def _nested_cv(self, run_id, model, input_dataset, output_dataset, db): - """Execute nested cross-validation with inner HPO loop. - - Nested CV implements a two-level validation scheme to prevent overfitting during - hyperparameter optimization and give an unbiased estimate of model performance. - - Outer loop: Standard k-fold CV for unbiased final performance estimation - - Inner loop: Separate CV fold within each outer fold for HPO - - This prevents "information leakage" where the data used to score a fold also - influences hyperparameter selection, which would overestimate true - generalization performance. - - Parameters - ---------- - run_id : int - The database run ID for metric storage. - model : BaseModel - The model instance to optimize and evaluate. - input_dataset : list of DatasetDict - List of outer fold data - {"train": X_train, "validation": X_validation}. - output_dataset : list of DatasetDict - List of outer fold labels - {"train": y_train, "validation": y_validation}. - db : Session - SQLAlchemy database session for persisting aggregated metrics. - """ - # Nested CV Outer Loop - # For each outer fold, optimize hyperparameters on inner folds - for i in range(len(input_dataset) - 1): - x_outer = input_dataset[i] - y_outer = output_dataset[i] - - # Create inner folds from outer fold's training data - # This ensures HPO validation data is independent of outer test fold - inner_x, inner_y, _ = self.inner_splitter.split( - x_outer["train"], y_outer["train"] - ) - - # Create evaluation strategy for this outer fold - # Passes fold_index so evaluate() knows it's in nested CV context - strategy_with_context = partial(self.evaluate, fold_index=i) - - # INNER LOOP: Optimize hyperparameters using inner CV - # The optimizer will iteratively: - # - Generate hyperparameter candidates - # - Train models on inner fold combinations - # - Evaluate using inner CV (calls strategy_with_context) - # - Select hyperparameters with best inner CV performance - self.optimizer.optimize( - model, - inner_x, - inner_y, - self.run_optimizable_parameters, - self.goal_metric, - strategy=strategy_with_context, - ) - outer_model = self.optimizer.get_model() - - # Set model's data references for outer fold evaluation - outer_model.x_data = x_outer - outer_model.y_data = y_outer - - # Train model on outer fold's training data with optimized hyperparameters - outer_model.train(x_outer["train"], y_outer["train"]) - - # Evaluate on outer fold's test data (this is OUTER_FOLD level metric) - outer_model.calculate_metrics( - split=SplitEnum.VALIDATION, level=LevelEnum.OUTER_FOLD, fold_index=i - ) - if SplitEnum.TRAIN in self.SCORED_SPLITS: - outer_model.calculate_metrics( - split=SplitEnum.TRAIN, level=LevelEnum.OUTER_FOLD, fold_index=i - ) - - # Aggregate outer fold metrics - # Compute mean and std of OUTER_FOLD metrics and store as LAST_OUTER level - self._aggregate_fold_metrics( - run_id=run_id, - db=db, - level_to_agg=LevelEnum.OUTER_FOLD, - level_to_save=LevelEnum.LAST_OUTER, - ) - - def _aggregate_fold_metrics( - self, - run_id: int, - db, - level_to_agg=LevelEnum.FOLD, - level_to_save=LevelEnum.LAST, - ): - """Aggregate and average fold metrics across cross-validation folds. - - This method computes the mean and standard deviation of metrics collected - at the fold level and stores the aggregated results at a higher level. - This is used to provide summary statistics for model performance. - - Typical usage patterns: - - Aggregate FOLD metrics -> store as LAST (simple CV summary) - - Aggregate OUTER_FOLD metrics -> store as LAST_OUTER (nested CV summary) - - Parameters - ---------- - run_id : int - The database run ID to aggregate metrics for. - db : Session - SQLAlchemy database session for querying and persisting metrics. - level_to_agg : LevelEnum, optional - The source metric level to aggregate. Default: FOLD. - level_to_save : LevelEnum, optional - The destination level for aggregated metrics. Default: LAST. - """ - # Query all metrics with level=level_to_agg for this run - fold_metrics = ( - db.query(Metric) - .filter(Metric.run_id == run_id, Metric.level == level_to_agg) - .all() - ) - - # If no metrics found, nothing to aggregate - if not fold_metrics: - return - - # Group metrics by (split, name) for aggregation - metrics_by_split_name = {} - for metric in fold_metrics: - key = (metric.split, metric.name) - if key not in metrics_by_split_name: - metrics_by_split_name[key] = [] - metrics_by_split_name[key].append(metric.value) - - # Aggregate and persist metrics - for (split, name), values in metrics_by_split_name.items(): - # Compute aggregation statistics - avg_value = np.mean(values) - std_value = np.std(values) if len(values) > 1 else 0.0 - - # Check if aggregated metric already exists for this split/name/run - existing = ( - db.query(Metric) - .filter_by(run_id=run_id, split=split, level=level_to_save, name=name) - .first() - ) - - if existing: - # Update existing metric with aggregated values - existing.value = avg_value - existing.std_value = std_value - else: - # Create new aggregated metric - db.add( - Metric( - run_id=run_id, - split=split, - level=level_to_save, - name=name, - value=avg_value, - std_value=std_value, - step=0, - ) - ) - - # Persist aggregated metrics to database - db.commit() - class CrossValidationEvaluationStrategy(FoldEvaluationStrategy): """Score a model across folds, recording train and validation for each. diff --git a/DashAI/back/evaluation/holdout.py b/DashAI/back/evaluation/holdout.py index 6fb602ccc..52b9b7e4c 100644 --- a/DashAI/back/evaluation/holdout.py +++ b/DashAI/back/evaluation/holdout.py @@ -1,171 +1,19 @@ -from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum -from DashAI.back.dependencies.database.models import Metric, Run +"""Holdout evaluation: one split into three partitions, scored once.""" + from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy class SinglePartitionEvaluationStrategy(BaseEvaluationStrategy): - """Evaluation strategy implementing holdout (train/validation/test split) - validation. - - This strategy divides the dataset into three mutually exclusive partitions: - training, validation, and test. The training set is used for model training, - the validation set for HPO, and the test set for final evaluation. + """Split once into train, validation and test. - The strategy handles metric aggregation at multiple levels: - - TRIAL level: Metrics during HPO trials on validation set - - LAST level: Final metrics computed on all three partitions after training + The training set fits the model, the validation set is what a + hyperparameter search is measured on, and the test set is scored once at + the end. A run carved this way is prepared by ``PrepareAndSplitUnit`` and + fitted by ``FitModelUnit``. """ KIND: str = "holdout" - def execute(self, x, y, run: Run, db): - """Execute holdout validation: train on training set, optimize with validation, - evaluate on test. - - Trains a model on the training partition and optionally performs hyperparameter - optimization using the validation set. Finally evaluates the trained model on - all three partitions (train, validation, test) and returns the model with plots. - - Parameters - ---------- - x : DatasetDict - DatasetDict with data partitions: - {"train": X_train, "validation": X_val, "test": X_test} - y : DatasetDict - DatasetDict with label partitions: - {"train": y_train, "validation": y_val, "test": y_test} - run : Run - Database run instance for storing results and configuration. - db : Session - SQLAlchemy database session for persisting metrics. - - Returns - ------- - tuple - (trained_model, plot_paths) where: - - trained_model : BaseModel - The trained model - - plot_paths : list[str] - Paths to HPO visualization plot files - """ - plot_paths = [] - model = self.model - - # set the data used for model training and evaluation - model.x_data = x - model.y_data = y - - # Execute HPO if optimizer and there are parameters to optimize - if self.optimizer and self.run_optimizable_parameters: - self._report_progress(0.2, "Hyperparameter optimization") - model = self._do_hpo(model, x, y, run, db) - plot_paths = self._generate_hpo_plots(run) - - # Train the model with the provided data and return it - self._report_progress(0.5, "Training") - self._fit_final_model(model, x, y) - - # Calculate metrics at the end of training if not done already - self._report_progress(0.85, "Computing metrics") - for split in self.SCORED_SPLITS: - self._calculate_metrics_if_missing(model, run, db, split) - - return model, plot_paths - - def _fit_final_model(self, model, x, y): - """Fit the model that gets kept. - - Separate from the trial fits so a strategy can differ on what the - kept model is allowed to learn from, which is the one thing - forecasting needs to change here. - - Parameters - ---------- - model : BaseModel - The model to fit. - x : DatasetDict - Input partitions. - y : DatasetDict - Target partitions. - """ - model.train(x["train"], y["train"], x["validation"], y["validation"]) - - def _calculate_metrics_if_missing(self, model, run: Run, db, split: SplitEnum): - """Compute and persist LAST-level metrics for a split, unless already saved. - - Parameters - ---------- - model : BaseModel - The trained model to compute metrics on. - run : Run - Database run instance the metrics belong to. - db : Session - SQLAlchemy database session used to check for existing metrics. - split : SplitEnum - The data split to compute metrics for. - """ - existing_metric = ( - db.query(Metric) - .filter_by(run_id=run.id, split=split, level=LevelEnum.LAST) - .first() - ) - if not existing_metric: - model.calculate_metrics(split=split, level=LevelEnum.LAST) - - def evaluate(self, model, input_dataset, output_dataset, metric): - """Evaluate model on validation set during HPO trials. - - Trains the model on the training set and computes the metric on the validation - set. Used as the objective function during hyperparameter optimization. - - Parameters - ---------- - model : BaseModel - The model instance to evaluate with specific hyperparameters. - input_dataset : DatasetDict - DatasetDict with data partitions - {"train": X_train, "validation": X_val, "test": X_test}. - output_dataset : DatasetDict - DatasetDict with label partitions - {"train": y_train, "validation": y_val, "test": y_test}. - metric : Metric - The metric function to compute on predictions. - - Returns - ------- - float - The metric score value for this hyperparameter combination. - """ - # Validation data is passed on purpose. Without it the epoch loops - # skip `calculate_metrics(split=VALIDATION, level=EPOCH)` entirely — - # they guard it behind `if x_validation is not None` — so during - # optimization the per-epoch validation score was never computed. - # - # That is the deeper reason pruning could not work here: the number a - # pruner needs to decide did not exist, independently of whether the - # pruner itself was an instance or a string. - model.train( - input_dataset["train"], - output_dataset["train"], - input_dataset["validation"], - output_dataset["validation"], - ) - - # Evaluate the model on the validation set - y_pred = model.predict(input_dataset["validation"]) - - output_dataset_transformed = model.prepare_output( - output_dataset["validation"], is_fit=False - ) - - # Calculate metric for train and validation data each trial. - if SplitEnum.TRAIN in self.SCORED_SPLITS: - model.calculate_metrics(split=SplitEnum.TRAIN, level=LevelEnum.TRIAL) - model.calculate_metrics(split=SplitEnum.VALIDATION, level=LevelEnum.TRIAL) - - # Compute the objective metric score on the validation set - score = metric.score(output_dataset_transformed, y_pred) - - return score - class HoldoutEvaluationStrategy(SinglePartitionEvaluationStrategy): """Split once into train, validation and test, and score all three. @@ -176,8 +24,10 @@ class HoldoutEvaluationStrategy(SinglePartitionEvaluationStrategy): not belong in the same results table as one. ``ForecastingHoldoutEvaluationStrategy`` records validation and test only. - The final fit is the same in both: the kept model is fitted on the training - partition alone, so it is the model the recorded metrics describe. + The kept model is fitted on the training partition alone in both, so it is + the model the recorded metrics describe. The validation partition is handed + to the fit, which models use to watch training and stop early -- being + given it to watch is not the same as being fitted on it. """ COMPATIBLE_COMPONENTS = [ diff --git a/DashAI/back/explainability/explainers/contrastive_shap.py b/DashAI/back/explainability/explainers/contrastive_shap.py index fc7141bf1..43c8bb853 100644 --- a/DashAI/back/explainability/explainers/contrastive_shap.py +++ b/DashAI/back/explainability/explainers/contrastive_shap.py @@ -239,7 +239,10 @@ def fit( """ import shap - from DashAI.back.explainability.model_input import prepare_model_input + from DashAI.back.explainability.model_input import ( + as_shap_predictor, + prepare_model_input, + ) x, y = background_dataset # SHAP calls the model with perturbed matrices, so the background must @@ -256,7 +259,7 @@ def fit( background_data = shap.sample(background_data, n_samples) self.explainer = shap.KernelExplainer( - model=self.model.predict_prepared, + model=as_shap_predictor(self.model), data=background_data, feature_names=feature_names, ) diff --git a/DashAI/back/explainability/explainers/kernel_shap.py b/DashAI/back/explainability/explainers/kernel_shap.py index 2248e00b4..e808b9b0f 100644 --- a/DashAI/back/explainability/explainers/kernel_shap.py +++ b/DashAI/back/explainability/explainers/kernel_shap.py @@ -335,7 +335,10 @@ def fit( """ sample_background_data = bool(sample_background_data) - from DashAI.back.explainability.model_input import prepare_model_input + from DashAI.back.explainability.model_input import ( + as_shap_predictor, + prepare_model_input, + ) x, y = background_dataset @@ -367,7 +370,7 @@ def fit( import shap self.explainer = shap.KernelExplainer( - model=self.model.predict_prepared, + model=as_shap_predictor(self.model), data=background_data, feature_names=feature_names, link=self.link, diff --git a/DashAI/back/explainability/explainers/regression_kernel_shap.py b/DashAI/back/explainability/explainers/regression_kernel_shap.py index cbc4471d3..5ebaa07f9 100644 --- a/DashAI/back/explainability/explainers/regression_kernel_shap.py +++ b/DashAI/back/explainability/explainers/regression_kernel_shap.py @@ -180,7 +180,10 @@ def fit( """ import shap - from DashAI.back.explainability.model_input import prepare_model_input + from DashAI.back.explainability.model_input import ( + as_shap_predictor, + prepare_model_input, + ) x, y = background_dataset # SHAP calls the model with perturbed matrices, so the background must @@ -197,7 +200,7 @@ def fit( background_data = shap.sample(background_data, n_samples) self.explainer = shap.KernelExplainer( - model=self.model.predict_prepared, + model=as_shap_predictor(self.model), data=background_data, feature_names=feature_names, ) diff --git a/DashAI/back/explainability/model_input.py b/DashAI/back/explainability/model_input.py index 3126ed3b5..c459f1018 100644 --- a/DashAI/back/explainability/model_input.py +++ b/DashAI/back/explainability/model_input.py @@ -17,12 +17,57 @@ the already prepared matrix a second time. """ -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Callable, List, Optional if TYPE_CHECKING: from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset +def as_shap_predictor(model: Any) -> Callable: + """Wrap the model's prepared-matrix prediction so SHAP gets a plain function. + + SHAP suppresses scikit-learn's "X does not have valid feature names" + warning by blanking ``feature_names_in_`` on whatever object the callable + is bound to (``shap.utils._legacy.convert_to_model``). It reaches that + object through ``__self__``, so it only does this when handed a *bound + method*, and it assumes the attribute is writable. + + That assumption does not hold for estimators that expose + ``feature_names_in_`` as a read-only ``property`` — a common shape among + third party wrappers — where the assignment raises ``AttributeError: + property 'feature_names_in_' ... has no setter`` and the explanation fails + before it starts. + + Handing over a plain closure instead leaves ``__self__`` absent, so SHAP + skips that step entirely — a function is SHAP's primary documented + interface for ``model``. The only thing lost is the suppression of a + cosmetic scikit-learn warning. + + It routes to ``predict_prepared``, not to ``predict``. Callers hand SHAP a + background already moved into the model's feature space with + ``prepare_model_input``, and SHAP then queries the model with perturbed + copies of *that* matrix. Going through ``predict`` would run the model's + input preparation a second time over an already prepared matrix — and SHAP + passes plain arrays, which the preparation cannot consume at all. + + Parameters + ---------- + model : Any + The trained model being explained. + + Returns + ------- + Callable + A one-argument function calling ``model.predict_prepared`` + positionally, the same way SHAP calls it. + """ + + def predict(x): + return model.predict_prepared(x) + + return predict + + def prepare_model_input(model: Any, dataset: "DashAIDataset") -> "DashAIDataset": """Apply the model's own input preprocessing to a dataset. diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 868c376d5..fa358c163 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -511,6 +511,49 @@ from DashAI.back.tasks.text_to_text_generation_task import TextToTextGenerationTask from DashAI.back.tasks.translation_task import TranslationTask +# Units +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.apply_dataset_schema_unit import ApplyDatasetSchemaUnit +from DashAI.back.units.build_global_explainer_unit import BuildGlobalExplainerUnit +from DashAI.back.units.build_local_explainer_unit import BuildLocalExplainerUnit +from DashAI.back.units.build_manual_input_unit import BuildManualInputUnit +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.compute_dataset_metadata_unit import ComputeDatasetMetadataUnit +from DashAI.back.units.evaluate_model_to_artifact_unit import ( + EvaluateModelToArtifactUnit, +) +from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit +from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit +from DashAI.back.units.fit_model_over_nested_folds_unit import ( + FitModelOverNestedFoldsUnit, +) +from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.generate_global_explanation_unit import ( + GenerateGlobalExplanationUnit, +) +from DashAI.back.units.generate_local_explanation_unit import ( + GenerateLocalExplanationUnit, +) +from DashAI.back.units.infer_dataset_types_unit import InferDatasetTypesUnit +from DashAI.back.units.load_datafile_dataset_unit import LoadDatafileDatasetUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_run_model_unit import LoadRunModelUnit +from DashAI.back.units.load_trained_model_unit import LoadTrainedModelUnit +from DashAI.back.units.load_training_dataset_unit import LoadTrainingDatasetUnit +from DashAI.back.units.load_uploaded_dataset_unit import LoadUploadedDatasetUnit +from DashAI.back.units.predict_unit import PredictUnit +from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.prepare_explanation_data_unit import PrepareExplanationDataUnit +from DashAI.back.units.run_exploration_unit import RunExplorationUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_exploration_unit import SaveExplorationUnit +from DashAI.back.units.save_model_unit import SaveModelUnit +from DashAI.back.units.save_prediction_unit import SavePredictionUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit + logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -687,6 +730,40 @@ def get_initial_components(): DatasetJob, GenerativeJob, PipelineJob, + # Units + LoadDatasetUnit, + PrepareAndSplitUnit, + PrepareAndFoldUnit, + BuildModelUnit, + FitModelUnit, + FitModelOverFoldsUnit, + FitModelOverNestedFoldsUnit, + EvaluateModelUnit, + EvaluateModelToArtifactUnit, + SaveModelUnit, + ApplyConverterUnit, + FitConverterUnit, + TransformDatasetUnit, + SaveDatasetUnit, + RunExplorationUnit, + SaveExplorationUnit, + LoadTrainedModelUnit, + LoadTrainingDatasetUnit, + BuildManualInputUnit, + PredictUnit, + SavePredictionUnit, + LoadRunModelUnit, + BuildGlobalExplainerUnit, + BuildLocalExplainerUnit, + PrepareExplanationDataUnit, + GenerateGlobalExplanationUnit, + GenerateLocalExplanationUnit, + LoadUploadedDatasetUnit, + LoadDatafileDatasetUnit, + InferDatasetTypesUnit, + ApplyDatasetSchemaUnit, + ComputeDatasetMetadataUnit, + SaveDatasetToPathUnit, RAGJob, # Reports ConfusionMatrix, diff --git a/DashAI/back/job/converter_job.py b/DashAI/back/job/converter_job.py index fcbfc4d68..f833dd3b7 100644 --- a/DashAI/back/job/converter_job.py +++ b/DashAI/back/job/converter_job.py @@ -1,123 +1,24 @@ import logging -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING from kink import inject from sqlalchemy import exc -from DashAI.back.api.api_v1.schemas.converter_params import ConverterParams from DashAI.back.dependencies.database.models import Converter from DashAI.back.dependencies.database.models import Dataset as DatasetModel from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker - from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset - logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) -def _rebuild_dataset_with_transformed_columns( - base: "DashAIDataset", - transformed: "DashAIDataset", - scope_column_names: List[str], - scope_column_indexes: List[int], -) -> "DashAIDataset": - """ - Replaces specific columns in the base dataset with columns from the transformed - dataset, preserving their original positions. Also appends any additional columns - that were generated by the transformer at the end. Keeps the features and metadata - consistent. - - Parameters - ---------- - base : DashAIDataset - The original dataset before transformation. - - transformed : DashAIDataset - The dataset resulting from applying a transformer, containing updated and/or - new columns. - - scope_column_names : List[str] - Names of the columns that were originally selected for transformation. - - scope_column_indexes : List[int] - The indices of the columns in the base dataset that were replaced. - Must match the order of scope_column_names. - - Returns - ------- - DashAIDataset - A new dataset with the specified columns replaced in place, new columns - appended, and original metadata and split information preserved. - """ - from DashAI.back.dataloaders.classes.dashai_dataset import modify_table - - original_columns = base.column_names - transformed_cols = transformed.column_names - - transformed_cols_set = set(transformed_cols) - scope_column_names_set = set(scope_column_names) - - removed_cols = [ - col for col in scope_column_names if col not in transformed_cols_set - ] - replacement_cols = [ - col for col in scope_column_names if col in transformed_cols_set - ] - new_cols = [col for col in transformed_cols if col not in scope_column_names_set] - - removed_cols_set = set(removed_cols) - - new_columns_order = [] - seen_cols = set() - for col in original_columns: - if col in removed_cols_set: - continue - if col not in seen_cols: - new_columns_order.append(col) - seen_cols.add(col) - - col_name_mapping = {} - for col in new_cols: - unique_col = col - counter = 1 - while unique_col in seen_cols: - unique_col = f"{col}_{counter}" - counter += 1 - new_columns_order.append(unique_col) - seen_cols.add(unique_col) - col_name_mapping[col] = unique_col - - updated_arrays = {} - for col in replacement_cols: - if col in transformed_cols_set: - updated_arrays[col] = transformed.arrow_table[col] - for col, unique_col in col_name_mapping.items(): - if col in transformed_cols_set: - updated_arrays[unique_col] = transformed.arrow_table[col] - - updated_types = base.types.copy() - - for col in removed_cols: - if col in updated_types: - del updated_types[col] - - for col in replacement_cols: - if col in transformed.types: - updated_types[col] = transformed.types[col] - for col, unique_col in col_name_mapping.items(): - if col in transformed.types: - updated_types[unique_col] = transformed.types[col] - - # Use existing modify_table (imported at module level) - modified_dataset = modify_table(base, updated_arrays, types=updated_types) - modified_dataset = modified_dataset.select_columns(new_columns_order) - - return modified_dataset - - class ConverterJob(BaseJob): """ConverterJob class to modify a dataset by applying a sequence of converters.""" @@ -198,31 +99,9 @@ def run( ) -> None: from kink import di - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - save_dataset, - ) - session_factory = di["session_factory"] - component_registry = di["component_registry"] - - def instantiate_converters( - converter_name: str, - converter_params: ConverterParams, - ) -> object: - # Import the converter - try: - converter_constructor = component_registry[converter_name]["class"] - except KeyError as e: - log.exception(e) - raise JobError( - f"Error importing converter {converter_name}: {e}" - ) from e - - # Get parameters or empty dict if none - converter_parameters = converter_params.get("params", {}) - return converter_constructor(**converter_parameters) + ctx = ExecutionContext() # Extract job parameters converter_id = self.kwargs["converter_id"] @@ -246,38 +125,48 @@ def instantiate_converters( # Get dataset try: + notebook_id = converter.notebook_id dataset_id = converter.notebook.dataset_id dataset = db.get(DatasetModel, dataset_id) - # dataset to edit - dataset_path = f"{converter.notebook.file_path}/dataset" - loaded_dataset = load_dataset(dataset_path) - params = converter.parameters or {} - target_column_index = ( - params["target"].get("idx") - if params.get("target") is not None - else None - ) + # dataset to edit: the notebook's own working copy + LoadDatasetUnit(notebook_id=notebook_id)(ctx) + dataset_path = ctx.require("dataset_path") - if not loaded_dataset: - raise JobError(f"Dataset with path {dataset_path} not found") + # How the converter configuration is stored on the row, not + # part of the transformation itself. + params = converter.parameters or {} except exc.SQLAlchemyError as e: log.exception(e) converter.set_status_as_error() db.commit() raise JobError("Error loading dataset info") from e - - # Load dataset + except Exception: + # Anything the load unit raises (missing notebook, unreadable + # dataset) also has to leave the row in ERROR. Nothing marks it + # otherwise: the Huey error signal only writes to its own + # task_copy table, never to the Converter row, so without this + # the converter would stay STARTED forever. Re-raised as-is so + # the unit's specific message survives. + converter.set_status_as_error() + db.commit() + raise + + apply_converter = ApplyConverterUnit( + converter={ + "component": converter.converter, + "params": params.get("params") or {}, + }, + scope=params.get("scope"), + target=params.get("target"), + ) + + # Validating before the work starts keeps an impossible target + # index reported as a dataset problem, which is where it was + # reported before the job was split into units. try: - # Validate target column index - if target_column_index is not None and ( - int(target_column_index) < 1 - or int(target_column_index) > len(loaded_dataset.features) - ): - raise JobError( - f"Target column index {target_column_index} is out of bounds" - ) + apply_converter.validate(ctx) except Exception as e: log.exception(e) converter.set_status_as_error() @@ -285,142 +174,11 @@ def instantiate_converters( raise JobError(f"Cannot load dataset from {dataset_path}") from e try: - # Get stored converter configurations - converters_stored_info = {converter.converter: converter.parameters} - dataset_original_columns = loaded_dataset.column_names - - # Sort converters by order - converters_sorted_list = sorted( - converters_stored_info.items(), key=lambda x: x[1]["order"] - ) - - i = 0 - converter_instances = [] - - while i < len(converters_sorted_list): - converter_name = converters_sorted_list[i][0] - converter_params = converters_sorted_list[i][1] - # Regular converter - converter_instance = instantiate_converters( - converter_name, - converter_params, - ) - - # Get scope or use default - scope = converter_params.get("scope", {"columns": [], "rows": []}) - - # Add to instances - converter_instances.append( - { - "name": converter_name, - "instance": converter_instance, - "scope": scope, - } - ) - i += 1 - - # Apply each converter in sequence - total_converters = len(converter_instances) - for converter_index, converter_info in enumerate(converter_instances): - converter_instance = converter_info["instance"] - converter_name = converter_info["name"] - converter_scope = converter_info["scope"] - - # Map converter progress onto the 0.2-0.9 band. - self.report_progress( - 0.2 + 0.7 * (converter_index / max(total_converters, 1)), - f"Applying {converter_name}", - ) - log.info(f"Applying converter: {converter_name}") - - columns_scope = [ - column["idx"] - 1 for column in converter_scope["columns"] - ] - scope_column_indexes = sorted(set(columns_scope)) - - if not scope_column_indexes: - scope_column_indexes = list(range(len(loaded_dataset.features))) - - scope_column_names = [ - dataset_original_columns[index] - for index in scope_column_indexes - ] - - rows_scope = [row - 1 for row in converter_scope["rows"]] - scope_rows_indexes = sorted(set(rows_scope)) - - y_dataset_fit = None - target_column_name = None - y_full_transform = None - if target_column_index is not None: - target_column_index_0based = int(target_column_index) - 1 - target_column_name = dataset_original_columns[ - target_column_index_0based - ] - y_dataset_fit = loaded_dataset.select_columns( - [target_column_name] - ) - if scope_rows_indexes: - y_dataset_fit = y_dataset_fit.select(scope_rows_indexes) - y_full_transform = loaded_dataset.select_columns( - [target_column_name] - ) - else: - y_full_transform = y_dataset_fit - - X_dataset_fit = loaded_dataset.select_columns(scope_column_names) - - if scope_rows_indexes: - X_dataset_fit = X_dataset_fit.select(scope_rows_indexes) - - try: - converter_instance = converter_instance.fit( - X_dataset_fit, y_dataset_fit - ) - except ValueError as e: - log.error(f"Validation error in {converter_name}: {e}") - raise JobError( - f"Validation error fitting {converter_name}: {e}" - ) from e - except Exception as e: - log.exception(e) - raise JobError( - f"Error fitting converter {converter_name}: {e}" - ) from e - - if scope_rows_indexes: - X_full_transform = loaded_dataset.select_columns( - scope_column_names - ) - else: - # Same reuse as above: no row-level fit scope means - # X_dataset_fit already covers the full transform scope. - X_full_transform = X_dataset_fit - - try: - transformed_dataset = converter_instance.transform( - X_full_transform, y_full_transform - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Error transforming data with {converter_name}: {e}" - ) from e - - if type(converter_instance).CHANGES_ROW_COUNT: - loaded_dataset = transformed_dataset - else: - loaded_dataset = _rebuild_dataset_with_transformed_columns( - loaded_dataset, - transformed_dataset, - scope_column_names, - scope_column_indexes, - ) - - dataset_original_columns = loaded_dataset.column_names + self.report_progress(0.2, f"Applying {converter.converter}") + apply_converter(ctx) self.report_progress(0.95, "Saving dataset") - save_dataset(loaded_dataset, f"{dataset_path}") + SaveDatasetUnit()(ctx) converter.set_status_as_finished() db.commit() db.refresh(dataset) @@ -432,3 +190,5 @@ def instantiate_converters( raise JobError( f"Error applying converters to dataset {dataset_id}: {e}" ) from e + finally: + ctx.clear_cache() diff --git a/DashAI/back/job/dataset_job.py b/DashAI/back/job/dataset_job.py index 43e25c54b..f1f851b0d 100644 --- a/DashAI/back/job/dataset_job.py +++ b/DashAI/back/job/dataset_job.py @@ -9,6 +9,14 @@ from DashAI.back.api.utils import parse_params from DashAI.back.dependencies.database.models import Converter, Dataset, Notebook from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.units.apply_dataset_schema_unit import ApplyDatasetSchemaUnit +from DashAI.back.units.compute_dataset_metadata_unit import ComputeDatasetMetadataUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.infer_dataset_types_unit import InferDatasetTypesUnit +from DashAI.back.units.load_datafile_dataset_unit import LoadDatafileDatasetUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_uploaded_dataset_unit import LoadUploadedDatasetUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -142,17 +150,10 @@ def run( import json import os import shutil + import tempfile import uuid from pathlib import Path - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - save_dataset, - transform_dataset_with_schema, - ) - from DashAI.back.types.inf.type_inference import infer_types - - component_registry = di["component_registry"] session_factory = di["session_factory"] config = di["config"] @@ -162,8 +163,26 @@ def run( n_sample = self.kwargs.get("n_sample", None) file_path = self.kwargs.get("file_path") temp_dir = self.kwargs.get("temp_dir") + if not temp_dir: + # The dataloaders forward this path to HuggingFace as ``cache_dir``. + # Passing it along unset would stringify to "None" and create a + # directory literally named "None" in the working directory. + temp_dir = tempfile.mkdtemp(prefix="dashai-dataset-") url = self.kwargs.get("url", "") + ctx = ExecutionContext() + + # Whether the destination folder is this job's to delete. Re-importing + # into an existing dataset writes over its current folder, and the + # failure paths below clean up by removing it — which would destroy data + # the surviving row still points at. Only a folder this run created may + # be removed. + # + # Bound before the try, not inside it: the handler at the bottom reads + # it, and a failure raised before the assignment would otherwise reach + # that handler with the name unbound. + folder_is_ours = False + try: with session_factory() as db: dataset = db.get(Dataset, dataset_id) @@ -190,6 +209,7 @@ def run( raise JobError( f"A dataset with the name {random_name} already exists." ) from e + folder_is_ours = True # Write folder_path to DB immediately so on_cancel can delete it # even if the job is killed before the final commit. @@ -209,6 +229,9 @@ def run( .filter(Notebook.id == notebook_id) .first() ) + # Checked here rather than left to the load unit, whose + # own wording for a missing notebook differs from this + # one. The message reaches the UI, so it is preserved. if not notebook_dataset: msg = ( "Notebook with ID " @@ -228,208 +251,168 @@ def run( is not None ) from_notebook_no_converters = not has_converters - new_dataset = load_dataset( - os.path.join(notebook_dataset.file_path, "dataset") - ) + + # ``LoadDatasetUnit`` also publishes ``dataset_path`` (the + # notebook's own copy) and ``dataset_id`` (the *source* + # dataset). Neither describes what is being created here: the + # save goes to a new folder, and this job's ``dataset_id`` is + # the destination row, read from kwargs. Nothing below reads + # them from the context, and nothing should start to. + LoadDatasetUnit(notebook_id=notebook_id)(ctx) + + # No schema is applied to a notebook copy: it is already a + # stored dataset, with its types settled when it was created. else: source_name = self.kwargs.get("source_name") if source_name: # --- Hub import path --- - from DashAI.back.core.enums.status import DatafileStatus - from DashAI.back.dependencies.database.models import ( - Datafile, - ) - + # The id is required, and it is the request that is + # malformed without it, so the check stays here rather + # than becoming a unit that cannot be configured. datafile_id = params.get("datafile_id") - selected_file = params.get("selected_file") - if datafile_id is None: raise JobError("datafile_id is required for hub imports.") - with session_factory() as db: - hub_row = db.get(Datafile, datafile_id) - if hub_row is None or hub_row.status != DatafileStatus.READY: - raise JobError(f"Datafile {datafile_id} is not ready.") - hub_work_dir = hub_row.local_path - if selected_file: - file_path_hub = str(Path(hub_work_dir) / selected_file) - else: - hub_base = Path(hub_work_dir) - files = sorted( - str(p) - for p in hub_base.rglob("*") - if p.is_file() - and not any( - part.startswith(".") - for part in p.relative_to(hub_base).parts - ) - ) - if not files: - raise JobError("Hub download directory is empty.") - file_path_hub = files[0] - - selected_dataloader = params.get("dataloader", "") - _reg = component_registry._registry - dl_registry = _reg.get("DataLoader", {}) - if selected_dataloader not in dl_registry: - raise JobError( - f"DataLoader '{selected_dataloader}'" - " not found in registry." - ) - dataloader = dl_registry[selected_dataloader]["class"]() - log.debug( - "Loading hub dataset from %s using %s", - file_path_hub, - selected_dataloader, - ) - hub_loader_params = params.get("dataloader_params", {}) - new_dataset = dataloader.load_data( - filepath_or_buffer=file_path_hub, - temp_path=hub_work_dir, - params=hub_loader_params, - n_sample=None, - ) + LoadDatafileDatasetUnit( + dataloader={ + "component": params.get("dataloader", ""), + "params": params.get("dataloader_params", {}), + }, + datafile_id=datafile_id, + selected_file=params.get("selected_file"), + )(ctx) else: - # --- File / URL upload path (unchanged) --- + # --- File / URL upload path --- + # Validating the request's params is unpacking of how the + # upload arrived, not part of reading the data: the unit + # gets the reader already picked and its params already + # checked. ``model_dump()`` keeps the payload the reader + # receives exactly as it was. parsed_params = parse_params(DatasetParams, json.dumps(params)) - dataloader = component_registry[parsed_params.dataloader][ - "class" - ]() log.debug("Storing dataset in %s", folder_path) - new_dataset = dataloader.load_data( - filepath_or_buffer=( - str(file_path) if file_path is not None else url - ), - temp_path=str(temp_dir) if temp_dir is not None else None, - params=parsed_params.model_dump(), + LoadUploadedDatasetUnit( + dataloader={ + "component": parsed_params.dataloader, + "params": parsed_params.model_dump(), + }, + source=str(file_path) if file_path is not None else url, + temp_path=str(temp_dir), n_sample=n_sample, - ) + )(ctx) + # The types either come with the request or are worked out + # from the data. Both paths end in the same context key, so + # the unit that applies them has a single input either way. if params.get("inferred_types"): - schema = params["inferred_types"] - elif new_dataset.types: - schema = { - col: typ.to_string() - for col, typ in new_dataset.types.items() - } + ctx.put_ref("inferred_types", params["inferred_types"]) else: - schema = infer_types( - new_dataset.to_pandas(), method="DashAIPtype" - ) - if "column_renames" in params: - renames = params["column_renames"] - original_names = new_dataset.arrow_table.schema.names - new_names = [renames.get(col, col) for col in original_names] - - if len(new_names) != len(set(new_names)): - duplicate_names = set() - seen = set() - for name in new_names: - if name in seen: - duplicate_names.add(name) - else: - seen.add(name) - msg = ( - "Invalid column_renames: resulting column names " - "contain duplicates: " - f"{sorted(duplicate_names)}" - ) - raise JobError(msg) - - arrow_table = new_dataset.arrow_table.rename_columns(new_names) - new_dataset = new_dataset.__class__( - arrow_table, - splits=new_dataset.splits, - types=new_dataset.types, - ) - schema = {renames.get(col, col): schema[col] for col in schema} + InferDatasetTypesUnit(method="DashAIPtype")(ctx) - new_dataset = transform_dataset_with_schema(new_dataset, schema) + ApplyDatasetSchemaUnit( + column_renames=params.get("column_renames"), + )(ctx) self.report_progress(0.5, "Computing metadata") - compute_meta = params.get("compute_metadata", True) - extended_keys = ( - "general_info", - "numeric_stats", - "categorical_stats", - "text_stats", - "quality_info", - "correlations", - ) - - if from_notebook_no_converters: - # No converters applied - saved data matches the source - # dataset byte-for-byte. Reuse the source's splits.json - # (already loaded into ``new_dataset.splits`` by - # ``load_dataset``) instead of recomputing. - has_extended = any(k in new_dataset.splits for k in extended_keys) - if compute_meta: - # Use source's full metadata if present, else compute. - if not has_extended: - new_dataset.compute_metadata() - else: - # Keep only base metadata; drop any inherited extended. - if "total_rows" not in new_dataset.splits: - new_dataset.compute_base_metadata() - for stale_key in extended_keys: - new_dataset.splits.pop(stale_key, None) - elif compute_meta: - new_dataset.compute_metadata() - else: - new_dataset.compute_base_metadata() - # Defensive: strip any extended keys that may have been - # inherited from a source dataset (e.g. notebook flow - # with converters that ran before this rule existed). - for stale_key in extended_keys: - new_dataset.splits.pop(stale_key, None) + # ``from_notebook_no_converters`` means the saved data matches + # the source dataset byte-for-byte, so the metadata it arrived + # with still describes it. That is a policy decision this job + # makes from the notebook's converter history; the unit only + # needs the answer. + ComputeDatasetMetadataUnit( + compute_metadata=params.get("compute_metadata", True), + trust_inherited_metadata=from_notebook_no_converters, + )(ctx) gc.collect() self.report_progress(0.8, "Saving dataset") dataset_save_path = folder_path / "dataset" log.debug("Saving dataset in %s", str(dataset_save_path)) - save_dataset(new_dataset, dataset_save_path) + SaveDatasetToPathUnit(path=str(dataset_save_path))(ctx) except Exception as e: log.exception(e) - shutil.rmtree(folder_path, ignore_errors=True) + if folder_is_ours: + shutil.rmtree(folder_path, ignore_errors=True) raise JobError(f"Error loading dataset: {str(e)}") from e - # Add dataset to database + # Add dataset to database. The counts are read back off the dataset + # rather than published by the metadata unit: they describe the + # dataset as it is right now, and a key holding them would go stale + # the moment anything else transformed it. + stored_metadata = ctx.require("dataset").splits with session_factory() as db: log.debug("Storing dataset metadata in database.") try: folder_path = os.path.realpath(folder_path) dataset = db.get(Dataset, dataset_id) + # Re-read, so it can be gone by now: the row is deletable + # through the API while the job runs. Without this guard the + # assignment below raises AttributeError on None, and the + # data just written to disk is orphaned. + if dataset is None: + raise JobError( + f"Dataset with ID {dataset_id} no longer exists." + ) dataset.file_path = folder_path - dataset.total_rows = new_dataset.splits.get("total_rows") - dataset.total_columns = len( - new_dataset.splits.get("column_names", []) - ) + dataset.total_rows = stored_metadata.get("total_rows") + dataset.total_columns = len(stored_metadata.get("column_names", [])) dataset.set_status_as_finished() db.commit() db.refresh(dataset) except exc.SQLAlchemyError as e: log.exception(e) - shutil.rmtree(folder_path, ignore_errors=True) + if folder_is_ours: + shutil.rmtree(folder_path, ignore_errors=True) raise JobError("Internal database error") from e + except Exception: + # Anything else here leaves a dataset on disk that no row + # points at, so it has to be cleaned up too. Re-raised + # unchanged; the handler below writes the status. + if folder_is_ours: + shutil.rmtree(folder_path, ignore_errors=True) + raise log.debug("Dataset creation successfully finished.") - except JobError as e: + except Exception as e: + # Every failure, not just ``JobError``. Nothing else moves this row + # out of STARTED: Huey's error signal writes to its own ``task_copy`` + # table and never touches ``Dataset``, so an exception this handler + # does not see leaves the dataset stuck as in-progress forever, and + # the UI keeps showing a spinner for work that already died. log.error(f"Dataset creation failed: {e}") - with session_factory() as db: - dataset = db.get(Dataset, dataset_id) - if dataset: - dataset.set_status_as_error() - db.commit() - db.refresh(dataset) - raise e + try: + with session_factory() as db: + dataset = db.get(Dataset, dataset_id) + if dataset: + dataset.set_status_as_error() + # The path is written to the row as soon as the folder + # is created, so a cancelled job can find it and delete + # it. The failure paths above already deleted it, so + # leaving the path behind would point the row at a + # folder that is gone -- and the next re-import would + # take it for a folder it may reuse. + # + # Only a folder this run created: a re-import that + # failed leaves the previous one in place, and the row + # still points at real data. + if folder_is_ours: + dataset.file_path = "" + db.commit() + db.refresh(dataset) + except Exception as bookkeeping_error: + # Never let the bookkeeping mask what actually went wrong. + log.exception(bookkeeping_error) + # Re-raised as-is: the message is the contract, and wrapping it here + # would change what every branch above reports. + raise finally: + ctx.clear_cache() gc.collect() if temp_dir and os.path.exists(temp_dir): try: diff --git a/DashAI/back/job/explainer_job.py b/DashAI/back/job/explainer_job.py index 2033353e1..adc0898df 100644 --- a/DashAI/back/job/explainer_job.py +++ b/DashAI/back/job/explainer_job.py @@ -1,5 +1,5 @@ import logging -from typing import TYPE_CHECKING, Any, Dict, Tuple +from typing import TYPE_CHECKING from kink import inject from sqlalchemy import exc @@ -11,18 +11,25 @@ ModelSession, Run, ) -from DashAI.back.explainability.global_explainer import BaseGlobalExplainer -from DashAI.back.explainability.local_explainer import BaseLocalExplainer from DashAI.back.job.base_job import BaseJob, JobError -from DashAI.back.models.base_model import BaseModel from DashAI.back.splitters.splits_payload import ( explainable_indexes, splitter_class_for, ) -from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_global_explainer_unit import BuildGlobalExplainerUnit +from DashAI.back.units.build_local_explainer_unit import BuildLocalExplainerUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.generate_global_explanation_unit import ( + GenerateGlobalExplanationUnit, +) +from DashAI.back.units.generate_local_explanation_unit import ( + GenerateLocalExplanationUnit, +) +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_run_model_unit import LoadRunModelUnit +from DashAI.back.units.prepare_explanation_data_unit import PrepareExplanationDataUnit if TYPE_CHECKING: - from datasets import DatasetDict from sqlalchemy.orm import sessionmaker logging.basicConfig(level=logging.DEBUG) @@ -120,282 +127,6 @@ def get_job_name(self) -> str: return f"{explainer_scope.capitalize()} Explanation ({explainer_id})" - @inject - def _generate_global_explanation( - self, - explainer: BaseGlobalExplainer, - dataset=Tuple["DatasetDict", "DatasetDict"], - ) -> None: - import os - import pickle - - from kink import di - - from DashAI.back.core.artifacts import normalize_artifacts - - explainer_id: int = self.kwargs["explainer_id"] - session_factory = di["session_factory"] - config = di["config"] - with session_factory() as db: - try: - explanation = explainer.explain(dataset) - plot = normalize_artifacts(explainer.plot(explanation)) - except Exception as e: - log.exception(e) - raise JobError( - "Failed to generate the explanation", - ) from e - try: - from pathlib import Path as _Path - - from DashAI.back.core.atomic import atomic_open - - explanation_filename = f"global_explanation_{explainer_id}.pickle" - explanation_path = os.path.join( - config["EXPLANATIONS_PATH"], explanation_filename - ) - with atomic_open(_Path(explanation_path), "wb") as file: - pickle.dump(explanation, file) - - plot_filename = f"global_explanation_plot_{explainer_id}.pickle" - plot_path = os.path.join(config["EXPLANATIONS_PATH"], plot_filename) - with atomic_open(_Path(plot_path), "wb") as file: - pickle.dump(plot, file) - - except Exception as e: - log.exception(e) - raise JobError( - "Explanation file saving failed", - ) from e - try: - self.explainer_db.explanation_path = explanation_path - self.explainer_db.plot_path = plot_path - self.explainer_db.plot_overrides = None - db.commit() - except Exception as e: - log.exception(e) - raise JobError( - "Explanation path saving failed", - ) from e - - @inject - def _generate_local_explanation( - self, - explainer: BaseLocalExplainer, - dataset: Tuple["DatasetDict", "DatasetDict"], - splits: Dict[str, Any], - task: BaseTask, - same_dataset: bool, - ) -> None: - import json - import os - import pickle - - from datasets import DatasetDict - from kink import di - - from DashAI.back.core.artifacts import normalize_artifacts - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - prepare_for_model_session, - save_dataset, - select_columns, - split_dataset, - ) - - explainer_id: int = self.kwargs["explainer_id"] - session_factory = di["session_factory"] - config = di["config"] - - explainer.fit(dataset, **self.explainer_db.fit_parameters) - instance_id = self.explainer_db.dataset_id - with session_factory() as db: - instance: Dataset = db.get(Dataset, instance_id) - if not instance: - raise JobError( - f"Dataset {instance_id} to be explained does not exist in DB." - ) - try: - loaded_instance = load_dataset(f"{instance.file_path}/dataset") - except Exception as e: - log.exception(e) - raise JobError( - f"Can not load instance from path {instance.file_path}", - ) from e - try: - # The data source is selected via scope["mode"]. It defaults to - # "split" so explainers created before this field existed keep - # their original split + percentage behavior. - mode = self.explainer_db.scope.get("mode", "split") - - if mode == "manual": - # Build the instances from values the user typed in by hand, - # reusing the same conversion the manual prediction flow uses. - # The rows (and any image files rewritten by the job endpoint) - # travel in the job kwargs, not in scope. - manual_input_data = self.kwargs.get("manual_input_data") or [] - if not manual_input_data: - raise JobError( - "No manual input data provided for the explanation" - ) - prepared_instance = task.process_manual_input( - manual_input_data, - f"{instance.file_path}/dataset", - ) - # Manual input carries only the input columns (no target), so - # keep just those instead of the standard input/output split. - # select_columns returns a DashAIDataset (same shape the - # split path produces), which is what the explainers expect. - X = prepared_instance.select_columns(self.input_columns) - else: - prepared_instance = task.prepare_for_task( - loaded_instance, - input_columns=self.input_columns, - output_columns=self.output_columns, - ) - - if mode == "rows": - # Explain a set of rows the user marked in the table. - # Indexes are over the whole dataset (the split does not - # apply in this mode). - row_indexes = self.explainer_db.scope.get("row_indexes") or [] - valid_indexes = [ - i - for i in row_indexes - if isinstance(i, int) - and 0 <= i < prepared_instance.num_rows - ] - if row_indexes and not valid_indexes: - raise JobError( - "No valid row indexes provided for the explanation" - ) - if valid_indexes: - prepared_instance = prepared_instance.select(valid_indexes) - else: - split = self.explainer_db.scope.get("split") - valid_splits = ["train", "val", "all", "test"] - if split not in valid_splits: - raise JobError(f"{split} is not a valid split") - - if split != "all": - if not same_dataset: - if isinstance(splits, str): - splits = json.loads(splits) - if "train" not in splits: - # A cross-validation session describes folds, - # not proportions, so there is no split of its - # own to apply to a dataset it never saw. - raise JobError( - "This run was cross-validated, so a split " - "of another dataset cannot be derived from " - "it. Explain specific rows, manual input, " - "or the whole dataset instead." - ) - ( - prepared_dataset_dict, - splits, - ) = prepare_for_model_session( - dataset=prepared_instance, - splits=splits, - output_columns=self.output_columns, - ) - split_key = "validation" if split == "val" else split - prepared_instance = prepared_dataset_dict[split_key] - else: - prepared_instance = split_dataset( - prepared_instance, - train_indexes=splits["train_indexes"], - test_indexes=splits["test_indexes"], - val_indexes=splits["val_indexes"], - ) - split_key = "validation" if split == "val" else split - prepared_instance = prepared_instance[split_key] - - n_rows = max( - 1, - int( - prepared_instance.num_rows - * self.explainer_db.scope.get("percentage") - / 100 - ), - ) - # When "shuffle" is set the percentage is taken as a random - # sample of the split; otherwise it is the leading rows. - if self.explainer_db.scope.get("shuffle"): - prepared_instance = prepared_instance.shuffle(seed=42) - prepared_instance = prepared_instance.select(range(n_rows)) - - prepared_instance = DatasetDict({"train": prepared_instance}) - X, _ = select_columns( - prepared_instance, - self.input_columns, - self.output_columns, - ) - # Persist the original selected rows (the model input for each - # explained instance) as a DashAIDataset before the model's own - # preprocessing runs, so the frontend can read them back with - # the existing dataset endpoints. - input_source = X["train"] if isinstance(X, DatasetDict) else X - input_dataset_path = os.path.join( - config["EXPLANATIONS_PATH"], - f"local_explanation_input_{explainer_id}", - ) - save_dataset(input_source, os.path.join(input_dataset_path, "dataset")) - # The instances are handed over unprepared, the same way the - # prediction job calls model.predict: the model applies its own - # preprocessing. Explainers that need the model feature space - # ask for it with prepare_model_input. - - except Exception as e: - log.exception(e) - raise JobError( - f"""Can not prepare Dataset with {instance_id} - to generate the local explanation.""", - ) from e - try: - explanation = explainer.explain_instance(X) - plots = normalize_artifacts( - explainer.plot(explanation), create_grouped=True - ) - except Exception as e: - log.exception(e) - raise JobError( - "Failed to generate the explanation", - ) from e - try: - from pathlib import Path as _Path - - from DashAI.back.core.atomic import atomic_open - - explanation_filename = f"local_explanation_{explainer_id}.pickle" - explanation_path = os.path.join( - config["EXPLANATIONS_PATH"], explanation_filename - ) - with atomic_open(_Path(explanation_path), "wb") as file: - pickle.dump(explanation, file) - - plots_filename = f"local_explanation_plots_{explainer_id}.pickle" - plots_path = os.path.join(config["EXPLANATIONS_PATH"], plots_filename) - with atomic_open(_Path(plots_path), "wb") as file: - pickle.dump(plots, file) - - except Exception as e: - log.exception(e) - raise JobError( - "Explanation file saving failed", - ) from e - try: - self.explainer_db.explanation_path = explanation_path - self.explainer_db.plots_path = plots_path - self.explainer_db.input_dataset_path = input_dataset_path - self.explainer_db.plot_overrides = None - db.commit() - except Exception as e: - log.exception(e) - raise JobError( - "Explanation path saving failed", - ) from e - @inject def run( self, @@ -404,17 +135,13 @@ def run( from kink import di - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - select_columns, - split_dataset, - ) - - component_registry = di["component_registry"] session_factory = di["session_factory"] explainer_id: int = self.kwargs["explainer_id"] explainer_scope: str = self.kwargs["explainer_scope"] + + ctx = ExecutionContext() + with session_factory() as db: if explainer_scope == "global": self.explainer_db: GlobalExplainer = db.get( @@ -425,6 +152,14 @@ def run( else: raise JobError(f"{explainer_scope} is an invalid explainer type") + if not self.explainer_db: + # Checked before the try below, whose handler would otherwise + # be the thing that crashes: it calls set_status_as_error on + # this very row. + raise JobError( + f"Explainer with id {explainer_id} does not exist in DB." + ) + try: run: Run = db.get(Run, self.explainer_db.run_id) if not run: @@ -438,138 +173,90 @@ def run( ) dataset: Dataset = db.get(Dataset, model_session.dataset_id) if not dataset: + # The id named here is the one that was looked up. It used + # to interpolate the explainer's own dataset_id, a column + # global explainers do not even have. raise JobError( - f"Dataset {self.explainer_db.dataset_id} does not exist in DB." + f"Dataset {model_session.dataset_id} does not exist in DB." ) - self.input_columns = model_session.input_columns - self.output_columns = model_session.output_columns + self.explainer_db.huey_id = self.kwargs.get("huey_id", None) + db.commit() - try: - run_model_class = component_registry[run.model_name]["class"] - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to find Model with name {run.model_name} in registry.", - ) from e - try: - model: BaseModel = run_model_class(**run.parameters) - except Exception as e: - log.exception(e) - raise JobError("Unable to instantiate model") from e - try: - trained_model = model.load(run.run_path) - except Exception as e: - log.exception(e) - raise JobError( - f"Can not load model from path {run.run_path}" - ) from e - try: - explainer_class = component_registry[ - self.explainer_db.explainer_name - ]["class"] - except Exception as e: - log.exception(e) - raise JobError( - f"""Unable to find the {explainer_scope} explainer with name - {self.explainer_db.explainer_name} in registry.""", - ) from e + input_columns = model_session.input_columns + output_columns = model_session.output_columns + + LoadRunModelUnit(run_id=run.id)(ctx) + + # How the explainer configuration is stored on the row — the + # component name and its parameters live in separate columns — + # rather than part of the explanation itself. + explainer_config = { + "component": self.explainer_db.explainer_name, + "params": self.explainer_db.parameters, + } + build_explainer = ( + BuildGlobalExplainerUnit + if explainer_scope == "global" + else BuildLocalExplainerUnit + ) + build_explainer(explainer=explainer_config)(ctx) + + LoadDatasetUnit(dataset_id=model_session.dataset_id)(ctx) + + prepare = PrepareExplanationDataUnit( + task_name=model_session.task_name, + input_columns=input_columns, + output_columns=output_columns, + ) + # Resolving the task outside the wrapper below keeps a missing + # task reported as a registry problem rather than a generic + # "cannot prepare" message. + prepare.validate(ctx) try: - explainer = explainer_class( - model=trained_model, **self.explainer_db.parameters - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to instantiate {explainer_scope} explainer.", - ) from e - try: - loaded_dataset: "DatasetDict" = load_dataset( - f"{dataset.file_path}/dataset" - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Can not load dataset from path {dataset.file_path}", - ) from e - try: - task: BaseTask = component_registry[model_session.task_name][ - "class" - ]() - except Exception as e: - log.exception(e) - raise JobError( - ( - f"Unable to find Task with name {model_session.task_name} " - "in registry" - ), - ) from e - try: - # The splitter that produced the run decides which - # partitions exist and what they are called, so it is asked - # rather than guessing from the payload's shape. - session_splits = model_session.splits - if isinstance(session_splits, str): - session_splits = json.loads(session_splits) + # Which partitions a run has, and what they are called, is + # the splitter's answer and not something to read off the + # shape of the payload: a holdout run stores one flat + # mapping and a cross-validated one stores an entry per + # fold plus the pooled rows. The splitter maps whichever it + # produced onto the three slots the explainers are built + # from, so a splitter added later needs no change here. + from kink import di + splitter_class = splitter_class_for( - session_splits, component_registry + json.loads(model_session.splits), di["component_registry"] ) - train_idx, test_idx, val_idx = explainable_indexes( + train_idx, evaluation_idx, val_idx = explainable_indexes( splitter_class, json.loads(run.split_indexes) ) - splits = { - "train_indexes": train_idx, - "test_indexes": test_idx, - "val_indexes": val_idx, - } except ValueError as e: + # Reported as itself: it says which of the two happened -- + # a payload that does not match its splitter, or a run that + # fitted on every row it had and so has nothing to explain. log.exception(e) raise JobError(str(e)) from e - try: - loaded_dataset = split_dataset( - loaded_dataset, - train_indexes=train_idx, - test_indexes=test_idx, - val_indexes=val_idx, - ) - - prepared_dataset = task.prepare_for_task( - dataset=loaded_dataset, - input_columns=self.input_columns, - output_columns=self.output_columns, - ) - data = select_columns( - prepared_dataset, - self.input_columns, - self.output_columns, - ) - data_x = split_dataset( - data[0], - train_indexes=train_idx, - test_indexes=test_idx, - val_indexes=val_idx, - ) - data_y = split_dataset( - data[1], - train_indexes=train_idx, - test_indexes=test_idx, - val_indexes=val_idx, + try: + # Unpacking the JSON column is an artifact of how the row + # stores it, but it stays inside this block because a + # malformed value has always been reported as a + # preparation failure. + ctx.put_ref( + "split_indexes", + { + "train_indexes": train_idx, + "test_indexes": evaluation_idx, + "val_indexes": val_idx, + }, ) - # Inputs stay unprepared (see the note in the local - # explanation path); targets are encoded because explainers - # compare them against the model's class indexes. - for split_name in data_y: - data_y[split_name] = trained_model.prepare_output( - data_y[split_name], is_fit=False - ) - + prepare(ctx) except Exception as e: log.exception(e) raise JobError( f"""Can not prepare dataset {dataset.id} for the explanation""", ) from e + try: self.explainer_db.set_status_as_started() db.commit() @@ -578,27 +265,44 @@ def run( raise JobError( "Connection with the database failed", ) from e - if explainer_scope == "global": - self._generate_global_explanation( - explainer=explainer, dataset=(data_x, data_y) - ) - elif explainer_scope == "local": + if explainer_scope == "global": + GenerateGlobalExplanationUnit(explainer_id=explainer_id)(ctx) + paths = { + "explanation_path": ctx.require("explanation_path"), + "plot_path": ctx.require("plot_path"), + } + else: same_dataset = ( model_session.dataset_id == self.explainer_db.dataset_id ) - if not same_dataset: - splits = model_session.splits - - self._generate_local_explanation( - explainer=explainer, - dataset=(data_x, data_y), - splits=splits, - task=task, + GenerateLocalExplanationUnit( + explainer_id=explainer_id, + instance_dataset_id=self.explainer_db.dataset_id, + scope=self.explainer_db.scope, + fit_parameters=self.explainer_db.fit_parameters, + input_columns=input_columns, + output_columns=output_columns, + manual_input_data=self.kwargs.get("manual_input_data"), same_dataset=same_dataset, - ) - else: - raise JobError(f"{explainer_scope} is an invalid explainer type") + session_splits=(None if same_dataset else model_session.splits), + )(ctx) + paths = { + "explanation_path": ctx.require("explanation_path"), + "plots_path": ctx.require("plots_path"), + "input_dataset_path": ctx.require("input_dataset_path"), + } + + try: + for column, value in paths.items(): + setattr(self.explainer_db, column, value) + self.explainer_db.plot_overrides = None + db.commit() + except Exception as e: + log.exception(e) + raise JobError( + "Explanation path saving failed", + ) from e self.explainer_db.set_status_as_finished() db.commit() @@ -607,3 +311,5 @@ def run( self.explainer_db.set_status_as_error() db.commit() raise e + finally: + ctx.clear_cache() diff --git a/DashAI/back/job/explorer_job.py b/DashAI/back/job/explorer_job.py index beb261e95..d4643bd0b 100644 --- a/DashAI/back/job/explorer_job.py +++ b/DashAI/back/job/explorer_job.py @@ -1,12 +1,15 @@ import logging -from typing import TYPE_CHECKING, Type +from typing import TYPE_CHECKING from kink import inject from sqlalchemy import exc from DashAI.back.dependencies.database.models import Explorer, Notebook -from DashAI.back.exploration.base_explorer import BaseExplorer from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.run_exploration_unit import RunExplorationUnit +from DashAI.back.units.save_exploration_unit import SaveExplorationUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -85,18 +88,15 @@ def get_job_name(self) -> str: def run( self, ) -> None: - import os - import pathlib - from kink import di - from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset from DashAI.back.exploration.artifact_store import store_artifacts - component_registry = di["component_registry"] session_factory = di["session_factory"] - config = di["config"] explorer_id: int = self.kwargs["explorer_id"] + + ctx = ExecutionContext() + with session_factory() as db: # Load the explorer information try: @@ -124,103 +124,61 @@ def run( explorer_info.set_status_as_error() db.commit() raise JobError("Error while loading the notebook info.") from e - - # Load the dataset from the notebook - try: - loaded_dataset = load_dataset(f"{notebook_info.file_path}/dataset") - except Exception as e: - log.exception(e) - explorer_info.set_status_as_error() - db.commit() - raise JobError( - f"Can not load dataset from path {notebook_info.file_path}", - ) from e - - # obtain the explorer component from the registry - try: - explorer_component_class: Type[BaseExplorer] = component_registry[ - explorer_info.exploration_type - ]["class"] - except KeyError as e: - log.exception(e) - explorer_info.set_status_as_error() - db.commit() - raise JobError( - ( - f"Explorer {explorer_info.exploration_type} " - "not found in the registry." - ) - ) from e - - # Instance the explorer (the explorer handles its validation) - try: - explorer_instance = explorer_component_class(**explorer_info.parameters) - assert isinstance(explorer_instance, BaseExplorer) - except Exception as e: - log.exception(e) + except Exception: + # A notebook that is simply not there used to escape the + # SQLAlchemyError handler above and leave the row STARTED + # forever, because nothing else marks it: the Huey error signal + # writes only to its own task_copy table, and + # _execute_base_job calls run() with no handler at all. + # Re-raised as-is so the "not found" message survives. explorer_info.set_status_as_error() db.commit() - raise JobError( - f"Error instancing the explorer {explorer_info.exploration_type}." - ) from e + raise - # prepare the dataset + # Load the dataset from the notebook: its own working copy, which + # is what the converters rewrite. try: - prepared_dataset = explorer_instance.prepare_dataset( - loaded_dataset, explorer_info.columns - ) + LoadDatasetUnit(notebook_id=notebook_info.id)(ctx) except Exception as e: + # Anything the load unit raises has to leave the row in ERROR. + # Nothing else marks it: the Huey error signal only writes to + # its own task_copy table, never to the Explorer row, so + # without this the exploration would stay STARTED forever. + # Re-raised as-is; the unit reports the same + # "Can not load dataset from path ..." message the job used to + # build here. log.exception(e) explorer_info.set_status_as_error() db.commit() - raise JobError( - ( - "Error preparing the dataset for the exploration " - f"{explorer_info.exploration_type}." - ) - ) from e - - # Launch the exploration + raise + + # How the exploration configuration is stored on the row — the + # component name and its parameters live in separate columns — + # rather than part of the exploration itself. + explorer = { + "component": explorer_info.exploration_type, + "params": explorer_info.parameters, + } + + # Run the exploration. The unit reports the registry, instancing, + # preparation and launch errors with the same texts the job used + # to build here; re-raised as-is so they reach the user intact. try: - result = explorer_instance.launch_exploration( - prepared_dataset, explorer_info - ) + RunExplorationUnit(explorer_id=explorer_id, explorer=explorer)(ctx) except Exception as e: log.exception(e) explorer_info.set_status_as_error() db.commit() - raise JobError( - f"Error launching the exploration {explorer_info.exploration_type}." - ) from e + raise # Save the result try: - # save in the notebook folder - save_path = pathlib.Path( - os.path.join( - config["NOTEBOOK_PATH"], - (f"{notebook_info.id}"), - ) - ) - if not save_path.exists(): - save_path.mkdir(parents=True) - - save_path = explorer_instance.save_notebook( - notebook_info, explorer_info, save_path, result - ) - if isinstance(save_path, str): - save_path = pathlib.Path(save_path) - if not isinstance(save_path, pathlib.Path): - raise JobError( - ( - f"Error while saving the exploration" - f" {explorer_info.exploration_type}" - f", save path is not a pathlib.Path." - ) - ) + SaveExplorationUnit(explorer_id=explorer_id)(ctx) - # Update the explorer info - explorer_info.exploration_path = save_path.as_posix() + # Update the explorer info. The status is not set to finished + # here: the artifacts below are part of the work, so the row + # only counts as done once they exist too. + explorer_info.exploration_path = ctx.require("exploration_path") db.commit() except Exception as e: log.exception(e) @@ -237,9 +195,17 @@ def run( # the explorer class is asked for its results: from here on the # stored artifacts are served as is, so the exploration keeps # rendering even if the explorer is removed from the registry. + # + # Both inputs come from the context rather than from local + # variables: the explorer instance is what ran the exploration + # (published by RunExplorationUnit, so the artifacts are built from + # the same object that produced the result, not a rebuilt one), and + # the path is where SaveExplorationUnit actually wrote it. try: explorer_info.artifacts_path = store_artifacts( - explorer_instance, save_path, explorer_info.id + ctx.require("explorer"), + ctx.require("exploration_path"), + explorer_info.id, ) explorer_info.set_status_as_finished() db.commit() @@ -253,3 +219,5 @@ def run( f"{explorer_info.exploration_type}." ) ) from e + finally: + ctx.clear_cache() diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index 7ae8528c7..0f16ede98 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -1,25 +1,36 @@ import logging -from typing import TYPE_CHECKING, Any, Dict, List +from typing import TYPE_CHECKING, Any, Dict from kink import inject from sqlalchemy import exc - -from DashAI.back.core.atomic import atomic_save_path -from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run -from DashAI.back.dependencies.downloads.nested import missing_downloads -from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy +from sqlalchemy.orm.attributes import flag_modified + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.dependencies.database.models import ( + Dataset, + Metric, + ModelSession, + Run, +) from DashAI.back.job.base_job import BaseJob, JobError -from DashAI.back.metrics.base_metric import BaseMetric -from DashAI.back.models.model_factory import ModelFactory from DashAI.back.optimizers.base_optimizer import BaseOptimizer -from DashAI.back.splitters.base_splitter import BaseSplitter from DashAI.back.splitters.splits_payload import normalize_splits_payload -from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit +from DashAI.back.units.fit_model_over_nested_folds_unit import ( + FitModelOverNestedFoldsUnit, +) +from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_model_unit import SaveModelUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker - from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -94,28 +105,53 @@ def run( ) -> None: import gc import json - import os from kink import di component_registry = di["component_registry"] session_factory = di["session_factory"] - config = di["config"] # Get the necessary parameters run_id: int = self.kwargs["run_id"] + ctx = ExecutionContext() + with session_factory() as db: run: Run = db.get(Run, run_id) + # Without this the next line raises AttributeError on None, which + # reaches the user as a stack trace rather than as the reason. + if not run: + raise JobError(f"Run {run_id} does not exist in DB.") run.huey_id = self.kwargs.get("huey_id", None) db.commit() self.report_progress(0.05, "Preparing data") try: try: - # Get the dataset and components prepared for the model training + # What is left in the helper is reading the configuration + # this run was created with off its rows. The work that + # configuration describes is done by the units below. preparation_results = self._prepare_dataset_and_components( run_id=run_id, db=db, component_registry=component_registry ) + model_session: ModelSession = preparation_results["model_session"] + prepare = preparation_results["prepare_unit"] + + LoadDatasetUnit(dataset_id=model_session.dataset_id)(ctx) + + build_model = BuildModelUnit( + model={ + "component": run.model_name, + "params": run.parameters, + }, + train_metrics=model_session.train_metrics, + validation_metrics=model_session.validation_metrics, + test_metrics=model_session.test_metrics, + run_id=run_id, + ) + # The download gate lives in validate(), and running it here + # keeps a model that cannot be trained from being reported + # as a splitting failure further down. + build_model.validate(ctx) except Exception as e: log.exception(e) raise JobError( @@ -123,23 +159,34 @@ def run( ) from e try: - # Get splits from the splitter - splitter: BaseSplitter = preparation_results["splitter"] - # Get the dataset splits between input columns and output column - X, Y = preparation_results["X"], preparation_results["Y"] - - # Get x,y but now splitted with train, validation and test indexes - # each one, and the indexes used for the splits - x, y, splits = splitter.split(X, Y) + # Preparing the dataset for the task and partitioning it are + # one step: how many partitions there are and what they are + # called is the splitter's answer, and the pair of units + # differ only in which family of splitters they offer and in + # the shape they publish for it. + prepare(ctx) + + # Only the partitions are needed here, and only to ask + # whether the session reserved any rows: the units read + # what they work on from the context themselves. + x = ctx.get("x") if ctx.has("x") else ctx.require("x_folds") # save the obtained splits into the database - run.split_indexes = json.dumps(splits) + run.split_indexes = json.dumps(ctx.require("split_indexes")) except Exception as e: log.exception(e) raise JobError( f"Error splitting the dataset for run {run_id}: {e}", ) from e + try: + build_model(ctx) + except Exception as e: + log.exception(e) + raise JobError( + f"Error preparing dataset and components for run {run_id}: {e}", + ) from e + try: run.set_status_as_started() db.commit() @@ -149,22 +196,106 @@ def run( "Connection with the database failed", ) from e + strategy_class = preparation_results["evaluation_strategy_class"] + # Which partitions a run records a score for is declared by the + # strategy the session chose: a forecaster is not judged on the + # dates it was fitted on, so its training partition is not in + # this list even though metrics are configured for it. + scored_splits = [split.name for split in strategy_class.SCORED_SPLITS] + self.report_progress(0.2, "Training") try: - # Hyperparameter Tunning plot_paths = [] - evaluation_estrategy: BaseEvaluationStrategy = preparation_results[ - "evaluation_strategy" - ] - - evaluation_estrategy.set_progress_reporter(self.report_progress) - model, plot_paths = evaluation_estrategy.execute( - x=x, - y=y, - run=run, - db=db, - ) + if getattr(strategy_class, "KIND", "holdout") == "holdout": + fit_model = FitModelUnit( + optimizer={ + "component": run.optimizer_name, + "params": run.optimizer_parameters, + }, + goal_metric=run.goal_metric, + run_id=run_id, + # The run names its own artifacts, which is what + # keeps the plot filenames of two runs apart inside + # the runs directory. + artifact_prefix=str(run_id), + # A trial never scores the test partition, so what + # it may record is whatever else the strategy scores. + trial_splits=[ + name for name in scored_splits if name != "TEST" + ], + ) + fit_model(ctx) + + plot_paths = ctx.require("plot_paths") + if ctx.has("best_parameters"): + run.parameters = ctx.get("best_parameters") + flag_modified(run, "parameters") + db.commit() + + self.report_progress(0.85, "Computing metrics") + EvaluateModelUnit(run_id=run_id, splits=scored_splits)(ctx) + else: + # Two units and not one with a flag: the nested one + # takes a required component field for its inner + # splitter, and a component field cannot be made + # optional without leaving the user without a selector. + fold_config = { + "optimizer": { + "component": run.optimizer_name, + "params": run.optimizer_parameters, + }, + "goal_metric": run.goal_metric, + "run_id": run_id, + "artifact_prefix": str(run_id), + "scored_splits": [ + name for name in scored_splits if name != "TEST" + ], + } + if run.nested: + fold_config["inner_splitter"] = { + "component": run.nested.get("splitter_name"), + "params": run.nested, + } + fit_folds = FitModelOverNestedFoldsUnit(**fold_config) + else: + fit_folds = FitModelOverFoldsUnit(**fold_config) + + fit_folds(ctx) + + plot_paths = ctx.require("plot_paths") + if ctx.has("best_parameters"): + run.parameters = ctx.get("best_parameters") + flag_modified(run, "parameters") + db.commit() + + self.report_progress(0.85, "Computing metrics") + if ctx.has("outer_fold_metrics"): + # Kept at its own level: it answers a different + # question from the ordinary summary -- how the + # procedure does, rather than how this model does -- + # and the two would be indistinguishable side by + # side. + self._aggregate_fold_metrics( + db, + run_id, + ctx.get("outer_fold_metrics"), + LevelEnum.LAST_OUTER, + ) + self._aggregate_fold_metrics( + db, run_id, ctx.require("fold_metrics"), LevelEnum.LAST + ) + + # The rows the session reserved are the only ones no + # fold and no trial ever saw, so they are the only + # honest estimate left once a model is picked out of a + # comparison table -- and scoring them is an ordinary + # LAST metric, so it is the same unit a holdout run + # uses. Whether there is anything to score is the + # caller's to know: a session that reserved nothing + # leaves that partition empty rather than absent. + if "TEST" in scored_splits and len(x[-1]["test"]) > 0: + EvaluateModelUnit(run_id=run_id, splits=["TEST"])(ctx) except Exception as e: log.exception(e) raise JobError( @@ -187,18 +318,10 @@ def run( ) from e self.report_progress(0.95, "Saving model") - try: - run_path = os.path.join(config["RUNS_PATH"], str(run.id)) - with atomic_save_path(run_path) as tmp_run_path: - model.save(str(tmp_run_path)) - except Exception as e: - log.exception(e) - raise JobError( - "Model saving failed", - ) from e + SaveModelUnit(artifact_prefix=str(run_id))(ctx) try: - run.run_path = run_path + run.run_path = ctx.require("model_path") db.commit() except exc.SQLAlchemyError as e: log.exception(e) @@ -221,53 +344,114 @@ def run( db.commit() raise e finally: + ctx.clear_cache() gc.collect() + @staticmethod + def _aggregate_fold_metrics( + db, run_id: int, fold_metrics: Dict[str, Any], level: LevelEnum + ) -> None: + """Summarise the per-fold scores into one row per split and metric. + + The unit that fitted the folds publishes their scores rather than + aggregating them, because a summary row carries a standard deviation + and a unit may not write domain rows -- the one sanctioned write in the + domain layer has nowhere to put one. So the arithmetic and the writing + happen here, where every other row this job persists is written. + + A single fold gets a deviation of zero rather than none: none is what + the reserved-rows measurement carries, and the two say different + things -- "one fold, so nothing varied" against "not a summary at all". + + Parameters + ---------- + db : Session + The session this job is already holding. + run_id : int + The run the rows belong to. + fold_metrics : dict + ``{split name: {metric name: [one score per fold]}}``. + level : LevelEnum + Where the summary goes. The ordinary fold scores summarise to + ``LAST``; the outer folds of a nested run summarise to + ``LAST_OUTER``, because they answer a different question and would + be indistinguishable from the first if they shared a level. + """ + import numpy as np + + for split_name, by_metric in fold_metrics.items(): + for metric_name, values in by_metric.items(): + if not values: + continue + existing = ( + db.query(Metric) + .filter_by( + run_id=run_id, + split=SplitEnum[split_name], + level=level, + name=metric_name, + ) + .first() + ) + mean = float(np.mean(values)) + deviation = float(np.std(values)) if len(values) > 1 else 0.0 + + if existing: + existing.value = mean + existing.std_value = deviation + else: + db.add( + Metric( + run_id=run_id, + split=SplitEnum[split_name], + level=level, + name=metric_name, + value=mean, + std_value=deviation, + step=0, + ) + ) + db.commit() + def _prepare_dataset_and_components( self, run_id: int, db, component_registry ) -> Dict[str, Any]: - """Prepare the dataset, task, splitter, metrics, model, and evaluation strategy. + """Read the configuration this run was created with, off its rows. - This helper resolves the persisted training configuration for a run, - loads the associated dataset from disk, prepares it for the selected - task, instantiates the required components from the component registry, - and builds the model factory together with the evaluation strategy. + What is resolved here is what the units cannot: rows, the JSON columns + stored on them, and the choice of which unit prepares the data. The + work those rows describe -- loading the dataset, validating it against + the task, separating features from targets, partitioning them and + building the model -- belongs to the units and happens in ``run``. Parameters ---------- run_id : int - Identifier of the training run whose configuration and artifacts must - be loaded. + Identifier of the training run whose configuration must be read. db : object - Database access object used to retrieve the run, model session, and - related persisted entities. + Database session used to retrieve the run and its model session. component_registry : object - Registry containing the available task, splitter, metric, model, - optimizer, and evaluation strategy implementations. + Registry used to resolve the splitter, the optimizer and the + evaluation strategy the session names. Returns ------- dict - A dictionary containing the prepared input and output datasets, the - instantiated splitter, and the evaluation strategy. + The model session, the unit that will prepare and partition the + dataset, the evaluation strategy class, and the optimizer and goal + metric it will be built with. Raises ------ JobError - If the run, model session, dataset, task, splitter, metrics, model, - optimizer, or evaluation strategy cannot be resolved or instantiated. + If the run's session, its splits payload, its splitter, its + optimizer or its evaluation strategy cannot be resolved. """ import json - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - select_columns, - ) - run: Run = db.get(Run, run_id) - # Get the model session and dataset from the database model_session: ModelSession = db.get(ModelSession, run.model_session_id) if not model_session: raise JobError( @@ -279,71 +463,13 @@ def _prepare_dataset_and_components( raise JobError(f"Dataset {model_session.dataset_id} does not exist in DB.") try: - # Load dataset from the file path - loaded_dataset: "DashAIDataset" = load_dataset( - f"{dataset.file_path}/dataset" - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Can not load dataset from path {dataset.file_path}", - ) from e - - try: - # Get task from model session - task: BaseTask = component_registry[model_session.task_name]["class"]() - except Exception as e: - log.exception(e) - raise JobError( - ( - f"Unable to find Task with name {model_session.task_name} " - "in registry" - ), - ) from e - - try: - # Prepare dataset for the task and get number of labels of the task - prepared_dataset = task.prepare_for_task( - dataset=loaded_dataset, - input_columns=model_session.input_columns, - output_columns=model_session.output_columns, - ) - n_labels = task.num_labels( - prepared_dataset, model_session.output_columns[0] - ) - except Exception as e: - log.exception(e) - raise JobError( - f"""Can not prepare Dataset {dataset.id} - for Task {model_session.task_name}""", - ) from e - - try: - # Divide the dataset into two datasets: - # one with the input columns and another with the output column. - # This reads the prepared dataset rather than the loaded one: a - # task may reorder or otherwise adjust the rows, and forecasting - # does, sorting them by date so the temporal splitter carves real - # periods of time. Selecting from the loaded dataset would drop - # that work on the floor. - X, Y = select_columns( - prepared_dataset, - model_session.input_columns, - model_session.output_columns, - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Error selecting input and output columns from dataset {dataset.id}", - ) from e - - try: - # Get splits data from model session + # Unpacking the JSON column is an artifact of how the row stores + # it, not part of the split. Sessions created before the payload + # followed the splitter schema use different keys for the seed and + # for manual indexes. splits_data = json.loads(model_session.splits) if run.split_indexes: splits_data["splitted_indexes"] = json.loads(run.split_indexes) - # Sessions created before the splits payload followed the splitter - # schema use different keys for the seed and for manual indexes. splits_data = normalize_splits_payload(splits_data) except Exception as e: log.exception(e) @@ -352,61 +478,28 @@ def _prepare_dataset_and_components( ) from e try: - # Get the splitter class from the registry and split the dataset splitter_name = splits_data.get("splitter_name", None) - splitter: BaseSplitter = component_registry[splitter_name]["class"]( - splits_data=splits_data, - ) + splitter_class = component_registry[splitter_name]["class"] except Exception as e: log.exception(e) raise JobError( - f"""Unable to find Splitter with name - {splitter_name} in registry.""", + f"Unable to find Splitter with name {splitter_name} in registry.", ) from e - try: - # Get metrics from model session - train_metrics: List[BaseMetric] = [ - component_registry[m]["class"] for m in model_session.train_metrics - ] - validation_metrics: List[BaseMetric] = [ - component_registry[m]["class"] for m in model_session.validation_metrics - ] - test_metrics: List[BaseMetric] = [ - component_registry[m]["class"] for m in model_session.test_metrics - ] - except Exception as e: - log.exception(e) - raise JobError( - "Unable to find metrics associated with" - f"Task {model_session.task_name} in registry", - ) from e - - try: - # Get the model class from the registry - run_model_class = component_registry[run.model_name]["class"] - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to find Model with name {run.model_name} in registry.", - ) from e - - # Make sure the model (and any nested components) are downloaded - # before attempting to train, otherwise fail fast with a clear error. - if getattr(run_model_class, "REQUIRES_DOWNLOAD", False) and not ( - run_model_class.is_downloaded() - ): - raise JobError( - f"Model {run.model_name} is not downloaded. " - "Download it before training." - ) - nested_missing = missing_downloads(run.parameters, component_registry) - if nested_missing: - names = ", ".join(m["name"] for m in nested_missing) - raise JobError( - "These components are not downloaded. " - f"Download them before training: {names}." - ) + # Which unit prepares the data follows from how the splitter carves it, + # which the splitter declares. The two units publish different shapes, + # so this is a choice of unit and not a flag on one. + prepare_class = ( + PrepareAndFoldUnit + if getattr(splitter_class, "PARTITIONING", "holdout") == "folds" + else PrepareAndSplitUnit + ) + prepare_unit = prepare_class( + task_name=model_session.task_name, + input_columns=model_session.input_columns, + output_columns=model_session.output_columns, + splitter={"component": splitter_name, "params": splits_data}, + ) try: # Get the optimizer if defined @@ -426,32 +519,9 @@ def _prepare_dataset_and_components( ) from e try: - # Instantiate the model using the ModelFactory - # and get the optimizable parameters - factory = ModelFactory( - model=run_model_class, - params=run.parameters, - run_id=run_id, - train_metrics=train_metrics, - validation_metrics=validation_metrics, - test_metrics=test_metrics, - n_labels=n_labels, - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to instantiate model using run {run_id}", - ) from e - - try: - # Get the evaluation strategy for the model session - evaluation_strategy: BaseEvaluationStrategy = component_registry[ + evaluation_strategy_class = component_registry[ model_session.evaluation_strategy - ]["class"]( - factory=factory, - optimizer=optimizer, - goal_metric=goal_metric, - ) + ]["class"] except Exception as e: log.exception(e) raise JobError( @@ -461,8 +531,9 @@ def _prepare_dataset_and_components( ) from e return { - "X": X, - "Y": Y, - "splitter": splitter, - "evaluation_strategy": evaluation_strategy, + "model_session": model_session, + "prepare_unit": prepare_unit, + "evaluation_strategy_class": evaluation_strategy_class, + "optimizer": optimizer, + "goal_metric": goal_metric, } diff --git a/DashAI/back/job/pipeline_job.py b/DashAI/back/job/pipeline_job.py index 1e679b4f5..cd1dff625 100644 --- a/DashAI/back/job/pipeline_job.py +++ b/DashAI/back/job/pipeline_job.py @@ -1,118 +1,180 @@ +"""Job that runs a pipeline as a graph of units.""" + import logging -from typing import TYPE_CHECKING, Any, Dict, List +from typing import TYPE_CHECKING -from kink import di +from kink import inject -from DashAI.back.dependencies.database.models import Pipeline +from DashAI.back.dag.engine import run as run_graph +from DashAI.back.dag.expand import expand +from DashAI.back.dag.graph import GraphError +from DashAI.back.dag.tracking import DatabaseSink from DashAI.back.job.base_job import BaseJob, JobError if TYPE_CHECKING: - from sqlalchemy.orm import Session - - from DashAI.back.dependencies.registry import ComponentRegistry + from sqlalchemy.orm import sessionmaker log = logging.getLogger(__name__) class PipelineJob(BaseJob): - ISOLATED = ( - False # async run() + live SQLAlchemy Session in kwargs — not serializable - ) + """Run a pipeline: expand its blocks into units and execute the graph. + + The job owns what a job owns -- the database session, the run row, its + status transitions -- and the engine owns execution. Between them sits the + expansion, which needs the run's id: naming a node's artifacts has to + happen before the graph is built, so the row is created here first and the + engine never creates anything. + + ``kwargs`` is a single id, as every job's is: the whole job is serialized + with dill to reach the worker process, which rebuilds its dependencies from + a fresh container, so nothing but plain data can travel in it. The graph is + read from the database inside ``run``. + """ + + @staticmethod + def _pipeline_id(kwargs) -> int: + """The id of the pipeline to run, under either name it arrives as. + + The front sends ``{"id": …}`` -- the wire contract predates this job -- + while every sibling job names its own subject (``run_id``, + ``converter_id``). Both are accepted so the existing caller keeps + working and a new one can use the name that matches the others; the + fallback goes when the endpoint is rewritten. + """ + pipeline_id = kwargs.get("pipeline_id", kwargs.get("id")) + if pipeline_id is None: + raise JobError("No pipeline id was given to run. Send it as 'pipeline_id'.") + return pipeline_id def set_status_as_delivered(self) -> None: - pass + """Nothing to mark: the run row does not exist until the job starts. - async def run( - self, - component_registry: "ComponentRegistry" = lambda di: di["component_registry"], - ) -> None: - db: "Session" = self.kwargs["db"] - id: int = self.kwargs.get("id", None) - pipeline: Pipeline = db.get(Pipeline, id) - steps: List[Dict[str, Any]] = self.kwargs.get("steps", []) or pipeline.steps - - if not id: - raise JobError("No id provided to execute the pipeline.") - if not steps: - raise JobError("No steps provided to execute the pipeline.") - - if not steps: - raise JobError("Pipeline has no steps to execute") - - log.info(f"Starting pipeline execution for pipeline {id}...") - - context: Dict[str, Any] = {"pipeline_id": id} - - for idx, step in enumerate(steps): - node_id = step.get("id") - node_type = step.get("type") - node_config = step.get("config", {}) - - log.debug( - f"Pipeline {id}: Executing step {idx + 1}/{len(steps)} - " - f"{node_type} ({node_id})" - ) - - try: - node_class = component_registry(di)[node_type]["class"] - except KeyError as e: - error_msg = f"Component type {node_type} not found in registry" - raise JobError( - f"Error in node {node_id} ({node_type}): {error_msg}" - ) from e - - try: - node_instance = node_class(**node_config) - except Exception as e: - error_msg = f"Error in node {node_id} ({node_type}): {str(e)}" - log.exception(error_msg) - raise JobError(error_msg) from e - - try: - output = await node_instance.run(context=context) - self._update_context(context, pipeline, node_type, node_id, output) - log.debug(f"Node {node_id} executed successfully.") - - except Exception as e: - error_msg = f"Error in node {node_id} ({node_type}): {str(e)}" - log.exception(error_msg) - raise JobError(error_msg) from e - - log.info(f"Pipeline {id} execution completed successfully.") - db.add(pipeline) - db.commit() - self.set_status_as_delivered() - - def _update_context( + A pipeline has no per-pipeline status, and the run row is created in + ``run`` because expanding the graph needs its id. Recording the + delivered state would mean creating the row when the job is enqueued + instead, which belongs with the endpoint that enqueues it. + """ + log.debug("PipelineJob delivered; no row to mark yet.") + + def set_status_as_error(self) -> None: + """Nothing to mark. + + This is called on a job deleted while still queued, before ``run`` has + created anything. A run that has started and then failed is marked by + the sink, from inside ``run``. + """ + log.debug("PipelineJob errored before starting; no row to mark.") + + @inject + def get_job_name(self) -> str: + """Get a descriptive name for the job.""" + from kink import di + + from DashAI.back.dependencies.database.models import Pipeline + + try: + pipeline_id = self._pipeline_id(self.kwargs) + except JobError: + return "Pipeline" + + try: + with di["session_factory"]() as db: + pipeline = db.get(Pipeline, pipeline_id) + if pipeline and pipeline.name: + return f"Pipeline: {pipeline.name}" + except Exception: + pass + + return f"Pipeline ({pipeline_id})" + + @inject + def run( self, - context: Dict[str, Any], - pipeline: Pipeline, - node_type: str, - node_id: str, - output: Dict[str, Any], + session_factory: "sessionmaker" = lambda di: di["session_factory"], ) -> None: + import gc + + from DashAI.back.dependencies.database.models import Pipeline, PipelineRun + + pipeline_id: int = self._pipeline_id(self.kwargs) + + with session_factory() as db: + pipeline: Pipeline = db.get(Pipeline, pipeline_id) + if not pipeline: + raise JobError(f"Pipeline {pipeline_id} does not exist in DB.") + + steps = pipeline.steps or [] + edges = pipeline.edges or [] + if not steps: + raise JobError(f"Pipeline {pipeline_id} has no steps to run.") + + # Created before the graph is built, because naming a node's + # artifacts needs the run's id: two executions of one pipeline that + # shared a name would have the second write over the first's model. + pipeline_run = PipelineRun(pipeline_id=pipeline_id) + db.add(pipeline_run) + db.commit() + pipeline_run_id = pipeline_run.id + + self.report_progress(0.05, "Preparing the graph") + + try: + graph = expand(steps, edges, pipeline_run_id) + except GraphError as e: + self._fail(session_factory, pipeline_run_id, str(e)) + raise JobError(str(e)) from e + + sink = DatabaseSink(pipeline_run_id, graph) + + try: + contexts = run_graph(graph, sink) + except GraphError as e: + # Validation happens before the engine reports anything, so the run + # is still untouched and saying why is this job's to do. + self._fail(session_factory, pipeline_run_id, str(e)) + raise JobError(str(e)) from e + except Exception as e: + # A node failed, and the sink normally recorded which one. But the + # sink's own calls are outside the engine's try, so one of them + # failing -- a locked database, a run row deleted underneath -- is + # also how we get here, and then nothing was recorded at all. _fail + # only writes when the run is not already in a terminal state, so + # calling it either way cannot overwrite what the sink said. + self._fail(session_factory, pipeline_run_id, str(e)) + raise JobError(str(e)) from e + finally: + gc.collect() + + # The contexts of the leaves are all that is still held. Their live + # objects are datasets and models, which nothing will read again. + for ctx in contexts.values(): + ctx.clear_cache() + gc.collect() + + self.report_progress(1.0, "Finished") + + @staticmethod + def _fail(session_factory, pipeline_run_id: int, message: str) -> None: + """Mark the run as failed, unless something already settled it. + + Never overwrites a terminal status. The sink records a node failure + with the message the unit produced, which is the better one; this is + the backstop for the paths where nothing recorded anything, and a run + left in STARTED with no error is indistinguishable from one still + going. """ - Update the pipeline context and database object based on node type. - - Args: - context: The pipeline context dictionary - pipeline: The pipeline database object - node_type: The type of node that was executed - node_id: The ID of the node that was executed - output: The output from the node execution - """ - if node_type == "DataSelector": - context["dataset"] = output.get("dataset") - elif node_type == "DataExploration": - context["exploration"] = output.get("exploration") - pipeline.exploration = context["exploration"] - elif node_type == "Train": - context["train"] = output.get("train") - pipeline.train = context["train"] - elif node_type == "RetrieveModel": - context["retrieve"] = output.get("retrieve") - elif node_type == "Prediction": - context["prediction"] = output.get("prediction") - pipeline.prediction = context["prediction"] - else: - context[node_id] = output + from DashAI.back.core.enums.status import PipelineRunStatus + from DashAI.back.dependencies.database.models import PipelineRun + + settled = { + PipelineRunStatus.FINISHED, + PipelineRunStatus.ERROR, + } + + with session_factory() as db: + pipeline_run = db.get(PipelineRun, pipeline_run_id) + if pipeline_run is None or pipeline_run.status in settled: + return + pipeline_run.set_status_as_error(message) + db.commit() diff --git a/DashAI/back/job/predict_job.py b/DashAI/back/job/predict_job.py index f509c101d..044761b45 100644 --- a/DashAI/back/job/predict_job.py +++ b/DashAI/back/job/predict_job.py @@ -1,5 +1,4 @@ import logging -from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Tuple from fastapi import status @@ -10,10 +9,13 @@ from DashAI.back.dependencies.database.models import Dataset, ModelSession, Prediction from DashAI.back.job.base_job import BaseJob, JobError -from DashAI.back.models.base_model import BaseModel -from DashAI.back.splitters.splits_payload import run_split_indexes -from DashAI.back.tasks.base_task import BaseTask -from DashAI.back.tasks.regression_task import RegressionTask +from DashAI.back.units.build_manual_input_unit import BuildManualInputUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_trained_model_unit import LoadTrainedModelUnit +from DashAI.back.units.load_training_dataset_unit import LoadTrainingDatasetUnit +from DashAI.back.units.predict_unit import PredictUnit +from DashAI.back.units.save_prediction_unit import SavePredictionUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -24,43 +26,6 @@ log = logging.getLogger(__name__) -def _run_prediction_pipeline( - task: BaseTask, - trained_model: BaseModel, - train_dataset: "DashAIDataset", - loaded_dataset: "DashAIDataset", - model_session: ModelSession, -) -> Tuple["DashAIDataset", Any]: - """Run shared prediction steps from prepared input data to final predictions. - - Parameters - ---------- - task : BaseTask - The task the run belongs to. - trained_model : BaseModel - The model loaded from the run. - train_dataset : DashAIDataset - The dataset the model was trained on, for type and label information. - loaded_dataset : DashAIDataset - The rows to predict. - model_session : ModelSession - The session declaring the input and output columns. - - Returns - ------- - tuple - The prepared inputs and the predictions for them. - """ - import numpy as np - - prepared_dataset = loaded_dataset.select_columns(model_session.input_columns) - y_pred_proba = np.array(trained_model.predict(prepared_dataset)) - y_pred = task.process_predictions( - train_dataset, y_pred_proba, model_session.output_columns[0] - ) - return prepared_dataset, y_pred - - def _build_preview_rows( prepared_dataset: "DashAIDataset", input_columns: List[str], @@ -99,21 +64,31 @@ def _to_native(v: Any) -> Any: def run_manual_prediction( run_id: int, manual_input_data: List[Dict], - component_registry: Any, session_factory: "sessionmaker", ) -> Tuple[List[str], List[List]]: """Execute a manual prediction synchronously without persisting results. + Composes the same units ``PredictJob`` does, so there is one definition of + what predicting means. The difference is entirely in the orchestration: no + state row to advance, nothing written to disk, and failures reported as + ``HTTPException`` because this runs inside a request instead of a worker. + + Each unit call sits in its own ``try`` so the HTTP response is decided by + *which step* failed, never by matching on an error message. That is what + keeps the endpoint's contract — eleven distinct responses across four status + codes — independent of how the units happen to word their errors. + Parameters ---------- run_id : int The ID of the trained run. manual_input_data : List[Dict] List of row dicts keyed by input column name. - component_registry : Any - The DashAI component registry. session_factory : sessionmaker - SQLAlchemy session factory. + SQLAlchemy session factory, used for this function's own row reads. It + must be the container's: the units resolve ``session_factory`` and + ``component_registry`` from the DI container themselves, so a different + one passed here would leave the two halves reading different databases. Returns ------- @@ -126,10 +101,11 @@ def run_manual_prediction( HTTPException On missing run, model session, or prediction failure. """ - with session_factory() as db: - from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset - from DashAI.back.dependencies.database.models import Run + from DashAI.back.dependencies.database.models import Run + # Read everything this function needs off the rows, then let the session go: + # nothing here writes, and the units open their own sessions. + with session_factory() as db: run = db.get(Run, run_id) if not run: raise HTTPException( @@ -163,54 +139,65 @@ def run_manual_prediction( detail="Model session has no output columns configured", ) - try: - task: BaseTask = component_registry[model_session.task_name]["class"]() - except Exception as e: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Task {model_session.task_name} not found in the registry", - ) from e + task_name = model_session.task_name + input_columns = list(model_session.input_columns) + output_columns = list(model_session.output_columns) + train_dataset_file_path = dataset_trained.file_path + + ctx = ExecutionContext() + + build_input = BuildManualInputUnit( + task_name=task_name, + train_dataset_file_path=train_dataset_file_path, + manual_input_data=manual_input_data, + ) + predict = PredictUnit( + task_name=task_name, + input_columns=input_columns, + output_columns=output_columns, + ) + try: + # Both units resolve the task; validating up front keeps a missing task + # reported as a task problem, ahead of the model, the way it always was. try: - model_cls = component_registry[run.model_name]["class"] - except KeyError as e: + build_input.validate(ctx) + predict.validate(ctx) + except JobError as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Model {run.model_name} not found in the registry", + detail=f"Task {task_name} not found in the registry", ) from e try: - trained_model: BaseModel = model_cls.load(run.run_path) - except Exception as e: + LoadTrainedModelUnit(run_id=run_id)(ctx) + except JobError as e: + # The unit distinguishes "not in the registry" from "cannot be read + # from disk" with the same two texts this endpoint has always + # returned, so its message is forwarded rather than rebuilt. raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=( - f"Failed to load model {run.model_name} from path {run.run_path}" - ), + detail=str(e), ) from e try: - train_dataset: "DashAIDataset" = load_dataset( - str(Path(f"{dataset_trained.file_path}/dataset/")) + LoadTrainingDatasetUnit(train_dataset_file_path=train_dataset_file_path)( + ctx ) - except Exception as e: + except JobError as e: + # Not forwarded: the unit names the path, this endpoint never has. raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Cannot load training dataset", ) from e try: - dataset_trained_path = str(Path(f"{dataset_trained.file_path}/dataset/")) - loaded_dataset: "DashAIDataset" = task.process_manual_input( - manual_input_data, dataset_trained_path - ) - prepared_dataset, y_pred = _run_prediction_pipeline( - task=task, - trained_model=trained_model, - train_dataset=train_dataset, - loaded_dataset=loaded_dataset, - model_session=model_session, - ) + build_input(ctx) + predict(ctx) + # Re-derived here rather than published by the unit: a narrowed view + # of the dataset is exactly the kind of derived value that must not + # cross a unit boundary, since anything upstream may reshape it. + prepared_dataset = ctx.require("dataset").select_columns(input_columns) except (ValueError, TypeError) as e: logging.exception("Manual prediction input error: %s", e) raise HTTPException( @@ -223,13 +210,14 @@ def run_manual_prediction( detail="Model prediction failed", ) from e - output_col = model_session.output_columns[0] return _build_preview_rows( prepared_dataset=prepared_dataset, - input_columns=list(model_session.input_columns), - output_col=output_col, - y_pred=y_pred, + input_columns=input_columns, + output_col=output_columns[0], + y_pred=ctx.require("y_pred"), ) + finally: + ctx.clear_cache() class PredictJob(BaseJob): @@ -295,18 +283,9 @@ def get_job_name(self) -> str: def run( self, ) -> List[Any]: - import uuid - from pathlib import Path - - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - save_dataset, - to_dashai_dataset, - ) - - component_registry = di["component_registry"] session_factory = di["session_factory"] - config = di["config"] + + ctx = ExecutionContext() prediction_id: int = self.kwargs["prediction_id"] manual_input_data: List[dict] = self.kwargs.get("manual_input_data", []) @@ -350,18 +329,17 @@ def run( detail="Model session not found", ) - # Retrieve Dataset if dataset_id is provided - dataset: Dataset = None + # The dataset the model was trained on. The one to predict on, + # when there is one, is resolved by the unit that loads it. dataset_trained: Dataset = db.get(Dataset, model_session.dataset_id) if not dataset_trained: + prediction.set_status_as_error() + db.commit() raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Training dataset not found", ) - if dataset_id: - dataset: Dataset = db.get(Dataset, dataset_id) - if not model_session.input_columns: prediction.set_status_as_error() db.commit() @@ -385,100 +363,81 @@ def run( detail="Internal database error", ) from e - # Retrieve Task + # The prediction step owns the task, and resolving it here — before + # the model is even looked up — is what keeps a missing task + # reported as a task problem instead of being overtaken by + # whatever fails next. Same shape as ModelJob validating the fit + # unit ahead of the status change. + predict = PredictUnit( + task_name=model_session.task_name, + input_columns=model_session.input_columns, + output_columns=model_session.output_columns, + ) try: - task: BaseTask = component_registry[model_session.task_name]["class"]() + predict.validate(ctx) except Exception as e: prediction.set_status_as_error() db.commit() log.exception(e) - raise JobError( - f"Task {model_session.task_name} not found in the registry", - ) from e + raise - # Load Model + # Load Model. The unit reports both the registry miss and the + # unreadable artifact with the same texts the job used to build + # here; re-raised as-is so they reach the user intact. try: - model = component_registry[prediction.run.model_name]["class"] - except KeyError as e: + LoadTrainedModelUnit(run_id=prediction.run_id)(ctx) + except Exception as e: prediction.set_status_as_error() db.commit() log.exception(e) - raise JobError( - f"Model {prediction.run.model_name} not found in the registry" - ) from e + raise + # Load training dataset for type info and label processing. Loaded + # before the dataset to predict on, which is the order the error + # messages depend on when both are unreadable. try: - trained_model: BaseModel = model.load(prediction.run.run_path) + LoadTrainingDatasetUnit( + train_dataset_file_path=dataset_trained.file_path + )(ctx) except Exception as e: + # This branch used to skip set_status_as_error, unlike every + # one around it, leaving the prediction STARTED forever. + # Re-raised as-is so the unit's specific message survives. prediction.set_status_as_error() db.commit() log.exception(e) - raise JobError( - f"Failed to load model {prediction.run.model_name} " - f"from path {prediction.run.run_path}" - ) from e + raise - # Load Dataset and make Predictions try: - # Load training dataset for type info and label processing - train_dataset: "DashAIDataset" = load_dataset( - str(Path(f"{dataset_trained.file_path}/dataset/")) - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Cannot load training dataset from " - f"{dataset_trained.file_path}/dataset/" - ) from e - - row_indexes = None - if prediction.split and dataset_id == model_session.dataset_id: - try: - row_indexes = run_split_indexes( - model_session.splits, - prediction.run.split_indexes, - component_registry, - prediction.split, - ) - except ValueError as e: - prediction.set_status_as_error() - db.commit() - log.exception(e) - raise JobError( - f"Cannot predict on the {prediction.split} split: {e}" - ) from e - - try: - # Load or create prediction dataset + # Load or create prediction dataset. Both branches publish the + # same "dataset" key, so the prediction below cannot tell a + # dataset read from disk from one typed in by hand. if dataset_id: - loaded_dataset: "DashAIDataset" = load_dataset( - str(Path(f"{dataset.file_path}/dataset/")) - ) - if row_indexes is not None: - loaded_dataset = loaded_dataset.select(row_indexes) + LoadDatasetUnit(dataset_id=dataset_id)(ctx) else: - dataset_trained_path = str( - Path(f"{dataset_trained.file_path}/dataset/") - ) - loaded_dataset = task.process_manual_input( - manual_input_data, dataset_trained_path - ) + BuildManualInputUnit( + task_name=model_session.task_name, + train_dataset_file_path=dataset_trained.file_path, + manual_input_data=manual_input_data, + )(ctx) self.report_progress(0.4, "Running prediction") - _, y_pred = _run_prediction_pipeline( - task=task, - trained_model=trained_model, - train_dataset=train_dataset, - loaded_dataset=loaded_dataset, - model_session=model_session, - ) + predict(ctx) except ValueError as ve: prediction.set_status_as_error() db.commit() log.error(f"Validation Error: {ve}") + # JobError, not HTTPException: this runs in the Huey worker + # and its exception travels back through dill. HTTPException + # stores nothing in ``args``, so unpickling calls it with no + # status code and the model's own complaint is replaced by a + # deserialisation failure. raise JobError(f"Invalid input data: {ve}") from ve except TypeError as te: + # Marked as failed like its ValueError neighbour: this branch + # used to return 400 without touching the row, which left the + # prediction STARTED forever. prediction.set_status_as_error() db.commit() log.error(f"Type Error: {te}") @@ -495,49 +454,13 @@ def run( # Save Predictions to Arrow file try: - # Create unique folder for predictions - path = str(Path(f"{config['DATASETS_PATH']}/predictions/")) - folder_name = str(uuid.uuid4()) - full_path = Path(path) / folder_name - full_path.mkdir(parents=True, exist_ok=True) - - output_col = model_session.output_columns[0] - base_columns = [ - col for col in loaded_dataset.column_names if col != output_col - ] - output_dataset = loaded_dataset.select_columns(base_columns) - dataset_with_prediction = to_dashai_dataset( - output_dataset.add_column(output_col, y_pred) - ) - - # Filter schema from trained dataset - trained_schema = train_dataset.types - filtered_schema = { - key: value.to_string() - for key, value in trained_schema.items() - if key in model_session.input_columns + model_session.output_columns - } - - # Regression models predict continuous values regardless of - # the training target's original dtype (e.g. a target column - # that happened to hold only integer-looking values), so the - # output column's saved schema must reflect that instead of - # inheriting the training dataset's type. - if isinstance(task, RegressionTask): - filtered_schema[output_col] = {"type": "Float", "dtype": "float64"} - - # Store num of rows, columns, and column names - dataset_with_prediction.compute_base_metadata() - - # Save dataset with predictions - save_dataset( - dataset_with_prediction, - str(full_path / "dataset"), - filtered_schema, - ) + SavePredictionUnit( + input_columns=model_session.input_columns, + output_columns=model_session.output_columns, + )(ctx) # Update Prediction record - prediction.results_path = str(full_path) + prediction.results_path = ctx.require("results_path") prediction.set_status_as_finished() db.commit() except Exception as e: @@ -547,3 +470,5 @@ def run( raise JobError( "Can not save prediction to json file", ) from e + finally: + ctx.clear_cache() diff --git a/DashAI/back/models/base_model.py b/DashAI/back/models/base_model.py index f53be2c6d..144f44ce5 100644 --- a/DashAI/back/models/base_model.py +++ b/DashAI/back/models/base_model.py @@ -243,57 +243,56 @@ def _save_metrics( db.commit() @final - def calculate_metrics( + def compute_metrics( self, split: SplitEnum = SplitEnum.VALIDATION, - level: LevelEnum = LevelEnum.LAST, - log_index: int = None, x_data: "DashAIDataset" = None, y_data: "DashAIDataset" = None, - fold_index: int = None, - inner_fold_index: int = None, - ): - """Calculate and save metrics for a given data split and level. + ) -> Dict[str, float]: + """Score a data split with this model's metrics, without persisting. + + Which metrics are computed is decided by the model rather than by the + caller: ``ModelFactory`` attaches the metric classes and the data + splits to the instance, and this reads them off it. + + Separate from :meth:`calculate_metrics` so a caller with no run to log + against can still have the numbers. A pipeline is that caller: it has + no ``Run`` row, so there is no foreign key for a ``Metric`` row to + point at, and its results are recorded as an artifact instead. Both + paths score through here, so the two cannot disagree. Parameters ---------- split : SplitEnum - The data split to evaluate (TRAIN, VALIDATION, - or TEST). Defaults to SplitEnum.VALIDATION. - level : LevelEnum - The metric granularity level (LAST, TRIAL, - STEP, or BATCH). Defaults to LevelEnum.LAST. - log_index : int, optional - Explicit step index for the metric - entry. If None, the next step index is computed automatically. - Defaults to None. + The data split to evaluate. Defaults to ``SplitEnum.VALIDATION``. x_data : DashAIDataset, optional - Input features. If None, the - dataset stored in the model for the given split is used. - Defaults to None. + Input features. Defaults to the split stored on the model. y_data : DashAIDataset, optional - Target labels. If None, the - labels stored in the model for the given split are used. - Defaults to None. + Target labels. Defaults to the split stored on the model. + + Returns + ------- + Dict[str, float] + Metric name to score. Empty when every metric returned a + non-finite value. ``None`` when there is nothing to score at all: + no metrics configured, or no data for this split. """ - # Get the appropriate metrics based on split metrics_attr = f"{split.value}_metrics" metrics = getattr(self, metrics_attr, None) - # If no metrics or run_id, skip calculation - if not metrics or not self.run_id: - return + if not metrics: + return None # Load data if not provided if x_data is None or y_data is None: if self.x_data is None or self.y_data is None: - return + return None x_data = self.x_data[split.value] y_data = self.y_data[split.value] # If data is empty after retrieval, skip calculation if x_data is None or y_data is None: - return + return None # Make predictions and transform outputs y_pred = self.predict(x_data) @@ -314,41 +313,33 @@ def calculate_metrics( continue results[metric.__name__] = score - # Save to database - self._save_metrics( - split=split, - level=level, - results=results, - log_index=log_index, - fold_index=fold_index, - inner_fold_index=inner_fold_index, - ) - - # Report the epoch to whoever is watching, AFTER persisting: the reporter - # is allowed to raise (Optuna prunes that way), and the metrics of the - # epoch that triggered the stop should survive it. - if ( - self._epoch_reporter is not None - and level is LevelEnum.EPOCH - and split is SplitEnum.VALIDATION - ): - self._epoch_reporter(results, log_index) + return results - # Create a function similar to calculate_metrics that returns the scores - # instead of saving them to the database, to be used in the CV evaluation loop - def compute_metrics( + @final + def calculate_metrics( self, - split: SplitEnum = SplitEnum.TEST, + split: SplitEnum = SplitEnum.VALIDATION, + level: LevelEnum = LevelEnum.LAST, + log_index: int = None, x_data: "DashAIDataset" = None, y_data: "DashAIDataset" = None, - ) -> Dict[str, float]: - """Calculate and return metric scores for a given data split. + fold_index: int = None, + inner_fold_index: int = None, + ): + """Calculate metrics for a data split and save them to the database. Parameters ---------- split : SplitEnum The data split to evaluate (TRAIN, VALIDATION, or TEST). Defaults to SplitEnum.VALIDATION. + level : LevelEnum + The metric granularity level (LAST, TRIAL, + STEP, or BATCH). Defaults to LevelEnum.LAST. + log_index : int, optional + Explicit step index for the metric + entry. If None, the next step index is computed automatically. + Defaults to None. x_data : DashAIDataset, optional Input features. If None, the dataset stored in the model for the given split is used. @@ -360,37 +351,53 @@ def compute_metrics( Returns ------- - Dict[str, float] - A dictionary mapping metric names to their computed scores. + Dict[str, float] or None + What was written, so a caller that also wants the numbers does not + have to score the same split twice. ``None`` when nothing was + written: no run to write against, or nothing to score. + + Notes + ----- + A metric row is keyed by the run it belongs to, so a model with no run + has nowhere to write and this returns without scoring anything. That is + what lets a caller with no ``Run`` row -- a pipeline -- train a model + that logs nothing at all, during training or after it, without having + to intercept anything. Such a caller uses :meth:`compute_metrics` and + keeps the numbers itself. """ - # Get the appropriate metrics based on split - metrics_attr = f"{split.value}_metrics" - metrics = getattr(self, metrics_attr, None) - - # If no metrics, return empty dict - if not metrics: - return {} - - # Load data if not provided - if x_data is None or y_data is None: - if self.x_data is None or self.y_data is None: - return {} - x_data = self.x_data[split.value] - y_data = self.y_data[split.value] - - # If data is empty after retrieval, return empty dict - if x_data is None or y_data is None: - return {} + # No run means no foreign key for a metric row to point at. + # + # getattr rather than self.run_id: ModelFactory attaches the attribute, + # but a model built directly never had it, and the guard this replaced + # happened to never reach the attribute for such a model because it + # checked for metrics first. Treating "no attribute" as "no run" keeps + # that path working and is the same answer for any caller that has one. + if not getattr(self, "run_id", None): + return None + + results = self.compute_metrics(split=split, x_data=x_data, y_data=y_data) + if results is None: + return None - # Make predictions and transform outputs - y_pred = self.predict(x_data) - y_transformed = self.prepare_output(y_data, is_fit=False) + # Save to database + self._save_metrics( + split=split, + level=level, + results=results, + log_index=log_index, + fold_index=fold_index, + inner_fold_index=inner_fold_index, + ) - # Calculate metric scores - results = {} - for metric in metrics: - score = metric.score(y_transformed, y_pred) - results[metric.__name__] = score + # Report the epoch to whoever is watching, AFTER persisting: the reporter + # is allowed to raise (Optuna prunes that way), and the metrics of the + # epoch that triggered the stop should survive it. + if ( + self._epoch_reporter is not None + and level is LevelEnum.EPOCH + and split is SplitEnum.VALIDATION + ): + self._epoch_reporter(results, log_index) return results diff --git a/DashAI/back/optimizers/base_optimizer.py b/DashAI/back/optimizers/base_optimizer.py index ce6da365e..b22a0ac7b 100644 --- a/DashAI/back/optimizers/base_optimizer.py +++ b/DashAI/back/optimizers/base_optimizer.py @@ -388,7 +388,7 @@ def importance_plot(self, trials, goal_metric): return PlotlyArtifact(payload=fig, title=title) - def create_plots(self, trials, run_id, n_params, goal_metric): + def create_plots(self, trials, run_id, n_params, goal_metric, artifact_prefix=None): """ List of available plots. @@ -400,18 +400,31 @@ def create_plots(self, trials, run_id, n_params, goal_metric): n_params (int): Number of the different hyperparameters involved in the process of hyperparameter optimization goal_metric (dict): Metric optimized in the process. + artifact_prefix (str, optional): Name to build the filenames from, + for callers that are not a run. Defaults to None, + which keeps deriving them from ``run_id``. Returns ------- plots_filenames (list): Filenames to persist each plot under. plots_list (list): The matching list of PlotlyArtifact instances. + + Notes + ----- + The filenames land in a directory shared with every other run's + artifacts, so whatever names them has to be unique per caller. + ``run_id`` is unique among runs but a pipeline has no run: it passes + ``artifact_prefix`` instead, and two pipeline executions that would + otherwise both write ``..._None.pickle`` stay apart. """ + tag = artifact_prefix if artifact_prefix is not None else run_id + if n_params >= 2: plots_filenames = [ - f"history_objective_plot_{run_id}.pickle", - f"slice_plot_{run_id}.pickle", - f"contour_plot_{run_id}.pickle", - f"importance_plot_{run_id}.pickle", + f"history_objective_plot_{tag}.pickle", + f"slice_plot_{tag}.pickle", + f"contour_plot_{tag}.pickle", + f"importance_plot_{tag}.pickle", ] plots_list = [ self.history_objective_plot(trials, goal_metric), @@ -422,8 +435,8 @@ def create_plots(self, trials, run_id, n_params, goal_metric): return plots_filenames, plots_list else: plots_filenames = [ - f"history_objective_plot_{run_id}.pickle", - f"slice_plot_{run_id}.pickle", + f"history_objective_plot_{tag}.pickle", + f"slice_plot_{tag}.pickle", ] plots_list = [ self.history_objective_plot(trials, goal_metric), diff --git a/DashAI/back/units/__init__.py b/DashAI/back/units/__init__.py new file mode 100644 index 000000000..f9cc3fccc --- /dev/null +++ b/DashAI/back/units/__init__.py @@ -0,0 +1,3 @@ +# flake8: noqa +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError diff --git a/DashAI/back/units/apply_converter_unit.py b/DashAI/back/units/apply_converter_unit.py new file mode 100644 index 000000000..da4bc636d --- /dev/null +++ b/DashAI/back/units/apply_converter_unit.py @@ -0,0 +1,145 @@ +"""Unit that fits a converter and transforms one dataset with it.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.converter_scope import ( + ConverterScopeMixin, + converter_field, + scope_field, + target_field, +) + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class ApplyConverterSchema(BaseSchema): + converter: converter_field() # type: ignore + scope: scope_field() # type: ignore + target: target_field() # type: ignore + + +class ApplyConverterUnit(BaseUnit, ConverterScopeMixin): + """Fit one converter on the scoped data and transform the same dataset. + + The single-dataset case, which is what applying a converter to a notebook + means. Equivalent to :class:`FitConverterUnit` followed by + :class:`TransformDatasetUnit` on the same dataset, kept as one unit for two + reasons: it is the common case, and when there is no row scope it hands + ``transform`` the very object ``fit`` saw, which some converters rely on to + skip recomputing (see below). + + Use the split pair instead whenever the converter has to be learned from one + dataset and applied to another — fit on train, transform on test. + + Reads ``dataset`` and writes ``dataset``: the same key on both sides, so any + number of these can be chained in one context, each one seeing what the + previous one produced. + + That is also why nothing about column identity ever crosses the context + boundary. The scope is expressed as 1-based column and row indexes, and + those indexes are resolved against ``dataset.column_names`` read at the top + of ``execute`` — the dataset as it is *right now*. A converter that renames, + drops or adds columns changes what index 3 means, so a resolved column list + published to the context would be stale for the very next converter. + """ + + SCHEMA = ApplyConverterSchema + + REQUIRES = ("dataset",) + PROVIDES = ("dataset", "fitted_converter") + + def __init__(self, **config) -> None: + super().__init__(**config) + # Memoized on the instance, never in the context: two units of this + # class can live in the same context and they are different converters. + self._converter_class = None + + @property + def _converter_name(self) -> str: + return self.config["converter"]["component"] + + def _resolve_converter(self): + """Look the converter class up in the registry, once per instance.""" + if self._converter_class is not None: + return self._converter_class + + from kink import di + + component_registry = di["component_registry"] + converter_name = self._converter_name + + try: + self._converter_class = component_registry[converter_name]["class"] + except KeyError as e: + log.exception(e) + raise JobError(f"Error importing converter {converter_name}: {e}") from e + + return self._converter_class + + def validate(self, ctx: ExecutionContext) -> None: + """Reject an out-of-bounds target before any work is done.""" + self._check_target_bounds(ctx.require("dataset")) + + def execute(self, ctx: ExecutionContext) -> None: + loaded_dataset: "DashAIDataset" = ctx.require("dataset") + converter_name = self._converter_name + + converter_constructor = self._resolve_converter() + converter_instance = converter_constructor( + **(self.config["converter"].get("params") or {}) + ) + + self._check_target_bounds(loaded_dataset) + ( + scope_column_names, + scope_rows_indexes, + target_column_name, + ) = self._resolve_scope(loaded_dataset) + + log.info(f"Applying converter: {converter_name}") + + x_dataset_fit, y_dataset_fit = self._slice_for_fit( + loaded_dataset, + scope_column_names, + scope_rows_indexes, + target_column_name, + ) + + converter_instance = self._fit( + converter_instance, x_dataset_fit, y_dataset_fit, converter_name + ) + + if scope_rows_indexes: + x_full_transform, y_full_transform = self._slice_for_transform( + loaded_dataset, scope_column_names, target_column_name + ) + else: + # Deliberately the *same objects* as the ones passed to fit, not + # equal ones: converters such as TypeCast key a cache off the + # identity of the dataset they were fitted on and skip recomputing + # when transform receives it again. Without a row scope the fit + # slice already covers every row, so reusing it is also correct. + x_full_transform, y_full_transform = x_dataset_fit, y_dataset_fit + + transformed_dataset = self._transform( + converter_instance, x_full_transform, y_full_transform, converter_name + ) + + ctx.put( + "dataset", + self._merge_transformed( + converter_instance, + loaded_dataset, + transformed_dataset, + scope_column_names, + ), + ) + ctx.put("fitted_converter", converter_instance) diff --git a/DashAI/back/units/apply_dataset_schema_unit.py b/DashAI/back/units/apply_dataset_schema_unit.py new file mode 100644 index 000000000..42f60af28 --- /dev/null +++ b/DashAI/back/units/apply_dataset_schema_unit.py @@ -0,0 +1,129 @@ +"""Unit that renames columns and casts them to their declared types.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + none_type, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class ApplyDatasetSchemaSchema(BaseSchema): + column_renames: schema_field( + none_type(dict), + placeholder=None, + description=MultilingualString( + en="Mapping of current column name to new name, applied before the " + "types are cast. Columns not mentioned keep their name. The " + "resulting names have to stay unique.", + es="Mapa de nombre de columna actual a nombre nuevo, aplicado antes " + "de convertir los tipos. Las columnas no mencionadas conservan su " + "nombre. Los nombres resultantes tienen que seguir siendo únicos.", + pt="Mapa de nome de coluna atual para nome novo, aplicado antes de " + "converter os tipos. As colunas não mencionadas mantêm o seu nome. " + "Os nomes resultantes têm de continuar únicos.", + de="Zuordnung von aktuellem zu neuem Spaltennamen, angewendet vor " + "der Typumwandlung. Nicht genannte Spalten behalten ihren Namen. " + "Die resultierenden Namen müssen eindeutig bleiben.", + zh="当前列名到新列名的映射，在类型转换前应用。未提及的列保留原名。" + "结果列名必须保持唯一。", + ), + alias=MultilingualString( + en="Column renames", + es="Renombres de columnas", + pt="Renomeações de colunas", + de="Spaltenumbenennungen", + zh="列重命名", + ), + ) # type: ignore + + +class ApplyDatasetSchemaUnit(BaseUnit): + """Rename the dataset's columns and cast them to their declared types. + + Takes the type declaration from the context (``inferred_types``) rather than + from its own configuration, so a single path feeds it whether the types were + inferred upstream or handed in by whoever set up the run. Renaming and + casting are one unit because a rename has to carry the declared types with + it: the type declared for ``Species`` has to end up on ``Variety``, not be + re-inferred from the data. + + ``validate`` rejects a declaration that names columns the dataset does not + have. Without that check the mismatch is silent — the underlying transform + passes unknown columns through untouched — so a stale declaration would + quietly leave columns with inferred types instead of the requested ones. + + Note the deliberate asymmetry: the unit does **not** republish a renamed + ``inferred_types``. After a rename that key describes column names that no + longer exist, and republishing it would invite a second consumer to trust a + declaration that no longer matches the dataset. Whoever needs types after + this unit reads them off the dataset. + """ + + SCHEMA = ApplyDatasetSchemaSchema + + REQUIRES = ("dataset", "inferred_types") + PROVIDES = ("dataset",) + + def validate(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") + schema = ctx.require("inferred_types") + + unknown = sorted(set(schema) - set(dataset.column_names)) + if unknown: + raise JobError( + f"The declared types name columns the dataset does not have: {unknown}" + ) + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import ( + transform_dataset_with_schema, + ) + + dataset = ctx.require("dataset") + schema = ctx.require("inferred_types") + renames = self.config.get("column_renames") + + if renames: + dataset, schema = _rename_columns(dataset, schema, renames) + + ctx.put("dataset", transform_dataset_with_schema(dataset, schema)) + + +def _rename_columns(dataset, schema: dict, renames: dict): + """Return the dataset with renamed columns and the schema remapped to match. + + A plain helper: it takes and returns values and never touches the context, so + the unit's declared contract cannot hide inside it. + """ + original_names = dataset.arrow_table.schema.names + new_names = [renames.get(column, column) for column in original_names] + + if len(new_names) != len(set(new_names)): + duplicate_names = set() + seen = set() + for name in new_names: + if name in seen: + duplicate_names.add(name) + else: + seen.add(name) + raise JobError( + "Invalid column_renames: resulting column names contain duplicates: " + f"{sorted(duplicate_names)}" + ) + + arrow_table = dataset.arrow_table.rename_columns(new_names) + renamed = dataset.__class__( + arrow_table, + splits=dataset.splits, + types=dataset.types, + ) + remapped_schema = {renames.get(column, column): schema[column] for column in schema} + return renamed, remapped_schema diff --git a/DashAI/back/units/base_unit.py b/DashAI/back/units/base_unit.py new file mode 100644 index 000000000..351f0fc3d --- /dev/null +++ b/DashAI/back/units/base_unit.py @@ -0,0 +1,123 @@ +"""Base class for the atomic units a job is composed of.""" + +import logging +from abc import ABCMeta, abstractmethod +from typing import Final, Tuple, final + +from DashAI.back.config_object import ConfigObject +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.units.context import ExecutionContext + +logger = logging.getLogger(__name__) + + +class BaseUnit(ConfigObject, metaclass=ABCMeta): + """Abstract class for all atomic units. + + A unit is the smallest reusable piece of a job: it declares the context + keys it needs, the keys it produces, and does one thing. Jobs compose units + into a sequence; the orchestration around them (database transactions, + status transitions, progress reporting) stays in the job. + + Units never read nor mutate the ``Run`` row. ``run_id`` travels through the + context as an opaque correlation id because ``ModelFactory`` needs it for + ``BaseModel.calculate_metrics`` to work, but the ownership of the row + belongs to the job. + + ``BaseUnit`` deliberately does not inherit from ``BaseJob``: the registry + derives a component's type by walking the MRO for ancestors whose name + contains "Base" and that declare a ``TYPE``, and it rejects components with + more than one candidate. Inheriting from both would make registration fail. + """ + + TYPE: Final[str] = "Unit" + + #: Context keys that must be present before the unit runs. + REQUIRES: Tuple[str, ...] = () + #: Context keys the unit guarantees after it runs. + PROVIDES: Tuple[str, ...] = () + + #: Configuration a user fills in. This is what the front renders. + SCHEMA: BaseSchema = BaseSchema + #: Configuration names supplied by whatever runs the unit, and never by a + #: user. Deliberately kept out of ``SCHEMA``: a value here is not known when + #: a form would be filled in, so there is nothing for a user to answer. + #: ``artifact_prefix`` is the clearest case — it is built from the id of a + #: run that has not started. + #: + #: Being a separate declaration rather than a flag inside the schema is the + #: whole point. A flag would still travel to the front, where every renderer + #: would have to remember to skip it and a new one would leak by default. + #: These names never reach the schema the front receives, so there is + #: nothing to filter and nothing to forget. + #: + #: They still arrive through ``**config`` and are read as ``self.config[…]`` + #: like any other, so declaring one here changes nothing about how a unit is + #: written — only about who is expected to supply it. + RUNTIME_PARAMS: Tuple[str, ...] = () + + def __init__(self, **config) -> None: + """Store the unit configuration. + + Parameters + ---------- + config : dict + Configuration of the unit, as declared by its schema. + """ + self.config = config + + def validate(self, ctx: ExecutionContext) -> None: + """Check preconditions without executing the unit. + + Runs before the job commits to any observable state change, so a unit + can reject an impossible configuration early. No-op by default. + + Parameters + ---------- + ctx : ExecutionContext + The shared execution context. + """ + + @abstractmethod + def execute(self, ctx: ExecutionContext) -> None: + """Do the unit's work, reading from and writing to the context. + + Parameters + ---------- + ctx : ExecutionContext + The shared execution context. + """ + raise NotImplementedError + + @final + def __call__(self, ctx: ExecutionContext) -> None: + """Run the unit, enforcing its declared contract. + + Calls ``validate`` before ``execute`` so a caller that just does + ``unit(ctx)`` — the sanctioned way to run a unit — always gets its + precondition checks (e.g. a download gate) for free. An orchestrator + that needs ``validate`` to run earlier, ahead of some other state + change, is still free to call ``unit.validate(ctx)`` directly first; + ``validate`` runs again here, which is redundant but harmless. + + Parameters + ---------- + ctx : ExecutionContext + The shared execution context. + + Raises + ------ + UnitContractError + If a required key is missing before execution or a promised key is + missing after it. + """ + for key in self.REQUIRES: + ctx.require(key) + + self.validate(ctx) + + logger.debug("Running unit %s", type(self).__name__) + self.execute(ctx) + + for key in self.PROVIDES: + ctx.require(key) diff --git a/DashAI/back/units/build_global_explainer_unit.py b/DashAI/back/units/build_global_explainer_unit.py new file mode 100644 index 000000000..2f6e6c717 --- /dev/null +++ b/DashAI/back/units/build_global_explainer_unit.py @@ -0,0 +1,63 @@ +"""Unit that instantiates a global explainer bound to a trained model.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import build_explainer + +log = logging.getLogger(__name__) + + +class BuildGlobalExplainerSchema(BaseSchema): + explainer: schema_field( + component_field(parent="BaseGlobalExplainer"), + placeholder={"component": "PermutationFeatureImportance", "params": {}}, + description=MultilingualString( + en="Explainer for the model as a whole, together with its own " + "configuration.", + es="Explicador para el modelo completo, junto con su propia configuración.", + pt="Explicador para o modelo como um todo, junto com a sua própria " + "configuração.", + de="Erklärer für das gesamte Modell samt eigener Konfiguration.", + zh="针对整个模型的解释器及其自身配置。", + ), + alias=MultilingualString( + en="Global explainer", + es="Explicador global", + pt="Explicador global", + de="Globaler Erklärer", + zh="全局解释器", + ), + ) # type: ignore + + +class BuildGlobalExplainerUnit(BaseUnit): + """Instantiate a global explainer over an already trained model. + + Sibling of ``BuildLocalExplainerUnit`` rather than one unit with a scope + flag, even though the building step itself is identical — the two share it + through a helper. Global and local explainers are separate registries with + separate base classes, and a component field carries a single ``parent`` + hint that the front reads straight off the property to list the candidates. + One field covering both scopes would have to be optional, and an optional + component field is emitted as an ``anyOf``, which buries the hint where the + front does not look and leaves the user with no picker at all. + """ + + SCHEMA = BuildGlobalExplainerSchema + + REQUIRES = ("model",) + PROVIDES = ("explainer",) + + def execute(self, ctx: ExecutionContext) -> None: + explainer = build_explainer( + "global", self.config["explainer"], ctx.require("model") + ) + ctx.put("explainer", explainer) diff --git a/DashAI/back/units/build_local_explainer_unit.py b/DashAI/back/units/build_local_explainer_unit.py new file mode 100644 index 000000000..b303da357 --- /dev/null +++ b/DashAI/back/units/build_local_explainer_unit.py @@ -0,0 +1,58 @@ +"""Unit that instantiates a local explainer bound to a trained model.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import build_explainer + +log = logging.getLogger(__name__) + + +class BuildLocalExplainerSchema(BaseSchema): + explainer: schema_field( + component_field(parent="BaseLocalExplainer"), + placeholder={"component": "KernelShap", "params": {}}, + description=MultilingualString( + en="Explainer for individual instances, together with its own " + "configuration.", + es="Explicador para instancias individuales, junto con su propia " + "configuración.", + pt="Explicador para instâncias individuais, junto com a sua própria " + "configuração.", + de="Erklärer für einzelne Instanzen samt eigener Konfiguration.", + zh="针对单个实例的解释器及其自身配置。", + ), + alias=MultilingualString( + en="Local explainer", + es="Explicador local", + pt="Explicador local", + de="Lokaler Erklärer", + zh="局部解释器", + ), + ) # type: ignore + + +class BuildLocalExplainerUnit(BaseUnit): + """Instantiate a local explainer over an already trained model. + + See ``BuildGlobalExplainerUnit`` for why the two scopes are two units even + though they share their whole implementation. + """ + + SCHEMA = BuildLocalExplainerSchema + + REQUIRES = ("model",) + PROVIDES = ("explainer",) + + def execute(self, ctx: ExecutionContext) -> None: + explainer = build_explainer( + "local", self.config["explainer"], ctx.require("model") + ) + ctx.put("explainer", explainer) diff --git a/DashAI/back/units/build_manual_input_unit.py b/DashAI/back/units/build_manual_input_unit.py new file mode 100644 index 000000000..4a0c3ac71 --- /dev/null +++ b/DashAI/back/units/build_manual_input_unit.py @@ -0,0 +1,119 @@ +"""Unit that turns hand-typed rows into a dataset to predict on.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +class BuildManualInputSchema(BaseSchema): + task_name: schema_field( + string_field(), + placeholder="TabularClassificationTask", + description=MultilingualString( + en="Name of the task that validates and types the hand-typed rows.", + es="Nombre de la tarea que valida y tipa las filas ingresadas a mano.", + pt="Nome da tarefa que valida e tipa as linhas introduzidas à mão.", + de="Name der Aufgabe, die die manuell eingegebenen Zeilen prüft " + "und typisiert.", + zh="用于校验并确定手工输入行类型的任务名称。", + ), + alias=MultilingualString( + en="Task", es="Tarea", pt="Tarefa", de="Aufgabe", zh="任务" + ), + ) # type: ignore + manual_input_data: schema_field( + list, + placeholder=[], + description=MultilingualString( + en="Rows to predict on, each a mapping from input column name to " + "value. Uploaded files arrive as a path reference instead of bytes.", + es="Filas a predecir, cada una un mapeo de nombre de columna de " + "entrada a valor. Los archivos subidos llegan como una referencia " + "a una ruta en vez de bytes.", + pt="Linhas a prever, cada uma um mapeamento de nome de coluna de " + "entrada para valor. Os ficheiros carregados chegam como uma " + "referência a um caminho em vez de bytes.", + de="Zu prognostizierende Zeilen, je eine Zuordnung von " + "Eingabespaltenname zu Wert. Hochgeladene Dateien kommen als " + "Pfadverweis statt als Bytes an.", + zh="要预测的行，每行是输入列名到值的映射。上传的文件以路径引用而非字节形式传入。", + ), + alias=MultilingualString( + en="Manual input", + es="Entrada manual", + pt="Entrada manual", + de="Manuelle Eingabe", + zh="手动输入", + ), + ) # type: ignore + + +class BuildManualInputUnit(BaseUnit): + """Build the dataset to predict on from values the user typed in. + + The counterpart of loading one from disk: it produces the same ``dataset`` + key, so whatever runs next cannot tell the two apart. That is what lets the + prediction step be written once for both sources. + + The task does the work — it is the task that knows the expected column + types and how to turn an uploaded file into a cell — so this unit only + resolves it and hands over the rows. + """ + + SCHEMA = BuildManualInputSchema + + PROVIDES = ("dataset",) + RUNTIME_PARAMS = ("train_dataset_file_path",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._task = None + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task from the registry, memoized on this unit.""" + if self._task is not None: + return self._task + + from kink import di + + component_registry = di["component_registry"] + task_name = self.config["task_name"] + + try: + task: "BaseTask" = component_registry[task_name]["class"]() + except Exception as e: + log.exception(e) + raise JobError(f"Task {task_name} not found in the registry") from e + + self._task = task + return task + + def validate(self, ctx: ExecutionContext) -> None: + """Resolve the task before anything observable happens.""" + self._resolve_task() + + def execute(self, ctx: ExecutionContext) -> None: + from pathlib import Path + + task = self._resolve_task() + + train_dataset_path = str( + Path(f"{self.config['train_dataset_file_path']}/dataset/") + ) + rows = self.config["manual_input_data"] + + ctx.put("dataset", task.process_manual_input(rows, train_dataset_path)) diff --git a/DashAI/back/units/build_model_unit.py b/DashAI/back/units/build_model_unit.py new file mode 100644 index 000000000..e0f5bc0b4 --- /dev/null +++ b/DashAI/back/units/build_model_unit.py @@ -0,0 +1,261 @@ +"""Unit that instantiates a model with its parameters, data and metrics.""" + +import logging +from typing import TYPE_CHECKING, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.nested import missing_downloads +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.metrics.base_metric import BaseMetric + from DashAI.back.models.base_model import BaseModel + +log = logging.getLogger(__name__) + + +def _metrics_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field()), + placeholder=[], + description=description, + alias=alias, + ) + + +class BuildModelSchema(BaseSchema): + model: schema_field( + component_field(parent="BaseModel"), + placeholder={"component": "SVC", "params": {}}, + description=MultilingualString( + en="Model to instantiate, along with its own configuration.", + es="Modelo a instanciar, junto con su propia configuración.", + pt="Modelo a instanciar, junto com a sua própria configuração.", + de="Zu instanziierendes Modell samt seiner eigenen Konfiguration.", + zh="要实例化的模型及其自身配置。", + ), + alias=MultilingualString( + en="Model", es="Modelo", pt="Modelo", de="Modell", zh="模型" + ), + ) # type: ignore + train_metrics: _metrics_field( + alias=MultilingualString( + en="Train metrics", + es="Métricas de entrenamiento", + pt="Métricas de treino", + de="Trainingsmetriken", + zh="训练指标", + ), + description=MultilingualString( + en="Metrics evaluated on the train split.", + es="Métricas evaluadas sobre la partición de entrenamiento.", + pt="Métricas avaliadas na partição de treino.", + de="Auf der Trainingsteilmenge ausgewertete Metriken.", + zh="在训练集上评估的指标。", + ), + ) # type: ignore + validation_metrics: _metrics_field( + alias=MultilingualString( + en="Validation metrics", + es="Métricas de validación", + pt="Métricas de validação", + de="Validierungsmetriken", + zh="验证指标", + ), + description=MultilingualString( + en="Metrics evaluated on the validation split.", + es="Métricas evaluadas sobre la partición de validación.", + pt="Métricas avaliadas na partição de validação.", + de="Auf der Validierungsteilmenge ausgewertete Metriken.", + zh="在验证集上评估的指标。", + ), + ) # type: ignore + test_metrics: _metrics_field( + alias=MultilingualString( + en="Test metrics", + es="Métricas de prueba", + pt="Métricas de teste", + de="Testmetriken", + zh="测试指标", + ), + description=MultilingualString( + en="Metrics evaluated on the test split.", + es="Métricas evaluadas sobre la partición de prueba.", + pt="Métricas avaliadas na partição de teste.", + de="Auf der Testteilmenge ausgewertete Metriken.", + zh="在测试集上评估的指标。", + ), + ) # type: ignore + + +class BuildModelUnit(BaseUnit): + """Instantiate an untrained model bound to its data and metrics. + + ``ModelFactory`` attaches the run id and the metric classes to the model + instance, which is what later lets the model log metrics on its own during + and after training. The metrics are configured here rather than in the + evaluation unit because models use them *while* training to log at the step + and epoch levels. + + **The data is not attached here.** It used to be, and that only worked + because a model was fitted once on one split. A model fitted over folds + sees different data on every iteration, so binding one partition at + construction would leave the metrics describing whichever fold happened to + be built with. Whoever fits the model points it at the data it is being + fitted on, and this unit no longer needs ``x`` or ``y`` at all -- only the + label count, which is a property of the dataset rather than of a split. + + ``validate`` checks that the model and every component nested in its + parameters have been downloaded, so an impossible run is rejected before + anything observable happens. + + The model is configured as a component field, the same way a model picks + its own sub-components: the value is ``{"component": , "params": + {...}}``. The parameter tree stays an opaque ``dict`` in the schema and the + front resolves it recursively, fetching the chosen component's schema to + render the nested form. ``ModelFactory`` walks the same shape to build the + object graph. + """ + + SCHEMA = BuildModelSchema + + # task_name only appears in the ModelFactory call and in error messages, + # but it is declared all the same: a key read without being declared is + # invisible to any caller — and to the DAG validator — that inspects + # REQUIRES instead of running the unit. + # + # run_id is configuration, not context: no unit publishes it, so nothing + # upstream could ever satisfy it. It is read without a default on purpose + # — a run_id nobody passed would read as "this model has no run", and a + # model with no run logs no metrics at all (see BaseModel). + REQUIRES = ("n_labels", "task_name") + PROVIDES = ("model", "factory", "optimizable_parameters", "model_parameters") + RUNTIME_PARAMS = ("run_id",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._model_class = None + + @property + def model_name(self) -> str: + return self.config["model"]["component"] + + @property + def model_parameters(self) -> dict: + return self.config["model"]["params"] + + def _resolve_model_class(self) -> type: + """Resolve the model class from the registry, memoized on this unit. + + Memoized on the unit instance, not in the shared context: a context + can outlive a single ``BuildModelUnit`` (a future DAG could have more + than one build-model node feeding the same run), and a context-global + cache key would make the second instance silently reuse the first + one's model class. + """ + if self._model_class is not None: + return self._model_class + + from kink import di + + component_registry = di["component_registry"] + model_name: str = self.model_name + + try: + model_class = component_registry[model_name]["class"] + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find Model with name {model_name} in registry.", + ) from e + + self._model_class = model_class + return model_class + + def validate(self, ctx: ExecutionContext) -> None: + from kink import di + + component_registry = di["component_registry"] + model_name: str = self.model_name + parameters = self.model_parameters + + model_class = self._resolve_model_class() + + if getattr(model_class, "REQUIRES_DOWNLOAD", False) and not ( + model_class.is_downloaded() + ): + raise JobError( + f"Model {model_name} is not downloaded. Download it before training." + ) + + nested_missing = missing_downloads(parameters, component_registry) + if nested_missing: + names = ", ".join(m["name"] for m in nested_missing) + raise JobError( + "These components are not downloaded. " + f"Download them before training: {names}." + ) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.models.model_factory import ModelFactory + + component_registry = di["component_registry"] + + parameters = self.model_parameters + run_id = self.config["run_id"] + task_name = ctx.require("task_name") + + model_class = self._resolve_model_class() + + try: + train_metrics: List["BaseMetric"] = [ + component_registry[m]["class"] for m in self.config["train_metrics"] + ] + validation_metrics: List["BaseMetric"] = [ + component_registry[m]["class"] + for m in self.config["validation_metrics"] + ] + test_metrics: List["BaseMetric"] = [ + component_registry[m]["class"] for m in self.config["test_metrics"] + ] + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find metrics associated with Task {task_name} in registry", + ) from e + + try: + factory = ModelFactory( + model_class, + parameters, + run_id, + train_metrics=train_metrics, + validation_metrics=validation_metrics, + test_metrics=test_metrics, + n_labels=ctx.require("n_labels"), + ) + model: "BaseModel" = factory.model + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to instantiate model using run {run_id}", + ) from e + + # The original tree is what the search unit rewrites with the best + # values found, so it travels as a reference instead of being read + # from this unit's configuration again. + ctx.put_ref("model_parameters", parameters) + ctx.put("factory", factory) + ctx.put("model", model) + ctx.put("optimizable_parameters", factory.optimizable_parameters) diff --git a/DashAI/back/units/compute_dataset_metadata_unit.py b/DashAI/back/units/compute_dataset_metadata_unit.py new file mode 100644 index 000000000..124c28e5b --- /dev/null +++ b/DashAI/back/units/compute_dataset_metadata_unit.py @@ -0,0 +1,108 @@ +"""Unit that computes the metadata stored alongside a dataset.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema, bool_field, schema_field +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + +#: Metadata keys produced by ``compute_metadata`` on top of the base ones. +#: +#: They are expensive (correlations and per-column statistics over the whole +#: dataset) and optional, so the unit both skips computing them and strips any it +#: finds already present when they were not asked for. +EXTENDED_METADATA_KEYS = ( + "general_info", + "numeric_stats", + "categorical_stats", + "text_stats", + "quality_info", + "correlations", +) + + +class ComputeDatasetMetadataSchema(BaseSchema): + compute_metadata: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether to compute the extended exploratory metadata " + "(per-column statistics, quality report, correlations) on top of " + "the base one. Disabling it makes large datasets much faster to " + "store, at the cost of the exploration views.", + es="Si se debe calcular la metadata exploratoria extendida " + "(estadísticas por columna, informe de calidad, correlaciones) " + "además de la básica. Desactivarla hace mucho más rápido el " + "guardado de conjuntos de datos grandes, a costa de las vistas de " + "exploración.", + pt="Se deve calcular a metadata exploratória estendida " + "(estatísticas por coluna, relatório de qualidade, correlações) " + "além da básica. Desativá-la torna o armazenamento de grandes " + "conjuntos de dados muito mais rápido, ao custo das vistas de " + "exploração.", + de="Ob die erweiterten explorativen Metadaten (Spaltenstatistiken, " + "Qualitätsbericht, Korrelationen) zusätzlich zu den Basisdaten " + "berechnet werden. Das Deaktivieren beschleunigt das Speichern " + "großer Datensätze erheblich, auf Kosten der Explorationsansichten.", + zh="是否在基础元数据之外计算扩展的探索性元数据（各列统计、质量报告、" + "相关性）。禁用后可大幅加快大型数据集的存储速度，但会失去探索视图。", + ), + alias=MultilingualString( + en="Compute extended metadata", + es="Calcular metadata extendida", + pt="Calcular metadata estendida", + de="Erweiterte Metadaten berechnen", + zh="计算扩展元数据", + ), + ) # type: ignore + + +class ComputeDatasetMetadataUnit(BaseUnit): + """Fill in the metadata a dataset carries in its ``splits`` mapping. + + Two independent knobs, because "how much metadata" and "can the metadata + already there be trusted" are different questions: + + * ``compute_metadata`` picks the depth: base only (column names, row count, + NaN counts) or base plus the extended exploratory fields. + * ``trust_inherited_metadata`` says the dataset arrived carrying metadata + that still describes it — the case of a copy whose bytes never changed. + Then the unit fills only what is missing instead of recomputing. + + Either way the requested depth is what ends up stored: asking for base only + always strips the extended keys, even ones inherited from elsewhere, so the + result never depends on where the dataset came from. + + Re-publishes ``dataset`` although it mutates it in place: the key is what + makes the unit chainable, and the contract audit reads the ``ctx.put`` call. + """ + + SCHEMA = ComputeDatasetMetadataSchema + + REQUIRES = ("dataset",) + PROVIDES = ("dataset",) + RUNTIME_PARAMS = ("trust_inherited_metadata",) + + def execute(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") + + compute_extended = self.config.get("compute_metadata", True) + trust_inherited = self.config.get("trust_inherited_metadata", False) + + if compute_extended: + has_extended = any(key in dataset.splits for key in EXTENDED_METADATA_KEYS) + if not (trust_inherited and has_extended): + dataset.compute_metadata() + else: + if not (trust_inherited and "total_rows" in dataset.splits): + dataset.compute_base_metadata() + # Strip extended keys the dataset may have arrived with, so the + # stored metadata matches what was asked for rather than the + # history of the file. + for stale_key in EXTENDED_METADATA_KEYS: + dataset.splits.pop(stale_key, None) + + ctx.put("dataset", dataset) diff --git a/DashAI/back/units/context.py b/DashAI/back/units/context.py new file mode 100644 index 000000000..daa1ef6a0 --- /dev/null +++ b/DashAI/back/units/context.py @@ -0,0 +1,218 @@ +"""Execution context shared between atomic units.""" + +import copy +import json +from typing import Any, Dict, Optional + + +class UnitContractError(Exception): + """Raised when a unit's declared contract is violated. + + Either a required key is missing from the context before execution, or a + promised key is missing after it. + """ + + +class ExecutionContext: + """State that crosses the boundary between atomic units. + + The context deliberately splits its state in two halves with different + rules: + + * ``refs``: JSON serializable references (ids, paths, index lists). This is + the only half that can cross a process boundary. Jobs are shipped to the + Huey worker with dill and rebuild their dependencies from a fresh + container, so anything that must survive that trip has to be expressible + as plain data. + * ``cache``: live objects (datasets, models, tasks). Never serialized; + always derivable again from the refs. + + Keeping the halves apart is what allows the same unit to run in-process + (cache hit, nothing is reloaded) or, in the future, as an independently + enqueued job (refs travel, the heavy objects are derived again). + + Parameters + ---------- + refs : Dict[str, Any], optional + Initial JSON serializable references. + """ + + def __init__(self, refs: Optional[Dict[str, Any]] = None) -> None: + self._refs: Dict[str, Any] = {} + self._cache: Dict[str, Any] = {} + + for key, value in (refs or {}).items(): + self.put_ref(key, value) + + @property + def refs(self) -> Dict[str, Any]: + """A deep copy of the serializable references held by the context.""" + return copy.deepcopy(self._refs) + + def put_ref(self, key: str, value: Any) -> None: + """Store a durable, JSON serializable reference. + + The value is deep-copied on the way in, so a mutable dict handed in + by the caller (e.g. an ORM-attached ``dict`` column) is never aliased. + Without this, a later in-place edit made through the context — such as + ``ModelFactory.update_parameters`` rewriting a nested ``fixed_value`` + — would silently write through to the caller's object. + + Parameters + ---------- + key : str + Name of the reference. + value : Any + Value to store. Must be JSON serializable. + + Raises + ------ + UnitContractError + If the value cannot be serialized to JSON. + """ + try: + json.dumps(value) + except (TypeError, ValueError) as e: + raise UnitContractError( + f"Context reference '{key}' is not JSON serializable " + f"({type(value).__name__}). Only ids, paths and plain data can " + "cross a unit boundary; store live objects with put() instead." + ) from e + + self._refs[key] = copy.deepcopy(value) + + def put(self, key: str, value: Any) -> None: + """Store a live object in the in-process cache. + + Parameters + ---------- + key : str + Name of the value. + value : Any + Any object. It is never serialized. + """ + self._cache[key] = value + + def get(self, key: str, default: Any = None) -> Any: + """Retrieve a value, looking in the cache before the references. + + A live object from the cache is returned by reference — that is the + point of the cache half. A reference is returned as a deep copy, so + mutating what comes back never reaches into the context's own state + or, transitively, into whatever object a ``put_ref`` call was given. + + Parameters + ---------- + key : str + Name of the value. + default : Any, optional + Returned when the key is absent from both halves. + + Returns + ------- + Any + The cached object, a copy of the stored reference, or ``default``. + """ + if key in self._cache: + return self._cache[key] + if key in self._refs: + return copy.deepcopy(self._refs[key]) + return default + + def require(self, key: str) -> Any: + """Retrieve a value, failing when it is absent. + + See :meth:`get` for the copy-on-read guarantee for references. + + Parameters + ---------- + key : str + Name of the value. + + Returns + ------- + Any + The cached object or a copy of the stored reference. + + Raises + ------ + UnitContractError + If the key is present in neither half of the context. + """ + if key in self._cache: + return self._cache[key] + if key in self._refs: + return copy.deepcopy(self._refs[key]) + + raise UnitContractError( + f"Context key '{key}' is not available. " + f"Present keys: {sorted(set(self._cache) | set(self._refs))}." + ) + + def has(self, key: str) -> bool: + """Return whether a key is present in either half of the context.""" + return key in self._cache or key in self._refs + + def origin(self, key: str) -> Optional[str]: + """Return which half holds a key, without copying anything. + + ``get`` and ``has`` merge the two halves on purpose, and ``refs`` + deep-copies the whole reference half, so neither can answer this + cheaply. Something that moves a value from one context to another has + to know: the two halves have incompatible rules, and picking the wrong + one fails in both directions. Handing a live dataset to ``put_ref`` + raises, and handing ``put`` a dict that was a reference silently drops + the copy-on-read guarantee the reference depended on. + + The cache is checked first, matching ``get`` and ``require``, so a key + present in both halves has one answer rather than two. + + Parameters + ---------- + key : str + Name of the value. + + Returns + ------- + Optional[str] + ``"cache"``, ``"ref"``, or ``None`` when the key is absent. + """ + if key in self._cache: + return "cache" + if key in self._refs: + return "ref" + return None + + def clear_cache(self) -> None: + """Drop every live object, keeping the references. + + Called by the orchestrator before ``gc.collect()`` so datasets, splits + and models become collectable as soon as the job is done. + """ + self._cache.clear() + + def to_dict(self) -> Dict[str, Any]: + """Serialize the context to the references that can cross a process. + + Returns + ------- + Dict[str, Any] + A deep copy of the JSON serializable half of the context. + """ + return self.refs + + @classmethod + def from_dict(cls, refs: Dict[str, Any]) -> "ExecutionContext": + """Rebuild a context from previously serialized references. + + Parameters + ---------- + refs : Dict[str, Any] + References as returned by :meth:`to_dict`. + + Returns + ------- + ExecutionContext + A context with an empty cache; live objects are derived on demand. + """ + return cls(refs=refs) diff --git a/DashAI/back/units/converter_scope.py b/DashAI/back/units/converter_scope.py new file mode 100644 index 000000000..73a10a712 --- /dev/null +++ b/DashAI/back/units/converter_scope.py @@ -0,0 +1,357 @@ +"""Shared scope resolution and slicing for the converter units. + +The three converter units — fit, transform, and the fused apply — all express +which part of a dataset they touch the same way, and all resolve it the same +way. That logic lives here so the fused path and the split path cannot drift +apart. +""" + +import logging +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +from DashAI.back.core.schema_fields import ( + component_field, + none_type, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + +EMPTY_SCOPE: Dict[str, List] = {"columns": [], "rows": []} + + +def converter_field(): + """Schema field for picking a converter and configuring it.""" + return schema_field( + component_field(parent="BaseConverter"), + placeholder={"component": "ColumnRemover", "params": {}}, + description=MultilingualString( + en="Converter to apply, together with its own configuration.", + es="Convertidor a aplicar, junto con su propia configuración.", + pt="Conversor a aplicar, junto com a sua própria configuração.", + de="Anzuwendender Konverter samt seiner eigenen Konfiguration.", + zh="要应用的转换器及其自身的配置。", + ), + alias=MultilingualString( + en="Converter", + es="Convertidor", + pt="Conversor", + de="Konverter", + zh="转换器", + ), + ) + + +def scope_field(): + """Schema field for the part of the dataset a converter touches.""" + return schema_field( + none_type(dict), + placeholder={"columns": [], "rows": []}, + description=MultilingualString( + en="Part of the dataset the converter applies to: 'columns' is a " + "list of {'idx': n} and 'rows' a list of n, both 1-based. An " + "empty column list means every column.", + es="Parte del conjunto de datos a la que se aplica el convertidor: " + "'columns' es una lista de {'idx': n} y 'rows' una lista de n, " + "ambas con base 1. Una lista de columnas vacía significa todas " + "las columnas.", + pt="Parte do conjunto de dados à qual o conversor se aplica: " + "'columns' é uma lista de {'idx': n} e 'rows' uma lista de n, " + "ambas com base 1. Uma lista de colunas vazia significa todas " + "as colunas.", + de="Teil des Datensatzes, auf den der Konverter angewendet wird: " + "'columns' ist eine Liste von {'idx': n} und 'rows' eine Liste " + "von n, beide 1-basiert. Eine leere Spaltenliste bedeutet alle " + "Spalten.", + zh="转换器作用的数据集范围：'columns' 是 {'idx': n} 的列表，" + "'rows' 是 n 的列表，均从 1 开始。空的列列表表示所有列。", + ), + alias=MultilingualString( + en="Scope", + es="Alcance", + pt="Escopo", + de="Geltungsbereich", + zh="范围", + ), + ) + + +def target_field(): + """Schema field for the column handed to the converter as ``y``.""" + return schema_field( + none_type(dict), + placeholder=None, + description=MultilingualString( + en="Target column handed to the converter as y, as {'idx': n} with " + "n 1-based. Leave empty for unsupervised converters.", + es="Columna objetivo entregada al convertidor como y, como " + "{'idx': n} con n en base 1. Dejar vacío para convertidores no " + "supervisados.", + pt="Coluna alvo entregue ao conversor como y, como {'idx': n} com " + "n com base 1. Deixe vazio para conversores não supervisionados.", + de="Zielspalte, die dem Konverter als y übergeben wird, als " + "{'idx': n} mit 1-basiertem n. Für unüberwachte Konverter leer " + "lassen.", + zh="作为 y 传给转换器的目标列，格式为 {'idx': n}，n 从 1 开始。" + "无监督转换器留空。", + ), + alias=MultilingualString( + en="Target column", + es="Columna objetivo", + pt="Coluna alvo", + de="Zielspalte", + zh="目标列", + ), + ) + + +def rebuild_dataset_with_transformed_columns( + base: "DashAIDataset", + transformed: "DashAIDataset", + scope_column_names: List[str], +) -> "DashAIDataset": + """ + Replaces specific columns in the base dataset with columns from the transformed + dataset, preserving their original positions. Also appends any additional columns + that were generated by the transformer at the end. Keeps the features and metadata + consistent. + + Parameters + ---------- + base : DashAIDataset + The original dataset before transformation. + + transformed : DashAIDataset + The dataset resulting from applying a transformer, containing updated and/or + new columns. + + scope_column_names : List[str] + Names of the columns that were originally selected for transformation. + + Returns + ------- + DashAIDataset + A new dataset with the specified columns replaced in place, new columns + appended, and original metadata and split information preserved. + """ + from DashAI.back.dataloaders.classes.dashai_dataset import modify_table + + original_columns = base.column_names + transformed_cols = transformed.column_names + + transformed_cols_set = set(transformed_cols) + scope_column_names_set = set(scope_column_names) + + removed_cols = [ + col for col in scope_column_names if col not in transformed_cols_set + ] + replacement_cols = [ + col for col in scope_column_names if col in transformed_cols_set + ] + new_cols = [col for col in transformed_cols if col not in scope_column_names_set] + + removed_cols_set = set(removed_cols) + + new_columns_order = [] + seen_cols = set() + for col in original_columns: + if col in removed_cols_set: + continue + if col not in seen_cols: + new_columns_order.append(col) + seen_cols.add(col) + + col_name_mapping = {} + for col in new_cols: + unique_col = col + counter = 1 + while unique_col in seen_cols: + unique_col = f"{col}_{counter}" + counter += 1 + new_columns_order.append(unique_col) + seen_cols.add(unique_col) + col_name_mapping[col] = unique_col + + updated_arrays = {} + for col in replacement_cols: + if col in transformed_cols_set: + updated_arrays[col] = transformed.arrow_table[col] + for col, unique_col in col_name_mapping.items(): + if col in transformed_cols_set: + updated_arrays[unique_col] = transformed.arrow_table[col] + + updated_types = base.types.copy() + + for col in removed_cols: + if col in updated_types: + del updated_types[col] + + for col in replacement_cols: + if col in transformed.types: + updated_types[col] = transformed.types[col] + for col, unique_col in col_name_mapping.items(): + if col in transformed.types: + updated_types[unique_col] = transformed.types[col] + + modified_dataset = modify_table(base, updated_arrays, types=updated_types) + modified_dataset = modified_dataset.select_columns(new_columns_order) + + return modified_dataset + + +class ConverterScopeMixin: + """Scope handling shared by the converter units. + + Deliberately **not** named ``Base*`` and deliberately not inheriting from + ``BaseUnit``. ``ComponentRegistry._get_base_type`` walks the MRO for + ancestors whose name contains "Base" and that declare a ``TYPE``, and + rejects a component with more than one candidate — an intermediate class + called ``BaseConverterUnit`` would be a second candidate and break the + registration of every unit inheriting from it. + """ + + def _scope(self) -> Dict[str, List]: + """The configured scope, tolerating an explicit ``None``. + + The API schema always writes the ``scope`` key but allows it to be null, + so a converter saved without a scope arrives here as ``None`` rather + than as a missing key. + """ + return self.config.get("scope") or EMPTY_SCOPE + + def _target_index(self) -> Optional[int]: + """The 1-based target column index, or None when there is no target.""" + target = self.config.get("target") + if target is None: + return None + return target.get("idx") + + def _check_target_bounds(self, dataset: "DashAIDataset") -> None: + """Reject a target index that does not point at a column. + + Checked against the dataset in hand rather than once up front: when + several converters are chained the column count changes underneath, and + an unchecked index would surface as a bare IndexError. + """ + target_column_index = self._target_index() + if target_column_index is None: + return + + if int(target_column_index) < 1 or int(target_column_index) > len( + dataset.features + ): + raise JobError( + f"Target column index {target_column_index} is out of bounds" + ) + + def _resolve_scope( + self, dataset: "DashAIDataset" + ) -> Tuple[List[str], List[int], Optional[str]]: + """Turn the 1-based scope into names and indexes for **this** dataset. + + Always resolved against the dataset in hand, never against a list + carried through the context: a converter that renames, drops or adds + columns changes what index 3 means for whatever runs next. + """ + scope = self._scope() + column_names = dataset.column_names + + columns_scope = [column["idx"] - 1 for column in scope["columns"]] + scope_column_indexes = sorted(set(columns_scope)) + + if not scope_column_indexes: + scope_column_indexes = list(range(len(dataset.features))) + + scope_column_names = [column_names[index] for index in scope_column_indexes] + + rows_scope = [row - 1 for row in scope["rows"]] + scope_rows_indexes = sorted(set(rows_scope)) + + target_column_name = None + target_column_index = self._target_index() + if target_column_index is not None: + target_column_name = column_names[int(target_column_index) - 1] + + return scope_column_names, scope_rows_indexes, target_column_name + + def _slice_for_fit( + self, + dataset: "DashAIDataset", + scope_column_names: List[str], + scope_rows_indexes: List[int], + target_column_name: Optional[str], + ) -> Tuple["DashAIDataset", Optional["DashAIDataset"]]: + """The X and y a converter is fitted on: scoped columns, scoped rows.""" + y_dataset = None + if target_column_name is not None: + y_dataset = dataset.select_columns([target_column_name]) + if scope_rows_indexes: + y_dataset = y_dataset.select(scope_rows_indexes) + + x_dataset = dataset.select_columns(scope_column_names) + if scope_rows_indexes: + x_dataset = x_dataset.select(scope_rows_indexes) + + return x_dataset, y_dataset + + def _slice_for_transform( + self, + dataset: "DashAIDataset", + scope_column_names: List[str], + target_column_name: Optional[str], + ) -> Tuple["DashAIDataset", Optional["DashAIDataset"]]: + """The X and y a converter transforms: scoped columns, **every** row. + + A row scope narrows what the converter learns from, never what it is + applied to. + """ + y_dataset = None + if target_column_name is not None: + y_dataset = dataset.select_columns([target_column_name]) + + return dataset.select_columns(scope_column_names), y_dataset + + def _merge_transformed( + self, + converter_instance, + dataset: "DashAIDataset", + transformed_dataset: "DashAIDataset", + scope_column_names: List[str], + ) -> "DashAIDataset": + """Fold the transform output back into the dataset. + + A converter that changes the row count cannot have its columns merged + back position by position, so its output replaces the dataset outright. + """ + if type(converter_instance).CHANGES_ROW_COUNT: + return transformed_dataset + + return rebuild_dataset_with_transformed_columns( + dataset, + transformed_dataset, + scope_column_names, + ) + + def _fit(self, converter_instance, x_dataset, y_dataset, converter_name): + """Fit a converter, preserving the job's two error messages.""" + try: + return converter_instance.fit(x_dataset, y_dataset) + except ValueError as e: + log.error(f"Validation error in {converter_name}: {e}") + raise JobError(f"Validation error fitting {converter_name}: {e}") from e + except Exception as e: + log.exception(e) + raise JobError(f"Error fitting converter {converter_name}: {e}") from e + + def _transform(self, converter_instance, x_dataset, y_dataset, converter_name): + """Transform with a fitted converter, preserving the job's message.""" + try: + return converter_instance.transform(x_dataset, y_dataset) + except Exception as e: + log.exception(e) + raise JobError(f"Error transforming data with {converter_name}: {e}") from e diff --git a/DashAI/back/units/evaluate_model_to_artifact_unit.py b/DashAI/back/units/evaluate_model_to_artifact_unit.py new file mode 100644 index 000000000..14e73ceb1 --- /dev/null +++ b/DashAI/back/units/evaluate_model_to_artifact_unit.py @@ -0,0 +1,90 @@ +"""Unit that computes a trained model's metrics and publishes them.""" + +import logging + +from DashAI.back.core.enums.metrics import SplitEnum +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + list_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + +DEFAULT_SPLITS = ["TRAIN", "VALIDATION", "TEST"] + + +class EvaluateModelToArtifactSchema(BaseSchema): + splits: schema_field( + list_field(enum_field(enum=DEFAULT_SPLITS)), + placeholder=DEFAULT_SPLITS, + description=MultilingualString( + en="Data splits the model is evaluated on.", + es="Particiones de datos sobre las que se evalúa el modelo.", + pt="Partições de dados sobre as quais o modelo é avaliado.", + de="Datenteilmengen, auf denen das Modell ausgewertet wird.", + zh="用于评估模型的数据划分。", + ), + alias=MultilingualString( + en="Splits", + es="Particiones", + pt="Partições", + de="Teilmengen", + zh="数据划分", + ), + ) # type: ignore + + +class EvaluateModelToArtifactUnit(BaseUnit): + """Compute a trained model's metrics and publish them as data. + + The sibling of ``EvaluateModelUnit``, for a caller that has no ``Run`` row. + That unit writes ``Metric`` rows against a foreign key to ``run.id``, so it + cannot work without one; this returns the same numbers instead of + persisting them, and whoever asked for them decides where they go. + + Which metrics are computed is decided by the model, not by this unit, and + not configured here either: ``ModelFactory`` attaches the metric classes + and the data splits to the instance, so the metrics stay configured in one + place — where the model is built — rather than in two nodes that could + disagree. + + Both this and ``EvaluateModelUnit`` score through + ``BaseModel.compute_metrics``, which is what keeps them from drifting apart + into two answers for the same model and split. + """ + + SCHEMA = EvaluateModelToArtifactSchema + + REQUIRES = ("model",) + PROVIDES = ("metrics",) + + def execute(self, ctx: ExecutionContext) -> None: + model = ctx.require("model") + splits = [SplitEnum[name] for name in self.config.get("splits", DEFAULT_SPLITS)] + + metrics = {} + try: + for split in splits: + scores = model.compute_metrics(split=split) + # None means there was nothing to score -- no metrics + # configured, or no data for this split. Recording an empty + # entry would claim the split was evaluated and scored zero + # metrics, which is a different statement. + if scores is None: + continue + metrics[split.value] = scores + except Exception as e: + log.exception(e) + raise JobError( + f"Metric calculation failed {e}", + ) from e + + # A ref, not a cached object: plain numbers are exactly what survives + # leaving the process, and what a caller records. + ctx.put_ref("metrics", metrics) diff --git a/DashAI/back/units/evaluate_model_unit.py b/DashAI/back/units/evaluate_model_unit.py new file mode 100644 index 000000000..edcfa0c9e --- /dev/null +++ b/DashAI/back/units/evaluate_model_unit.py @@ -0,0 +1,111 @@ +"""Unit that computes and stores the final metrics of a trained model.""" + +import logging + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + list_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Metric +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + +DEFAULT_SPLITS = ["TRAIN", "VALIDATION", "TEST"] + + +class EvaluateModelSchema(BaseSchema): + splits: schema_field( + list_field(enum_field(enum=DEFAULT_SPLITS)), + placeholder=DEFAULT_SPLITS, + description=MultilingualString( + en="Data splits the model is evaluated on.", + es="Particiones de datos sobre las que se evalúa el modelo.", + pt="Partições de dados sobre as quais o modelo é avaliado.", + de="Datenteilmengen, auf denen das Modell ausgewertet wird.", + zh="用于评估模型的数据划分。", + ), + alias=MultilingualString( + en="Splits", + es="Particiones", + pt="Partições", + de="Teilmengen", + zh="数据划分", + ), + ) # type: ignore + + +class EvaluateModelUnit(BaseUnit): + """Compute the final metrics of a trained model, once per split. + + The unit is idempotent: a split whose final metrics were already logged + (for instance by a model that evaluates itself while training) is skipped. + + Which metrics are computed is decided by the model, not by this unit: + ``ModelFactory`` attaches the metric classes to the model instance, and + ``BaseModel.calculate_metrics`` is ``final``. That method persists the + rows through a session of its own, so these writes are not part of the + transaction the calling job controls. + + **This unit needs a real ``Run`` row and cannot be used without one.** + Everything it does is write ``Metric`` rows against a foreign key to + ``run.id``, so there is nothing left for it to do when there is no run. + ``validate`` refuses a missing run id rather than letting it through, + because the failure would otherwise be silent in both directions: the + idempotency query below would match no row whatever was already logged, + and ``calculate_metrics`` no-ops on a model with no run — so the unit + would report success having written nothing at all. + + A caller that wants a model's metrics *without* a run wants a different + unit, one that returns them instead of persisting them. + """ + + SCHEMA = EvaluateModelSchema + + # run_id is configuration, not context: no unit publishes it. + REQUIRES = ("model",) + RUNTIME_PARAMS = ("run_id",) + + def validate(self, ctx: ExecutionContext) -> None: + if self.config["run_id"] is None: + raise JobError( + "Metrics can only be logged against a run, and this one has no " + "run id. Use a unit that returns the metrics instead of writing " + "them if there is no run to attach them to." + ) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + session_factory = di["session_factory"] + + model = ctx.require("model") + # validate() already refused a None here, which is what keeps the + # idempotency query below from matching no row regardless of what was + # actually logged. + run_id = self.config["run_id"] + splits = [SplitEnum[name] for name in self.config.get("splits", DEFAULT_SPLITS)] + + try: + for split in splits: + with session_factory() as db: + already_logged = ( + db.query(Metric) + .filter_by(run_id=run_id, split=split, level=LevelEnum.LAST) + .first() + ) + if already_logged: + continue + + model.calculate_metrics(split=split, level=LevelEnum.LAST) + except Exception as e: + log.exception(e) + raise JobError( + f"Metric calculation failed {e}", + ) from e diff --git a/DashAI/back/units/explanation_artifacts.py b/DashAI/back/units/explanation_artifacts.py new file mode 100644 index 000000000..e8ee8416a --- /dev/null +++ b/DashAI/back/units/explanation_artifacts.py @@ -0,0 +1,128 @@ +"""Shared helpers for the two explanation-generating units. + +Not a unit: no configuration, no context, nothing to declare. It lives here +rather than in ``job/`` because importing from a job into a unit would invert +the dependency. +""" + +import logging +from typing import Any, Tuple + +from DashAI.back.job.base_job import JobError + +log = logging.getLogger(__name__) + + +def build_explainer(scope: str, selected: dict, trained_model: Any) -> Any: + """Resolve an explainer component and bind it to a trained model. + + Shared by the two scope-specific build units: building is identical either + way, only the registry the component comes from differs. + + Takes and returns plain values instead of touching the context. That keeps + every context write inside the unit itself, where the contract audit can + see it — a ``ctx.put`` hidden in a helper is invisible to the static check + and would let a broken ``PROVIDES`` through. + + Parameters + ---------- + scope : str + ``"global"`` or ``"local"``. Only decorates the error messages, which + are user-visible and worded per scope. + selected : dict + The ``{"component": ..., "params": ...}`` value of the unit's field. + trained_model : Any + The model the explainer explains. + + Returns + ------- + Any + The instantiated explainer. + + Raises + ------ + JobError + If the component is not registered or cannot be instantiated. + """ + from kink import di + + component_registry = di["component_registry"] + + explainer_name = selected["component"] + + try: + explainer_class = component_registry[explainer_name]["class"] + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find the {scope} explainer with name " + f"{explainer_name} in registry.", + ) from e + + try: + return explainer_class(model=trained_model, **(selected.get("params") or {})) + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to instantiate {scope} explainer.", + ) from e + + +def dump_explanation(explanation: Any, plots: Any, prefix: str, key: int) -> Tuple: + """Pickle an explanation and its plots under the explanations directory. + + Both files are named after the explanation they belong to, so re-running + one overwrites its own artifacts and never another's. + + Parameters + ---------- + explanation : Any + Whatever the explainer's ``explain``/``explain_instance`` returned. + plots : Any + The normalized artifacts produced from that explanation. + prefix : str + ``"global"`` or ``"local"``: the two scopes keep separate file names + because their ids come from separate tables and would otherwise clash. + key : int + Identifier of the explanation row. + + Returns + ------- + Tuple[str, str] + The explanation path and the plot path. + + Raises + ------ + JobError + If either file cannot be written. + """ + import os + import pickle + + from kink import di + + config = di["config"] + + plots_name = "plot" if prefix == "global" else "plots" + + try: + explanation_path = os.path.join( + config["EXPLANATIONS_PATH"], f"{prefix}_explanation_{key}.pickle" + ) + with open(explanation_path, "wb") as file: + pickle.dump(explanation, file) + + plot_path = os.path.join( + config["EXPLANATIONS_PATH"], + f"{prefix}_explanation_{plots_name}_{key}.pickle", + ) + with open(plot_path, "wb") as file: + pickle.dump(plots, file) + + except Exception as e: + log.exception(e) + raise JobError( + "Explanation file saving failed", + ) from e + + return explanation_path, plot_path diff --git a/DashAI/back/units/fit_converter_unit.py b/DashAI/back/units/fit_converter_unit.py new file mode 100644 index 000000000..645f10d52 --- /dev/null +++ b/DashAI/back/units/fit_converter_unit.py @@ -0,0 +1,107 @@ +"""Unit that fits a converter on a dataset without transforming anything.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.converter_scope import ( + ConverterScopeMixin, + converter_field, + scope_field, + target_field, +) + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class FitConverterSchema(BaseSchema): + converter: converter_field() # type: ignore + scope: scope_field() # type: ignore + target: target_field() # type: ignore + + +class FitConverterUnit(BaseUnit, ConverterScopeMixin): + """Fit a converter on the dataset in the context and publish it fitted. + + Splits the "fit" half out of :class:`ApplyConverterUnit` so a converter can + be learned from one dataset and applied to another. That is the standard + train/test discipline: a scaler, an encoder or an imputer must learn its + statistics from the training data only, and then be applied unchanged to the + test data — refitting on test would leak information into the evaluation. + + The dataset is left untouched: this unit produces a fitted converter, not + data. Pair it with :class:`TransformDatasetUnit`, once per dataset the + fitted converter has to be applied to. + """ + + SCHEMA = FitConverterSchema + + REQUIRES = ("dataset",) + PROVIDES = ("fitted_converter",) + + def __init__(self, **config) -> None: + super().__init__(**config) + # Memoized on the instance, never in the context: two units of this + # class can live in the same context and they are different converters. + self._converter_class = None + + @property + def _converter_name(self) -> str: + return self.config["converter"]["component"] + + def _resolve_converter(self): + """Look the converter class up in the registry, once per instance.""" + if self._converter_class is not None: + return self._converter_class + + from kink import di + + component_registry = di["component_registry"] + converter_name = self._converter_name + + try: + self._converter_class = component_registry[converter_name]["class"] + except KeyError as e: + log.exception(e) + raise JobError(f"Error importing converter {converter_name}: {e}") from e + + return self._converter_class + + def validate(self, ctx: ExecutionContext) -> None: + """Reject an out-of-bounds target before any work is done.""" + self._check_target_bounds(ctx.require("dataset")) + + def execute(self, ctx: ExecutionContext) -> None: + dataset: "DashAIDataset" = ctx.require("dataset") + converter_name = self._converter_name + + converter_constructor = self._resolve_converter() + converter_instance = converter_constructor( + **(self.config["converter"].get("params") or {}) + ) + + self._check_target_bounds(dataset) + ( + scope_column_names, + scope_rows_indexes, + target_column_name, + ) = self._resolve_scope(dataset) + + log.info(f"Fitting converter: {converter_name}") + + x_dataset, y_dataset = self._slice_for_fit( + dataset, + scope_column_names, + scope_rows_indexes, + target_column_name, + ) + + fitted = self._fit(converter_instance, x_dataset, y_dataset, converter_name) + + ctx.put("fitted_converter", fitted) diff --git a/DashAI/back/units/fit_model_over_folds_unit.py b/DashAI/back/units/fit_model_over_folds_unit.py new file mode 100644 index 000000000..b530df41b --- /dev/null +++ b/DashAI/back/units/fit_model_over_folds_unit.py @@ -0,0 +1,281 @@ +"""Unit that fits a model across cross-validation folds.""" + +import logging + +import numpy as np + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.fit_scope import ( + TRIAL_SPLITS, + ModelFitScopeMixin, + goal_metric_field, + optimizer_field, + trial_splits_field, +) + +log = logging.getLogger(__name__) + + +class FitModelOverFoldsSchema(BaseSchema): + optimizer: optimizer_field() # type: ignore + goal_metric: goal_metric_field() # type: ignore + scored_splits: trial_splits_field() # type: ignore + + +class FitModelOverFoldsUnit(BaseUnit, ModelFitScopeMixin): + """Fit a model once per fold, then once more on everything the folds used. + + The sibling of ``FitModelUnit``. It takes a list of partition sets rather + than one, which is a different ``REQUIRES`` and therefore a different unit; + everything around the fit is shared between them. + + Three things happen here, in this order, and the order is the point: + + 1. If the model declares optimizable parameters, the search runs first. + Its objective is the whole fold loop, so one trial costs k fits. That is + what cross-validation buys and what it costs, and it is why the + objective had to become something the caller supplies rather than + something the optimizer does. + 2. Every fold is fitted and scored, and a ``FOLD`` metric row is written + for it. Those rows are what the fold charts and the statistical tests + read, and they are indexed from zero without gaps -- repeated + cross-validation buckets them by integer division, so a gap silently + moves a fold into the wrong repetition. + 3. The model that gets kept is fitted on every row the folds could use, + and left pointing at that partition set. Scoring it against the rows the + session reserved is not done here: that is a ``LAST`` metric like any + other, so it is ``EvaluateModelUnit``'s, the same unit a holdout run + uses -- and the caller is the one that knows whether anything was + reserved to score against. + + **The per-fold scores are published rather than aggregated here.** Turning + them into a mean and a deviation means writing ``Metric`` rows that carry a + ``std_value``, and a unit may not write domain rows -- the one sanctioned + write in the domain layer, ``BaseModel._save_metrics``, has nowhere to put + a deviation. Whoever asked for the fit aggregates and persists. + + Nested cross-validation is not here either: its inner splitter is a + required component field, and a component field cannot be made optional + without leaving the user without a selector, so it is a further sibling + rather than a flag on this one. + """ + + SCHEMA = FitModelOverFoldsSchema + + REQUIRES = ( + "model", + "factory", + "optimizable_parameters", + "model_parameters", + "x_folds", + "y_folds", + ) + PROVIDES = ("model", "plot_paths", "fold_metrics") + RUNTIME_PARAMS = ("run_id", "artifact_prefix") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._optimizer = None + self._goal_metric = None + + def validate(self, ctx: ExecutionContext) -> None: + # ctx.require, not ctx.get: an absent key means BuildModelUnit has not + # run yet, which is a call-order mistake rather than a model with + # nothing to optimize. + self._validate_search(ctx.require("optimizable_parameters")) + + def execute(self, ctx: ExecutionContext) -> None: + model = ctx.require("model") + x_folds = ctx.require("x_folds") + y_folds = ctx.require("y_folds") + optimizable_parameters = ctx.require("optimizable_parameters") + + plot_paths = [] + try: + if self._will_search(optimizable_parameters): + # Every read of the context happens in this file rather than in + # the shared helper: the contract audit parses it, so a require + # moved out makes a declared key look unread. + model, best_parameters, plot_paths = self._search( + model, + x_folds, + y_folds, + optimizable_parameters, + ctx.require("factory"), + ctx.require("model_parameters"), + self.config["run_id"], + self.config["artifact_prefix"], + self._score_one_trial, + ) + ctx.put_ref("best_parameters", best_parameters) + + fold_metrics = self._score_every_fold(model, x_folds, y_folds) + + # The model that gets kept is fitted last, on every row the folds + # could use, and is left pointing at that partition set -- which is + # what whoever scores it afterwards will be scoring. + model.x_data = x_folds[-1] + model.y_data = y_folds[-1] + self._fit_kept_model(model, x_folds[-1], y_folds[-1]) + except Exception as e: + log.exception(e) + raise JobError( + f"Model training failed {e}", + ) from e + + ctx.put("model", model) + ctx.put_ref("plot_paths", plot_paths) + ctx.put_ref("fold_metrics", fold_metrics) + + # ----------------------------------------------------------------- # + # The three things, one method each + # ----------------------------------------------------------------- # + + def _score_one_trial(self, model, x_folds, y_folds, metric) -> float: + return self._score_folds(model, x_folds, y_folds, metric, record=True) + + def _score_folds(self, model, x_folds, y_folds, metric, record: bool) -> float: + """Fit and score every fold, and return the mean, for one trial. + + This is what the optimizer measures, and it is the whole difference + between this unit and its sibling: there, a trial is one fit; here it + is k of them. The optimizer passes the partitions through without + looking at them, so it does not have to know which. + + One row per split is written for the trial, holding the mean over its + folds -- not one row per fold. The folds of a trial measure a + hyperparameter setting rather than the model that gets kept, so + recording each of them would bury the rows that describe it, and the + live chart that watches a search wants one point per trial anyway. + + The objective is scored directly rather than read out of those means, + because the goal metric is not necessarily one of the metrics the + session configured for the validation partition. + + Returns + ------- + float + The mean of the goal metric over the folds, which is the objective. + """ + scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) + scores = [] + accumulated: dict = {} + + for x_fold, y_fold in self._folds(x_folds, y_folds): + # Pointed at the fold it is about to be fitted on, the same as in + # the scoring loop. Nothing here reads it back, but the model is + # expected to carry the data it was last fitted on -- that is what + # the check after the search asserts, and what anything scoring it + # afterwards relies on. + model.x_data = x_fold + model.y_data = y_fold + self._fit_kept_model(model, x_fold, y_fold) + predictions = model.predict(x_fold["validation"]) + expected = model.prepare_output(y_fold["validation"], is_fit=False) + scores.append(metric.score(expected, predictions)) + + for name in scored_splits: + results = model.compute_metrics(split=SplitEnum[name]) + for metric_name, value in (results or {}).items(): + accumulated.setdefault(name, {}).setdefault(metric_name, []).append( + value + ) + + if record: + self._record_trial(model, accumulated) + return float(np.mean(scores)) + + @staticmethod + def _record_trial(model, accumulated: dict) -> None: + """Write one metric row per split for this trial, holding its mean. + + Guarded on the run, which ``_save_metrics`` does not do for itself: a + caller with no run -- a pipeline -- would otherwise write rows against + a foreign key pointing at nothing, and they insert without complaint + because the database has no enforcement to refuse them. That is the + failure mode a metrics unit already exists to avoid, so it is refused + here too rather than repeated. + """ + if not getattr(model, "run_id", None): + return + + for split_name, by_metric in accumulated.items(): + averaged = { + metric_name: float(np.mean(values)) + for metric_name, values in by_metric.items() + } + if averaged: + model._save_metrics( + results=averaged, + split=SplitEnum[split_name], + level=LevelEnum.TRIAL, + ) + + def _score_every_fold(self, model, x_folds, y_folds) -> dict: + """Fit and score each fold, writing its rows and keeping its numbers. + + Returns + ------- + dict + ``{split name: {metric name: [one score per fold]}}``, in fold + order, for whoever aggregates them. A split with no metrics + configured is absent rather than present and empty: the two are + different statements. + """ + scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) + fold_metrics: dict = {} + + for index, (x_fold, y_fold) in enumerate(self._folds(x_folds, y_folds)): + # The metric methods read the data off the instance to decide what + # they are scoring, so it follows the iteration. + model.x_data = x_fold + model.y_data = y_fold + self._fit_kept_model(model, x_fold, y_fold) + + for name in scored_splits: + split = SplitEnum[name] + # calculate_metrics hands back what it wrote, so the split is + # scored once rather than once for the row and once for the + # number. It writes nothing, and returns nothing, when there is + # no run to write against -- a pipeline -- and the numbers are + # wanted there too, hence the fallback. + results = model.calculate_metrics( + split=split, level=LevelEnum.FOLD, fold_index=index + ) + if results is None: + results = model.compute_metrics(split=split) + if not results: + continue + for metric_name, value in results.items(): + fold_metrics.setdefault(name, {}).setdefault( + metric_name, [] + ).append(value) + + return fold_metrics + + def _fit_kept_model(self, model, x, y) -> None: + """Fit on the training partition, never handing over the validation one. + + A fold is scored on the rows it held back from its own training, so + letting the fit watch them would measure it on data it was allowed to + see -- and nothing raises when that happens, the score simply comes out + better than the model deserves. Unlike the holdout sibling this is not + configurable, because there is no reading of a fold under which it + would be right. + """ + self._fit(model, x, y, with_validation=False) + + @staticmethod + def _folds(x_folds, y_folds): + """Pair up the folds, leaving out the entry that is not one. + + A fold splitter returns one entry per fold plus a trailing entry + holding every row the folds could use and the rows reserved from them. + That last one fits the model that gets kept; it is not a fold and is + never scored as one. + """ + return zip(x_folds[:-1], y_folds[:-1], strict=True) diff --git a/DashAI/back/units/fit_model_over_nested_folds_unit.py b/DashAI/back/units/fit_model_over_nested_folds_unit.py new file mode 100644 index 000000000..a97a99f83 --- /dev/null +++ b/DashAI/back/units/fit_model_over_nested_folds_unit.py @@ -0,0 +1,197 @@ +"""Unit that measures a fold run with a search of its own inside each fold.""" + +import logging + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit +from DashAI.back.units.fit_scope import ( + TRIAL_SPLITS, + goal_metric_field, + optimizer_field, + trial_splits_field, +) +from DashAI.back.units.splitter_scope import _splitter_field + +log = logging.getLogger(__name__) + + +def inner_splitter_field(): + """The splitter that carves an outer fold into folds of its own. + + Required rather than optional, which is what makes this a separate unit. + A component field wrapped in ``none_type`` is emitted as ``anyOf`` and the + front reads ``parent`` directly off the property, so an optional one leaves + the user with no selector at all -- the same wall that made the two + explainer units siblings. + """ + return _splitter_field( + parent="FoldSplitter", + placeholder={ + "component": "KFoldSplitter", + "params": {"n_splits": 3, "shuffle": True, "random_state": 42}, + }, + ) + + +class FitModelOverNestedFoldsSchema(BaseSchema): + optimizer: optimizer_field() # type: ignore + goal_metric: goal_metric_field() # type: ignore + scored_splits: trial_splits_field() # type: ignore + inner_splitter: inner_splitter_field() # type: ignore + + +class FitModelOverNestedFoldsUnit(FitModelOverFoldsUnit): + """Score every outer fold with a search that never saw it. + + ``FitModelOverFoldsUnit`` plus one measurement taken before it. The + relationship is inheritance rather than a shared mixin because that is what + it is: everything the sibling does still happens, and this adds a step in + front of it. + + **What the extra step is for.** In an ordinary cross-validated search, the + same folds choose the hyperparameters and report the score, so the reported + score is optimistic by however much the search managed to fit them. The + nested loop measures that honestly: for each outer fold, a search is run + from scratch on folds carved out of *that fold's training rows only*, and + the chosen model is then scored on the outer fold's validation rows, which + that search never saw. Those are the ``OUTER_FOLD`` rows. + + What it does **not** do is choose the hyperparameters. Each outer fold + picks its own, and they generally differ; there is no single model to keep + out of that loop. So the ordinary search still runs afterwards, over all + the folds, and it is what produces the model that gets saved. The nested + numbers are a statement about the procedure, not about the artifact. + + That also means the inner trials record nothing: their partitions belong to + one outer fold, and rows from them would sit alongside rows describing the + kept model as if they were comparable. + """ + + SCHEMA = FitModelOverNestedFoldsSchema + + PROVIDES = ("model", "plot_paths", "fold_metrics", "outer_fold_metrics") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._inner_splitter = None + + def _resolve_inner_splitter(self): + """Build the splitter that carves an outer fold, memoized on this unit.""" + if self._inner_splitter is not None: + return self._inner_splitter + + from kink import di + + chosen = self.config["inner_splitter"] + try: + splitter_class = di["component_registry"][chosen["component"]]["class"] + self._inner_splitter = splitter_class(splits_data=dict(chosen["params"])) + except Exception as e: + log.exception(e) + raise JobError( + f"Error configuring inner splitter for nested CV: {e}", + ) from e + return self._inner_splitter + + def execute(self, ctx: ExecutionContext) -> None: + model = ctx.require("model") + x_folds = ctx.require("x_folds") + y_folds = ctx.require("y_folds") + optimizable_parameters = ctx.require("optimizable_parameters") + + if self._will_search(optimizable_parameters): + # Resolved here and not at the top: a run with nothing to search + # never carves an outer fold, so a session left carrying a nested + # configuration it no longer uses must not fail on it. And resolved + # outside the wrapper below, so a splitter that cannot be built is + # reported as that rather than as a training failure. + inner_splitter = self._resolve_inner_splitter() + try: + outer_fold_metrics = self._measure_every_outer_fold( + model, + x_folds, + y_folds, + inner_splitter, + optimizable_parameters, + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Model training failed {e}", + ) from e + else: + # Nothing to search means nothing for the nested loop to measure: + # every outer fold would choose the same parameters, which is what + # the ordinary loop already reports. + outer_fold_metrics = {} + + ctx.put_ref("outer_fold_metrics", outer_fold_metrics) + + # And then the ordinary run, which is what produces the kept model. + super().execute(ctx) + + def _measure_every_outer_fold( + self, model, x_folds, y_folds, inner_splitter, optimizable_parameters + ) -> dict: + """Search inside each outer fold, then score that fold with what it chose. + + Returns + ------- + dict + ``{split name: {metric name: [one score per outer fold]}}``, for + whoever aggregates them -- the same shape and the same reason as + the ordinary fold scores. + """ + optimizer, goal_metric = self._resolve_search() + scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) + outer_fold_metrics: dict = {} + + for index, (x_outer, y_outer) in enumerate(self._folds(x_folds, y_folds)): + # Carved out of this fold's training rows alone. Its validation + # rows are what the search will be judged on, so nothing drawn from + # them may reach it. + inner_x, inner_y, _ = inner_splitter.split( + x_outer["train"], y_outer["train"] + ) + + optimizer.optimize( + model, + inner_x, + inner_y, + optimizable_parameters, + goal_metric, + self._score_one_inner_trial, + ) + outer_model = optimizer.get_model() + + outer_model.x_data = x_outer + outer_model.y_data = y_outer + self._fit_kept_model(outer_model, x_outer, y_outer) + + for name in scored_splits: + split = SplitEnum[name] + results = outer_model.calculate_metrics( + split=split, level=LevelEnum.OUTER_FOLD, fold_index=index + ) + if results is None: + results = outer_model.compute_metrics(split=split) + if not results: + continue + for metric_name, value in results.items(): + outer_fold_metrics.setdefault(name, {}).setdefault( + metric_name, [] + ).append(value) + + return outer_fold_metrics + + def _score_one_inner_trial(self, model, x_folds, y_folds, metric) -> float: + """The objective of an inner search: the fold loop, recording nothing. + + The partitions of an inner trial belong to one outer fold. Rows written + from them would sit in the same table as the rows describing the model + that gets kept, indistinguishable from them and far more numerous. + """ + return self._score_folds(model, x_folds, y_folds, metric, record=False) diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py new file mode 100644 index 000000000..3e3d69b10 --- /dev/null +++ b/DashAI/back/units/fit_model_unit.py @@ -0,0 +1,163 @@ +"""Unit that fits a model on one set of partitions, optionally searching.""" + +import logging + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.fit_scope import ( + TRIAL_SPLITS, + ModelFitScopeMixin, + goal_metric_field, + optimizer_field, + trial_splits_field, + validation_during_fit_field, +) + +log = logging.getLogger(__name__) + + +class FitModelSchema(BaseSchema): + optimizer: optimizer_field() # type: ignore + goal_metric: goal_metric_field() # type: ignore + trial_splits: trial_splits_field() # type: ignore + validation_during_fit: validation_during_fit_field() # type: ignore + + +class FitModelUnit(BaseUnit, ModelFitScopeMixin): + """Train a model, running a hyperparameter search when there is one to run. + + Hyperparameter optimization is a fitting strategy rather than a separate + step: it returns a fitted model, and the trial plots are a by-product only + that branch produces. Both paths therefore live in this unit. + + ``validate`` resolves the optimizer and the goal metric so an impossible + configuration is rejected before the job reports that training started. + + The sibling for cross-validation is ``FitModelOverFoldsUnit``: it takes a + list of partition sets rather than one, which is a different ``REQUIRES`` + and therefore a different unit. What surrounds the fit is shared between + them; only the objective the search measures, and what happens once the + search is over, is written here. + """ + + SCHEMA = FitModelSchema + + # run_id and artifact_prefix are configuration, not context: no unit + # publishes them, so nothing upstream could ever satisfy them as REQUIRES. + # See the artifact_prefix section of DAG_ENGINE.md. + REQUIRES = ( + "model", + "factory", + "optimizable_parameters", + "model_parameters", + "x", + "y", + ) + PROVIDES = ("model", "plot_paths") + RUNTIME_PARAMS = ("run_id", "artifact_prefix") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._optimizer = None + self._goal_metric = None + + def validate(self, ctx: ExecutionContext) -> None: + # ctx.require, not ctx.get: an absent key means BuildModelUnit has not + # run yet, which is a call-order mistake rather than a model with + # nothing to optimize. + self._validate_search(ctx.require("optimizable_parameters")) + + def execute(self, ctx: ExecutionContext) -> None: + model = ctx.require("model") + x = ctx.require("x") + y = ctx.require("y") + optimizable_parameters = ctx.require("optimizable_parameters") + + # The model is pointed at the data it is about to be fitted on, here + # rather than where it was built: the metric methods read these + # attributes off the instance to decide what they are scoring. + model.x_data = x + model.y_data = y + + plot_paths = [] + try: + if not self._will_search(optimizable_parameters): + self._fit_kept_model(model, x, y) + else: + # Every read of the context happens here rather than in the + # shared helper: the contract audit parses this file, so a + # require moved out of it makes a declared key look unread. + model, best_parameters, plot_paths = self._search( + model, + x, + y, + optimizable_parameters, + ctx.require("factory"), + # ctx.require already hands back an isolated copy of the + # stored parameter tree, so writing the best values into it + # cannot touch the Run row it came from. + ctx.require("model_parameters"), + self.config["run_id"], + self.config["artifact_prefix"], + self._score_one_trial, + ) + ctx.put_ref("best_parameters", best_parameters) + except Exception as e: + log.exception(e) + raise JobError( + f"Model training failed {e}", + ) from e + + ctx.put("model", model) + ctx.put_ref("plot_paths", plot_paths) + + def _fit_kept_model(self, model, x, y) -> None: + """Fit the model that gets kept, on the partitions it was given.""" + self._fit(model, x, y, self.config.get("validation_during_fit", True)) + + def _score_one_trial(self, model, x, y, metric) -> float: + """Fit the model once and score it, for one point of the search. + + This is what the optimizer measures. It used to be the optimizer's own + business: the objective fitted and scored inline, and the sixth + argument of ``optimize`` was the task. It is now a callable the caller + supplies, which is the change that lets the search be reused over + anything that can be fitted and scored -- a single split here, a whole + set of folds in the sibling -- without the optimizer knowing which. + + The trial metrics are written here rather than by the optimizer for the + same reason: what counts as a scored partition is a property of the + thing being fitted, not of the search. Which ones those are is + configured, because it is not always the same as which ones have + metrics: a forecaster has training metrics and still must not be judged + on the dates it was fitted on, since an in-sample fit statistic is not + comparable with a forecast. A partition left in the list but with no + metrics configured writes nothing anyway -- ``calculate_metrics`` finds + nothing to score and returns. + + Parameters + ---------- + model : BaseModel + The instance the optimizer has just set this trial's parameters on. + x, y : DatasetDict + The partitions this fit may use. + metric : BaseMetric + The metric class the search is optimizing, already unwrapped from + its registry entry by the optimizer. + + Returns + ------- + float + The score on the validation partition, which is the objective. + """ + self._fit_kept_model(model, x, y) + + for name in self.config.get("trial_splits", TRIAL_SPLITS): + model.calculate_metrics(split=SplitEnum[name], level=LevelEnum.TRIAL) + + predictions = model.predict(x["validation"]) + expected = model.prepare_output(y["validation"], is_fit=False) + return metric.score(expected, predictions) diff --git a/DashAI/back/units/fit_scope.py b/DashAI/back/units/fit_scope.py new file mode 100644 index 000000000..e484cc26c --- /dev/null +++ b/DashAI/back/units/fit_scope.py @@ -0,0 +1,368 @@ +"""Shared body of the units that fit a model, with or without a search. + +``FitModelUnit`` fits one set of partitions; ``FitModelOverFoldsUnit`` fits a +list of them. Everything around the fit is the same for both -- resolving the +optimizer and the goal metric, running the search, checking that the optimizer +gave back the model it was handed, recording the best parameters, and writing +the plots the search produces -- and it is exactly the sort of thing that ends +up written twice and then drifting. + +What differs is only the objective the search measures and what happens once it +is over, so those stay in the units. + +Named ``ModelFitScopeMixin`` rather than ``BaseSomething``: the registry derives +a component's type by walking its ``__mro__`` for a class whose name contains +"Base" and that declares ``TYPE``, and demands exactly one. A shared parent +called ``Base*`` would be a second candidate and would break the registration of +every unit that inherited it. +""" + +import logging +from typing import TYPE_CHECKING, List, Tuple + +from DashAI.back.core.schema_fields import ( + bool_field, + component_field, + enum_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError + +if TYPE_CHECKING: + from DashAI.back.optimizers.base_optimizer import BaseOptimizer + +log = logging.getLogger(__name__) + +#: A trial may score the partition it fitted on and the one it is measured +#: against, and nothing else. The test partition is deliberately absent: +#: scoring it once per trial would let the search see it, and a model chosen +#: with the test set in view has no honest score left to report. +TRIAL_SPLITS = ["TRAIN", "VALIDATION"] + + +def optimizer_field(): + return schema_field( + component_field(parent="BaseOptimizer"), + placeholder={"component": "OptunaOptimizer", "params": {}}, + description=MultilingualString( + en="Optimizer used to search for hyperparameters, along with its own " + "configuration. Only used when the model declares optimizable " + "parameters.", + es="Optimizador usado para buscar hiperparámetros, junto con su propia " + "configuración. Solo se usa cuando el modelo declara parámetros " + "optimizables.", + pt="Otimizador usado para procurar hiperparâmetros, junto com a sua " + "própria configuração. Só é usado quando o modelo declara parâmetros " + "otimizáveis.", + de="Optimierer für die Hyperparametersuche samt seiner eigenen " + "Konfiguration. Wird nur verwendet, wenn das Modell optimierbare " + "Parameter deklariert.", + zh="用于搜索超参数的优化器及其自身配置。仅当模型声明了可优化参数时使用。", + ), + alias=MultilingualString( + en="Optimizer", + es="Optimizador", + pt="Otimizador", + de="Optimierer", + zh="优化器", + ), + ) + + +def goal_metric_field(): + return schema_field( + string_field(), + placeholder="Accuracy", + description=MultilingualString( + en="Metric the hyperparameter search optimizes.", + es="Métrica que optimiza la búsqueda de hiperparámetros.", + pt="Métrica que a procura de hiperparâmetros otimiza.", + de="Metrik, die die Hyperparametersuche optimiert.", + zh="超参数搜索所优化的指标。", + ), + alias=MultilingualString( + en="Goal metric", + es="Métrica objetivo", + pt="Métrica objetivo", + de="Zielmetrik", + zh="目标指标", + ), + ) + + +def trial_splits_field(): + return schema_field( + list_field(enum_field(enum=TRIAL_SPLITS)), + placeholder=TRIAL_SPLITS, + description=MultilingualString( + en="Partitions each trial of the search records a score for. A " + "partition the model is not meant to be judged on belongs out of " + "this list even when metrics are configured for it.", + es="Particiones para las que cada intento de la búsqueda registra " + "un puntaje. Una partición sobre la que el modelo no debe juzgarse " + "no va en esta lista aunque tenga métricas configuradas.", + pt="Partições para as quais cada tentativa da procura regista uma " + "pontuação. Uma partição sobre a qual o modelo não deve ser julgado " + "fica fora desta lista mesmo que tenha métricas configuradas.", + de="Partitionen, für die jeder Versuch der Suche einen Wert " + "festhält. Eine Partition, nach der das Modell nicht beurteilt " + "werden soll, gehört nicht in diese Liste, auch wenn Metriken für " + "sie konfiguriert sind.", + zh="搜索的每次试验为其记录分数的分区。不应据以评判模型的分区不列入此处，" + "即使已为其配置了指标。", + ), + alias=MultilingualString( + en="Trial splits", + es="Particiones por intento", + pt="Partições por tentativa", + de="Versuchspartitionen", + zh="试验分区", + ), + ) + + +def validation_during_fit_field(): + return schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether the validation partition is handed to the model while " + "fitting. Models use it to watch training and stop early. Turn it " + "off when the same partition is what the fit will be scored on.", + es="Si la particion de validacion se entrega al modelo durante el " + "ajuste. Los modelos la usan para vigilar el entrenamiento y " + "detenerse antes. Desactivar cuando esa misma particion es sobre " + "la que se va a evaluar el ajuste.", + pt="Se a particao de validacao e entregue ao modelo durante o " + "ajuste. Os modelos usam-na para acompanhar o treino e parar mais " + "cedo. Desative quando essa mesma particao for aquela sobre a qual " + "o ajuste sera avaliado.", + de="Ob die Validierungspartition dem Modell beim Fitten uebergeben " + "wird. Modelle nutzen sie, um das Training zu beobachten und frueh " + "abzubrechen. Abschalten, wenn genau diese Partition den Fit " + "bewerten soll.", + zh="拟合时是否将验证分区交给模型。", + ), + alias=MultilingualString( + en="Validate while fitting", + es="Validar durante el ajuste", + pt="Validar durante o ajuste", + de="Beim Fitten validieren", + zh="拟合时验证", + ), + ) + + +class ModelFitScopeMixin: + """Everything around a fit that does not depend on how the data is shaped. + + Takes and returns plain values and never touches the execution context. + That is not tidiness: the contract audit parses each unit's own source, so + a ``ctx.require`` moved in here makes a declared key look unread and a + ``ctx.put`` makes a broken ``PROVIDES`` pass. Reading and publishing stay + in the units, and so does reading the runtime parameters, which the audit + checks are used by whoever declares them. + """ + + #: Declared here so a reader can see what state an instance carries; each + #: unit sets them in its own ``__init__``, because ``BaseUnit.__init__`` + #: comes first in the MRO and does not chain. + _optimizer = None + _goal_metric = None + + def _resolve_search(self): + """Resolve the optimizer and the goal metric, memoized on this unit. + + Kept on the instance rather than in the context on purpose. These are + this unit's own state, not something it hands to another unit: two + fitting units sharing a context -- a graph with two training nodes -- + would otherwise overwrite each other's optimizer, and the second one + would silently run the first one's. + """ + if self._optimizer is not None: + return self._optimizer, self._goal_metric + + from kink import di + + component_registry = di["component_registry"] + goal_metric_name: str = self.config["goal_metric"] + optimizer_name: str = self.config["optimizer"]["component"] + + try: + # The whole registry entry, not the class: the optimizer reads + # metadata["maximize"] from it to pick a direction. + goal_metric = component_registry[goal_metric_name] + except Exception as e: + log.exception(e) + raise JobError( + f"Metric is not compatible with the Task. {e}", + ) from e + + try: + optimizer_class = component_registry[optimizer_name]["class"] + optimizer: "BaseOptimizer" = optimizer_class( + **self.config["optimizer"]["params"] + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Error instantiating optimizer {optimizer_name}, {e}", + ) from e + + self._goal_metric = goal_metric + self._optimizer = optimizer + return optimizer, goal_metric + + def _will_search(self, optimizable_parameters) -> bool: + """Whether there is a search to run: something to tune, and a tuner. + + Both halves are needed. A run can name no optimizer at all -- the + column is a plain string and the wizard leaves it empty when the user + does not ask for a search -- while the model still declares a parameter + as optimizable, and that combination has always meant "fit it once with + the values given". Checking only the parameters turns it into a lookup + of the empty string in the registry, which fails with a message about + the metric being incompatible with the task. + """ + return bool(optimizable_parameters) and bool( + self.config["optimizer"]["component"] + ) + + def _validate_search(self, optimizable_parameters) -> None: + """Refuse an impossible search before anything observable happens. + + Handed the value rather than the context: the caller reads it with + ``ctx.require`` and not ``ctx.get``, because an absent key means the + model has not been built yet -- a call-order mistake, not "there is + nothing to optimize". A run with nothing to search skips the checks + below, so no registry lookup is needed either. + """ + if not self._will_search(optimizable_parameters): + return + + self._resolve_search() + + def _fit(self, model, x, y, with_validation: bool) -> None: + """Fit the model on the training partition of the data it was given. + + Whether the validation partition goes with it is a policy and not a + shape: a model uses it to watch the fit and stop early, which is what + an ordinary holdout run wants, and which is exactly wrong when that + same partition is what the fit will be scored on -- a fold is scored on + the rows it held back, so handing them over would be measuring the fit + on data it was allowed to watch. + + The key may also simply not be there. The trailing entry a fold + splitter produces holds the pooled rows and the reserved ones and has + no validation partition at all, so there is nothing to hand over even + where the policy would allow it. + """ + if with_validation and "validation" in x: + model.train(x["train"], y["train"], x["validation"], y["validation"]) + else: + model.train(x["train"], y["train"]) + + def _search( + self, + model, + x, + y, + optimizable_parameters, + factory, + old_parameters, + run_id, + artifact_prefix, + objective, + ) -> Tuple[object, dict, List[str]]: + """Run the hyperparameter search and report what it produced. + + Parameters + ---------- + model : BaseModel + The instance the search sets its parameters on. + x, y : object + Whatever the objective knows how to fit and score -- one set of + partitions, or a list of them. The optimizer passes it through + without looking at it, which is what lets one search serve both. + optimizable_parameters : list + The search space, as ModelFactory built it. + factory : ModelFactory + Used to write the values found back into the parameter tree. + old_parameters : dict + The tree they are written into. The caller is expected to hand over + a copy it owns; this mutates nothing else. + run_id, artifact_prefix : object + What the plots are named after. + objective : callable + ``(model, x, y, metric) -> float``: what the search measures. + + Returns + ------- + tuple + The fitted model, the parameter tree with the best values in it, + and the paths of the plots the search produced. + """ + import os + import pickle + + from kink import di + + from DashAI.back.core.artifacts import normalize_artifacts + + # Memoized: validate() resolved these already, and resolving again + # here would be the same lookup. + optimizer, goal_metric = self._resolve_search() + + optimizer.optimize(model, x, y, optimizable_parameters, goal_metric, objective) + model = optimizer.get_model() + best_params = optimizer.get_best_params() + + self._assert_model_keeps_its_runtime_state(model) + + best_parameters = factory.update_parameters(old_parameters, best_params) + + config = di["config"] + plot_paths: List[str] = [] + trials = optimizer.get_trials_values() + plot_filenames, plots = optimizer.create_plots( + trials, + run_id, + n_params=len(optimizable_parameters), + goal_metric=goal_metric, + artifact_prefix=artifact_prefix, + ) + normalized_plots = normalize_artifacts(plots) + for filename, plot in zip(plot_filenames, normalized_plots, strict=False): + plot_path = os.path.join(config["RUNS_PATH"], filename) + with open(plot_path, "wb") as file: + pickle.dump(plot, file) + plot_paths.append(plot_path) + + return model, best_parameters, plot_paths + + @staticmethod + def _assert_model_keeps_its_runtime_state(model) -> None: + """Fail loudly if the optimizer returned a model that cannot be scored. + + ``ModelFactory`` attaches the metric classes to the model instance and + whoever fits it points it at its data, and optimizers are expected to + return that same instance. If one ever returns a fresh object instead, + scoring it finds nothing to score and the caller ends up with no + metrics rather than an error. + + The check is on the data, not on the run id. Keying it to ``run_id`` + made it a no-op for every caller that has no run -- a pipeline, where + ``run_id`` is always None -- which is exactly the caller with no other + signal that anything went wrong: it would finish with an empty metrics + artifact. What both callers need is the same, so this asks for that + instead. + """ + if getattr(model, "x_data", None) is None: + raise JobError( + "The optimizer returned a model detached from its data: metrics " + "could not be computed for it. Optimizers must return the same " + "model instance they received." + ) diff --git a/DashAI/back/units/generate_global_explanation_unit.py b/DashAI/back/units/generate_global_explanation_unit.py new file mode 100644 index 000000000..677929b19 --- /dev/null +++ b/DashAI/back/units/generate_global_explanation_unit.py @@ -0,0 +1,88 @@ +"""Unit that explains a model as a whole and stores the result.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import dump_explanation + +log = logging.getLogger(__name__) + + +class GenerateGlobalExplanationSchema(BaseSchema): + explainer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the global explanation being produced. It names " + "the files on disk, so a re-run overwrites its own artifacts and " + "never another explanation's.", + es="Identificador de la explicación global que se produce. Da " + "nombre a los archivos en disco, de modo que volver a ejecutarla " + "sobrescribe sus propios artefactos y nunca los de otra.", + pt="Identificador da explicação global a ser produzida. Dá nome aos " + "ficheiros em disco, pelo que uma nova execução substitui os seus " + "próprios artefactos e nunca os de outra.", + de="Kennung der erzeugten globalen Erklärung. Sie benennt die " + "Dateien auf der Festplatte, sodass ein erneuter Lauf nur die " + "eigenen Artefakte überschreibt.", + zh="所生成全局解释的标识符。它命名磁盘上的文件，因此重新运行只会覆盖自身产物。", + ), + alias=MultilingualString( + en="Explanation", + es="Explicación", + pt="Explicação", + de="Erklärung", + zh="解释", + ), + ) # type: ignore + + +class GenerateGlobalExplanationUnit(BaseUnit): + """Explain the model over the whole dataset and pickle the result. + + Sibling of ``GenerateLocalExplanationUnit`` rather than one unit with a + branch, for three reasons that all point the same way: the two produce + different outputs (a single plot here, a set of plots plus the explained + instances there), which a single ``PROVIDES`` could not describe since it + is checked unconditionally; the local path has steps this one does not + (fitting, selecting instances); and their configurations point at two + different component registries. + + The unit never touches the explanation row: it publishes where it wrote, + and the job owns the columns. + """ + + SCHEMA = GenerateGlobalExplanationSchema + + REQUIRES = ("explainer", "data_x", "data_y") + PROVIDES = ("explanation_path", "plot_path") + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.core.artifacts import normalize_artifacts + + explainer = ctx.require("explainer") + dataset = (ctx.require("data_x"), ctx.require("data_y")) + + try: + explanation = explainer.explain(dataset) + plot = normalize_artifacts(explainer.plot(explanation)) + except Exception as e: + log.exception(e) + raise JobError( + "Failed to generate the explanation", + ) from e + + explanation_path, plot_path = dump_explanation( + explanation, plot, "global", self.config["explainer_id"] + ) + + ctx.put_ref("explanation_path", explanation_path) + ctx.put_ref("plot_path", plot_path) diff --git a/DashAI/back/units/generate_local_explanation_unit.py b/DashAI/back/units/generate_local_explanation_unit.py new file mode 100644 index 000000000..32cca16fd --- /dev/null +++ b/DashAI/back/units/generate_local_explanation_unit.py @@ -0,0 +1,420 @@ +"""Unit that explains individual instances and stores the result.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + bool_field, + int_field, + list_field, + none_type, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import dump_explanation + +log = logging.getLogger(__name__) + + +def _columns_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=description, + alias=alias, + ) + + +class GenerateLocalExplanationSchema(BaseSchema): + explainer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the local explanation being produced. It names " + "the files on disk, so a re-run overwrites its own artifacts.", + es="Identificador de la explicación local que se produce. Da nombre " + "a los archivos en disco, de modo que volver a ejecutarla " + "sobrescribe sus propios artefactos.", + pt="Identificador da explicação local a ser produzida. Dá nome aos " + "ficheiros em disco, pelo que uma nova execução substitui os seus " + "próprios artefactos.", + de="Kennung der erzeugten lokalen Erklärung. Sie benennt die " + "Dateien auf der Festplatte, sodass ein erneuter Lauf nur die " + "eigenen Artefakte überschreibt.", + zh="所生成局部解释的标识符。它命名磁盘上的文件，因此重新运行只会覆盖自身产物。", + ), + alias=MultilingualString( + en="Explanation", + es="Explicación", + pt="Explicação", + de="Erklärung", + zh="解释", + ), + ) # type: ignore + instance_dataset_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the dataset the explained instances come from. " + "It may differ from the one the model was trained on.", + es="Identificador del conjunto de datos del que provienen las " + "instancias explicadas. Puede diferir de aquel con el que se " + "entrenó el modelo.", + pt="Identificador do conjunto de dados de onde vêm as instâncias " + "explicadas. Pode diferir daquele com que o modelo foi treinado.", + de="Kennung des Datensatzes, aus dem die erklärten Instanzen " + "stammen. Er kann sich von dem des Trainings unterscheiden.", + zh="被解释实例所属数据集的标识符。它可能与模型训练所用的数据集不同。", + ), + alias=MultilingualString( + en="Instance dataset", + es="Conjunto de instancias", + pt="Conjunto de instâncias", + de="Instanzdatensatz", + zh="实例数据集", + ), + ) # type: ignore + scope: schema_field( + dict, + placeholder={"mode": "split", "split": "test", "percentage": 20}, + description=MultilingualString( + en="Which instances to explain. A 'mode' of 'split' takes a share " + "of one data split, 'rows' takes the row indexes the user marked, " + "and 'manual' takes hand-typed values. Defaults to 'split'.", + es="Qué instancias explicar. Un 'mode' de 'split' toma una " + "proporción de una partición, 'rows' toma los índices de fila que " + "marcó el usuario, y 'manual' toma valores ingresados a mano. Por " + "defecto es 'split'.", + pt="Que instâncias explicar. Um 'mode' de 'split' toma uma parte de " + "uma partição, 'rows' toma os índices de linha que o utilizador " + "marcou, e 'manual' toma valores introduzidos à mão. Por omissão é " + "'split'.", + de="Welche Instanzen erklärt werden. 'mode' 'split' nimmt einen " + "Anteil einer Teilmenge, 'rows' die vom Benutzer markierten " + "Zeilenindizes und 'manual' manuell eingegebene Werte. Standard " + "ist 'split'.", + zh="要解释哪些实例。'mode' 为 'split' 时取某个划分的一部分，" + "'rows' 取用户标记的行索引，'manual' 取手工输入的值。默认为 'split'。", + ), + alias=MultilingualString( + en="Scope", es="Alcance", pt="Âmbito", de="Umfang", zh="范围" + ), + ) # type: ignore + fit_parameters: schema_field( + dict, + placeholder={}, + description=MultilingualString( + en="Extra arguments handed to the explainer's fit step.", + es="Argumentos adicionales entregados al paso de ajuste del explicador.", + pt="Argumentos adicionais entregues ao passo de ajuste do explicador.", + de="Zusätzliche Argumente für den Fit-Schritt des Erklärers.", + zh="传递给解释器拟合步骤的额外参数。", + ), + alias=MultilingualString( + en="Fit parameters", + es="Parámetros de ajuste", + pt="Parâmetros de ajuste", + de="Fit-Parameter", + zh="拟合参数", + ), + ) # type: ignore + input_columns: _columns_field( + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + description=MultilingualString( + en="Names of the columns used as model input.", + es="Nombres de las columnas usadas como entrada del modelo.", + pt="Nomes das colunas usadas como entrada do modelo.", + de="Namen der als Modelleingabe verwendeten Spalten.", + zh="用作模型输入的列名。", + ), + ) # type: ignore + output_columns: _columns_field( + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + description=MultilingualString( + en="Names of the columns the model predicts.", + es="Nombres de las columnas que el modelo predice.", + pt="Nomes das colunas que o modelo prevê.", + de="Namen der Spalten, die das Modell vorhersagt.", + zh="模型需要预测的列名。", + ), + ) # type: ignore + manual_input_data: schema_field( + none_type(list), + placeholder=None, + description=MultilingualString( + en="Rows to explain when the scope mode is 'manual'. Ignored otherwise.", + es="Filas a explicar cuando el modo del alcance es 'manual'. Se " + "ignora en otro caso.", + pt="Linhas a explicar quando o modo do âmbito é 'manual'. Ignorado " + "caso contrário.", + de="Zu erklärende Zeilen, wenn der Modus 'manual' ist. Sonst ignoriert.", + zh="范围模式为 'manual' 时要解释的行。其他情况下忽略。", + ), + alias=MultilingualString( + en="Manual input", + es="Entrada manual", + pt="Entrada manual", + de="Manuelle Eingabe", + zh="手动输入", + ), + ) # type: ignore + same_dataset: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether the instances come from the very dataset the model was " + "trained on. When they do not, the run's row indexes mean nothing " + "here and the split has to be recomputed.", + es="Si las instancias provienen del mismo conjunto de datos con el " + "que se entrenó el modelo. Si no, los índices de fila de la " + "ejecución no significan nada acá y la partición se recalcula.", + pt="Se as instâncias vêm do mesmo conjunto de dados com que o " + "modelo foi treinado. Se não, os índices de linha da execução não " + "significam nada aqui e a divisão tem de ser recalculada.", + de="Ob die Instanzen aus genau dem Datensatz stammen, mit dem das " + "Modell trainiert wurde. Andernfalls sind die Zeilenindizes des " + "Laufs hier bedeutungslos und der Split wird neu berechnet.", + zh="实例是否来自模型训练所用的同一数据集。若不是，则运行记录的行索引在此无意义，" + "需要重新计算划分。", + ), + alias=MultilingualString( + en="Same dataset", + es="Mismo conjunto", + pt="Mesmo conjunto", + de="Gleicher Datensatz", + zh="同一数据集", + ), + ) # type: ignore + + +class GenerateLocalExplanationUnit(BaseUnit): + """Explain a selection of instances and pickle the result. + + Sibling of ``GenerateGlobalExplanationUnit``; see that class for why the + two scopes are two units and not one with a branch. + + The three ways of choosing instances — a share of a split, the rows the + user marked, or hand-typed values — stay one unit on purpose. They are + three branches of a single decision with a single output; as separate + nodes exactly one could ever run, which the contract cannot express. + + Row indexes are never taken on trust across datasets: when the instances + come from a dataset other than the one the run was trained on, the run's + indexes address rows that do not correspond, so the split is recomputed + from the session's ratios instead. That derived state is resolved here, + inside ``execute``, and never published. + """ + + SCHEMA = GenerateLocalExplanationSchema + + REQUIRES = ("explainer", "data_x", "data_y", "task", "split_indexes") + PROVIDES = ("explanation_path", "plots_path", "input_dataset_path") + RUNTIME_PARAMS = ("session_splits",) + + def _select_instances(self, prepared_instance, splits, instance, task): + """Narrow the loaded dataset down to the instances to explain.""" + import json + + from datasets import DatasetDict + + from DashAI.back.dataloaders.classes.dashai_dataset import ( + prepare_for_model_session, + select_columns, + split_dataset, + ) + + scope = self.config["scope"] or {} + input_columns = self.config["input_columns"] + output_columns = self.config["output_columns"] + + # The data source is selected via scope["mode"]. It defaults to + # "split" so explainers created before this field existed keep + # their original split + percentage behavior. + mode = scope.get("mode", "split") + + if mode == "manual": + # Build the instances from values the user typed in by hand, + # reusing the same conversion the manual prediction flow uses. + # The rows (and any image files rewritten by the job endpoint) + # travel in the job kwargs, not in scope. + manual_input_data = self.config.get("manual_input_data") or [] + if not manual_input_data: + raise JobError("No manual input data provided for the explanation") + prepared_instance = task.process_manual_input( + manual_input_data, + f"{instance.file_path}/dataset", + ) + # Manual input carries only the input columns (no target), so + # keep just those instead of the standard input/output split. + # select_columns returns a DashAIDataset (same shape the + # split path produces), which is what the explainers expect. + return prepared_instance.select_columns(input_columns) + + prepared_instance = task.prepare_for_task( + prepared_instance, + input_columns=input_columns, + output_columns=output_columns, + ) + + if mode == "rows": + # Explain a set of rows the user marked in the table. + # Indexes are over the whole dataset (the split does not + # apply in this mode). + row_indexes = scope.get("row_indexes") or [] + valid_indexes = [ + i + for i in row_indexes + if isinstance(i, int) and 0 <= i < prepared_instance.num_rows + ] + if row_indexes and not valid_indexes: + raise JobError("No valid row indexes provided for the explanation") + if valid_indexes: + prepared_instance = prepared_instance.select(valid_indexes) + else: + split = scope.get("split") + if split not in ["train", "test", "val", "all"]: + raise JobError(f"{split} is not a valid split") + + if split != "all": + if not self.config["same_dataset"]: + if isinstance(splits, str): + splits = json.loads(splits) + ( + prepared_dataset_dict, + splits, + ) = prepare_for_model_session( + dataset=prepared_instance, + splits=splits, + output_columns=output_columns, + ) + split_key = "validation" if split == "val" else split + prepared_instance = prepared_dataset_dict[split_key] + else: + prepared_instance = split_dataset( + prepared_instance, + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + split_key = "validation" if split == "val" else split + prepared_instance = prepared_instance[split_key] + + n_rows = max( + 1, + int(prepared_instance.num_rows * scope.get("percentage") / 100), + ) + # When "shuffle" is set the percentage is taken as a random + # sample of the split; otherwise it is the leading rows. + if scope.get("shuffle"): + prepared_instance = prepared_instance.shuffle(seed=42) + prepared_instance = prepared_instance.select(range(n_rows)) + + prepared_instance = DatasetDict({"train": prepared_instance}) + x, _ = select_columns(prepared_instance, input_columns, output_columns) + return x + + def execute(self, ctx: ExecutionContext) -> None: + import os + + from datasets import DatasetDict + from kink import di + + from DashAI.back.core.artifacts import normalize_artifacts + from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + ) + + config = di["config"] + session_factory = di["session_factory"] + + explainer = ctx.require("explainer") + task = ctx.require("task") + dataset = (ctx.require("data_x"), ctx.require("data_y")) + splits = ctx.require("split_indexes") + + explainer_id = self.config["explainer_id"] + instance_id = self.config["instance_dataset_id"] + + # Fitting happens before the instances are even looked up, and is left + # unwrapped on purpose: the explainer's own error is what the user gets. + explainer.fit(dataset, **(self.config["fit_parameters"] or {})) + + if not self.config["same_dataset"]: + splits = self.config["session_splits"] + + with session_factory() as db: + instance: Dataset = db.get(Dataset, instance_id) + if not instance: + raise JobError( + f"Dataset {instance_id} to be explained does not exist in DB." + ) + + try: + loaded_instance = load_dataset(f"{instance.file_path}/dataset") + except Exception as e: + log.exception(e) + raise JobError( + f"Can not load instance from path {instance.file_path}", + ) from e + + try: + x = self._select_instances(loaded_instance, splits, instance, task) + + # Persist the original selected rows (the model input for each + # explained instance) as a DashAIDataset before the model's own + # preprocessing runs, so the frontend can read them back with + # the existing dataset endpoints. + input_source = x["train"] if isinstance(x, DatasetDict) else x + input_dataset_path = os.path.join( + config["EXPLANATIONS_PATH"], + f"local_explanation_input_{explainer_id}", + ) + save_dataset(input_source, os.path.join(input_dataset_path, "dataset")) + # The instances are handed over unprepared, the same way the + # prediction job calls model.predict: the model applies its own + # preprocessing. Explainers that need the model feature space + # ask for it with prepare_model_input. + + except Exception as e: + log.exception(e) + raise JobError( + f"""Can not prepare Dataset with {instance_id} + to generate the local explanation.""", + ) from e + + try: + explanation = explainer.explain_instance(x) + plots = normalize_artifacts( + explainer.plot(explanation), create_grouped=True + ) + except Exception as e: + log.exception(e) + raise JobError( + "Failed to generate the explanation", + ) from e + + explanation_path, plots_path = dump_explanation( + explanation, plots, "local", explainer_id + ) + + ctx.put_ref("explanation_path", explanation_path) + ctx.put_ref("plots_path", plots_path) + ctx.put_ref("input_dataset_path", input_dataset_path) diff --git a/DashAI/back/units/infer_dataset_types_unit.py b/DashAI/back/units/infer_dataset_types_unit.py new file mode 100644 index 000000000..58f3785c3 --- /dev/null +++ b/DashAI/back/units/infer_dataset_types_unit.py @@ -0,0 +1,76 @@ +"""Unit that works out what type each column of a dataset holds.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema, schema_field, string_field +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class InferDatasetTypesSchema(BaseSchema): + method: schema_field( + string_field(), + placeholder="DashAIPtype", + description=MultilingualString( + en="Inference method used when the dataset does not already carry " + "types of its own.", + es="Método de inferencia que se usa cuando el conjunto de datos no " + "trae tipos propios.", + pt="Método de inferência usado quando o conjunto de dados não traz " + "tipos próprios.", + de="Inferenzmethode, die verwendet wird, wenn der Datensatz keine " + "eigenen Typen mitbringt.", + zh="当数据集本身未携带类型时使用的推断方法。", + ), + alias=MultilingualString( + en="Inference method", + es="Método de inferencia", + pt="Método de inferência", + de="Inferenzmethode", + zh="推断方法", + ), + ) # type: ignore + + +class InferDatasetTypesUnit(BaseUnit): + """Publish a type declaration for every column of the dataset. + + Prefers the types the dataset already carries: a dataloader that can read + them from the source (a Parquet schema, a typed database column) knows better + than any inference over the values, so re-inferring would throw that away. + Only when there are none does it fall back to inferring from the data. + + Publishes the declaration as a plain reference rather than applying it, so + the decision and the transformation stay separable — the same declaration can + be shown to a user for review before anything is cast. + + Note the reference names columns, so it goes stale the moment something + renames or drops one. It is meant to be consumed by the next step, not + carried across a chain of transformations. + """ + + SCHEMA = InferDatasetTypesSchema + + REQUIRES = ("dataset",) + PROVIDES = ("inferred_types",) + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.types.inf.type_inference import infer_types + + dataset = ctx.require("dataset") + + if dataset.types: + schema = { + column: dashai_type.to_string() + for column, dashai_type in dataset.types.items() + } + else: + schema = infer_types( + dataset.to_pandas(), + method=self.config.get("method", "DashAIPtype"), + ) + + ctx.put_ref("inferred_types", schema) diff --git a/DashAI/back/units/load_datafile_dataset_unit.py b/DashAI/back/units/load_datafile_dataset_unit.py new file mode 100644 index 000000000..b73631a9a --- /dev/null +++ b/DashAI/back/units/load_datafile_dataset_unit.py @@ -0,0 +1,182 @@ +"""Unit that reads a dataset out of an already downloaded datafile.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + int_field, + none_type, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class LoadDatafileDatasetSchema(BaseSchema): + dataloader: schema_field( + component_field(parent="BaseDataLoader"), + placeholder={"component": "CSVDataLoader", "params": {"separator": ","}}, + description=MultilingualString( + en="Reader used to parse the downloaded file, together with its own " + "configuration.", + es="Lector usado para interpretar el archivo descargado, junto con " + "su propia configuración.", + pt="Leitor usado para interpretar o ficheiro descarregado, junto com " + "a sua própria configuração.", + de="Leser zum Parsen der heruntergeladenen Datei samt eigener " + "Konfiguration.", + zh="用于解析已下载文件的读取器及其自身配置。", + ), + alias=MultilingualString( + en="Data loader", + es="Cargador de datos", + pt="Carregador de dados", + de="Datenlader", + zh="数据加载器", + ), + ) # type: ignore + datafile_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the completed download to read from. It has to have " + "finished successfully; a download still running has no files yet.", + es="Identificador de la descarga completada desde la que leer. Tiene " + "que haber terminado con éxito; una descarga en curso todavía no " + "tiene archivos.", + pt="Identificador da descarga concluída de onde ler. Tem de ter " + "terminado com sucesso; uma descarga em curso ainda não tem " + "ficheiros.", + de="Kennung des abgeschlossenen Downloads, aus dem gelesen wird. Er " + "muss erfolgreich beendet sein; ein laufender Download hat noch " + "keine Dateien.", + zh="要读取的已完成下载的标识符。必须已成功完成；仍在进行的下载尚无文件。", + ), + alias=MultilingualString( + en="Downloaded file", + es="Archivo descargado", + pt="Ficheiro descarregado", + de="Heruntergeladene Datei", + zh="已下载文件", + ), + ) # type: ignore + selected_file: schema_field( + none_type(string_field()), + placeholder=None, + description=MultilingualString( + en="Which file inside the download to read, relative to its root. " + "Leave empty to take the first one, ignoring hidden files.", + es="Qué archivo dentro de la descarga leer, relativo a su raíz. " + "Dejar vacío para tomar el primero, ignorando archivos ocultos.", + pt="Qual ficheiro dentro da descarga ler, relativo à sua raiz. " + "Deixar vazio para tomar o primeiro, ignorando ficheiros ocultos.", + de="Welche Datei innerhalb des Downloads gelesen wird, relativ zu " + "dessen Wurzel. Leer lassen, um die erste zu nehmen; versteckte " + "Dateien werden übersprungen.", + zh="读取下载内容中的哪个文件（相对于其根目录）。留空则取第一个，忽略" + "隐藏文件。", + ), + alias=MultilingualString( + en="File", + es="Archivo", + pt="Ficheiro", + de="Datei", + zh="文件", + ), + ) # type: ignore + + +class LoadDatafileDatasetUnit(BaseUnit): + """Parse a dataset out of a download that already completed. + + Separate from ``LoadUploadedDatasetUnit`` even though both end in the same + reader call, because finding *what* to read is the whole job here: the + download is a directory tree recorded in the database, not a file the caller + hands over. The unit re-reads that row in its own read-only session, the same + way ``LoadDatasetUnit`` does, and never writes to it. + + Publishes only ``dataset``, for the same reason as its sibling: an imported + file is not a stored dataset yet, so there is no id to correlate and no path + to save back to. + """ + + SCHEMA = LoadDatafileDatasetSchema + + PROVIDES = ("dataset",) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.core.enums.status import DatafileStatus + from DashAI.back.dependencies.database.models import Datafile + + component_registry = di["component_registry"] + session_factory = di["session_factory"] + + datafile_id = self.config["datafile_id"] + selected_file = self.config.get("selected_file") + dataloader_config = self.config["dataloader"] + + with session_factory() as db: + datafile = db.get(Datafile, datafile_id) + # A missing row and an unfinished download are the same problem from + # here: there are no files to read either way. + if datafile is None or datafile.status != DatafileStatus.READY: + raise JobError(f"Datafile {datafile_id} is not ready.") + work_dir = datafile.local_path + + source = _resolve_source_file(work_dir, selected_file) + + # Looked up among the readers specifically, not with the registry's + # global ``registry[name]``: that one walks every component type, so a + # metric or a model whose name happened to be passed here would be + # instantiated as if it could read a file instead of being rejected. + dataloader_name = dataloader_config["component"] + readers = component_registry.registry.get("DataLoader", {}) + if dataloader_name not in readers: + raise JobError(f"DataLoader '{dataloader_name}' not found in registry.") + dataloader = readers[dataloader_name]["class"]() + + log.debug("Loading hub dataset from %s using %s", source, dataloader_name) + ctx.put( + "dataset", + dataloader.load_data( + filepath_or_buffer=source, + temp_path=work_dir, + params=dataloader_config.get("params") or {}, + n_sample=None, + ), + ) + + +def _resolve_source_file(work_dir: str, selected_file) -> str: + """Pick the file to read inside a completed download. + + A plain helper: takes and returns values, never touches the context. + + With no explicit choice it walks the tree and takes the first file in sorted + order, skipping anything under a dot-prefixed path component — download tools + leave metadata directories (``.cache``, ``.git``) behind that sort before the + real data and would otherwise win. + """ + from pathlib import Path + + if selected_file: + return str(Path(work_dir) / selected_file) + + base = Path(work_dir) + files = sorted( + str(path) + for path in base.rglob("*") + if path.is_file() + and not any(part.startswith(".") for part in path.relative_to(base).parts) + ) + if not files: + raise JobError("Hub download directory is empty.") + return files[0] diff --git a/DashAI/back/units/load_dataset_unit.py b/DashAI/back/units/load_dataset_unit.py new file mode 100644 index 000000000..cac25e15c --- /dev/null +++ b/DashAI/back/units/load_dataset_unit.py @@ -0,0 +1,138 @@ +"""Unit that loads a stored dataset into memory.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + none_type, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Dataset, Notebook +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class LoadDatasetSchema(BaseSchema): + dataset_id: schema_field( + none_type(int_field(gt=0)), + placeholder=None, + description=MultilingualString( + en="Identifier of the stored dataset to load. Mutually exclusive " + "with the notebook identifier.", + es="Identificador del conjunto de datos almacenado a cargar. " + "Excluyente con el identificador del cuaderno.", + pt="Identificador do conjunto de dados armazenado a carregar. " + "Mutuamente exclusivo com o identificador do caderno.", + de="Kennung des zu ladenden gespeicherten Datensatzes. Schließt " + "die Notebook-Kennung aus.", + zh="要加载的已存储数据集的标识符。与笔记本标识符互斥。", + ), + alias=MultilingualString( + en="Dataset", + es="Conjunto de datos", + pt="Conjunto de dados", + de="Datensatz", + zh="数据集", + ), + ) # type: ignore + notebook_id: schema_field( + none_type(int_field(gt=0)), + placeholder=None, + description=MultilingualString( + en="Identifier of the notebook whose working copy of the dataset " + "should be loaded. Mutually exclusive with the dataset identifier.", + es="Identificador del cuaderno cuya copia de trabajo del conjunto " + "de datos se debe cargar. Excluyente con el identificador del " + "conjunto de datos.", + pt="Identificador do caderno cuja cópia de trabalho do conjunto de " + "dados deve ser carregada. Mutuamente exclusivo com o " + "identificador do conjunto de dados.", + de="Kennung des Notebooks, dessen Arbeitskopie des Datensatzes " + "geladen werden soll. Schließt die Datensatz-Kennung aus.", + zh="要加载其数据集工作副本的笔记本标识符。与数据集标识符互斥。", + ), + alias=MultilingualString( + en="Notebook", + es="Cuaderno", + pt="Caderno", + de="Notebook", + zh="笔记本", + ), + ) # type: ignore + + +class LoadDatasetUnit(BaseUnit): + """Load a dataset from disk into the execution context. + + Takes exactly one of two starting points and materialises the dataset each + one points at, so downstream units receive a dataset instead of an + identifier: + + * ``dataset_id``: the stored dataset itself, read from ``Dataset.file_path``. + * ``notebook_id``: the notebook's own working copy, read from + ``Notebook.file_path``. A notebook holds a private copy precisely so + converters can rewrite it without touching the source dataset. + + Either way the unit publishes ``dataset_path``, which is where a later unit + has to write the dataset back for the change to be visible. + """ + + SCHEMA = LoadDatasetSchema + + PROVIDES = ("dataset", "dataset_id", "dataset_path") + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + session_factory = di["session_factory"] + + dataset_id = self.config.get("dataset_id") + notebook_id = self.config.get("notebook_id") + + if (dataset_id is None) == (notebook_id is None): + raise JobError( + "LoadDatasetUnit needs exactly one of dataset_id or notebook_id." + ) + + with session_factory() as db: + if dataset_id is not None: + dataset: Dataset = db.get(Dataset, dataset_id) + if not dataset: + raise JobError(f"Dataset {dataset_id} does not exist in DB.") + file_path = dataset.file_path + else: + notebook: Notebook = db.get(Notebook, notebook_id) + if not notebook: + raise JobError(f"Notebook {notebook_id} does not exist in DB.") + file_path = notebook.file_path + # The notebook's copy still belongs to a source dataset, and + # downstream error messages identify the work by that id. + dataset_id = notebook.dataset_id + + dataset_path = f"{file_path}/dataset" + + try: + loaded_dataset: "DashAIDataset" = load_dataset(dataset_path) + except Exception as e: + log.exception(e) + raise JobError( + f"Can not load dataset from path {file_path}", + ) from e + + if not loaded_dataset: + raise JobError(f"Dataset with path {dataset_path} not found") + + ctx.put("dataset", loaded_dataset) + ctx.put_ref("dataset_id", dataset_id) + ctx.put_ref("dataset_path", dataset_path) diff --git a/DashAI/back/units/load_run_model_unit.py b/DashAI/back/units/load_run_model_unit.py new file mode 100644 index 000000000..3ed3970d7 --- /dev/null +++ b/DashAI/back/units/load_run_model_unit.py @@ -0,0 +1,124 @@ +"""Unit that restores a run's model the way the explanation flow expects.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Run +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class LoadRunModelSchema(BaseSchema): + run_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the run whose trained model is restored. The " + "model component, its parameters and the artifact path all come " + "from that row.", + es="Identificador de la ejecución cuyo modelo entrenado se " + "restaura. El componente del modelo, sus parámetros y la ruta del " + "artefacto salen de esa fila.", + pt="Identificador da execução cujo modelo treinado é restaurado. O " + "componente do modelo, os seus parâmetros e o caminho do artefacto " + "vêm todos dessa linha.", + de="Kennung des Laufs, dessen trainiertes Modell wiederhergestellt " + "wird. Modellkomponente, Parameter und Artefaktpfad stammen alle " + "aus dieser Zeile.", + zh="要恢复其已训练模型的运行标识符。模型组件、参数和产物路径都来自该行。", + ), + alias=MultilingualString( + en="Run", es="Ejecución", pt="Execução", de="Lauf", zh="运行" + ), + ) # type: ignore + + +class LoadRunModelUnit(BaseUnit): + """Restore a run's trained model for an explanation. + + Deliberately **not** ``LoadTrainedModelUnit``, and the difference is not + cosmetic to preserve even though it is very likely accidental: + + * this unit builds an instance with ``model_class(**run.parameters)`` and + only then calls ``load`` on it, the way the explanation flow always has; + * ``LoadTrainedModelUnit`` calls ``load`` straight on the class. + + Every concrete model in this codebase declares ``load`` as a + ``staticmethod`` or a ``classmethod`` that rebuilds the object from the + file, so the instance built here is thrown away and the extra step changes + nothing — and no explainer reads anything that ``__init__`` sets: the only + model attributes any of them touch (``one_hot_encoder``, + ``categorical_columns``, ``label_encoder``) are set during training and + restored from the artifact. + + The two units are kept apart because merging them would have to unify their + error messages, which are user-visible and differ word for word. Do not + collapse them into one with a flag without deciding that first. + """ + + SCHEMA = LoadRunModelSchema + + PROVIDES = ("model",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._model_class = None + + def _resolve_model_class(self, model_name: str) -> type: + """Resolve the model class from the registry, memoized on this unit.""" + if self._model_class is not None: + return self._model_class + + from kink import di + + component_registry = di["component_registry"] + + try: + model_class = component_registry[model_name]["class"] + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find Model with name {model_name} in registry.", + ) from e + + self._model_class = model_class + return model_class + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + session_factory = di["session_factory"] + + run_id = self.config["run_id"] + + with session_factory() as db: + run: Run = db.get(Run, run_id) + if not run: + raise JobError(f"Run {run_id} does not exist in DB.") + model_name = run.model_name + run_path = run.run_path + parameters = dict(run.parameters or {}) + + run_model_class = self._resolve_model_class(model_name) + + try: + model = run_model_class(**parameters) + except Exception as e: + log.exception(e) + raise JobError("Unable to instantiate model") from e + + try: + trained_model = model.load(run_path) + except Exception as e: + log.exception(e) + raise JobError(f"Can not load model from path {run_path}") from e + + ctx.put("model", trained_model) diff --git a/DashAI/back/units/load_trained_model_unit.py b/DashAI/back/units/load_trained_model_unit.py new file mode 100644 index 000000000..c8e194028 --- /dev/null +++ b/DashAI/back/units/load_trained_model_unit.py @@ -0,0 +1,114 @@ +"""Unit that restores a trained model from the artifact a run left on disk.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Run +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class LoadTrainedModelSchema(BaseSchema): + run_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the run whose trained model is restored. Both " + "the model component and the artifact path come from that row.", + es="Identificador de la ejecución cuyo modelo entrenado se " + "restaura. Tanto el componente del modelo como la ruta del " + "artefacto salen de esa fila.", + pt="Identificador da execução cujo modelo treinado é restaurado. " + "Tanto o componente do modelo como o caminho do artefacto vêm " + "dessa linha.", + de="Kennung des Laufs, dessen trainiertes Modell wiederhergestellt " + "wird. Sowohl die Modellkomponente als auch der Artefaktpfad " + "stammen aus dieser Zeile.", + zh="要恢复其已训练模型的运行标识符。模型组件和产物路径都来自该行。", + ), + alias=MultilingualString( + en="Run", es="Ejecución", pt="Execução", de="Lauf", zh="运行" + ), + ) # type: ignore + + +class LoadTrainedModelUnit(BaseUnit): + """Rebuild a trained model from the run that produced it. + + Reads the model component name and the artifact path off the ``Run`` row + rather than taking them as configuration, so the model that is restored is + always the one that run actually saved. + + ``load`` is invoked on the model *class*, not on an instance, which is what + every concrete model in this codebase expects: they all declare it as a + ``staticmethod`` or a ``classmethod`` that rebuilds the object from the + file. ``ExplainerJob`` instead instantiates the class before calling + ``load``; that extra step has no effect for those models, and preserving + the difference is why ``LoadRunModelUnit`` exists separately instead of + this unit growing a flag. + """ + + SCHEMA = LoadTrainedModelSchema + + PROVIDES = ("model",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._model_class = None + + def _resolve_model_class(self, model_name: str) -> type: + """Resolve the model class from the registry, memoized on this unit. + + Memoized on the instance, not in the shared context: a context can hold + more than one model-loading node, and a context-global cache key would + make the second one silently reuse the first one's class. + """ + if self._model_class is not None: + return self._model_class + + from kink import di + + component_registry = di["component_registry"] + + try: + model_class = component_registry[model_name]["class"] + except KeyError as e: + log.exception(e) + raise JobError(f"Model {model_name} not found in the registry") from e + + self._model_class = model_class + return model_class + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + session_factory = di["session_factory"] + + run_id = self.config["run_id"] + + with session_factory() as db: + run: Run = db.get(Run, run_id) + if not run: + raise JobError(f"Run {run_id} does not exist in DB.") + model_name = run.model_name + run_path = run.run_path + + model_class = self._resolve_model_class(model_name) + + try: + trained_model = model_class.load(run_path) + except Exception as e: + log.exception(e) + raise JobError( + f"Failed to load model {model_name} from path {run_path}" + ) from e + + ctx.put("model", trained_model) diff --git a/DashAI/back/units/load_training_dataset_unit.py b/DashAI/back/units/load_training_dataset_unit.py new file mode 100644 index 000000000..8d92280c3 --- /dev/null +++ b/DashAI/back/units/load_training_dataset_unit.py @@ -0,0 +1,56 @@ +"""Unit that loads the dataset a model was trained on.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class LoadTrainingDatasetUnit(BaseUnit): + """Load the dataset a model was trained on, under a key of its own. + + Deliberately not ``LoadDatasetUnit``: this dataset is not the one being + transformed, it is a *reference* the prediction needs — the task decodes + predicted class indexes against its labels, and its declared types become + the schema of the saved result. Publishing it as ``dataset`` would collide + with the dataset actually being predicted on, since ``PROVIDES`` is fixed + per class and both would want the same key. + + Two outputs, with different rules on purpose: the live dataset is cached + for the prediction step, while the types travel as a plain JSON-serializable + mapping so the saving step never has to reopen the file. Nothing derived + from the dataset *being predicted on* crosses this boundary. + """ + + PROVIDES = ("train_dataset", "train_dataset_types") + RUNTIME_PARAMS = ("train_dataset_file_path",) + + def execute(self, ctx: ExecutionContext) -> None: + from pathlib import Path + + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + file_path = self.config["train_dataset_file_path"] + + try: + train_dataset: "DashAIDataset" = load_dataset( + str(Path(f"{file_path}/dataset/")) + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Cannot load training dataset from {file_path}/dataset/" + ) from e + + ctx.put("train_dataset", train_dataset) + ctx.put_ref( + "train_dataset_types", + {name: kind.to_string() for name, kind in train_dataset.types.items()}, + ) diff --git a/DashAI/back/units/load_uploaded_dataset_unit.py b/DashAI/back/units/load_uploaded_dataset_unit.py new file mode 100644 index 000000000..a287b8f17 --- /dev/null +++ b/DashAI/back/units/load_uploaded_dataset_unit.py @@ -0,0 +1,133 @@ +"""Unit that reads a dataset from an uploaded file or a URL.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + int_field, + none_type, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class LoadUploadedDatasetSchema(BaseSchema): + dataloader: schema_field( + component_field(parent="BaseDataLoader"), + placeholder={"component": "CSVDataLoader", "params": {"separator": ","}}, + description=MultilingualString( + en="Reader used to parse the source, together with its own " + "configuration (delimiter, sheet, encoding, and so on).", + es="Lector usado para interpretar el origen, junto con su propia " + "configuración (delimitador, hoja, codificación, etc.).", + pt="Leitor usado para interpretar a origem, junto com a sua própria " + "configuração (delimitador, folha, codificação, etc.).", + de="Leser zum Parsen der Quelle samt eigener Konfiguration " + "(Trennzeichen, Blatt, Kodierung usw.).", + zh="用于解析数据源的读取器及其自身配置（分隔符、工作表、编码等）。", + ), + alias=MultilingualString( + en="Data loader", + es="Cargador de datos", + pt="Carregador de dados", + de="Datenlader", + zh="数据加载器", + ), + ) # type: ignore + source: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="Where to read from: a path to a local file or a URL. Archives " + "are downloaded and extracted before being parsed.", + es="Desde dónde leer: una ruta a un archivo local o una URL. Los " + "archivos comprimidos se descargan y extraen antes de interpretarse.", + pt="De onde ler: um caminho para um ficheiro local ou um URL. Os " + "arquivos comprimidos são descarregados e extraídos antes de serem " + "interpretados.", + de="Woher gelesen wird: ein Pfad zu einer lokalen Datei oder eine " + "URL. Archive werden vor dem Parsen heruntergeladen und entpackt.", + zh="读取来源：本地文件路径或 URL。压缩包会先下载并解压再解析。", + ), + alias=MultilingualString( + en="Source", + es="Origen", + pt="Origem", + de="Quelle", + zh="来源", + ), + ) # type: ignore + n_sample: schema_field( + none_type(int_field(gt=0)), + placeholder=None, + description=MultilingualString( + en="Read only this many rows instead of the whole source. Leave " + "empty to read everything.", + es="Leer solo esta cantidad de filas en vez de todo el origen. " + "Dejar vacío para leer todo.", + pt="Ler apenas esta quantidade de linhas em vez de toda a origem. " + "Deixar vazio para ler tudo.", + de="Nur so viele Zeilen lesen statt der gesamten Quelle. Leer " + "lassen, um alles zu lesen.", + zh="仅读取这么多行而非整个数据源。留空表示全部读取。", + ), + alias=MultilingualString( + en="Row sample", + es="Muestra de filas", + pt="Amostra de linhas", + de="Zeilenstichprobe", + zh="行采样", + ), + ) # type: ignore + + +class LoadUploadedDatasetUnit(BaseUnit): + """Parse a file or URL into a dataset with the chosen reader. + + The counterpart of ``LoadDatasetUnit``: that one materialises something DashAI + already stores, this one brings in data that is not a dataset yet. So it + publishes only ``dataset`` — there is no stored id to correlate against and no + path to save back to, because nothing decided yet where the result belongs. + + Types are not applied here. What the reader produces is whatever the source + suggests; declaring and casting types is a separate step, so the same load can + be reviewed before anything is committed to. + """ + + SCHEMA = LoadUploadedDatasetSchema + + PROVIDES = ("dataset",) + RUNTIME_PARAMS = ("temp_path",) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + component_registry = di["component_registry"] + + dataloader_config = self.config["dataloader"] + source = self.config["source"] + + # Raises KeyError with the registry's own wording when the reader does + # not exist. That message reaches the user, so it is not reworded here. + dataloader = component_registry[dataloader_config["component"]]["class"]() + + log.debug( + "Loading dataset from %s using %s", + source, + dataloader_config["component"], + ) + ctx.put( + "dataset", + dataloader.load_data( + filepath_or_buffer=source, + temp_path=self.config["temp_path"], + params=dataloader_config.get("params") or {}, + n_sample=self.config.get("n_sample"), + ), + ) diff --git a/DashAI/back/units/predict_unit.py b/DashAI/back/units/predict_unit.py new file mode 100644 index 000000000..4a6542fff --- /dev/null +++ b/DashAI/back/units/predict_unit.py @@ -0,0 +1,148 @@ +"""Unit that runs a trained model over a dataset and decodes its output.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +class PredictSchema(BaseSchema): + task_name: schema_field( + string_field(), + placeholder="TabularClassificationTask", + description=MultilingualString( + en="Name of the task that turns raw model output into labels.", + es="Nombre de la tarea que convierte la salida cruda del modelo en " + "etiquetas.", + pt="Nome da tarefa que converte a saída bruta do modelo em rótulos.", + de="Name der Aufgabe, die die Rohausgabe des Modells in Labels umwandelt.", + zh="将模型原始输出转换为标签的任务名称。", + ), + alias=MultilingualString( + en="Task", es="Tarea", pt="Tarefa", de="Aufgabe", zh="任务" + ), + ) # type: ignore + input_columns: schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=MultilingualString( + en="Names of the columns handed to the model as input.", + es="Nombres de las columnas entregadas al modelo como entrada.", + pt="Nomes das colunas entregues ao modelo como entrada.", + de="Namen der Spalten, die dem Modell als Eingabe übergeben werden.", + zh="作为输入交给模型的列名。", + ), + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + ) # type: ignore + output_columns: schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=MultilingualString( + en="Names of the columns the model predicts. Only the first one is " + "used: it names the column the predictions are written to.", + es="Nombres de las columnas que el modelo predice. Solo se usa la " + "primera: da nombre a la columna donde se escriben las predicciones.", + pt="Nomes das colunas que o modelo prevê. Apenas a primeira é " + "usada: dá nome à coluna onde as previsões são escritas.", + de="Namen der Spalten, die das Modell vorhersagt. Nur die erste " + "wird verwendet: sie benennt die Spalte für die Vorhersagen.", + zh="模型预测的列名。仅使用第一个：它命名写入预测结果的列。", + ), + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + ) # type: ignore + + +class PredictUnit(BaseUnit): + """Predict with a trained model and decode the result into labels. + + The input columns are selected against the dataset the context holds at + this moment, never against a column list captured earlier: whatever built + that dataset — a load from disk or hand-typed rows — is free to have + produced a different shape. + + The training dataset is required rather than reloaded because the task + decodes predicted class indexes against its labels. The model is handed the + selected columns unprepared: models apply their own preprocessing inside + ``predict``, and preparing beforehand would break the ones that replace + their input columns with derived features. + """ + + SCHEMA = PredictSchema + + REQUIRES = ("dataset", "model", "train_dataset") + PROVIDES = ("y_pred",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._task = None + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task from the registry, memoized on this unit.""" + if self._task is not None: + return self._task + + from kink import di + + component_registry = di["component_registry"] + task_name = self.config["task_name"] + + try: + task: "BaseTask" = component_registry[task_name]["class"]() + except Exception as e: + log.exception(e) + raise JobError(f"Task {task_name} not found in the registry") from e + + self._task = task + return task + + def validate(self, ctx: ExecutionContext) -> None: + """Resolve the task before anything observable happens. + + The orchestrator calls this ahead of loading the model, which is what + keeps a missing task reported as a task problem rather than being + overtaken by whatever fails next. + """ + self._resolve_task() + + def execute(self, ctx: ExecutionContext) -> None: + import numpy as np + + task = self._resolve_task() + + dataset = ctx.require("dataset") + model = ctx.require("model") + train_dataset = ctx.require("train_dataset") + + prepared_dataset = dataset.select_columns(self.config["input_columns"]) + y_pred_proba = np.array(model.predict(prepared_dataset)) + y_pred = task.process_predictions( + train_dataset, y_pred_proba, self.config["output_columns"][0] + ) + + ctx.put("y_pred", y_pred) diff --git a/DashAI/back/units/prepare_and_fold_unit.py b/DashAI/back/units/prepare_and_fold_unit.py new file mode 100644 index 000000000..c71842328 --- /dev/null +++ b/DashAI/back/units/prepare_and_fold_unit.py @@ -0,0 +1,73 @@ +"""Unit that prepares a dataset for a task and carves it into folds.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.splitter_scope import ( + SplitterScopeMixin, + fold_splitter_field, + input_columns_field, + output_columns_field, + task_name_field, +) + +log = logging.getLogger(__name__) + + +class PrepareAndFoldSchema(BaseSchema): + task_name: task_name_field() # type: ignore + input_columns: input_columns_field() # type: ignore + output_columns: output_columns_field() # type: ignore + splitter: fold_splitter_field() # type: ignore + + +class PrepareAndFoldUnit(BaseUnit, SplitterScopeMixin): + """Validate a dataset against a task and carve it into cross-validation folds. + + The sibling of ``PrepareAndSplitUnit``, sharing its whole body. The two + differ in which family of splitters they offer and, because of that, in the + shape of what they publish. + + **What comes back is a list, and that is why the keys are different.** A + fold splitter returns one entry per fold plus a trailing entry that is not + a fold: its train partition is every row the folds could use and its test + partition holds the rows reserved for scoring the model that gets kept, + empty when the session reserved none. Publishing that list as ``x`` would + give one key two possible types, which a contract comparing key names + cannot express -- a graph would validate and then fail at run time, or + worse, quietly train on the wrong thing. So it is ``x_folds`` and + ``y_folds``. + + ``split_indexes`` keeps its name because it keeps its meaning -- the rows + of every partition this run produced -- even though a fold payload is + shaped differently from a holdout one. Which shape it is, is answered by + asking the splitter that produced it, not by inspecting the payload. + """ + + SCHEMA = PrepareAndFoldSchema + + REQUIRES = ("dataset", "dataset_id") + PROVIDES = ("x_folds", "y_folds", "n_labels", "task", "split_indexes", "task_name") + + def __init__(self, **config) -> None: + super().__init__(**config) + # Kept on the instance, never in the context: two of these units in one + # context would otherwise overwrite each other's resolved task. + self._task = None + self._splitter_class = None + + def execute(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") + dataset_id = ctx.require("dataset_id") + + task, n_labels, x, y = self._prepare(dataset, dataset_id) + x_folds, y_folds, split_indexes = self._split(x, y) + + ctx.put_ref("task_name", self.config["task_name"]) + ctx.put_ref("split_indexes", split_indexes) + ctx.put("task", task) + ctx.put("n_labels", n_labels) + ctx.put("x_folds", x_folds) + ctx.put("y_folds", y_folds) diff --git a/DashAI/back/units/prepare_and_split_unit.py b/DashAI/back/units/prepare_and_split_unit.py new file mode 100644 index 000000000..26aeb9caf --- /dev/null +++ b/DashAI/back/units/prepare_and_split_unit.py @@ -0,0 +1,75 @@ +"""Unit that prepares a dataset for a task and splits it into train/val/test.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.splitter_scope import ( + SplitterScopeMixin, + input_columns_field, + output_columns_field, + partition_splitter_field, + task_name_field, +) + +log = logging.getLogger(__name__) + + +class PrepareAndSplitSchema(BaseSchema): + task_name: task_name_field() # type: ignore + input_columns: input_columns_field() # type: ignore + output_columns: output_columns_field() # type: ignore + splitter: partition_splitter_field() # type: ignore + + +class PrepareAndSplitUnit(BaseUnit, SplitterScopeMixin): + """Validate a dataset against a task and split it once into three partitions. + + Runs the task's own validation, counts the labels, separates features from + targets and hands the pair to a holdout splitter. + + The split configuration is a component field rather than an untyped + dictionary: which partitions exist, what they are called and what may be + configured are the splitter's own schema, so the front renders that form + instead of a free-form object nobody could validate. + + The sibling for cross-validation is ``PrepareAndFoldUnit``. They are two + units and not one with a flag for two reasons that both bite. A component + field carries a single ``parent`` and the front reads it directly off the + property, so offering two families from one field leaves the user with no + selector at all. And what comes back has a different *type* -- one + ``DatasetDict`` against a list of them -- which a contract that compares + key names cannot express: the same key holding two shapes would validate + statically and fail at run time. The fold unit therefore publishes + ``x_folds`` and ``y_folds`` rather than reusing ``x`` and ``y``. + """ + + SCHEMA = PrepareAndSplitSchema + + # dataset_id is declared even though it only decorates an error message: + # an undeclared read is a contract a DAG validator cannot see, and the + # loaders publish it precisely so downstream units can name their input. + REQUIRES = ("dataset", "dataset_id") + PROVIDES = ("x", "y", "n_labels", "task", "split_indexes", "task_name") + + def __init__(self, **config) -> None: + super().__init__(**config) + # Kept on the instance, never in the context: two of these units in one + # context would otherwise overwrite each other's resolved task. + self._task = None + self._splitter_class = None + + def execute(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") + dataset_id = ctx.require("dataset_id") + + task, n_labels, x, y = self._prepare(dataset, dataset_id) + x, y, split_indexes = self._split(x, y) + + ctx.put_ref("task_name", self.config["task_name"]) + ctx.put_ref("split_indexes", split_indexes) + ctx.put("task", task) + ctx.put("n_labels", n_labels) + ctx.put("x", x) + ctx.put("y", y) diff --git a/DashAI/back/units/prepare_explanation_data_unit.py b/DashAI/back/units/prepare_explanation_data_unit.py new file mode 100644 index 000000000..2ebc0ed20 --- /dev/null +++ b/DashAI/back/units/prepare_explanation_data_unit.py @@ -0,0 +1,188 @@ +"""Unit that rebuilds a run's train/test/val splits for an explanation.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +def _columns_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=description, + alias=alias, + ) + + +class PrepareExplanationDataSchema(BaseSchema): + task_name: schema_field( + string_field(), + placeholder="TabularClassificationTask", + description=MultilingualString( + en="Name of the task the dataset is prepared for.", + es="Nombre de la tarea para la que se prepara el conjunto de datos.", + pt="Nome da tarefa para a qual o conjunto de dados é preparado.", + de="Name der Aufgabe, für die der Datensatz vorbereitet wird.", + zh="数据集所准备的任务名称。", + ), + alias=MultilingualString( + en="Task", es="Tarea", pt="Tarefa", de="Aufgabe", zh="任务" + ), + ) # type: ignore + input_columns: _columns_field( + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + description=MultilingualString( + en="Names of the columns used as model input.", + es="Nombres de las columnas usadas como entrada del modelo.", + pt="Nomes das colunas usadas como entrada do modelo.", + de="Namen der als Modelleingabe verwendeten Spalten.", + zh="用作模型输入的列名。", + ), + ) # type: ignore + output_columns: _columns_field( + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + description=MultilingualString( + en="Names of the columns the model predicts.", + es="Nombres de las columnas que el modelo predice.", + pt="Nomes das colunas que o modelo prevê.", + de="Namen der Spalten, die das Modell vorhersagt.", + zh="模型需要预测的列名。", + ), + ) # type: ignore + + +class PrepareExplanationDataUnit(BaseUnit): + """Rebuild the exact train/test/val split the run was trained on. + + Deliberately not ``PrepareAndSplitUnit``: that one *computes* a split from + a ratio configuration, which would hand the explainer different rows than + the model ever saw. This one replays the row indexes the run recorded, so + the explanation is about the model that exists. + + Features stay unprepared — models apply their own preprocessing inside + ``predict``, and preparing beforehand would break the ones that replace + their input columns with derived features — while targets are encoded, + because explainers compare them against the model's class indexes. + """ + + SCHEMA = PrepareExplanationDataSchema + + # Exactly the keys ``execute`` reads, and no more. ``dataset_id`` is + # deliberately absent: the unit does not name it anywhere, and every key + # listed here is demanded unconditionally by ``__call__``, so declaring an + # unused one would reject any upstream that publishes a dataset without an + # id — ``BuildManualInputUnit``, for one. + REQUIRES = ("dataset", "model", "split_indexes") + PROVIDES = ("data_x", "data_y", "task") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._task = None + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task from the registry, memoized on this unit.""" + if self._task is not None: + return self._task + + from kink import di + + component_registry = di["component_registry"] + task_name = self.config["task_name"] + + try: + task: "BaseTask" = component_registry[task_name]["class"]() + except Exception as e: + log.exception(e) + raise JobError( + (f"Unable to find Task with name {task_name} in registry"), + ) from e + + self._task = task + return task + + def validate(self, ctx: ExecutionContext) -> None: + """Resolve the task before anything observable happens. + + The orchestrator calls this outside the block that wraps preparation + failures, which is what keeps a missing task reported as a registry + problem instead of a generic "cannot prepare" message. + """ + self._resolve_task() + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import ( + select_columns, + split_dataset, + ) + + task = self._resolve_task() + + loaded_dataset = ctx.require("dataset") + trained_model = ctx.require("model") + splits = ctx.require("split_indexes") + input_columns = self.config["input_columns"] + output_columns = self.config["output_columns"] + + loaded_dataset = split_dataset( + loaded_dataset, + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + + prepared_dataset = task.prepare_for_task( + dataset=loaded_dataset, + input_columns=input_columns, + output_columns=output_columns, + ) + data = select_columns(prepared_dataset, input_columns, output_columns) + + data_x = split_dataset( + data[0], + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + data_y = split_dataset( + data[1], + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + # Inputs stay unprepared (see the class docstring); targets are encoded + # because explainers compare them against the model's class indexes. + for split_name in data_y: + data_y[split_name] = trained_model.prepare_output( + data_y[split_name], is_fit=False + ) + + ctx.put("data_x", data_x) + ctx.put("data_y", data_y) + ctx.put("task", task) diff --git a/DashAI/back/units/run_exploration_unit.py b/DashAI/back/units/run_exploration_unit.py new file mode 100644 index 000000000..99f0ac5c2 --- /dev/null +++ b/DashAI/back/units/run_exploration_unit.py @@ -0,0 +1,187 @@ +"""Unit that runs one exploration over the dataset in the context.""" + +import logging +from typing import TYPE_CHECKING, Type + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Explorer +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.exploration.base_explorer import BaseExplorer + +log = logging.getLogger(__name__) + + +class RunExplorationSchema(BaseSchema): + explorer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the exploration whose selected columns and " + "display name the explorer component reads.", + es="Identificador de la exploración cuyas columnas seleccionadas y " + "nombre para mostrar lee el componente de exploración.", + pt="Identificador da exploração cujas colunas selecionadas e nome " + "de exibição o componente de exploração lê.", + de="Kennung der Exploration, deren ausgewählte Spalten und " + "Anzeigename die Explorer-Komponente liest.", + zh="探索的标识符，探索组件从中读取所选列和显示名称。", + ), + alias=MultilingualString( + en="Exploration", + es="Exploración", + pt="Exploração", + de="Exploration", + zh="探索", + ), + ) # type: ignore + explorer: schema_field( + component_field(parent="BaseExplorer"), + placeholder={ + "component": "DescribeExplorer", + "params": {"percentiles": "25, 50, 75", "include": "all", "exclude": None}, + }, + description=MultilingualString( + en="Exploration to run, together with its own configuration.", + es="Exploración a ejecutar, junto con su propia configuración.", + pt="Exploração a executar, junto com a sua própria configuração.", + de="Auszuführende Exploration samt ihrer eigenen Konfiguration.", + zh="要运行的探索及其自身配置。", + ), + alias=MultilingualString( + en="Explorer", + es="Explorador", + pt="Explorador", + de="Explorer", + zh="探索器", + ), + ) # type: ignore + + +class RunExplorationUnit(BaseUnit): + """Instantiate an explorer component and run it over the dataset. + + Preparing the dataset and launching the exploration are one unit, not two: + ``prepare_dataset`` is a hook on the explorer component itself + (``BaseExplorer.prepare_dataset``), so it does not exist without an + instantiated explorer. Splitting them would force the live explorer + instance through the context, which is instance state wearing a context + key's clothes. + + The unit re-reads the ``Explorer`` row because the component API takes it: + ``prepare_dataset`` needs its ``columns`` and ``launch_exploration`` + receives the row itself. The read is strictly read-only — the row's status + belongs to the job. + + The explorer instance is published alongside the result because saving is + also a method on it (``save_notebook``). Handing over the same object, + rather than letting the save unit build its own from the same + configuration, is what keeps a stateful explorer working: ``CorrMatrix`` + and ``CovMatrix`` read ``self.plot`` while saving. Same shape as + ``BuildModelUnit`` publishing ``model`` for ``SaveModelUnit``. + """ + + SCHEMA = RunExplorationSchema + + PROVIDES = ("exploration_result", "explorer") + REQUIRES = ("dataset",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._explorer_class = None + + @property + def exploration_type(self) -> str: + return self.config["explorer"]["component"] + + @property + def parameters(self) -> dict: + return self.config["explorer"]["params"] + + def _resolve_explorer_class(self) -> Type["BaseExplorer"]: + """Resolve the explorer class from the registry, memoized on this unit. + + Memoized on the instance rather than in the shared context: a context + can hold more than one exploration node, and a context-global cache key + would make the second one silently reuse the first one's class. + """ + if self._explorer_class is not None: + return self._explorer_class + + from kink import di + + component_registry = di["component_registry"] + exploration_type = self.exploration_type + + try: + explorer_class = component_registry[exploration_type]["class"] + except KeyError as e: + log.exception(e) + raise JobError( + (f"Explorer {exploration_type} not found in the registry.") + ) from e + + self._explorer_class = explorer_class + return explorer_class + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.exploration.base_explorer import BaseExplorer + + session_factory = di["session_factory"] + + explorer_id = self.config["explorer_id"] + exploration_type = self.exploration_type + loaded_dataset = ctx.require("dataset") + + explorer_component_class = self._resolve_explorer_class() + + with session_factory() as db: + explorer_info: Explorer = db.get(Explorer, explorer_id) + if explorer_info is None: + raise JobError(f"Explorer with id {explorer_id} not found.") + + try: + explorer_instance = explorer_component_class(**self.parameters) + assert isinstance(explorer_instance, BaseExplorer) + except Exception as e: + log.exception(e) + raise JobError( + f"Error instancing the explorer {exploration_type}." + ) from e + + try: + prepared_dataset = explorer_instance.prepare_dataset( + loaded_dataset, explorer_info.columns + ) + except Exception as e: + log.exception(e) + raise JobError( + ( + "Error preparing the dataset for the exploration " + f"{exploration_type}." + ) + ) from e + + try: + result = explorer_instance.launch_exploration( + prepared_dataset, explorer_info + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Error launching the exploration {exploration_type}." + ) from e + + ctx.put("exploration_result", result) + ctx.put("explorer", explorer_instance) diff --git a/DashAI/back/units/save_dataset_to_path_unit.py b/DashAI/back/units/save_dataset_to_path_unit.py new file mode 100644 index 000000000..6c5c80649 --- /dev/null +++ b/DashAI/back/units/save_dataset_to_path_unit.py @@ -0,0 +1,78 @@ +"""Unit that writes the dataset in the context to a destination of its own.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema, schema_field, string_field +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class SaveDatasetToPathSchema(BaseSchema): + path: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="Destination directory for the dataset. Unlike saving in place, " + "this writes wherever it is told, so it can store a dataset that " + "was loaded from somewhere else entirely.", + es="Directorio de destino del conjunto de datos. A diferencia de " + "guardar en su lugar, escribe donde se le indique, así que puede " + "almacenar un conjunto de datos cargado desde otro origen.", + pt="Diretório de destino do conjunto de dados. Ao contrário de " + "guardar no lugar, escreve onde lhe for indicado, pelo que pode " + "armazenar um conjunto de dados carregado de outra origem.", + de="Zielverzeichnis für den Datensatz. Anders als beim Speichern am " + "Ursprungsort schreibt diese Einheit dorthin, wo es ihr gesagt " + "wird, und kann so einen anderswo geladenen Datensatz ablegen.", + zh="数据集的目标目录。与原地保存不同，它会写入指定位置，因此可以存储" + "从其他来源加载的数据集。", + ), + alias=MultilingualString( + en="Destination path", + es="Ruta de destino", + pt="Caminho de destino", + de="Zielpfad", + zh="目标路径", + ), + ) # type: ignore + + +class SaveDatasetToPathUnit(BaseUnit): + """Write the dataset to a destination given in the configuration. + + The sibling of ``SaveDatasetUnit``, and deliberately a separate unit rather + than a flag on it. ``SaveDatasetUnit`` saves back to ``dataset_path`` — where + the load came from — which is what makes editing a dataset in place safe. This + one stores the dataset as something new, so it never reads ``dataset_path``: + if it did, a flow that loads a working copy and registers the result as a + fresh dataset would overwrite the copy it read. + + Declares no outputs — its result is on disk, not in the context. + + The error message keeps the original exception's text. A failure here is an + infrastructure failure (out of space, permissions, a path the filesystem + rejects), and only the message of the outermost error reaches the user: the + job queue stores ``str(exc)``, never the ``__cause__`` chain. Swallowing it + would drop the diagnosis exactly when it is needed. + """ + + SCHEMA = SaveDatasetToPathSchema + + REQUIRES = ("dataset",) + PROVIDES = () + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import save_dataset + + dataset = ctx.require("dataset") + path = self.config["path"] + + try: + save_dataset(dataset, path) + except Exception as e: + log.exception(e) + raise JobError(f"Can not save dataset to path {path}: {e}") from e diff --git a/DashAI/back/units/save_dataset_unit.py b/DashAI/back/units/save_dataset_unit.py new file mode 100644 index 000000000..519482edc --- /dev/null +++ b/DashAI/back/units/save_dataset_unit.py @@ -0,0 +1,40 @@ +"""Unit that persists the dataset in the context back to disk.""" + +import logging + +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class SaveDatasetUnit(BaseUnit): + """Write the dataset back to the path it was loaded from. + + Takes no configuration: the destination is ``dataset_path``, published by + whichever unit loaded the dataset, so a save can never land somewhere the + load did not come from. Declares no outputs — its result is on disk, not in + the context. + + The error message keeps the original exception's text. A failure here is an + infrastructure failure (out of space, permissions, a path the filesystem + rejects), and only the message of the outermost error reaches the user: the + job queue stores ``str(exc)``, never the ``__cause__`` chain. Swallowing it + would drop the diagnosis exactly when it is needed. + """ + + REQUIRES = ("dataset", "dataset_path") + PROVIDES = () + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import save_dataset + + dataset = ctx.require("dataset") + dataset_path = ctx.require("dataset_path") + + try: + save_dataset(dataset, dataset_path) + except Exception as e: + log.exception(e) + raise JobError(f"Can not save dataset to path {dataset_path}: {e}") from e diff --git a/DashAI/back/units/save_exploration_unit.py b/DashAI/back/units/save_exploration_unit.py new file mode 100644 index 000000000..71146ac76 --- /dev/null +++ b/DashAI/back/units/save_exploration_unit.py @@ -0,0 +1,122 @@ +"""Unit that persists an exploration result under its notebook's folder.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Explorer, Notebook +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class SaveExplorationSchema(BaseSchema): + explorer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the exploration being saved. It decides the " + "destination folder and the file name, so a re-run overwrites its " + "own artifact and never another exploration's.", + es="Identificador de la exploración que se guarda. Determina la " + "carpeta de destino y el nombre del archivo, de modo que volver a " + "ejecutarla sobrescribe su propio artefacto y nunca el de otra.", + pt="Identificador da exploração que está a ser guardada. Determina " + "a pasta de destino e o nome do ficheiro, pelo que uma nova " + "execução substitui o seu próprio artefacto e nunca o de outra.", + de="Kennung der zu speichernden Exploration. Sie bestimmt " + "Zielordner und Dateinamen, sodass ein erneuter Lauf nur das " + "eigene Artefakt überschreibt und nie das einer anderen.", + zh="要保存的探索的标识符。它决定目标文件夹和文件名，因此重新运行只会覆盖" + "自身的产物，而不会覆盖其他探索的产物。", + ), + alias=MultilingualString( + en="Exploration", + es="Exploración", + pt="Exploração", + de="Exploration", + zh="探索", + ), + ) # type: ignore + + +class SaveExplorationUnit(BaseUnit): + """Write an exploration result to disk and publish where it landed. + + How the result is serialised is the explorer component's own business — + a Plotly figure as JSON, a DataFrame as JSON, a word cloud as PNG — so the + unit delegates to ``save_notebook`` and only owns the destination: the + notebook's folder under ``NOTEBOOK_PATH``, keyed by the notebook id. + + The explorer arrives through the context rather than being rebuilt from a + configuration of its own, so the object that saves is the object that ran. + + Declares only ``exploration_path``: the result itself is on disk, and the + row that records the path belongs to the job. + """ + + SCHEMA = SaveExplorationSchema + + REQUIRES = ("exploration_result", "explorer") + PROVIDES = ("exploration_path",) + + def execute(self, ctx: ExecutionContext) -> None: + import os + import pathlib + + from kink import di + + config = di["config"] + session_factory = di["session_factory"] + + explorer_id = self.config["explorer_id"] + explorer_instance = ctx.require("explorer") + result = ctx.require("exploration_result") + + with session_factory() as db: + explorer_info: Explorer = db.get(Explorer, explorer_id) + if explorer_info is None: + raise JobError(f"Explorer with id {explorer_id} not found.") + + notebook_info: Notebook = db.get(Notebook, explorer_info.notebook_id) + if notebook_info is None: + raise JobError( + f"Notebook with id {explorer_info.notebook_id} not found." + ) + + # Read while the row is still attached: the error message below is + # built after the session is gone. + exploration_type = explorer_info.exploration_type + + # save in the notebook folder + save_path = pathlib.Path( + os.path.join( + config["NOTEBOOK_PATH"], + (f"{notebook_info.id}"), + ) + ) + if not save_path.exists(): + save_path.mkdir(parents=True) + + save_path = explorer_instance.save_notebook( + notebook_info, explorer_info, save_path, result + ) + + if isinstance(save_path, str): + save_path = pathlib.Path(save_path) + if not isinstance(save_path, pathlib.Path): + raise JobError( + ( + f"Error while saving the exploration" + f" {exploration_type}" + f", save path is not a pathlib.Path." + ) + ) + + ctx.put_ref("exploration_path", save_path.as_posix()) diff --git a/DashAI/back/units/save_model_unit.py b/DashAI/back/units/save_model_unit.py new file mode 100644 index 000000000..97ce74c6a --- /dev/null +++ b/DashAI/back/units/save_model_unit.py @@ -0,0 +1,74 @@ +"""Unit that persists a trained model to disk.""" + +import logging +import re + +from DashAI.back.core.atomic import atomic_save_path +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + +#: A prefix names a directory directly under RUNS_PATH, so anything that could +#: be read as a path — a separator, a parent reference, a drive letter — has to +#: be refused rather than cleaned up: a caller that meant one destination and +#: silently got another is the failure this guards against. +_SAFE_PREFIX = re.compile(r"^[A-Za-z0-9_-]+$") + + +class SaveModelUnit(BaseUnit): + """Write a trained model under the runs directory, in its own subdirectory. + + The destination is named by ``artifact_prefix``, which the caller chooses: + a run passes its own id, so a re-run overwrites its own artifact and never + another's. A caller that is not a run passes something unique to itself. + + The prefix is configuration rather than a context key on purpose: no unit + publishes it, so nothing upstream could ever satisfy it as a requirement. + It is the same split ``SaveDatasetUnit`` and ``SaveDatasetToPathUnit`` + already draw — a destination an upstream unit produced against one the + caller names. + """ + + REQUIRES = ("model",) + PROVIDES = ("model_path",) + RUNTIME_PARAMS = ("artifact_prefix",) + + def validate(self, ctx: ExecutionContext) -> None: + prefix = self.config["artifact_prefix"] + + if not isinstance(prefix, str) or not _SAFE_PREFIX.match(prefix): + raise JobError( + "The artifact prefix names a directory under the runs " + "directory, so it can only contain letters, digits, hyphens " + f"and underscores. Got: {prefix!r}" + ) + + def execute(self, ctx: ExecutionContext) -> None: + import os + + from kink import di + + config = di["config"] + + model = ctx.require("model") + + try: + model_path = os.path.join( + config["RUNS_PATH"], self.config["artifact_prefix"] + ) + # Written aside and moved into place, so a save that dies halfway + # leaves the previous artifact intact instead of a truncated one + # that the row still points at. The temporary path is handed to the + # model rather than derived here because only the model knows + # whether it writes a file or a directory of weights. + with atomic_save_path(model_path) as tmp_path: + model.save(str(tmp_path)) + except Exception as e: + log.exception(e) + raise JobError( + "Model saving failed", + ) from e + + ctx.put_ref("model_path", model_path) diff --git a/DashAI/back/units/save_prediction_unit.py b/DashAI/back/units/save_prediction_unit.py new file mode 100644 index 000000000..71e9625a7 --- /dev/null +++ b/DashAI/back/units/save_prediction_unit.py @@ -0,0 +1,140 @@ +"""Unit that stores a prediction alongside the data it was made on.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +def _columns_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=description, + alias=alias, + ) + + +class SavePredictionSchema(BaseSchema): + input_columns: _columns_field( + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + description=MultilingualString( + en="Names of the model input columns. Together with the output " + "column they decide which declared types are kept in the schema " + "written next to the result.", + es="Nombres de las columnas de entrada del modelo. Junto con la " + "columna de salida deciden qué tipos declarados se conservan en el " + "esquema que se escribe junto al resultado.", + pt="Nomes das colunas de entrada do modelo. Juntamente com a coluna " + "de saída decidem que tipos declarados são mantidos no esquema " + "escrito junto ao resultado.", + de="Namen der Modelleingabespalten. Zusammen mit der Ausgabespalte " + "bestimmen sie, welche deklarierten Typen im Schema neben dem " + "Ergebnis erhalten bleiben.", + zh="模型输入列的列名。它们与输出列共同决定结果旁写入的模式中保留哪些声明类型。", + ), + ) # type: ignore + output_columns: _columns_field( + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + description=MultilingualString( + en="Names of the predicted columns. Only the first one is used: it " + "names the column the predictions are written to.", + es="Nombres de las columnas predichas. Solo se usa la primera: da " + "nombre a la columna donde se escriben las predicciones.", + pt="Nomes das colunas previstas. Apenas a primeira é usada: dá nome " + "à coluna onde as previsões são escritas.", + de="Namen der vorhergesagten Spalten. Nur die erste wird verwendet: " + "sie benennt die Spalte für die Vorhersagen.", + zh="预测列的列名。仅使用第一个：它命名写入预测结果的列。", + ), + ) # type: ignore + + +class SavePredictionUnit(BaseUnit): + """Write the predicted column next to the data it was predicted from. + + The destination is a fresh folder under the datasets directory, named by a + generated identifier: a prediction has no natural key to overwrite, so + every run gets its own and no result can clobber another's. + + The columns to keep are resolved against the dataset the context holds + right now, at the top of this method. Publishing a column list earlier + would go stale the moment anything upstream renamed, added or dropped a + column. + """ + + SCHEMA = SavePredictionSchema + + REQUIRES = ("dataset", "y_pred", "train_dataset_types") + PROVIDES = ("results_path",) + + def execute(self, ctx: ExecutionContext) -> None: + import uuid + from pathlib import Path + + from kink import di + + from DashAI.back.dataloaders.classes.dashai_dataset import ( + save_dataset, + to_dashai_dataset, + ) + + config = di["config"] + + dataset = ctx.require("dataset") + y_pred = ctx.require("y_pred") + train_dataset_types = ctx.require("train_dataset_types") + + input_columns = self.config["input_columns"] + output_col = self.config["output_columns"][0] + + path = str(Path(f"{config['DATASETS_PATH']}/predictions/")) + folder_name = str(uuid.uuid4()) + full_path = Path(path) / folder_name + full_path.mkdir(parents=True, exist_ok=True) + + base_columns = [col for col in dataset.column_names if col != output_col] + output_dataset = dataset.select_columns(base_columns) + dataset_with_prediction = to_dashai_dataset( + output_dataset.add_column(output_col, y_pred) + ) + + # Only the columns the model session declares carry a type; anything + # the input dataset happened to bring along is left untyped. + filtered_schema = { + name: kind + for name, kind in train_dataset_types.items() + if name in input_columns + self.config["output_columns"] + } + + # Store num of rows, columns, and column names + dataset_with_prediction.compute_base_metadata() + + save_dataset( + dataset_with_prediction, + str(full_path / "dataset"), + filtered_schema, + ) + + ctx.put_ref("results_path", str(full_path)) diff --git a/DashAI/back/units/splitter_scope.py b/DashAI/back/units/splitter_scope.py new file mode 100644 index 000000000..1954ba823 --- /dev/null +++ b/DashAI/back/units/splitter_scope.py @@ -0,0 +1,303 @@ +"""Shared body of the two units that prepare a dataset and partition it. + +``PrepareAndSplitUnit`` and ``PrepareAndFoldUnit`` do the same four things -- +resolve the task, prepare the dataset for it, separate features from targets, +and hand the pair to a splitter -- and differ only in which family of splitters +they offer and in the shape of what comes back. Keeping the body here is what +stops the two from drifting into two answers for the same dataset. + +Named ``SplitterScopeMixin`` rather than ``BaseSomething`` on purpose: the +registry derives a component's type by walking its ``__mro__`` for a class whose +name contains "Base" and that declares ``TYPE``, and demands exactly one. A +shared parent called ``Base*`` would be a second candidate and would break the +registration of every unit that inherited it. +""" + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Tuple + +from DashAI.back.core.schema_fields import ( + component_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + from DashAI.back.splitters.base_splitter import BaseSplitter + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +def task_name_field(): + return schema_field( + string_field(), + placeholder="TabularClassificationTask", + description=MultilingualString( + en="Name of the task the dataset is prepared for.", + es="Nombre de la tarea para la que se prepara el conjunto de datos.", + pt="Nome da tarefa para a qual o conjunto de dados é preparado.", + de="Name der Aufgabe, für die der Datensatz vorbereitet wird.", + zh="数据集所准备的任务名称。", + ), + alias=MultilingualString( + en="Task", es="Tarea", pt="Tarefa", de="Aufgabe", zh="任务" + ), + ) + + +def input_columns_field(): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=MultilingualString( + en="Names of the columns used as model input.", + es="Nombres de las columnas usadas como entrada del modelo.", + pt="Nomes das colunas usadas como entrada do modelo.", + de="Namen der als Modelleingabe verwendeten Spalten.", + zh="用作模型输入的列名。", + ), + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + ) + + +def output_columns_field(): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=MultilingualString( + en="Names of the columns the model has to predict.", + es="Nombres de las columnas que el modelo debe predecir.", + pt="Nomes das colunas que o modelo deve prever.", + de="Namen der Spalten, die das Modell vorhersagen soll.", + zh="模型需要预测的列名。", + ), + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + ) + + +def _splitter_field(parent: str, placeholder: Dict[str, Any]): + """A splitter chosen from one family, with its own configuration. + + Parameters + ---------- + parent : str + Class name every offered splitter has in its ``__mro__``. The two + families are told apart here rather than by a flag, because the front + resolves the choices with ``component_parent``, which matches any + ancestor by name: ``PartitionSplitter`` offers the two holdout + splitters and ``FoldSplitter`` the eight fold ones, with no renaming + and no change to the registry. + placeholder : dict + The ``{"component": …, "params": {…}}`` value the form starts on. + """ + return schema_field( + component_field(parent=parent), + placeholder=placeholder, + description=MultilingualString( + en="Splitter that decides how the dataset is partitioned, along " + "with its own configuration.", + es="Particionador que decide cómo se divide el conjunto de datos, " + "junto con su propia configuración.", + pt="Divisor que decide como o conjunto de dados é particionado, " + "junto com a sua própria configuração.", + de="Splitter, der über die Aufteilung des Datensatzes entscheidet, " + "samt seiner eigenen Konfiguration.", + zh="决定数据集如何划分的划分器及其自身配置。", + ), + alias=MultilingualString( + en="Splitter", + es="Particionador", + pt="Divisor", + de="Splitter", + zh="划分器", + ), + ) + + +def partition_splitter_field(): + """The field of the unit that splits once, offering the holdout family.""" + return _splitter_field( + parent="PartitionSplitter", + placeholder={ + "component": "HoldoutSplitter", + "params": { + "train": 0.6, + "test": 0.2, + "validation": 0.2, + "stratify": False, + "shuffle": True, + "random_state": 42, + }, + }, + ) + + +def fold_splitter_field(): + """The field of the unit that splits into folds, offering the fold family.""" + return _splitter_field( + parent="FoldSplitter", + placeholder={ + "component": "KFoldSplitter", + "params": { + "n_splits": 5, + "test_size": 0.1, + "shuffle": True, + "random_state": 42, + }, + }, + ) + + +class SplitterScopeMixin: + """Prepare a dataset for a task and hand it to a splitter. + + Every method here takes and returns plain values and never touches the + execution context. A ``ctx.put`` hidden in a shared helper is invisible to + the audit that parses each unit's own source, so a broken ``PROVIDES`` + would pass it. + """ + + #: Declared here so both the memoizing helpers and a reader can see what + #: state an instance carries; each unit sets them in its own ``__init__``, + #: because ``BaseUnit.__init__`` comes first in the MRO and does not chain. + _task = None + _splitter_class = None + + @property + def splitter_name(self) -> str: + return self.config["splitter"]["component"] + + @property + def splitter_params(self) -> Dict[str, Any]: + return self.config["splitter"]["params"] + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task, memoized on this unit. + + On the instance rather than in the context: a context can hold two of + these units, and a context-global cache would silently give the second + one the first one's task. + """ + if self._task is not None: + return self._task + + from kink import di + + task_name: str = self.config["task_name"] + try: + self._task = di["component_registry"][task_name]["class"]() + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find Task with name {task_name} in registry", + ) from e + return self._task + + def _resolve_splitter(self) -> "BaseSplitter": + """Build the configured splitter. + + ``BaseSplitter.__init__`` takes a single ``splits_data`` mapping rather + than keyword arguments, so this is the one component field in the units + that is not expanded with ``**params``. The shape is the splitter's own + schema either way; only how it is handed over differs. + """ + from kink import di + + splitter_name = self.splitter_name + if self._splitter_class is None: + try: + self._splitter_class = di["component_registry"][splitter_name]["class"] + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find Splitter with name {splitter_name} in registry.", + ) from e + + try: + return self._splitter_class(splits_data=dict(self.splitter_params)) + except Exception as e: + log.exception(e) + raise JobError( + f"Error instantiating splitter {splitter_name}, {e}", + ) from e + + def _prepare( + self, dataset: "DashAIDataset", dataset_id: Any + ) -> Tuple["BaseTask", int, "DashAIDataset", "DashAIDataset"]: + """Validate the dataset against the task and separate x from y. + + Returns + ------- + tuple + The task, the number of labels, and the input and output datasets, + in that order. Nothing here is written to the context: the unit + that called this decides what it promises. + """ + from DashAI.back.dataloaders.classes.dashai_dataset import select_columns + + task = self._resolve_task() + task_name: str = self.config["task_name"] + input_columns: List[str] = self.config["input_columns"] + output_columns: List[str] = self.config["output_columns"] + + try: + prepared_dataset = task.prepare_for_task( + dataset=dataset, + input_columns=input_columns, + output_columns=output_columns, + ) + n_labels = task.num_labels(prepared_dataset, output_columns[0]) + except Exception as e: + log.exception(e) + raise JobError( + f"Can not prepare Dataset {dataset_id} for Task {task_name}", + ) from e + + try: + # Read from the prepared dataset rather than the loaded one: a task + # may reorder the rows, and forecasting does, sorting them by date + # so a temporal splitter carves real periods of time. Selecting + # from the loaded dataset would drop that work on the floor. + x, y = select_columns(prepared_dataset, input_columns, output_columns) + except Exception as e: + log.exception(e) + raise JobError( + f"Error selecting input and output columns from dataset {dataset_id}", + ) from e + + return task, n_labels, x, y + + def _split(self, x: "DashAIDataset", y: "DashAIDataset"): + """Partition the pair with the configured splitter. + + The splitter's own complaint is passed through as the whole message, + undecorated. It already names the numbers that explain the refusal -- + how many folds against how many rows -- and a caller that wants to say + which run this was frames it from outside, which is how the message + the user reads is built today. + """ + splitter = self._resolve_splitter() + try: + return splitter.split(x, y) + except JobError: + raise + except Exception as e: + log.exception(e) + raise JobError(str(e)) from e diff --git a/DashAI/back/units/transform_dataset_unit.py b/DashAI/back/units/transform_dataset_unit.py new file mode 100644 index 000000000..cba697af7 --- /dev/null +++ b/DashAI/back/units/transform_dataset_unit.py @@ -0,0 +1,82 @@ +"""Unit that transforms a dataset with an already fitted converter.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.converter_scope import ( + ConverterScopeMixin, + scope_field, + target_field, +) + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class TransformDatasetSchema(BaseSchema): + scope: scope_field() # type: ignore + target: target_field() # type: ignore + + +class TransformDatasetUnit(BaseUnit, ConverterScopeMixin): + """Apply an already fitted converter to the dataset in the context. + + Takes no converter configuration: the converter arrives fitted through the + context, from :class:`FitConverterUnit` or :class:`ApplyConverterUnit`. That + is what makes "fit on train, apply to test" possible without refitting — + run this once per dataset, and the converter keeps the state it learned. + + It does carry its own ``scope``, on purpose: the scope is a list of 1-based + indexes, and they are resolved against **this** unit's dataset. Reusing the + column names resolved during fit would break the moment the two datasets + order their columns differently, and would carry stale column identity + across the context boundary. + """ + + SCHEMA = TransformDatasetSchema + + REQUIRES = ("dataset", "fitted_converter") + PROVIDES = ("dataset",) + + def validate(self, ctx: ExecutionContext) -> None: + """Reject an out-of-bounds target before any work is done.""" + self._check_target_bounds(ctx.require("dataset")) + + def execute(self, ctx: ExecutionContext) -> None: + dataset: "DashAIDataset" = ctx.require("dataset") + converter_instance = ctx.require("fitted_converter") + converter_name = type(converter_instance).__name__ + + self._check_target_bounds(dataset) + ( + scope_column_names, + _scope_rows_indexes, + target_column_name, + ) = self._resolve_scope(dataset) + + log.info(f"Transforming with fitted converter: {converter_name}") + + x_dataset, y_dataset = self._slice_for_transform( + dataset, + scope_column_names, + target_column_name, + ) + + transformed_dataset = self._transform( + converter_instance, x_dataset, y_dataset, converter_name + ) + + ctx.put( + "dataset", + self._merge_transformed( + converter_instance, + dataset, + transformed_dataset, + scope_column_names, + ), + ) diff --git a/tests/back/api/test_converter_job.py b/tests/back/api/test_converter_job.py new file mode 100644 index 000000000..925113b65 --- /dev/null +++ b/tests/back/api/test_converter_job.py @@ -0,0 +1,336 @@ +"""End-to-end regression net for ``ConverterJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so that the refactor has something to be measured +against. The assertions are deliberately explicit — exact status values, exact +column names on disk, exact error message fragments — instead of the looser +``status in ["finished", "error"]`` style used elsewhere in this suite, which +cannot tell a unit that silently stopped doing part of its work from one that +did it. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import ConverterStatus +from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset +from DashAI.back.dependencies.database.models import Converter +from DashAI.back.job.base_job import JobError +from DashAI.back.job.converter_job import ConverterJob + +IRIS_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", +] +IRIS_ROWS = 150 + + +@pytest.fixture(name="notebook") +def create_notebook(client: TestClient, dataset_1): + """A notebook holding its own copy of the iris dataset. + + ``POST /notebook/`` copies the dataset folder, so every converter run + mutates the notebook's copy and never the source dataset. + """ + response = client.post( + "/api/v1/notebook/", + json={"dataset_id": dataset_1.id, "name": "converter job test"}, + ) + assert response.status_code == 201, response.text + return response.json() + + +def _create_converter(client, notebook_id, converter_name, scope=None, target=None): + """Create a Converter row through the API and return its id.""" + response = client.post( + "/api/v1/converter/", + json={ + "notebook_id": notebook_id, + "converter": converter_name, + "parameters": { + "order": 0, + "params": {}, + "scope": scope if scope is not None else {"columns": [], "rows": []}, + "target": target, + }, + }, + ) + assert response.status_code == 201, response.text + return response.json()["id"] + + +def _stored_converter(client, converter_id): + """Read the Converter row straight from the database.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + converter = db.get(Converter, converter_id) + return { + "status": converter.status, + "start_time": converter.start_time, + "end_time": converter.end_time, + } + + +def _notebook_dataset(notebook): + return load_dataset(f"{notebook['file_path']}/dataset") + + +def test_the_notebook_starts_as_an_untouched_copy_of_the_dataset(notebook): + """Guards the fixture itself: the assertions below mean nothing if the + notebook copy does not start out as the full iris dataset.""" + dataset = _notebook_dataset(notebook) + + assert dataset.column_names == IRIS_COLUMNS + assert len(dataset) == IRIS_ROWS + + +def test_converter_job_removes_the_scoped_column_and_finishes(client, notebook): + """The happy path, end to end: status transitions and the dataset on disk. + + ``ColumnRemover`` deletes whatever is in scope, so a one-column scope is + directly observable in the saved dataset. + """ + converter_id = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 2}], "rows": []}, + ) + + ConverterJob(converter_id=converter_id).run() + + stored = _stored_converter(client, converter_id) + assert stored["status"] == ConverterStatus.FINISHED + assert stored["start_time"] is not None + assert stored["end_time"] is not None + + dataset = _notebook_dataset(notebook) + assert dataset.column_names == [ + "SepalLengthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", + ] + assert len(dataset) == IRIS_ROWS + + +def test_two_converters_chained_see_the_columns_the_previous_one_left(client, notebook): + """Column scope is resolved by index against the *current* dataset. + + Each converter is its own row and its own job invocation, and the second + one's ``idx`` must be read against the four columns the first one left + behind, not against the original five. This is the behaviour any atomic + decomposition has to keep: nothing may cache the original column list. + """ + first = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + ) + ConverterJob(converter_id=first).run() + + assert _notebook_dataset(notebook).column_names == [ + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", + ] + + # idx 1 now means SepalWidthCm, not SepalLengthCm. + second = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + ) + ConverterJob(converter_id=second).run() + + assert _stored_converter(client, second)["status"] == ConverterStatus.FINISHED + assert _notebook_dataset(notebook).column_names == [ + "PetalLengthCm", + "PetalWidthCm", + "Species", + ] + + +def test_a_changes_row_count_converter_replaces_the_whole_dataset(client, notebook): + """``CHANGES_ROW_COUNT`` takes the transform output as the new dataset. + + ``NanRemover`` returns only the scoped columns, so the columns outside the + scope are dropped — the documented consequence of replacing the dataset + instead of merging the transformed columns back in. + """ + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [{"idx": 1}, {"idx": 5}], "rows": []}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + + dataset = _notebook_dataset(notebook) + assert dataset.column_names == ["SepalLengthCm", "Species"] + # iris has no missing values, so no row is dropped. + assert len(dataset) == IRIS_ROWS + + +def test_an_empty_column_scope_means_every_column(client, notebook): + """An empty scope is not "no columns"; it is "all of them".""" + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [], "rows": []}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS + + +def test_a_row_scope_fits_on_the_subset_and_transforms_the_whole_dataset( + client, notebook +): + """A row scope narrows ``fit`` only; ``transform`` still sees every row.""" + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [{"idx": 1}], "rows": [1, 2, 3]}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + + dataset = _notebook_dataset(notebook) + assert dataset.column_names == ["SepalLengthCm"] + assert len(dataset) == IRIS_ROWS + + +def test_a_target_column_is_resolved_and_passed_to_the_converter(client, notebook): + """The supervised path builds y and hands it to fit/transform.""" + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 5}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + assert _notebook_dataset(notebook).column_names == ["SepalLengthCm"] + + +def test_a_missing_converter_row_reports_it_by_id(client, notebook): + with pytest.raises(JobError, match="Converter with id 999999 not found"): + ConverterJob(converter_id=999999).run() + + +def test_an_unknown_converter_name_fails_with_the_wrapped_message(client, notebook): + """The registry lookup error is reported inside the outer wrapper. + + Both halves matter: the import error names the culprit, the wrapper is what + the jobs UI actually shows. + """ + converter_id = _create_converter( + client, notebook["id"], "ThisConverterDoesNotExist" + ) + + with pytest.raises(JobError) as excinfo: + ConverterJob(converter_id=converter_id).run() + + message = str(excinfo.value) + assert "Error applying converters to dataset" in message + assert "Error importing converter ThisConverterDoesNotExist" in message + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + # The dataset must be left untouched when the converter never ran. + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS + + +def test_an_out_of_bounds_target_index_reports_cannot_load_dataset(client, notebook): + """The "out of bounds" text is swallowed; only the wrapper reaches the user. + + The inner JobError is re-caught by the surrounding ``except Exception`` and + replaced, surviving only as ``__cause__``. Locking this in because it is an + easy detail to "fix" by accident while refactoring. + """ + converter_id = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 99}, + ) + + with pytest.raises(JobError, match="Cannot load dataset from") as excinfo: + ConverterJob(converter_id=converter_id).run() + + assert "Target column index 99 is out of bounds" in str(excinfo.value.__cause__) + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS + + +def test_a_dataset_that_cannot_be_loaded_still_leaves_the_row_in_error( + client, notebook +): + """A load failure must not leave the converter stuck in STARTED. + + Nothing else would fix it: the Huey error signal writes only to its own + ``task_copy`` table and never touches the ``Converter`` row, and the job + runs with no outer handler. Before this, only ``SQLAlchemyError`` was + caught here, so an unreadable dataset left the row STARTED forever and the + UI showed the converter as still running. + """ + import shutil + + converter_id = _create_converter(client, notebook["id"], "ColumnRemover") + shutil.rmtree(f"{notebook['file_path']}/dataset") + + with pytest.raises(JobError, match="Can not load dataset from path"): + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + + +def test_a_failing_converter_leaves_the_dataset_untouched(client, notebook): + """``ColumnRemover`` raises when asked for a column that is not there. + + Reaching that requires a scope index past the end of the dataset, which is + exactly what a stale column list would produce in a chained run. + """ + converter_id = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 99}], "rows": []}, + ) + + with pytest.raises(JobError, match="Error applying converters to dataset"): + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS diff --git a/tests/back/api/test_dataset_job.py b/tests/back/api/test_dataset_job.py new file mode 100644 index 000000000..f20e7caa7 --- /dev/null +++ b/tests/back/api/test_dataset_job.py @@ -0,0 +1,1094 @@ +"""End-to-end regression net for ``DatasetJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so the refactor has something to be measured against. +The assertions are deliberately explicit — exact status values, exact row/column +counts, exact keys in ``splits.json``, exact error message text — instead of the +looser ``status in ["finished", "error"]`` style used elsewhere in this suite, +which cannot tell a unit that silently stopped doing part of its work from one +that did it. + +``DatasetJob`` has the widest blast radius of any job in the repo: it is the +bootstrap fixture of half the suite (``tests/back/api/conftest.py``, +``tests/back/models/conftest.py``, ``tests/back/api/test_predict_api.py``, +``tests/back/types/load_preview_test.py``). Everything here has to keep passing +through every slice of the refactor. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import os +import shutil +from pathlib import Path +from typing import Optional + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import DatafileStatus, DatasetStatus +from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset +from DashAI.back.dependencies.database.models import Datafile, Dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.job.dataset_job import DatasetJob + +IRIS_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", +] +IRIS_ROWS = 150 + +IRIS_SCHEMA = { + "SepalLengthCm": {"type": "Float", "dtype": "float64"}, + "SepalWidthCm": {"type": "Float", "dtype": "float64"}, + "PetalLengthCm": {"type": "Float", "dtype": "float64"}, + "PetalWidthCm": {"type": "Float", "dtype": "float64"}, + "Species": {"type": "Categorical", "dtype": "string"}, +} + +#: The extended EDA keys ``compute_metadata=False`` has to leave out. +EXTENDED_KEYS = ( + "general_info", + "numeric_stats", + "categorical_stats", + "text_stats", + "quality_info", + "correlations", +) + +IRIS_CSV = Path(__file__).parent / "iris.csv" + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # + + +def _session_factory(client): + return client.app.container["session_factory"] + + +def _new_dataset_row(client, name: str) -> int: + """Create a ``Dataset`` row in DELIVERED, the state the job expects.""" + with _session_factory(client)() as db: + entry = Dataset(name=name, file_path="") + entry.set_status_as_delivered() + db.add(entry) + db.commit() + db.refresh(entry) + return entry.id + + +def _stored(client, dataset_id: int) -> Optional[dict]: + """Read the ``Dataset`` row straight from the database. + + No ``start_time`` / ``end_time``: unlike ``Run`` or ``Explorer``, this table + declares no such columns. The status setters used to assign them anyway, + which only set attributes that were dropped at commit. + """ + with _session_factory(client)() as db: + row = db.get(Dataset, dataset_id) + if row is None: + return None + return { + "status": row.status, + "file_path": row.file_path, + "total_rows": row.total_rows, + "total_columns": row.total_columns, + "last_modified": row.last_modified, + } + + +def _run_job(client, dataset_id: int, params: dict, **extra) -> None: + """Run ``DatasetJob`` synchronously, the way every fixture in the suite does.""" + kwargs = { + "dataset_id": dataset_id, + "url": "", + "params": params, + "file_path": IRIS_CSV, + **extra, + } + DatasetJob(job_type="DatasetJob", kwargs=kwargs).run() + + +def _csv_params(name: str, **overrides) -> dict: + params = { + "dataloader": "CSVDataLoader", + "separator": ",", + "name": name, + "schema": IRIS_SCHEMA, + } + params.update(overrides) + return params + + +def _splits(file_path: str) -> dict: + with open(Path(file_path) / "dataset" / "splits.json", encoding="utf-8") as f: + return json.load(f) + + +def _cleanup(client, dataset_id: int) -> None: + row = _stored(client, dataset_id) + if row and row["file_path"]: + shutil.rmtree(row["file_path"], ignore_errors=True) + + +# --------------------------------------------------------------------------- # +# status bookkeeping +# --------------------------------------------------------------------------- # + + +class TestTheRowIsNeverLeftStuckInStarted: + """Whatever goes wrong, the row must not stay in STARTED. + + Nothing else moves it: Huey's ``SIGNAL_ERROR`` writes to its own + ``task_copy`` table and never to the ``Dataset`` row + (``huey_job_queue.py:201-211``), so an exception the job's own handler does + not see leaves the dataset advertised as in-progress forever and the UI + spinning on work that already died. + + Both cases below used to escape — the handler only caught ``JobError``, and + the final database block only caught ``SQLAlchemyError``. They are the + reason the handler now catches ``Exception``. + """ + + def test_a_mkdir_failure_that_is_not_file_exists_still_marks_the_row( + self, client, monkeypatch + ): + """Creating the destination folder can fail with more than + ``FileExistsError``: no permission, no space, a path the filesystem + rejects. The status was already committed as STARTED by then.""" + dataset_id = _new_dataset_row(client, "stuck_on_mkdir") + + def deny(self, *args, **kwargs): + raise PermissionError("cannot create the dataset folder") + + monkeypatch.setattr(Path, "mkdir", deny) + + with pytest.raises(PermissionError): + _run_job(client, dataset_id, _csv_params("stuck_on_mkdir")) + + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + def test_a_row_deleted_mid_run_is_reported_instead_of_crashing(self, client): + """The row is deletable through the API while the job runs, so the final + block can find it gone. It used to assign to ``None`` and raise + ``AttributeError`` from inside a block that only caught database errors. + """ + dataset_id = _new_dataset_row(client, "stuck_on_deleted_row") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_save = dashai_dataset.save_dataset + + def save_then_delete(dataset, path): + real_save(dataset, path) + with _session_factory(client)() as db: + db.delete(db.get(Dataset, dataset_id)) + db.commit() + + dashai_dataset.save_dataset = save_then_delete + try: + with pytest.raises(JobError) as excinfo: + _run_job(client, dataset_id, _csv_params("stuck_on_deleted_row")) + finally: + dashai_dataset.save_dataset = real_save + + assert str(excinfo.value) == (f"Dataset with ID {dataset_id} no longer exists.") + # The row is gone, so there is no status to write — but the data written + # to disk has to go with it instead of being orphaned. + assert _stored(client, dataset_id) is None + assert _folders_in(datasets_dir) == before + + def test_a_non_database_failure_in_the_final_block_still_marks_the_row( + self, client + ): + """The mirror of the case above for a row that does still exist: the + block catches ``SQLAlchemyError`` by name, so anything else has to be + handled further out rather than escaping.""" + dataset_id = _new_dataset_row(client, "stuck_on_final_block") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + real_set_finished = Dataset.set_status_as_finished + + def boom(self): + raise RuntimeError("not a database error") + + Dataset.set_status_as_finished = boom + try: + with pytest.raises(RuntimeError): + _run_job(client, dataset_id, _csv_params("stuck_on_final_block")) + finally: + Dataset.set_status_as_finished = real_set_finished + + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + assert _folders_in(datasets_dir) == before + + +# --------------------------------------------------------------------------- # +# file / URL branch — the happy path +# --------------------------------------------------------------------------- # + + +def test_the_file_branch_finishes_and_writes_every_column_of_the_row(client): + """Status transitions, the row's columns, and the dataset on disk.""" + dataset_id = _new_dataset_row(client, "net_file_happy") + before = _stored(client, dataset_id) + assert before["status"] == DatasetStatus.DELIVERED + + _run_job(client, dataset_id, _csv_params("net_file_happy")) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + assert row["last_modified"] > before["last_modified"] + + # The path is a realpath under DATASETS_PATH, and the dataset is really there. + datasets_path = client.app.container["config"]["DATASETS_PATH"] + assert Path(row["file_path"]).parent == Path(os.path.realpath(datasets_path)) + + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + assert len(dataset) == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_declared_schema_lands_on_the_stored_types(client): + """``transform_dataset_with_schema`` is what makes Species categorical.""" + dataset_id = _new_dataset_row(client, "net_file_types") + + _run_job(client, dataset_id, _csv_params("net_file_types")) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + types = dataset.types + + assert type(types["Species"]).__name__ == "Categorical" + for column in IRIS_COLUMNS[:4]: + assert type(types[column]).__name__ == "Float" + + _cleanup(client, dataset_id) + + +def test_inferred_types_in_the_params_win_over_inference(client): + """``params["inferred_types"]`` is the first branch of the schema cascade.""" + dataset_id = _new_dataset_row(client, "net_file_inferred") + # Species declared as plain text instead of categorical: only an honoured + # override can produce this. + override = dict(IRIS_SCHEMA, Species={"type": "Text", "dtype": "string"}) + + _run_job( + client, + dataset_id, + _csv_params("net_file_inferred", inferred_types=override), + ) + + row = _stored(client, dataset_id) + types = load_dataset(f"{row['file_path']}/dataset").types + assert type(types["Species"]).__name__ != "Categorical" + + _cleanup(client, dataset_id) + + +def test_the_schema_is_inferred_when_the_params_declare_nothing(client): + """Third branch of the cascade: ``infer_types(..., "DashAIPtype")``.""" + dataset_id = _new_dataset_row(client, "net_file_no_schema") + + _run_job( + client, + dataset_id, + {"dataloader": "CSVDataLoader", "separator": ",", "name": "net_file_no_schema"}, + ) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + # Inference ran and produced a type for every column. + assert set(dataset.types) == set(IRIS_COLUMNS) + + _cleanup(client, dataset_id) + + +def test_column_renames_rewrite_the_columns_and_remap_the_schema(client): + """The rename remaps ``schema`` too (lines 291-297), so the renamed column + keeps the type its old name declared.""" + dataset_id = _new_dataset_row(client, "net_file_renames") + + _run_job( + client, + dataset_id, + _csv_params( + "net_file_renames", + inferred_types=IRIS_SCHEMA, + column_renames={"Species": "Variety", "SepalLengthCm": "SepalLength"}, + ), + ) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + + assert dataset.column_names == [ + "SepalLength", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Variety", + ] + # The type followed the rename rather than being re-inferred. + assert type(dataset.types["Variety"]).__name__ == "Categorical" + assert row["total_columns"] == len(IRIS_COLUMNS) + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# metadata policy +# --------------------------------------------------------------------------- # + + +def test_compute_metadata_true_writes_base_and_extended_metadata(client): + dataset_id = _new_dataset_row(client, "net_meta_full") + + _run_job(client, dataset_id, _csv_params("net_meta_full", compute_metadata=True)) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + assert splits["total_rows"] == IRIS_ROWS + assert splits["column_names"] == IRIS_COLUMNS + assert "nan" in splits + for key in EXTENDED_KEYS: + assert key in splits, f"missing {key} when compute_metadata=True" + + _cleanup(client, dataset_id) + + +def test_compute_metadata_false_writes_base_metadata_only(client): + dataset_id = _new_dataset_row(client, "net_meta_base") + + _run_job(client, dataset_id, _csv_params("net_meta_base", compute_metadata=False)) + + row = _stored(client, dataset_id) + splits = _splits(row["file_path"]) + assert splits["total_rows"] == IRIS_ROWS + assert splits["column_names"] == IRIS_COLUMNS + assert "nan" in splits + for key in EXTENDED_KEYS: + assert key not in splits, f"unexpected {key} when compute_metadata=False" + + # The row's columns come from splits, so they must survive the base-only path. + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + + _cleanup(client, dataset_id) + + +def test_omitting_the_flag_defaults_to_full_metadata(client): + """Backward compatibility: the 4 bootstrap fixtures never pass the flag.""" + dataset_id = _new_dataset_row(client, "net_meta_default") + + _run_job(client, dataset_id, _csv_params("net_meta_default")) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + for key in EXTENDED_KEYS: + assert key in splits + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# notebook branch +# --------------------------------------------------------------------------- # + + +@pytest.fixture(name="notebook") +def create_notebook(client: TestClient, dataset_1): + """A notebook holding its own copy of the iris dataset. + + ``POST /notebook/`` copies the dataset folder, so the notebook branch of the + job reads a private working copy and never the source dataset. + """ + response = client.post( + "/api/v1/notebook/", + json={"dataset_id": dataset_1.id, "name": "dataset job net"}, + ) + assert response.status_code == 201, response.text + return response.json() + + +def _add_converter(client, notebook_id: int) -> int: + """Register a converter against the notebook. + + The job only asks *whether any Converter row exists* for the notebook + (lines 174-180); it never runs it. So a row is enough to flip + ``from_notebook_no_converters`` to False. + """ + response = client.post( + "/api/v1/converter/", + json={ + "notebook_id": notebook_id, + "converter": "StandardScaler", + "parameters": { + "order": 0, + "params": {}, + "scope": {"columns": [], "rows": []}, + "target": None, + }, + }, + ) + assert response.status_code == 201, response.text + return response.json()["id"] + + +def test_the_notebook_branch_copies_the_working_copy_into_a_new_dataset( + client, notebook +): + dataset_id = _new_dataset_row(client, "net_notebook_happy") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_happy"}, + notebook_id=notebook["id"], + file_path=None, + ) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + + # A new folder, not the notebook's own copy. + assert Path(row["file_path"]) != Path(notebook["file_path"]) + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + assert len(dataset) == IRIS_ROWS + + # And the notebook is untouched. + assert load_dataset(f"{notebook['file_path']}/dataset").column_names == IRIS_COLUMNS + + _cleanup(client, dataset_id) + + +def test_a_notebook_without_converters_reuses_the_source_metadata(client, notebook): + """No converters means the bytes match the source, so ``splits.json`` is + inherited rather than recomputed (lines 313-328).""" + dataset_id = _new_dataset_row(client, "net_notebook_inherit") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_inherit", "compute_metadata": True}, + notebook_id=notebook["id"], + file_path=None, + ) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + # dataset_1 was built with the default (full metadata), so the inherited + # splits already carry the extended keys. + for key in EXTENDED_KEYS: + assert key in splits + assert splits["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_a_notebook_without_converters_still_drops_extended_when_asked( + client, notebook +): + """``compute_metadata=False`` purges the extended keys the source carried.""" + dataset_id = _new_dataset_row(client, "net_notebook_purge") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_purge", "compute_metadata": False}, + notebook_id=notebook["id"], + file_path=None, + ) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + assert splits["total_rows"] == IRIS_ROWS + for key in EXTENDED_KEYS: + assert key not in splits, f"inherited {key} was not purged" + + _cleanup(client, dataset_id) + + +def test_a_notebook_with_converters_recomputes_the_metadata(client, notebook): + """One Converter row is enough to stop trusting the source's splits.""" + _add_converter(client, notebook["id"]) + dataset_id = _new_dataset_row(client, "net_notebook_recompute") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_recompute", "compute_metadata": True}, + notebook_id=notebook["id"], + file_path=None, + ) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + for key in EXTENDED_KEYS: + assert key in splits + assert splits["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_notebook_branch_never_applies_a_schema_or_renames(client, notebook): + """Lines 260-299 sit in the ``else`` of the notebook check, so neither + ``inferred_types`` nor ``column_renames`` has any effect here.""" + dataset_id = _new_dataset_row(client, "net_notebook_no_schema") + + _run_job( + client, + dataset_id, + { + "name": "net_notebook_no_schema", + "column_renames": {"Species": "ShouldBeIgnored"}, + }, + notebook_id=notebook["id"], + file_path=None, + ) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# hub branch +# --------------------------------------------------------------------------- # + + +@pytest.fixture(name="datafile") +def create_datafile(client, tmp_path_factory, request): + """A READY ``Datafile`` row pointing at a directory holding iris.csv. + + Built by hand: the hub branch only reads the row and the files on disk, so + no network is involved. ``dataset_id`` carries the test name because the + table has a ``UNIQUE(source_name, dataset_id)`` constraint and each test + gets its own row. + """ + work_dir = tmp_path_factory.mktemp("hub_download") + shutil.copy(IRIS_CSV, work_dir / "iris.csv") + + with _session_factory(client)() as db: + row = Datafile( + source_name="HuggingFaceDatasetSource", + dataset_id=f"net/iris/{request.node.name}", + name="net iris", + local_path=str(work_dir), + status=DatafileStatus.READY, + ) + db.add(row) + db.commit() + db.refresh(row) + return {"id": row.id, "local_path": str(work_dir)} + + +def _hub_params(name: str, datafile_id: int, **overrides) -> dict: + params = { + "name": name, + "dataloader": "CSVDataLoader", + "dataloader_params": {"separator": ","}, + "datafile_id": datafile_id, + } + params.update(overrides) + return params + + +def _run_hub_job(client, dataset_id: int, params: dict) -> None: + DatasetJob( + job_type="DatasetJob", + kwargs={ + "dataset_id": dataset_id, + "source_name": "HuggingFaceDatasetSource", + "dataset_source_id": "net/iris", + "params": params, + }, + ).run() + + +def test_the_hub_branch_finishes_from_an_explicit_selected_file(client, datafile): + dataset_id = _new_dataset_row(client, "net_hub_selected") + + _run_hub_job( + client, + dataset_id, + _hub_params("net_hub_selected", datafile["id"], selected_file="iris.csv"), + ) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + assert load_dataset(f"{row['file_path']}/dataset").column_names == IRIS_COLUMNS + + _cleanup(client, dataset_id) + + +def test_the_hub_branch_picks_the_first_file_when_none_is_selected(client, datafile): + """No ``selected_file``: the first entry of a sorted ``rglob`` wins.""" + dataset_id = _new_dataset_row(client, "net_hub_first") + + _run_hub_job(client, dataset_id, _hub_params("net_hub_first", datafile["id"])) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_hub_branch_ignores_dotted_files_when_picking(client, datafile): + """Any path component starting with a dot is skipped (lines 213-217). + + ``.hidden`` sorts before ``iris.csv``, so without the filter it would win + and the job would fail on an unreadable file. + """ + hidden = Path(datafile["local_path"]) / ".hidden" + hidden.mkdir() + (hidden / "junk.csv").write_text("not,a,dataset\n", encoding="utf-8") + + dataset_id = _new_dataset_row(client, "net_hub_hidden") + + _run_hub_job(client, dataset_id, _hub_params("net_hub_hidden", datafile["id"])) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_hub_branch_applies_the_schema_and_the_renames(client, datafile): + """Unlike the notebook branch, hub imports do go through lines 260-299.""" + dataset_id = _new_dataset_row(client, "net_hub_renames") + + _run_hub_job( + client, + dataset_id, + _hub_params( + "net_hub_renames", + datafile["id"], + selected_file="iris.csv", + inferred_types=IRIS_SCHEMA, + column_renames={"Species": "Variety"}, + ), + ) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names[-1] == "Variety" + assert type(dataset.types["Variety"]).__name__ == "Categorical" + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# error messages, by exact text +# --------------------------------------------------------------------------- # +# +# These travel verbatim to ``task_copy.error_msg`` (``huey_job_queue.py:210`` +# stores ``str(exc)``) and from there to ``GET /job/status/{id}`` and the UI, so +# the text *is* the contract. Everything raised inside the try at line 153 comes +# back wrapped as ``Error loading dataset: ``. + + +def _datasets_dir(client) -> Path: + return Path(client.app.container["config"]["DATASETS_PATH"]) + + +def _folders_in(path: Path) -> set: + return {p.name for p in path.iterdir() if p.is_dir()} + + +def test_a_missing_dataset_row_is_reported_by_id(client): + """Raised before the status is touched, so nothing is left behind.""" + with pytest.raises(JobError) as excinfo: + _run_job(client, 987654, _csv_params("net_missing_row")) + + assert str(excinfo.value) == "Dataset with ID 987654 not found." + + +def test_a_file_exists_collision_names_the_generated_folder(client, monkeypatch): + dataset_id = _new_dataset_row(client, "net_collision") + + def collide(self, *args, **kwargs): + raise FileExistsError(self) + + monkeypatch.setattr(Path, "mkdir", collide) + + with pytest.raises(JobError) as excinfo: + _run_job(client, dataset_id, _csv_params("net_collision")) + + message = str(excinfo.value) + assert message.startswith("A dataset with the name ") + assert message.endswith(" already exists.") + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_missing_notebook_is_reported_with_the_original_wording(client): + """The wording says "has no associated dataset" even though what is missing + is the Notebook row itself. Preserved verbatim: it reaches the UI.""" + dataset_id = _new_dataset_row(client, "net_missing_notebook") + + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + {"name": "net_missing_notebook"}, + notebook_id=987654, + file_path=None, + ) + + assert str(excinfo.value) == ( + "Error loading dataset: Notebook with ID 987654 has no associated dataset." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_an_unreadable_notebook_copy_names_the_path_it_could_not_read(client, notebook): + """The notebook branch reads a stored dataset off disk. When that read fails + the message names the path instead of quoting the reader's exception. + + This is the one intentional wording change of the refactor: the branch now + goes through the same loading step as every other flow that materialises a + stored dataset, and that step reports failures by path. The path is the + actionable part — the reader's own text ("the arrow file is truncated") says + nothing about *which* file. + """ + dataset_id = _new_dataset_row(client, "net_unreadable_notebook") + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_load = dashai_dataset.load_dataset + + def fail_to_load(path): + raise OSError("the arrow file is truncated") + + dashai_dataset.load_dataset = fail_to_load + try: + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + {"name": "net_unreadable_notebook"}, + notebook_id=notebook["id"], + file_path=None, + ) + finally: + dashai_dataset.load_dataset = real_load + + assert str(excinfo.value) == ( + f"Error loading dataset: Can not load dataset from path {notebook['file_path']}" + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_an_unknown_dataloader_in_the_file_branch_reports_the_registry_key_error( + client, +): + """The file branch uses ``component_registry[name]``, whose ``KeyError`` + stringifies *with* the quotes.""" + dataset_id = _new_dataset_row(client, "net_bad_loader_file") + + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + _csv_params("net_bad_loader_file", dataloader="NoSuchDataLoader"), + ) + + assert str(excinfo.value) == ( + "Error loading dataset: \"Component 'NoSuchDataLoader' does not exists " + 'in the registry."' + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_an_unknown_dataloader_in_the_hub_branch_has_its_own_wording(client, datafile): + """The hub branch looks the loader up itself, with a different message than + the file branch. Both have to survive the refactor.""" + dataset_id = _new_dataset_row(client, "net_bad_loader_hub") + + with pytest.raises(JobError) as excinfo: + _run_hub_job( + client, + dataset_id, + _hub_params( + "net_bad_loader_hub", datafile["id"], dataloader="NoSuchDataLoader" + ), + ) + + assert str(excinfo.value) == ( + "Error loading dataset: DataLoader 'NoSuchDataLoader' not found in registry." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_hub_import_without_a_datafile_id_says_so(client): + dataset_id = _new_dataset_row(client, "net_hub_no_id") + + with pytest.raises(JobError) as excinfo: + _run_hub_job( + client, + dataset_id, + {"name": "net_hub_no_id", "dataloader": "CSVDataLoader"}, + ) + + assert str(excinfo.value) == ( + "Error loading dataset: datafile_id is required for hub imports." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_datafile_that_is_not_ready_is_reported_by_id(client, datafile): + dataset_id = _new_dataset_row(client, "net_hub_not_ready") + with _session_factory(client)() as db: + row = db.get(Datafile, datafile["id"]) + row.status = DatafileStatus.DOWNLOADING + db.commit() + + with pytest.raises(JobError) as excinfo: + _run_hub_job( + client, dataset_id, _hub_params("net_hub_not_ready", datafile["id"]) + ) + + assert str(excinfo.value) == ( + f"Error loading dataset: Datafile {datafile['id']} is not ready." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_missing_datafile_row_is_also_reported_as_not_ready(client): + """``hub_row is None`` and ``status != READY`` share one branch and one + message, so a nonexistent id reads as "not ready".""" + dataset_id = _new_dataset_row(client, "net_hub_missing_row") + + with pytest.raises(JobError) as excinfo: + _run_hub_job(client, dataset_id, _hub_params("net_hub_missing_row", 987654)) + + assert str(excinfo.value) == "Error loading dataset: Datafile 987654 is not ready." + + +def test_an_empty_hub_download_directory_says_so(client, tmp_path_factory): + dataset_id = _new_dataset_row(client, "net_hub_empty") + empty_dir = tmp_path_factory.mktemp("hub_empty") + with _session_factory(client)() as db: + row = Datafile( + source_name="HuggingFaceDatasetSource", + dataset_id="net/iris/empty", + name="empty", + local_path=str(empty_dir), + status=DatafileStatus.READY, + ) + db.add(row) + db.commit() + db.refresh(row) + datafile_id = row.id + + with pytest.raises(JobError) as excinfo: + _run_hub_job(client, dataset_id, _hub_params("net_hub_empty", datafile_id)) + + assert str(excinfo.value) == ( + "Error loading dataset: Hub download directory is empty." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_column_renames_that_collide_list_the_duplicates_sorted(client): + dataset_id = _new_dataset_row(client, "net_dup_renames") + + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + _csv_params( + "net_dup_renames", + inferred_types=IRIS_SCHEMA, + column_renames={ + "SepalWidthCm": "Same", + "PetalWidthCm": "Same", + "Species": "SepalLengthCm", + }, + ), + ) + + assert str(excinfo.value) == ( + "Error loading dataset: Invalid column_renames: resulting column names " + "contain duplicates: ['Same', 'SepalLengthCm']" + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_database_error_on_the_final_commit_is_reported_generically(client): + """The last block (lines 351-368) catches only ``SQLAlchemyError`` and + replaces it with a fixed message, dropping the database's own text. + + Note the asymmetry with everything above: this one is raised *outside* the + try at line 153, so it is not prefixed with "Error loading dataset: ". + """ + from sqlalchemy import exc as sa_exc + + dataset_id = _new_dataset_row(client, "net_db_error") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + real_set_finished = Dataset.set_status_as_finished + + def boom(self): + raise sa_exc.InvalidRequestError("the database said no") + + Dataset.set_status_as_finished = boom + try: + with pytest.raises(JobError) as excinfo: + _run_job(client, dataset_id, _csv_params("net_db_error")) + finally: + Dataset.set_status_as_finished = real_set_finished + + assert str(excinfo.value) == "Internal database error" + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + # This branch cleans up the folder too. + assert _folders_in(datasets_dir) == before + + +# --------------------------------------------------------------------------- # +# side effects of a failure +# --------------------------------------------------------------------------- # + + +def test_a_failure_after_the_folder_was_created_removes_it(client): + """The generated folder is created before loading, and the ``except`` at + line 345 has to clean it up — otherwise every failed import leaks a folder + under DATASETS_PATH.""" + dataset_id = _new_dataset_row(client, "net_leak_check") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + with pytest.raises(JobError): + _run_job( + client, + dataset_id, + _csv_params("net_leak_check", dataloader="NoSuchDataLoader"), + ) + + assert _folders_in(datasets_dir) == before + assert _stored(client, dataset_id)["file_path"] == "" + + +def test_a_failure_does_not_delete_a_folder_the_job_did_not_create(client, tmp_path): + """Re-importing into an existing dataset reuses that dataset's own folder + instead of creating one (the ``n_sample`` branch). + + The cleanup on failure removes the destination folder, which is right when + the job created it and destructive when it did not: the row survives the + failure pointing at a path whose contents would be gone. + """ + existing = tmp_path / "already_stored" + existing.mkdir() + (existing / "dataset").mkdir() + (existing / "dataset" / "data.arrow").write_text("payload", encoding="utf-8") + + with _session_factory(client)() as db: + entry = Dataset(name="net_reused_folder", file_path=str(existing)) + entry.set_status_as_delivered() + db.add(entry) + db.commit() + db.refresh(entry) + dataset_id = entry.id + + with pytest.raises(JobError): + _run_job( + client, + dataset_id, + _csv_params("net_reused_folder", dataloader="NoSuchDataLoader"), + n_sample=10, + ) + + assert (existing / "dataset" / "data.arrow").read_text( + encoding="utf-8" + ) == "payload" + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_the_temp_dir_is_removed_even_on_the_happy_path(client, tmp_path_factory): + """The ``finally`` removes ``temp_dir`` whether the job worked or not.""" + dataset_id = _new_dataset_row(client, "net_temp_happy") + temp_dir = tmp_path_factory.mktemp("net_temp_happy") + + _run_job(client, dataset_id, _csv_params("net_temp_happy"), temp_dir=str(temp_dir)) + + assert not temp_dir.exists() + assert _stored(client, dataset_id)["status"] == DatasetStatus.FINISHED + + _cleanup(client, dataset_id) + + +def test_the_temp_dir_is_removed_after_a_failure(client, tmp_path_factory): + dataset_id = _new_dataset_row(client, "net_temp_failure") + temp_dir = tmp_path_factory.mktemp("net_temp_failure") + + with pytest.raises(JobError): + _run_job( + client, + dataset_id, + _csv_params("net_temp_failure", dataloader="NoSuchDataLoader"), + temp_dir=str(temp_dir), + ) + + assert not temp_dir.exists() + + +def test_a_failed_notebook_import_leaves_the_notebook_intact(client, notebook): + """The cleanup ``rmtree`` targets the *generated* folder. If it ever pointed + at the loaded path instead, this is what would catch it: the notebook's own + working copy would be gone.""" + dataset_id = _new_dataset_row(client, "net_notebook_failure") + notebook_dataset = Path(notebook["file_path"]) / "dataset" + assert notebook_dataset.exists() + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_save = dashai_dataset.save_dataset + + def fail_to_save(dataset, path): + raise OSError("disk full") + + dashai_dataset.save_dataset = fail_to_save + try: + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + {"name": "net_notebook_failure"}, + notebook_id=notebook["id"], + file_path=None, + ) + finally: + dashai_dataset.save_dataset = real_save + + # The message gained the destination path, and deliberately kept the + # original exception's text: a save failure is an infrastructure failure and + # only the outermost message reaches the user (``huey_job_queue.py:210`` + # stores ``str(exc)``, never the ``__cause__`` chain), so swallowing it would + # drop the diagnosis exactly when it matters. + message = str(excinfo.value) + assert message.startswith("Error loading dataset: Can not save dataset to path ") + assert message.endswith(": disk full") + + assert notebook_dataset.exists() + assert load_dataset(str(notebook_dataset)).column_names == IRIS_COLUMNS + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR diff --git a/tests/back/api/test_explainer_job.py b/tests/back/api/test_explainer_job.py new file mode 100644 index 000000000..cb2485fdd --- /dev/null +++ b/tests/back/api/test_explainer_job.py @@ -0,0 +1,723 @@ +"""End-to-end regression net for ``ExplainerJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so that the refactor has something to be measured +against. The assertions are deliberately explicit — exact status values, exact +columns written on the row, exact error message fragments, and which message is +the one the user actually sees versus which survives only as ``__cause__`` — +instead of the looser ``status in [1, 3]`` style used by +``test_explainer_jobs.py``, which cannot tell a unit that silently stopped doing +part of its work from one that did it. + +Tests named ``test_currently_*`` pin behaviour that is known to be wrong. They +exist so the refactor can be proven behaviour-preserving first; the fix lands +afterwards as its own change, which flips the assertion and renames the test. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import shutil +from pathlib import Path + +import joblib +import pytest +from datasets import ClassLabel, Value +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import ExplainerStatus +from DashAI.back.dependencies.database.models import ( + Dataset, + GlobalExplainer, + LocalExplainer, + ModelSession, + Run, +) +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.explainability.global_explainer import BaseGlobalExplainer +from DashAI.back.explainability.local_explainer import BaseLocalExplainer +from DashAI.back.job.base_job import JobError +from DashAI.back.job.explainer_job import ExplainerJob +from DashAI.back.models.base_model import BaseModel +from DashAI.back.splitters.holdout import HoldoutSplitter +from DashAI.back.tasks.base_task import BaseTask + +INPUT_COLUMNS = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"] +OUTPUT_COLUMNS = ["Species"] + +SPLIT_INDEXES = json.dumps( + { + "train_indexes": [0, 1, 2, 3, 4], + "test_indexes": [5, 6, 7, 8], + "val_indexes": [9, 10, 11, 12], + } +) + +SPLITS = json.dumps( + { + # The session names the splitter that produced its partitions, which is + # how every consumer of Run.split_indexes tells a holdout payload from + # a fold one. + "splitter_name": "HoldoutSplitter", + "splitType": "random", + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } +) + + +class DummyTask(BaseTask): + name: str = "DummyTask" + metadata: dict = { + "inputs_types": [ClassLabel, Value], + "outputs_types": [ClassLabel], + "inputs_cardinality": "n", + "outputs_cardinality": 1, + } + + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + +class DummyModel(BaseModel): + COMPATIBLE_COMPONENTS = ["DummyTask"] + + @classmethod + def get_schema(cls): + return {} + + def save(self, filename): + joblib.dump(self, filename) + + @staticmethod + def load(filename): + return DummyModel() + + def predict(self, x): + return {} + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + return + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + def prepare_output(self, dataset, is_fit=False): + return dataset + + +class UninstantiableModel(DummyModel): + def __init__(self, *args, **kwargs): + raise RuntimeError("this model refuses to be built") + + +class UnloadableModel(DummyModel): + @staticmethod + def load(filename): + raise OSError("the artifact is not there") + + +class DummyGlobalExplainer(BaseGlobalExplainer): + COMPATIBLE_COMPONENTS = ["DummyTask"] + + def __init__(self, model: BaseModel) -> None: + self.model = model + self.explanation = None + + @classmethod + def get_schema(cls): + return {} + + def explain(self, dataset): + return {"importance": [1, 2, 3]} + + def plot(self, explanation): + return "a plot" + + +class ExplodingGlobalExplainer(DummyGlobalExplainer): + def explain(self, dataset): + raise RuntimeError("the explanation itself blew up") + + +class UninstantiableGlobalExplainer(DummyGlobalExplainer): + def __init__(self, model: BaseModel) -> None: + raise RuntimeError("this explainer refuses to be built") + + +class DummyLocalExplainer(BaseLocalExplainer): + COMPATIBLE_COMPONENTS = ["DummyTask"] + + def __init__(self, model: BaseModel) -> None: + self.model = model + self.explanation = None + + @classmethod + def get_schema(cls): + return {} + + def fit(self, dataset, **kwargs): + return self + + def explain_instance(self, instances): + return {"local": True} + + def plot(self, explanation): + return "a plot" + + +@pytest.fixture(autouse=True, name="test_registry") +def setup_test_registry(client, monkeypatch: pytest.MonkeyPatch): + container = client.app.container + + test_registry = ComponentRegistry( + initial_components=[ + DummyTask, + DummyModel, + UninstantiableModel, + UnloadableModel, + DummyGlobalExplainer, + ExplodingGlobalExplainer, + UninstantiableGlobalExplainer, + DummyLocalExplainer, + ExplainerJob, + # The job asks the splitter that produced the run which partitions + # it has, so the one the session names has to be resolvable. + HoldoutSplitter, + ] + ) + + monkeypatch.setitem(container._services, "component_registry", test_registry) + return test_registry + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session(client: TestClient, dataset_1: Dataset): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="ExplainerJobSession", + task_name="DummyTask", + input_columns=INPUT_COLUMNS, + output_columns=OUTPUT_COLUMNS, + evaluation_strategy="HoldoutEvaluationStrategy", + splits=SPLITS, + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + return model_session.id + + +@pytest.fixture(name="run_id") +def create_run(client: TestClient, model_session_id: int): + """Function scoped: the error-branch tests corrupt this row on purpose.""" + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + optimizer_name="OptunaOptimizer", + optimizer_parameters={}, + model_name="DummyModel", + parameters={}, + goal_metric="Accuracy", + name="ExplainerJobRun", + run_path="a/saved/model", + split_indexes=SPLIT_INDEXES, + ) + db.add(run) + db.commit() + db.refresh(run) + return run.id + + +def _create_global_explainer(client, run_id, explainer_name="DummyGlobalExplainer"): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explainer = GlobalExplainer( + run_id=run_id, + explainer_name=explainer_name, + parameters={}, + ) + db.add(explainer) + db.commit() + db.refresh(explainer) + return explainer.id + + +def _create_local_explainer(client, run_id, dataset_id, scope=None): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explainer = LocalExplainer( + run_id=run_id, + explainer_name="DummyLocalExplainer", + dataset_id=dataset_id, + scope=scope if scope is not None else {"split": "test", "percentage": 100}, + parameters={}, + fit_parameters={}, + ) + db.add(explainer) + db.commit() + db.refresh(explainer) + return explainer.id + + +def _stored(client, model, explainer_id): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + row = db.get(model, explainer_id) + stored = { + "status": row.status, + "explanation_path": row.explanation_path, + "plot_overrides": row.plot_overrides, + "huey_id": row.huey_id, + } + if model is GlobalExplainer: + stored["plot_path"] = row.plot_path + else: + stored["plots_path"] = row.plots_path + stored["input_dataset_path"] = row.input_dataset_path + return stored + + +# --- happy paths -------------------------------------------------------- + + +def test_a_global_explanation_writes_both_pickles_and_finishes(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + stored = _stored(client, GlobalExplainer, explainer_id) + assert stored["status"] == ExplainerStatus.FINISHED + assert Path(stored["explanation_path"]).name == ( + f"global_explanation_{explainer_id}.pickle" + ) + assert Path(stored["plot_path"]).name == ( + f"global_explanation_plot_{explainer_id}.pickle" + ) + assert Path(stored["explanation_path"]).exists() + assert Path(stored["plot_path"]).exists() + # Overrides belong to a previous result and must not survive a re-run. + assert stored["plot_overrides"] is None + + +def test_a_local_explanation_writes_its_three_paths_and_finishes( + client, run_id, dataset_1 +): + """The local row carries an extra artifact the global one does not: the + selected instances, saved so the frontend can read them back.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + explainer_id = _create_local_explainer(client, run_id, dataset_1.id) + + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + stored = _stored(client, LocalExplainer, explainer_id) + assert stored["status"] == ExplainerStatus.FINISHED + assert Path(stored["explanation_path"]).name == ( + f"local_explanation_{explainer_id}.pickle" + ) + assert Path(stored["plots_path"]).name == ( + f"local_explanation_plots_{explainer_id}.pickle" + ) + assert Path(stored["explanation_path"]).exists() + assert Path(stored["plots_path"]).exists() + assert stored["plot_overrides"] is None + + saved_input = load_dataset(str(Path(stored["input_dataset_path"]) / "dataset")) + assert saved_input.column_names == INPUT_COLUMNS + # scope percentage 100 over the four test indexes. + assert len(saved_input) == 4 + + +def test_a_rows_scope_explains_exactly_the_marked_rows(client, run_id, dataset_1): + """Row indexes address the whole dataset; the split does not apply.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + explainer_id = _create_local_explainer( + client, run_id, dataset_1.id, scope={"mode": "rows", "row_indexes": [0, 7, 42]} + ) + + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + stored = _stored(client, LocalExplainer, explainer_id) + assert stored["status"] == ExplainerStatus.FINISHED + + saved_input = load_dataset(str(Path(stored["input_dataset_path"]) / "dataset")) + assert len(saved_input) == 3 + + +# --- scope selection ---------------------------------------------------- + + +def test_an_invalid_scope_is_rejected_before_anything_is_touched(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + with pytest.raises(JobError, match="banana is an invalid explainer type"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="banana").run() + + # Nothing ran, so the row is untouched. + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.NOT_STARTED + ) + + +def test_a_missing_explainer_row_is_reported_by_id(client): + """The missing row must be named, not crash the handler meant to mark it. + + The row is looked up before the guarded block, because the outer + ``except Exception`` calls ``set_status_as_error`` on that very row — so + without the check a bad id used to surface as an ``AttributeError`` raised + by the error handler itself. + """ + with pytest.raises( + JobError, match="Explainer with id 999999 does not exist in DB." + ): + ExplainerJob(explainer_id=999999, explainer_scope="global").run() + + +def test_the_huey_id_is_recorded_on_the_row(client, run_id): + """Like the other three jobs, the queue task id is stored on the row. + + Without it the explanation cannot be matched back to its queue entry. + """ + explainer_id = _create_global_explainer(client, run_id) + + ExplainerJob( + explainer_id=explainer_id, explainer_scope="global", huey_id="task-abc" + ).run() + + assert _stored(client, GlobalExplainer, explainer_id)["huey_id"] == "task-abc" + + +# --- loading errors ----------------------------------------------------- + + +def test_a_missing_run_is_reported_and_the_row_goes_to_error(client): + explainer_id = _create_global_explainer(client, run_id=999999) + + with pytest.raises(JobError, match="Run 999999 does not exist in DB."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + # STARTED is set late, after everything is loaded, so a failure here takes + # the row straight from NOT_STARTED to ERROR. + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_missing_model_session_is_reported_by_id(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).model_session_id = 999999 + db.commit() + + with pytest.raises(JobError, match="Model session 999999 does not exist in DB."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_missing_training_dataset_names_the_id_that_was_looked_up( + client, run_id, model_session_id +): + """The message names the dataset the lookup actually used. + + It used to interpolate ``self.explainer_db.dataset_id`` while looking up + ``model_session.dataset_id`` — and that column exists on ``LocalExplainer`` + and *not* on ``GlobalExplainer``, so in the global scope the "does not + exist" error was never built at all: an ``AttributeError`` was raised while + formatting it. + """ + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = 999999 + db.commit() + try: + with pytest.raises(JobError, match="Dataset 999999 does not exist in DB."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + finally: + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = ( + db.query(Dataset).filter(Dataset.name == "test_csv_1").first().id + ) + db.commit() + + +def test_an_unknown_model_name_is_reported_by_name(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).model_name = "NoSuchModel" + db.commit() + + with pytest.raises( + JobError, match="Unable to find Model with name NoSuchModel in registry." + ): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_model_that_cannot_be_instantiated_is_reported(client, run_id): + """The job builds the model before loading it, unlike ``PredictJob``.""" + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).model_name = "UninstantiableModel" + db.commit() + + with pytest.raises(JobError, match="Unable to instantiate model"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + +def test_a_model_that_cannot_be_loaded_names_the_path(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = db.get(Run, run_id) + run.model_name = "UnloadableModel" + run.run_path = "gone/from/disk" + db.commit() + + with pytest.raises(JobError, match="Can not load model from path gone/from/disk"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + +def test_an_unknown_explainer_name_is_reported_with_the_multiline_message( + client, run_id +): + """One line, so it stays readable wherever it surfaces. + + It used to be a triple-quoted f-string, which put a newline and the source + file's indentation literally inside the text the user reads. + """ + explainer_id = _create_global_explainer( + client, run_id, explainer_name="NoSuchExplainer" + ) + + with pytest.raises(JobError) as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert str(excinfo.value) == ( + "Unable to find the global explainer with name NoSuchExplainer in registry." + ) + + +def test_an_explainer_that_cannot_be_instantiated_names_the_scope(client, run_id): + explainer_id = _create_global_explainer( + client, run_id, explainer_name="UninstantiableGlobalExplainer" + ) + + with pytest.raises(JobError, match="Unable to instantiate global explainer."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + +def test_a_dataset_that_cannot_be_loaded_names_the_path( + client, run_id, dataset_1, tmp_path +): + explainer_id = _create_global_explainer(client, run_id) + + stored_folder = Path(dataset_1.file_path) / "dataset" + backup = tmp_path / "explainer-dataset-backup" + shutil.copytree(stored_folder, backup) + shutil.rmtree(stored_folder) + try: + with pytest.raises(JobError, match="Can not load dataset from path"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + finally: + shutil.copytree(backup, stored_folder) + + +def test_an_unknown_task_name_is_reported_by_name(client, run_id, model_session_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "NoSuchTask" + db.commit() + try: + with pytest.raises( + JobError, match="Unable to find Task with name NoSuchTask in registry" + ): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + finally: + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "DummyTask" + db.commit() + + +def test_incomplete_split_indexes_name_the_splitter_they_disagree_with( + client, run_id, dataset_1 +): + """A payload that does not match its splitter says so, and stops there. + + It used to reach the generic preparation wrapper, which told the user that + the dataset could not be prepared -- true, but not the reason, and it reads + as a problem with the dataset rather than with the run's own record of how + it was split. The partitions are now resolved by asking the splitter that + produced the run, before anything is prepared, so the message is the + mismatch itself. + """ + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).split_indexes = json.dumps({"train_indexes": [0, 1]}) + db.commit() + + with pytest.raises(JobError) as error: + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert str(error.value) == ( + "The run's split indexes do not match the splitter that produced it, " + "so there is no data to explain." + ) + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +# --- generation errors -------------------------------------------------- + + +def test_a_failing_global_explanation_is_reported_and_errors(client, run_id): + explainer_id = _create_global_explainer( + client, run_id, explainer_name="ExplodingGlobalExplainer" + ) + + with pytest.raises(JobError, match="Failed to generate the explanation") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert "the explanation itself blew up" in str(excinfo.value.__cause__) + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_an_invalid_split_is_swallowed_by_the_preparation_wrapper( + client, run_id, dataset_1 +): + """The specific complaint never reaches the user. + + ``"notasplit is not a valid split"`` is raised inside the block whose + ``except Exception`` replaces it with the generic wrapper, so it survives + only as ``__cause__``. Locking this in because it is an easy detail to + "fix" by accident while refactoring. + """ + explainer_id = _create_local_explainer( + client, run_id, dataset_1.id, scope={"split": "notasplit", "percentage": 100} + ) + + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert "notasplit is not a valid split" in str(excinfo.value.__cause__) + assert _stored(client, LocalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_rows_scope_with_no_valid_index_is_swallowed_by_the_same_wrapper( + client, run_id, dataset_1 +): + explainer_id = _create_local_explainer( + client, + run_id, + dataset_1.id, + scope={"mode": "rows", "row_indexes": [10**9]}, + ) + + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert "No valid row indexes provided for the explanation" in str( + excinfo.value.__cause__ + ) + + +def test_a_manual_scope_with_no_rows_is_swallowed_by_the_same_wrapper( + client, run_id, dataset_1 +): + explainer_id = _create_local_explainer( + client, run_id, dataset_1.id, scope={"mode": "manual"} + ) + + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert "No manual input data provided for the explanation" in str( + excinfo.value.__cause__ + ) + + +def test_a_missing_instance_dataset_is_reported_by_id(client, run_id, dataset_1): + explainer_id = _create_local_explainer(client, run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(LocalExplainer, explainer_id).dataset_id = 999999 + db.commit() + + with pytest.raises( + JobError, match="Dataset 999999 to be explained does not exist in DB." + ): + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert _stored(client, LocalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +# --- delivery ----------------------------------------------------------- + + +def test_set_status_as_delivered_marks_the_right_row(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + ExplainerJob( + explainer_id=explainer_id, explainer_scope="global" + ).set_status_as_delivered() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.DELIVERED + ) + + +def test_set_status_as_delivered_rejects_an_invalid_scope(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + with pytest.raises(JobError, match="banana is an invalid explainer type"): + ExplainerJob( + explainer_id=explainer_id, explainer_scope="banana" + ).set_status_as_delivered() diff --git a/tests/back/api/test_explorer_job.py b/tests/back/api/test_explorer_job.py new file mode 100644 index 000000000..b591a06b3 --- /dev/null +++ b/tests/back/api/test_explorer_job.py @@ -0,0 +1,277 @@ +"""End-to-end regression net for ``ExplorerJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so that the refactor has something to be measured +against. The assertions are deliberately explicit — exact status values, exact +files on disk, exact error message fragments — instead of the looser +``status in ["finished", "error"]`` style used elsewhere in this suite, which +cannot tell a unit that silently stopped doing part of its work from one that +did it. + +``ExplorerJob`` had no tests at all before this file. + +Tests named ``test_currently_*`` pin behaviour that is known to be wrong. They +exist so the refactor can be proven behaviour-preserving first; the fix lands +afterwards as its own change, which flips the assertion and renames the test. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import pathlib +import shutil + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import ExplorerStatus +from DashAI.back.dependencies.database.models import Explorer +from DashAI.back.job.base_job import JobError +from DashAI.back.job.explorer_job import ExplorerJob + +#: ``DescribeExplorer.__init__`` reads these three keys unconditionally, and the +#: schema declares no defaults, so a valid configuration always carries all of +#: them. +DESCRIBE_PARAMETERS = {"percentiles": "25, 50, 75", "include": "all", "exclude": None} + +SEPAL_LENGTH = [{"columnName": "SepalLengthCm"}] + + +@pytest.fixture(name="notebook") +def create_notebook(client: TestClient, dataset_1): + """A notebook holding its own copy of the iris dataset. + + ``POST /notebook/`` copies the dataset folder, so an exploration always + reads the notebook's copy and never the source dataset. + """ + response = client.post( + "/api/v1/notebook/", + json={"dataset_id": dataset_1.id, "name": "explorer job test"}, + ) + assert response.status_code == 201, response.text + return response.json() + + +def _create_explorer(client, notebook_id, columns=None, parameters=None): + """Create an Explorer row through the API, which validates it.""" + response = client.post( + "/api/v1/explorer/", + json={ + "notebook_id": notebook_id, + "exploration_type": "DescribeExplorer", + "columns": columns if columns is not None else SEPAL_LENGTH, + "parameters": ( + parameters if parameters is not None else DESCRIBE_PARAMETERS + ), + }, + ) + assert response.status_code == 201, response.text + return response.json()["id"] + + +def _insert_explorer_row(client, **fields): + """Insert an Explorer row straight into the database. + + ``POST /explorer/`` validates the exploration type, the parameters and the + columns, so the branches of ``run()`` that react to an invalid row can only + be reached by writing the row directly. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explorer = Explorer(**fields) + db.add(explorer) + db.commit() + db.refresh(explorer) + return explorer.id + + +def _stored_explorer(client, explorer_id): + """Read the Explorer row straight from the database.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explorer = db.get(Explorer, explorer_id) + return { + "status": explorer.status, + "exploration_path": explorer.exploration_path, + "start_time": explorer.start_time, + "end_time": explorer.end_time, + } + + +def _notebook_folder(client, notebook_id): + return pathlib.Path(client.app.container["config"]["NOTEBOOK_PATH"]) / str( + notebook_id + ) + + +def test_the_notebook_starts_as_a_readable_copy_of_the_dataset(client, notebook): + """Guards the fixture itself: the assertions below mean nothing if the + notebook copy is not a loadable iris dataset.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + dataset = load_dataset(f"{notebook['file_path']}/dataset") + + assert "SepalLengthCm" in dataset.column_names + assert len(dataset) == 150 + + +def test_explorer_job_writes_the_result_and_finishes(client, notebook): + """The happy path, end to end: status transitions and the file on disk. + + ``DescribeExplorer`` writes ``{explorer_id}.json`` under the notebook's own + folder, and the row records that exact path. + """ + explorer_id = _create_explorer(client, notebook["id"]) + + ExplorerJob(explorer_id=explorer_id).run() + + stored = _stored_explorer(client, explorer_id) + assert stored["status"] == ExplorerStatus.FINISHED + assert stored["start_time"] is not None + assert stored["end_time"] is not None + + expected = _notebook_folder(client, notebook["id"]) / f"{explorer_id}.json" + assert expected.exists() + assert stored["exploration_path"] == expected.as_posix() + + +def test_two_explorations_on_one_notebook_keep_separate_files(client, notebook): + """The save path is keyed by explorer id, so runs never overwrite each + other's result. Any decomposition has to keep that key.""" + first = _create_explorer(client, notebook["id"]) + second = _create_explorer( + client, notebook["id"], columns=[{"columnName": "PetalWidthCm"}] + ) + + ExplorerJob(explorer_id=first).run() + ExplorerJob(explorer_id=second).run() + + first_path = _stored_explorer(client, first)["exploration_path"] + second_path = _stored_explorer(client, second)["exploration_path"] + + assert first_path != second_path + assert pathlib.Path(first_path).exists() + assert pathlib.Path(second_path).exists() + + +def test_a_missing_explorer_row_reports_it_by_id(client): + with pytest.raises(JobError, match="Explorer with id 999999 not found."): + ExplorerJob(explorer_id=999999).run() + + +def test_a_missing_notebook_leaves_the_row_in_error(client, notebook): + """A notebook that is gone must not leave the exploration stuck in STARTED. + + The "not found" error used to be raised inside a ``try`` whose only handler + was ``except exc.SQLAlchemyError``, so ``set_status_as_error`` never ran and + the UI showed the exploration as still running. Nothing else would have + fixed it: the Huey error signal writes only to its own ``task_copy`` table + and never touches the ``Explorer`` row, and ``_execute_base_job`` calls + ``job.run()`` with no handler at all. + + The message still has to be the specific one, not a generic wrapper. + """ + explorer_id = _create_explorer(client, notebook["id"]) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explorer = db.get(Explorer, explorer_id) + explorer.notebook_id = 999999 + db.commit() + + with pytest.raises(JobError, match="Notebook with id 999999 not found."): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_a_dataset_that_cannot_be_loaded_leaves_the_row_in_error(client, notebook): + """A load failure must not leave the explorer stuck in STARTED.""" + explorer_id = _create_explorer(client, notebook["id"]) + shutil.rmtree(f"{notebook['file_path']}/dataset") + + with pytest.raises(JobError, match="Can not load dataset from path"): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_an_unknown_exploration_type_reports_it_and_errors(client, notebook): + """The registry lookup error names the culprit and reaches the user intact. + + Unlike ``ConverterJob``, nothing wraps this message on the way out. + """ + explorer_id = _insert_explorer_row( + client, + notebook_id=notebook["id"], + exploration_type="ThisExplorerDoesNotExist", + columns=SEPAL_LENGTH, + parameters={}, + ) + + with pytest.raises( + JobError, + match="Explorer ThisExplorerDoesNotExist not found in the registry.", + ): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_parameters_the_explorer_cannot_accept_report_an_instancing_error( + client, notebook +): + """``DescribeExplorer.__init__`` reads its three keys unconditionally.""" + explorer_id = _insert_explorer_row( + client, + notebook_id=notebook["id"], + exploration_type="DescribeExplorer", + columns=SEPAL_LENGTH, + parameters={}, + ) + + with pytest.raises( + JobError, match="Error instancing the explorer DescribeExplorer." + ): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_a_column_absent_from_the_dataset_reports_a_preparation_error(client, notebook): + """``prepare_dataset`` selects the requested columns and fails loudly. + + The column list is resolved against the dataset the job just loaded, which + is the behaviour any decomposition has to keep: nothing may resolve these + names ahead of time against a different dataset. + """ + explorer_id = _insert_explorer_row( + client, + notebook_id=notebook["id"], + exploration_type="DescribeExplorer", + columns=[{"columnName": "ThisColumnDoesNotExist"}], + parameters=DESCRIBE_PARAMETERS, + ) + + with pytest.raises( + JobError, + match="Error preparing the dataset for the exploration DescribeExplorer.", + ): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_set_status_as_delivered_marks_the_row(client, notebook): + """The enqueue path marks the row before the worker ever picks it up.""" + explorer_id = _create_explorer(client, notebook["id"]) + + ExplorerJob(explorer_id=explorer_id).set_status_as_delivered() + + stored = _stored_explorer(client, explorer_id) + assert stored["status"] == ExplorerStatus.DELIVERED + + +def test_set_status_as_delivered_reports_a_missing_row(client): + with pytest.raises(JobError, match="Explorer with id 999999 not found."): + ExplorerJob(explorer_id=999999).set_status_as_delivered() diff --git a/tests/back/api/test_model_job_as_a_graph.py b/tests/back/api/test_model_job_as_a_graph.py new file mode 100644 index 000000000..7a6240f2c --- /dev/null +++ b/tests/back/api/test_model_job_as_a_graph.py @@ -0,0 +1,549 @@ +"""ModelJob's whole shape, run as a graph. The test that says the design works. + +The spike that established the engine never ran this, and the three problems it +left open all live here: where run_id goes, whether the bundling rule survives +fifteen edges, and whether a training pipeline can produce metrics with no Run +row to hang them on. + +Six nodes, and the pipeline is a sandbox: it creates no Run and no ModelSession, +writes no Metric rows, and keeps its model out of the way of every real run's. + + load --dataset,dataset_id--> prep --x,y,n_labels,task_name--> build + | | + +-------x,y,task-------------> fit <-+ + | + +-------model-------+-------+ + v v + eval save +""" + +import json +import os + +import joblib +import pytest +from datasets import ClassLabel, Value +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.enums.status import NodeRunStatus, PipelineRunStatus +from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader +from DashAI.back.dependencies.database.models import ( + Dataset, + Metric, + ModelSession, + NodeArtifact, + NodeRun, + Pipeline, + PipelineRun, + Run, +) +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.job.base_job import JobError +from DashAI.back.job.pipeline_job import PipelineJob +from DashAI.back.metrics.base_metric import BaseMetric +from DashAI.back.models.base_model import BaseModel +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.splitters.holdout import HoldoutSplitter +from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.evaluate_model_to_artifact_unit import ( + EvaluateModelToArtifactUnit, +) +from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_model_unit import SaveModelUnit + +SPLITTER = { + "component": "HoldoutSplitter", + "params": { + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "shuffle": True, + "stratify": False, + "random_state": 42, + }, +} + + +class GraphTask(BaseTask): + name: str = "GraphTask" + metadata: dict = { + "inputs_types": [ClassLabel, Value], + "outputs_types": [ClassLabel], + "inputs_cardinality": "n", + "outputs_cardinality": 1, + } + + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + def num_labels(self, dataset, output_column): + return 3 + + +#: How many times a model asked to log metrics while training, across the run. +#: A model instance does not survive to be inspected -- the graph releases it -- +#: and the count is what makes "no metric rows" a consequence rather than a +#: coincidence. +LOG_ATTEMPTS = [] + + +class GraphModel(BaseModel): + """Logs metrics while training, the way the torch-based models do. + + That matters for what this file is testing. A model that never asks to log + would leave the Metric table empty no matter what the pipeline did, so the + sandbox would look airtight without being tested at all. This one asks on + every train call, so an empty Metric table means the switch stopped it. + """ + + COMPATIBLE_COMPONENTS = ["GraphTask"] + + def __init__(self, **kwargs): + self.trained_with = None + + def save(self, filename): + joblib.dump({"trained_with": self.trained_with}, filename) + + def load(self, filename): + return joblib.load(filename) + + def predict(self, x): + return [0] * x.shape[0] + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.trained_with = {"train": x_train.shape[0]} + LOG_ATTEMPTS.append(getattr(self, "run_id", "missing")) + self.calculate_metrics(split=SplitEnum.TRAIN, level=LevelEnum.STEP) + return self + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + def prepare_output(self, dataset, is_fit=False): + return [0] * len(LOG_ATTEMPTS or [0]) + + +class GraphMetric(BaseMetric): + COMPATIBLE_COMPONENTS = ["GraphTask"] + MAXIMIZE = True + + @staticmethod + def score(true_labels, probs_pred_labels): + return 0.5 + + +@pytest.fixture(scope="module", name="graph_registry", autouse=True) +def setup_graph_registry(client): + services = client.app.container._services + sentinel = object() + old = services.get("component_registry", sentinel) + + services["component_registry"] = ComponentRegistry( + initial_components=[ + GraphTask, + GraphModel, + GraphMetric, + CSVDataLoader, + OptunaOptimizer, + PipelineJob, + LoadDatasetUnit, + PrepareAndSplitUnit, + HoldoutSplitter, + BuildModelUnit, + FitModelUnit, + EvaluateModelToArtifactUnit, + SaveModelUnit, + ] + ) + yield services["component_registry"] + if old is sentinel: + del services["component_registry"] + else: + services["component_registry"] = old + + +def _model_job_blocks(dataset_id: int): + """The six blocks, and the six drawn edges that stand for fifteen wires.""" + steps = [ + { + "id": "load", + "units": [ + { + "id": "load", + "unit": "LoadDatasetUnit", + "config": {"dataset_id": dataset_id}, + } + ], + }, + { + "id": "prep", + "units": [ + { + "id": "prep", + "unit": "PrepareAndSplitUnit", + "config": { + "task_name": "GraphTask", + "input_columns": ["SepalLengthCm", "SepalWidthCm"], + "output_columns": ["Species"], + "splitter": SPLITTER, + }, + } + ], + }, + { + "id": "build", + "units": [ + { + "id": "build", + "unit": "BuildModelUnit", + "config": { + "model": {"component": "GraphModel", "params": {}}, + "train_metrics": ["GraphMetric"], + "validation_metrics": ["GraphMetric"], + "test_metrics": ["GraphMetric"], + }, + } + ], + }, + { + "id": "fit", + "units": [ + { + "id": "fit", + "unit": "FitModelUnit", + "config": { + "optimizer": { + "component": "OptunaOptimizer", + "params": {}, + }, + "goal_metric": "GraphMetric", + }, + } + ], + }, + { + "id": "eval", + "units": [ + { + "id": "eval", + "unit": "EvaluateModelToArtifactUnit", + "config": {"splits": ["TRAIN", "VALIDATION", "TEST"]}, + } + ], + }, + { + "id": "save", + "units": [{"id": "save", "unit": "SaveModelUnit", "config": {}}], + }, + ] + edges = [ + {"source": "load", "target": "prep"}, + {"source": "prep", "target": "build"}, + {"source": "prep", "target": "fit"}, + {"source": "build", "target": "fit"}, + {"source": "fit", "target": "eval"}, + {"source": "fit", "target": "save"}, + ] + return steps, edges + + +@pytest.fixture(name="finished_pipeline_run", scope="module") +def run_the_graph(client: TestClient, dataset_1: Dataset, graph_registry): + session_factory = client.app.container["session_factory"] + steps, edges = _model_job_blocks(dataset_1.id) + + with session_factory() as db: + pipeline = Pipeline(name="ModelJob as a graph", steps=steps, edges=edges) + db.add(pipeline) + db.commit() + pipeline_id = pipeline.id + + PipelineJob(pipeline_id=pipeline_id).run() + + with session_factory() as db: + pipeline_run = db.query(PipelineRun).filter_by(pipeline_id=pipeline_id).one() + db.refresh(pipeline_run) + # Read everything now: the session closes with the fixture. + yield { + "id": pipeline_run.id, + "status": pipeline_run.status, + "error_message": pipeline_run.error_message, + "steps": pipeline_run.steps, + "edges": pipeline_run.edges, + "nodes": { + row.node_id: { + "status": row.status, + "start_time": row.start_time, + "end_time": row.end_time, + "block_id": row.block_id, + "node_type": row.node_type, + "config": row.config, + } + for row in pipeline_run.node_runs + }, + "artifacts": { + (row.node_run.node_id, row.key): row.value + for row in db.query(NodeArtifact).all() + }, + } + + +def test_the_whole_graph_runs_to_completion(finished_pipeline_run): + assert finished_pipeline_run["status"] == PipelineRunStatus.FINISHED + assert finished_pipeline_run["error_message"] is None + + nodes = finished_pipeline_run["nodes"] + assert set(nodes) == {"load", "prep", "build", "fit", "eval", "save"} + for node_id, node in nodes.items(): + assert node["status"] == NodeRunStatus.FINISHED, node_id + assert node["start_time"] is not None, node_id + assert node["end_time"] is not None, node_id + + +def test_six_drawn_edges_expand_into_twelve_wires(finished_pipeline_run): + """The granularity problem, measured. + + The unit contract is finer than a canvas can draw: FitModelUnit alone + requires seven keys. One drawn edge carries every key its two units agree + on, which is what makes six shapes on a canvas enough for this graph. + + It was fifteen while BuildModelUnit still took the data. It stopped taking + it once a model could be fitted over folds, where the data changes on every + iteration and binding one partition at construction would leave the metrics + describing whichever fold happened to be built with. The two wires that + disappeared are ``x`` and ``y`` from prep to build; whoever fits the model + points it at the data now. + + The twelfth went when the hyperparameter search stopped being handed the + task. It was the sixth argument of ``optimize`` while the optimizer fitted + and scored inline; the objective is now a callable the fitting unit builds, + so the task is no longer anything the search needs to be told. + """ + edges = finished_pipeline_run["edges"] + assert len(edges) == 12 + + carried = {} + for edge in edges: + carried.setdefault((edge["src"], edge["dst"]), set()).add(edge["src_key"]) + + assert carried == { + ("load", "prep"): {"dataset", "dataset_id"}, + ("prep", "build"): {"n_labels", "task_name"}, + ("prep", "fit"): {"x", "y"}, + ("build", "fit"): { + "model", + "factory", + "optimizable_parameters", + "model_parameters", + }, + ("fit", "eval"): {"model"}, + ("fit", "save"): {"model"}, + } + + +def test_the_sandbox_writes_no_metric_rows(client, finished_pipeline_run): + """The assertion that decides whether the sandbox holds. + + Metric.run_id is a foreign key to run.id, and SQLite is not enforcing it + here -- there is no PRAGMA foreign_keys=ON anywhere. So a row written with + a run id that matches nothing inserts happily, and pipeline_run.id and + run.id are independent sequences that both start at 1: a pipeline metric + would land in the metric list of the real run with that id and show up in + its live chart. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + assert db.query(Metric).count() == 0 + + +def test_the_model_did_ask_to_log_and_was_stopped(finished_pipeline_run): + """What makes the empty Metric table a consequence, not a coincidence. + + GraphModel calls calculate_metrics on every train call, the way the + torch-based models do. So the run above genuinely tried to write metric + rows and wrote none, and it tried with a model that had no run -- which is + the only thing standing between it and a row keyed by nothing. + """ + assert LOG_ATTEMPTS, "the model never tried to log, so nothing was tested" + assert set(LOG_ATTEMPTS) == {None} + + +def test_the_sandbox_creates_no_run_and_no_model_session(client, finished_pipeline_run): + """No application entity is manufactured to make a pipeline work. + + A Run needs a ModelSession, which needs a globally unique name, a dataset, + a task, input and output columns and splits, all NOT NULL -- two entities + per execution, visible in the Models UI as sessions nobody created, with a + deletion lifecycle nobody owns. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + assert db.query(Run).count() == 0 + assert db.query(ModelSession).count() == 0 + + +def test_the_metrics_of_the_pipeline_live_in_an_artifact(finished_pipeline_run): + """Where a pipeline's metrics go, now that they cannot go to Metric.""" + metrics = finished_pipeline_run["artifacts"][("eval", "metrics")] + + assert metrics == { + "train": {"GraphMetric": 0.5}, + "validation": {"GraphMetric": 0.5}, + "test": {"GraphMetric": 0.5}, + } + + +def test_the_model_is_saved_where_no_real_run_could_collide_with_it( + client, finished_pipeline_run +): + """The artifact name, which is the other half of what run_id used to do.""" + model_path = finished_pipeline_run["artifacts"][("save", "model_path")] + runs_path = str(client.app.container["config"]["RUNS_PATH"]) + + assert model_path.startswith(runs_path) + assert os.path.exists(model_path) + + directory = os.path.basename(model_path) + # Not a bare integer, which is what every real run's directory is called. + assert not directory.isdigit() + assert directory == f"pipeline-{finished_pipeline_run['id']}-save" + + # And the model that landed there is the one that was trained. + assert joblib.load(model_path)["trained_with"]["train"] > 0 + + +def test_the_model_never_belonged_to_a_run(finished_pipeline_run): + """The switch that turns metric persistence off, as it reaches the node. + + ModelFactory hangs run_id on the model instance, and + BaseModel.calculate_metrics returns early for a model with no run. That is + what makes the sandbox airtight without the pipeline having to intercept + anything: not even a model that logs while training can write a row. + """ + assert finished_pipeline_run["nodes"]["build"]["config"]["run_id"] is None + assert finished_pipeline_run["nodes"]["fit"]["config"]["run_id"] is None + + +def test_the_run_records_units_not_canvas_shapes(finished_pipeline_run): + nodes = finished_pipeline_run["nodes"] + + assert nodes["build"]["node_type"] == "BuildModelUnit" + assert nodes["eval"]["node_type"] == "EvaluateModelToArtifactUnit" + # One unit per block for now, so each node is its own block. + for node_id, node in nodes.items(): + assert node["block_id"] == node_id + + +def test_every_serializable_output_is_recorded(finished_pipeline_run): + """Artifacts are named by the unit's own PROVIDES keys.""" + artifacts = finished_pipeline_run["artifacts"] + + assert ("load", "dataset_id") in artifacts + assert ("load", "dataset_path") in artifacts + assert ("prep", "split_indexes") in artifacts + assert ("prep", "task_name") in artifacts + assert ("build", "model_parameters") in artifacts + assert ("fit", "plot_paths") in artifacts + assert ("eval", "metrics") in artifacts + assert ("save", "model_path") in artifacts + + # The live objects are not: x, y, the task, the model and the factory are + # all cache-half values, derivable again and never serialized. + for key in ("dataset", "x", "y", "task", "model", "factory"): + assert not any(recorded == key for _, recorded in artifacts) + + +def test_a_node_failing_halfway_cancels_the_rest( + client: TestClient, dataset_1: Dataset, graph_registry +): + """A run that died halfway must not look like one still in flight.""" + session_factory = client.app.container["session_factory"] + steps, edges = _model_job_blocks(dataset_1.id) + + # A model that is not in the registry: build fails, and everything the + # order puts after it never runs. + for step in steps: + if step["id"] == "build": + step["units"][0]["config"]["model"] = { + "component": "NoSuchModel", + "params": {}, + } + + with session_factory() as db: + pipeline = Pipeline(name="A graph that fails", steps=steps, edges=edges) + db.add(pipeline) + db.commit() + pipeline_id = pipeline.id + + with pytest.raises(JobError): + PipelineJob(pipeline_id=pipeline_id).run() + + with session_factory() as db: + pipeline_run = db.query(PipelineRun).filter_by(pipeline_id=pipeline_id).one() + statuses = {row.node_id: row.status for row in pipeline_run.node_runs} + assert pipeline_run.status == PipelineRunStatus.ERROR + + assert statuses["load"] == NodeRunStatus.FINISHED + assert statuses["prep"] == NodeRunStatus.FINISHED + assert statuses["build"] == NodeRunStatus.ERROR + assert statuses["fit"] == NodeRunStatus.CANCELLED + assert statuses["eval"] == NodeRunStatus.CANCELLED + assert statuses["save"] == NodeRunStatus.CANCELLED + + +def test_two_runs_of_the_same_graph_do_not_overwrite_each_others_model( + client: TestClient, dataset_1: Dataset, graph_registry +): + """Each execution names its own artifacts. + + Without the run id in the name, the second execution would write over the + first one's model, silently. + """ + session_factory = client.app.container["session_factory"] + steps, edges = _model_job_blocks(dataset_1.id) + + with session_factory() as db: + pipeline = Pipeline(name="A graph run twice", steps=steps, edges=edges) + db.add(pipeline) + db.commit() + pipeline_id = pipeline.id + + PipelineJob(pipeline_id=pipeline_id).run() + PipelineJob(pipeline_id=pipeline_id).run() + + with session_factory() as db: + runs = ( + db.query(PipelineRun) + .filter_by(pipeline_id=pipeline_id) + .order_by(PipelineRun.id) + .all() + ) + paths = [ + db.query(NodeArtifact) + .join(NodeRun) + .filter( + NodeRun.pipeline_run_id == pipeline_run.id, + NodeArtifact.key == "model_path", + ) + .one() + .value + for pipeline_run in runs + ] + + assert len(paths) == 2 + assert paths[0] != paths[1] + assert all(os.path.exists(path) for path in paths) + + +def test_the_graph_definition_survives_being_stored_as_json( + client: TestClient, dataset_1: Dataset, graph_registry +): + """The blocks are plain data, which is what a canvas can save and reload.""" + steps, edges = _model_job_blocks(dataset_1.id) + + assert json.loads(json.dumps(steps)) == steps + assert json.loads(json.dumps(edges)) == edges diff --git a/tests/back/api/test_model_job_cross_validation.py b/tests/back/api/test_model_job_cross_validation.py new file mode 100644 index 000000000..a86dd861c --- /dev/null +++ b/tests/back/api/test_model_job_cross_validation.py @@ -0,0 +1,892 @@ +"""End to end regression net for the ModelJob cross-validation paths. + +``test_model_job_orchestration.py`` pins the holdout path. Nothing pins the +fold paths: the only end-to-end cross-validated run in the suite is a fixture +in ``test_reports_api.py`` that asserts the run reached ``FINISHED`` and +nothing else, and it configures no optimizer, so cross-validation combined with +a hyperparameter search never executes anywhere. + +These tests pin the observable contract of a cross-validated run: the shape of +the persisted split indexes, which rows each partition holds, how many times +the model is actually fitted and on what, the fold level ``Metric`` rows and +their aggregation, the single saved artifact, and the verbatim text of every +error branch. + +They deliberately assert exact values rather than ``status in [...]``: a run +that quietly stops doing part of its work still finishes. +""" + +import json +import os + +import joblib +import numpy as np +import pytest +from datasets import ClassLabel, Value +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.enums.status import RunStatus +from DashAI.back.core.schema_fields import BaseSchema, int_field, schema_field +from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader +from DashAI.back.dependencies.database.models import Dataset, Metric, ModelSession, Run +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.evaluation.cv import CrossValidationEvaluationStrategy +from DashAI.back.evaluation.holdout import HoldoutEvaluationStrategy +from DashAI.back.job.base_job import JobError +from DashAI.back.job.model_job import ModelJob +from DashAI.back.metrics.base_metric import BaseMetric +from DashAI.back.models.base_model import BaseModel +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.splitters.holdout import HoldoutSplitter +from DashAI.back.splitters.k_fold import KFoldSplitter +from DashAI.back.splitters.stratified_k_fold import StratifiedKFoldSplitter +from DashAI.back.tasks.base_task import BaseTask + +INPUT_COLUMNS = ["SepalLengthCm", "SepalWidthCm"] +OUTPUT_COLUMNS = ["Species"] + +#: Rows in ``tests/back/api/iris.csv``. +DATASET_ROWS = 150 + +#: Every ``train`` call any model instance makes, in order. A module level list +#: rather than instance state: cross-validation reuses one model instance across +#: folds, so counting on the instance could not tell several fits apart from one. +FIT_LOG: list = [] + + +def _squash(text: str) -> str: + """Collapse whitespace, so a triple quoted message can be compared whole. + + Several messages on this path are written as triple quoted f-strings, so + they carry the newline and the source indentation literally. Comparing the + squashed form pins the whole sentence instead of a fragment of it, without + pinning the indentation of the file it came from. + """ + return " ".join(text.split()) + + +class CvTask(BaseTask): + name: str = "CvTask" + metadata: dict = { + "inputs_types": [ClassLabel, Value], + "outputs_types": [ClassLabel], + "inputs_cardinality": "n", + "outputs_cardinality": 1, + } + + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + def num_labels(self, dataset, output_column): + return 3 + + +class CountingModel(BaseModel): + """Model that records every fit and predicts from its own input rows. + + ``predict`` returns the first input column verbatim, so the paired metric + below scores every partition differently. That matters for the aggregation + assertions: with a constant score the standard deviation is zero whether it + was computed or hardcoded. + """ + + COMPATIBLE_COMPONENTS = ["CvTask"] + + def __init__(self, **kwargs): + self.trained_with = None + + def save(self, filename): + joblib.dump({"trained_with": self.trained_with}, filename) + + def load(self, filename): + return joblib.load(filename) + + def predict(self, x): + column = x.column_names[0] + return [float(value) for value in x[column]] + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.trained_with = { + "train": x_train.shape[0], + "validation": None if x_validation is None else x_validation.shape[0], + } + FIT_LOG.append(dict(self.trained_with)) + return self + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + +class TunableSchema(BaseSchema): + n_estimators: schema_field( + int_field(gt=0), + placeholder=2, + description="Number of estimators.", + ) # type: ignore + + +class TunableCvModel(CountingModel): + """A ``CountingModel`` with one optimizable parameter, for the search branch.""" + + COMPATIBLE_COMPONENTS = ["CvTask"] + SCHEMA = TunableSchema + + def __init__(self, n_estimators=2, **kwargs): + super().__init__(**kwargs) + self.n_estimators = n_estimators + + +class MeanPredictionMetric(BaseMetric): + """Mean of the predicted values. + + Deterministic and different for every partition, without depending on how + the target column happens to be encoded. + """ + + COMPATIBLE_COMPONENTS = ["CvTask"] + MAXIMIZE = True + + @staticmethod + def score(true_labels, probs_pred_labels): + return float(np.mean(np.asarray(probs_pred_labels, dtype=float))) + + +@pytest.fixture(scope="module", name="cv_registry", autouse=True) +def setup_cv_registry(client): + container = client.app.container + sentinel = object() + services = container._services + old = services.get("component_registry", sentinel) + + services["component_registry"] = ComponentRegistry( + initial_components=[ + CvTask, + CountingModel, + TunableCvModel, + MeanPredictionMetric, + CSVDataLoader, + ModelJob, + OptunaOptimizer, + HoldoutSplitter, + HoldoutEvaluationStrategy, + KFoldSplitter, + StratifiedKFoldSplitter, + CrossValidationEvaluationStrategy, + ] + ) + yield services["component_registry"] + if old is sentinel: + del services["component_registry"] + else: + services["component_registry"] = old + + +# --------------------------------------------------------------------------- # +# Helpers +# --------------------------------------------------------------------------- # + + +def _make_session( + client: TestClient, + dataset_id: int, + name: str, + splits: dict, + strategy: str = "CrossValidationEvaluationStrategy", +) -> int: + session_factory = client.app.container["session_factory"] + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_id, + name=name, + task_name="CvTask", + input_columns=INPUT_COLUMNS, + output_columns=OUTPUT_COLUMNS, + train_metrics=["MeanPredictionMetric"], + validation_metrics=["MeanPredictionMetric"], + test_metrics=["MeanPredictionMetric"], + evaluation_strategy=strategy, + splits=json.dumps(splits), + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + return model_session.id + + +def _kfold_splits(n_splits: int = 3, test_size: float = 0.2, **extra) -> dict: + payload = { + "splitter_name": "KFoldSplitter", + "splitType": "cv", + "n_splits": n_splits, + "shuffle": True, + "random_state": 42, + "test_size": test_size, + } + payload.update(extra) + return payload + + +def _make_run(client: TestClient, model_session_id: int, **overrides) -> int: + session_factory = client.app.container["session_factory"] + fields = { + "model_name": "CountingModel", + "parameters": {}, + "optimizer_name": "", + "optimizer_parameters": {}, + "goal_metric": "", + "name": f"CvRun-{model_session_id}", + } + fields.update(overrides) + with session_factory() as db: + run = Run(model_session_id=model_session_id, **fields) + db.add(run) + db.commit() + db.refresh(run) + return run.id + + +def _run_job(run_id: int) -> list: + """Run the job and return the fits it performed, in order.""" + FIT_LOG.clear() + ModelJob(run_id=run_id).run() + return list(FIT_LOG) + + +def _reload(client: TestClient, run_id: int) -> Run: + session_factory = client.app.container["session_factory"] + with session_factory() as db: + return db.get(Run, run_id) + + +def _metrics(client: TestClient, run_id: int, **filters) -> list: + session_factory = client.app.container["session_factory"] + with session_factory() as db: + return db.query(Metric).filter_by(run_id=run_id, **filters).all() + + +def _fold_values(client: TestClient, run_id: int, level, split) -> list: + """Fold metric values for one level and split, ordered by fold index.""" + rows = _metrics(client, run_id, level=level, split=split) + return [row.value for row in sorted(rows, key=lambda row: row.fold_index)] + + +# --------------------------------------------------------------------------- # +# A plain cross-validated run: 3 folds over 120 rows, 30 reserved +# --------------------------------------------------------------------------- # + + +@pytest.fixture(scope="module", name="cv_run") +def run_a_cross_validated_job(client: TestClient, dataset_1: Dataset) -> dict: + session_id = _make_session( + client, dataset_1.id, "CvSession", _kfold_splits(n_splits=3, test_size=0.2) + ) + run_id = _make_run(client, session_id, name="PlainCvRun") + fits = _run_job(run_id) + return {"run_id": run_id, "fits": fits} + + +def test_a_cross_validated_run_reaches_finished(client: TestClient, cv_run: dict): + assert _reload(client, cv_run["run_id"]).status == RunStatus.FINISHED + + +def test_a_cross_validated_run_stamps_its_timestamps(client: TestClient, cv_run: dict): + run = _reload(client, cv_run["run_id"]) + + assert run.start_time is not None + assert run.end_time is not None + assert run.end_time >= run.start_time + + +def test_the_split_indexes_name_every_fold_and_the_full_dataset( + client: TestClient, cv_run: dict +): + """The shape a fold splitter persists, which is not the holdout shape. + + A holdout run stores one flat ``{train,test,val}_indexes`` mapping. Every + consumer of ``Run.split_indexes`` has to tell the two apart, which is what + ``BaseSplitter.explainable_partitions`` exists for. + """ + split_indexes = json.loads(_reload(client, cv_run["run_id"]).split_indexes) + + assert set(split_indexes) == {"fold_0", "fold_1", "fold_2", "full_dataset"} + + +def test_every_fold_names_a_train_and_a_validation_partition( + client: TestClient, cv_run: dict +): + split_indexes = json.loads(_reload(client, cv_run["run_id"]).split_indexes) + + for name in ("fold_0", "fold_1", "fold_2"): + assert set(split_indexes[name]) == {"train_indexes", "validation_indexes"}, name + + +def test_the_full_dataset_entry_names_a_train_and_a_test_partition( + client: TestClient, cv_run: dict +): + """The trailing entry is not a fold: it fits the kept model and scores it.""" + split_indexes = json.loads(_reload(client, cv_run["run_id"]).split_indexes) + + assert set(split_indexes["full_dataset"]) == {"train_indexes", "test_indexes"} + assert len(split_indexes["full_dataset"]["train_indexes"]) == 120 + assert len(split_indexes["full_dataset"]["test_indexes"]) == 30 + + +def test_the_reserved_rows_reach_no_fold(client: TestClient, cv_run: dict): + """The assertion the whole reserved-rows design exists for. + + If a reserved row leaked into any fold, the score reported on the reserved + rows would be measured on data the folds had already seen, and nothing + about that failure raises. + """ + split_indexes = json.loads(_reload(client, cv_run["run_id"]).split_indexes) + reserved = set(split_indexes["full_dataset"]["test_indexes"]) + + for name in ("fold_0", "fold_1", "fold_2"): + fold = split_indexes[name] + seen = set(fold["train_indexes"]) | set(fold["validation_indexes"]) + assert reserved.isdisjoint(seen), name + + +def test_every_pooled_row_is_validated_exactly_once(client: TestClient, cv_run: dict): + """K-fold's defining property, checked on the rows it actually produced.""" + split_indexes = json.loads(_reload(client, cv_run["run_id"]).split_indexes) + pool = split_indexes["full_dataset"]["train_indexes"] + + validated = [ + index + for name in ("fold_0", "fold_1", "fold_2") + for index in split_indexes[name]["validation_indexes"] + ] + + assert sorted(validated) == sorted(pool) + assert len(validated) == len(set(validated)) + + +def test_the_pool_and_the_reserved_rows_cover_the_dataset( + client: TestClient, cv_run: dict +): + split_indexes = json.loads(_reload(client, cv_run["run_id"]).split_indexes) + pool = set(split_indexes["full_dataset"]["train_indexes"]) + reserved = set(split_indexes["full_dataset"]["test_indexes"]) + + assert pool.isdisjoint(reserved) + assert pool | reserved == set(range(DATASET_ROWS)) + + +def test_the_model_is_fitted_once_per_fold_and_once_more_on_the_whole_pool( + cv_run: dict, +): + """Three fold fits and a final refit, and only the refit is kept.""" + fits = cv_run["fits"] + + assert len(fits) == 4 + assert [fit["train"] for fit in fits] == [80, 80, 80, 120] + + +def test_no_cross_validated_fit_receives_validation_data(cv_run: dict): + """A fold is scored on its validation partition, so it is not fitted on it. + + The holdout strategy does hand validation data to ``train`` (models use it + for early stopping); the fold strategy deliberately does not, and the + trailing entry has no validation partition to hand over at all. + """ + assert [fit["validation"] for fit in cv_run["fits"]] == [None, None, None, None] + + +def test_the_saved_artifact_is_the_refitted_model_and_not_a_fold( + client: TestClient, cv_run: dict +): + """One model is saved, and it is the one fitted on every pooled row.""" + run = _reload(client, cv_run["run_id"]) + + assert run.run_path is not None + assert run.run_path.endswith(str(run.id)) + assert os.path.exists(run.run_path) + assert joblib.load(run.run_path)["trained_with"]["train"] == 120 + + +def test_a_fold_metric_is_written_for_every_fold_and_every_scored_split( + client: TestClient, cv_run: dict +): + rows = _metrics(client, cv_run["run_id"], level=LevelEnum.FOLD) + + by_split = {} + for row in rows: + by_split.setdefault(row.split, []).append(row) + + assert set(by_split) == {SplitEnum.TRAIN, SplitEnum.VALIDATION} + for split, split_rows in by_split.items(): + indexes = sorted(row.fold_index for row in split_rows) + # Contiguous from zero: the fold charts bucket repeated runs with + # ``fold_index // n_splits``, so a gap silently reassigns a fold to the + # wrong repetition instead of raising. + assert indexes == [0, 1, 2], split + assert all(row.name == "MeanPredictionMetric" for row in split_rows) + assert all(row.std_value is None for row in split_rows) + + +def test_the_last_metric_is_the_mean_and_deviation_of_the_fold_metrics( + client: TestClient, cv_run: dict +): + """The aggregation, checked against the fold rows it aggregated.""" + for split in (SplitEnum.TRAIN, SplitEnum.VALIDATION): + folds = _fold_values(client, cv_run["run_id"], LevelEnum.FOLD, split) + # Otherwise the deviation below is zero whether it was computed or not. + assert len(set(folds)) > 1, split + + aggregated = _metrics( + client, cv_run["run_id"], level=LevelEnum.LAST, split=split + ) + assert len(aggregated) == 1, split + assert aggregated[0].value == pytest.approx(float(np.mean(folds))) + assert aggregated[0].std_value == pytest.approx(float(np.std(folds))) + + +def test_the_reserved_rows_are_scored_once_and_carry_no_deviation( + client: TestClient, cv_run: dict +): + """A single measurement, so there is nothing to take a deviation over.""" + rows = _metrics( + client, cv_run["run_id"], level=LevelEnum.LAST, split=SplitEnum.TEST + ) + + assert len(rows) == 1 + assert rows[0].std_value is None + assert rows[0].fold_index is None + + +def test_the_aggregated_rows_are_written_at_step_zero(client: TestClient, cv_run: dict): + """Pins an inconsistency rather than endorsing it. + + Fold rows get an auto-incremented step, and so does the reserved-rows + measurement, because both go through ``calculate_metrics``. The aggregated + rows are written directly with ``step=0``. Anything that orders metrics by + step sees the summary before the folds it summarises. + """ + aggregated = _metrics( + client, cv_run["run_id"], level=LevelEnum.LAST, split=SplitEnum.VALIDATION + ) + folds = _metrics( + client, cv_run["run_id"], level=LevelEnum.FOLD, split=SplitEnum.VALIDATION + ) + + assert [row.step for row in aggregated] == [0] + assert sorted(row.step for row in folds) == [1, 2, 3] + + +def test_the_inner_fold_index_is_never_written(client: TestClient, cv_run: dict): + """``Metric.inner_fold_index`` was added with cross-validation and is unused. + + Pinned so that a reconciliation either starts populating it deliberately or + drops it deliberately, rather than either happening by accident. + """ + rows = _metrics(client, cv_run["run_id"]) + + assert rows + assert all(row.inner_fold_index is None for row in rows) + + +# --------------------------------------------------------------------------- # +# A session that reserves nothing +# --------------------------------------------------------------------------- # + + +@pytest.fixture(scope="module", name="unreserved_run") +def run_without_reserved_rows(client: TestClient, dataset_1: Dataset) -> dict: + session_id = _make_session( + client, + dataset_1.id, + "CvSessionNoTest", + _kfold_splits(n_splits=3, test_size=0), + ) + run_id = _make_run(client, session_id, name="UnreservedCvRun") + fits = _run_job(run_id) + return {"run_id": run_id, "fits": fits} + + +def test_a_session_that_reserves_nothing_still_finishes( + client: TestClient, unreserved_run: dict +): + run = _reload(client, unreserved_run["run_id"]) + split_indexes = json.loads(run.split_indexes) + + assert run.status == RunStatus.FINISHED + assert split_indexes["full_dataset"]["test_indexes"] == [] + assert len(split_indexes["full_dataset"]["train_indexes"]) == DATASET_ROWS + + +def test_a_session_that_reserves_nothing_writes_no_test_metric( + client: TestClient, unreserved_run: dict +): + """With nothing held out there is no honest test score, so none is written.""" + rows = _metrics(client, unreserved_run["run_id"], split=SplitEnum.TEST) + + assert rows == [] + + +def test_a_session_that_reserves_nothing_refits_on_every_row(unreserved_run: dict): + fits = unreserved_run["fits"] + + assert [fit["train"] for fit in fits] == [100, 100, 100, DATASET_ROWS] + + +# --------------------------------------------------------------------------- # +# Cross-validation inside a hyperparameter search +# --------------------------------------------------------------------------- # + +HPO_FOLDS = 2 +HPO_TRIALS = 3 + + +@pytest.fixture(scope="module", name="hpo_run") +def run_a_search_over_folds(client: TestClient, dataset_1: Dataset) -> dict: + session_id = _make_session( + client, + dataset_1.id, + "CvSessionHpo", + _kfold_splits(n_splits=HPO_FOLDS, test_size=0.2), + ) + run_id = _make_run( + client, + session_id, + model_name="TunableCvModel", + parameters={ + "n_estimators": { + "optimize": True, + "lower_bound": 1, + "upper_bound": 5, + "fixed_value": 2, + } + }, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": HPO_TRIALS, + "sampler": "TPESampler", + "pruner": "None", + }, + goal_metric="MeanPredictionMetric", + name="HpoCvRun", + ) + fits = _run_job(run_id) + return {"run_id": run_id, "fits": fits} + + +def test_a_search_over_folds_reaches_finished(client: TestClient, hpo_run: dict): + assert _reload(client, hpo_run["run_id"]).status == RunStatus.FINISHED + + +def test_every_trial_fits_every_fold(hpo_run: dict): + """Cross-validation runs *inside* each trial, which is what it costs. + + The optimizer's objective is the whole fold loop, so one trial is k fits. + On top of the search come the scoring loop's k fits and the final refit. + """ + fits = hpo_run["fits"] + pooled = [fit for fit in fits if fit["train"] == 60] + refits = [fit for fit in fits if fit["train"] == 120] + + assert len(pooled) == HPO_TRIALS * HPO_FOLDS + HPO_FOLDS + assert len(refits) == 1 + assert len(fits) == len(pooled) + len(refits) + + +def test_a_trial_metric_is_written_once_per_trial_for_each_scored_split( + client: TestClient, hpo_run: dict +): + """One row per trial, holding the mean over that trial's folds.""" + rows = _metrics(client, hpo_run["run_id"], level=LevelEnum.TRIAL) + + by_split = {} + for row in rows: + by_split.setdefault(row.split, []).append(row) + + assert set(by_split) == {SplitEnum.TRAIN, SplitEnum.VALIDATION} + for split, split_rows in by_split.items(): + assert len(split_rows) == HPO_TRIALS, split + assert all(row.fold_index is None for row in split_rows), split + + +def test_a_search_over_folds_persists_the_best_parameters( + client: TestClient, hpo_run: dict +): + run = _reload(client, hpo_run["run_id"]) + best = run.parameters["n_estimators"]["fixed_value"] + + assert 1 <= best <= 5 + + +def test_a_search_over_folds_saves_two_plots_named_after_the_run( + client: TestClient, hpo_run: dict +): + run = _reload(client, hpo_run["run_id"]) + + assert run.plot_history_path is not None + assert run.plot_slice_path is not None + assert os.path.exists(run.plot_history_path) + assert os.path.exists(run.plot_slice_path) + assert os.path.basename(run.plot_history_path) == ( + f"history_objective_plot_{run.id}.pickle" + ) + # The contour and importance plots need more than one searched parameter. + assert run.plot_contour_path is None + assert run.plot_importance_path is None + + +def test_a_search_over_folds_still_aggregates_its_folds( + client: TestClient, hpo_run: dict +): + """The scoring loop runs after the search, so the summary is still written.""" + for split in (SplitEnum.TRAIN, SplitEnum.VALIDATION): + folds = _fold_values(client, hpo_run["run_id"], LevelEnum.FOLD, split) + aggregated = _metrics( + client, hpo_run["run_id"], level=LevelEnum.LAST, split=split + ) + + assert len(folds) == HPO_FOLDS, split + assert len(aggregated) == 1, split + assert aggregated[0].value == pytest.approx(float(np.mean(folds))) + assert aggregated[0].std_value == pytest.approx(float(np.std(folds))) + + +# --------------------------------------------------------------------------- # +# Nested cross-validation +# --------------------------------------------------------------------------- # + +OUTER_FOLDS = 2 +INNER_FOLDS = 2 +NESTED_TRIALS = 2 + + +@pytest.fixture(scope="module", name="nested_run") +def run_a_nested_search(client: TestClient, dataset_1: Dataset) -> dict: + session_id = _make_session( + client, + dataset_1.id, + "CvSessionNested", + _kfold_splits(n_splits=OUTER_FOLDS, test_size=0.2), + ) + run_id = _make_run( + client, + session_id, + model_name="TunableCvModel", + parameters={ + "n_estimators": { + "optimize": True, + "lower_bound": 1, + "upper_bound": 5, + "fixed_value": 2, + } + }, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": NESTED_TRIALS, + "sampler": "TPESampler", + "pruner": "None", + }, + goal_metric="MeanPredictionMetric", + nested={ + "splitter_name": "KFoldSplitter", + "n_splits": INNER_FOLDS, + "shuffle": True, + "random_state": 42, + }, + name="NestedCvRun", + ) + fits = _run_job(run_id) + return {"run_id": run_id, "fits": fits} + + +def test_a_nested_run_reaches_finished(client: TestClient, nested_run: dict): + assert _reload(client, nested_run["run_id"]).status == RunStatus.FINISHED + + +def test_a_nested_run_scores_every_outer_fold(client: TestClient, nested_run: dict): + rows = _metrics(client, nested_run["run_id"], level=LevelEnum.OUTER_FOLD) + + by_split = {} + for row in rows: + by_split.setdefault(row.split, []).append(row) + + assert set(by_split) == {SplitEnum.TRAIN, SplitEnum.VALIDATION} + for split, split_rows in by_split.items(): + assert sorted(row.fold_index for row in split_rows) == [0, 1], split + + +def test_a_nested_run_aggregates_its_outer_folds_apart_from_the_inner_ones( + client: TestClient, nested_run: dict +): + """``LAST_OUTER`` is the unbiased estimate; ``LAST`` summarises the plain folds.""" + for split in (SplitEnum.TRAIN, SplitEnum.VALIDATION): + outer = _fold_values(client, nested_run["run_id"], LevelEnum.OUTER_FOLD, split) + aggregated = _metrics( + client, nested_run["run_id"], level=LevelEnum.LAST_OUTER, split=split + ) + + assert len(aggregated) == 1, split + assert aggregated[0].value == pytest.approx(float(np.mean(outer))) + assert aggregated[0].std_value == pytest.approx(float(np.std(outer))) + + # The ordinary aggregation still happens, at its own level: train and + # validation from the fold loop, plus the reserved rows. + assert len(_metrics(client, nested_run["run_id"], level=LevelEnum.LAST)) == 3 + + +def test_a_nested_run_searches_inside_every_outer_fold(nested_run: dict): + """Three nested loops, counted by the size of the partition each fit saw. + + Outer folds hold 60 pooled rows; splitting one into inner folds leaves 30 + rows to fit on. Those 30 row fits only exist because the search is run + again, from scratch, inside each outer fold. + """ + fits = nested_run["fits"] + inner = [fit for fit in fits if fit["train"] == 30] + outer = [fit for fit in fits if fit["train"] == 60] + refits = [fit for fit in fits if fit["train"] == 120] + + assert len(inner) == OUTER_FOLDS * NESTED_TRIALS * INNER_FOLDS + # One fit per outer fold with the parameters its inner search chose, then + # the ordinary search and the scoring loop over the outer folds again. + assert len(outer) == OUTER_FOLDS + NESTED_TRIALS * OUTER_FOLDS + OUTER_FOLDS + assert len(refits) == 1 + assert len(fits) == len(inner) + len(outer) + len(refits) + + +# --------------------------------------------------------------------------- # +# Error branches, by the text the user is shown +# --------------------------------------------------------------------------- # + + +def _failing_run(client: TestClient, dataset_1: Dataset, name: str, **kwargs) -> int: + session_id = _make_session( + client, + dataset_1.id, + name, + kwargs.pop("splits"), + strategy=kwargs.pop("strategy", "CrossValidationEvaluationStrategy"), + ) + return _make_run(client, session_id, name=name + "Run", **kwargs) + + +def test_an_unknown_splitter_is_reported_with_the_run_it_belongs_to( + client: TestClient, dataset_1: Dataset +): + run_id = _failing_run( + client, + dataset_1, + "UnknownSplitter", + splits=_kfold_splits(splitter_name="ThereIsNoSuchSplitter"), + ) + + with pytest.raises(JobError) as error: + ModelJob(run_id=run_id).run() + + assert _squash(str(error.value)) == _squash( + f"Error preparing dataset and components for run {run_id}: " + "Unable to find Splitter with name ThereIsNoSuchSplitter in registry." + ) + assert _reload(client, run_id).status == RunStatus.ERROR + + +def test_an_unknown_evaluation_strategy_is_reported( + client: TestClient, dataset_1: Dataset +): + run_id = _failing_run( + client, + dataset_1, + "UnknownStrategy", + splits=_kfold_splits(), + strategy="ThereIsNoSuchStrategy", + ) + + with pytest.raises(JobError) as error: + ModelJob(run_id=run_id).run() + + assert _squash(str(error.value)) == _squash( + f"Error preparing dataset and components for run {run_id}: " + "Unable to find Evaluation Strategy with name " + "ThereIsNoSuchStrategy in registry." + ) + assert _reload(client, run_id).status == RunStatus.ERROR + + +def test_more_folds_than_rows_is_reported_as_a_splitting_failure( + client: TestClient, dataset_1: Dataset +): + run_id = _failing_run( + client, + dataset_1, + "TooManyFolds", + splits=_kfold_splits(n_splits=DATASET_ROWS + 1, test_size=0), + ) + + with pytest.raises(JobError) as error: + ModelJob(run_id=run_id).run() + + assert _squash(str(error.value)) == _squash( + f"Error splitting the dataset for run {run_id}: " + f"Number of splits (n_splits={DATASET_ROWS + 1}) cannot be greater " + f"than the number of samples ({DATASET_ROWS})." + ) + assert _reload(client, run_id).status == RunStatus.ERROR + + +def test_reserving_almost_everything_is_reported_as_a_splitting_failure( + client: TestClient, dataset_1: Dataset +): + run_id = _failing_run( + client, + dataset_1, + "TooLargeTestSet", + splits=_kfold_splits(n_splits=3, test_size=0.99), + ) + + with pytest.raises(JobError) as error: + ModelJob(run_id=run_id).run() + + message = _squash(str(error.value)) + assert message.startswith(f"Error splitting the dataset for run {run_id}: ") + assert "leaves too few rows" in message + assert _reload(client, run_id).status == RunStatus.ERROR + + +def test_an_unknown_inner_splitter_is_reported_as_a_training_failure( + client: TestClient, dataset_1: Dataset +): + """Nested cross-validation resolves its inner splitter while already training. + + So a typo there is not caught with the rest of the configuration: the run + has already reported that training started, and the message the user sees + comes from a different branch. + """ + run_id = _failing_run( + client, + dataset_1, + "UnknownInnerSplitter", + splits=_kfold_splits(n_splits=2), + model_name="TunableCvModel", + parameters={ + "n_estimators": { + "optimize": True, + "lower_bound": 1, + "upper_bound": 5, + "fixed_value": 2, + } + }, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": 1, + "sampler": "TPESampler", + "pruner": "None", + }, + goal_metric="MeanPredictionMetric", + nested={"splitter_name": "ThereIsNoSuchSplitter", "n_splits": 2}, + ) + + with pytest.raises(JobError) as error: + ModelJob(run_id=run_id).run() + + message = _squash(str(error.value)) + assert message.startswith("Model training and evaluation failed ") + assert "Error configuring inner splitter for nested CV" in message + assert _reload(client, run_id).status == RunStatus.ERROR diff --git a/tests/back/api/test_model_job_orchestration.py b/tests/back/api/test_model_job_orchestration.py new file mode 100644 index 000000000..fa962c1d9 --- /dev/null +++ b/tests/back/api/test_model_job_orchestration.py @@ -0,0 +1,378 @@ +"""End to end regression net for the ModelJob orchestration. + +``test_jobs.py`` accepts either ``finished`` or ``error`` as a job outcome, so +it cannot catch a job that silently stops doing part of its work. These tests +pin the observable contract of a successful run: status transitions, the +columns written on ``Run``, the ``Metric`` rows and the saved model artifact. +""" + +import json +import os + +import joblib +import pytest +from datasets import ClassLabel, Value +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.enums.status import RunStatus +from DashAI.back.core.schema_fields import BaseSchema, int_field, schema_field +from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader +from DashAI.back.dependencies.database.models import ( + Dataset, + Metric, + ModelSession, + Run, +) +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.evaluation.holdout import HoldoutEvaluationStrategy +from DashAI.back.job.base_job import JobError +from DashAI.back.job.model_job import ModelJob +from DashAI.back.metrics.base_metric import BaseMetric +from DashAI.back.models.base_model import BaseModel +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.splitters.holdout import HoldoutSplitter +from DashAI.back.tasks.base_task import BaseTask + + +class OrchestrationTask(BaseTask): + name: str = "OrchestrationTask" + metadata: dict = { + "inputs_types": [ClassLabel, Value], + "outputs_types": [ClassLabel], + "inputs_cardinality": "n", + "outputs_cardinality": 1, + } + + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + def num_labels(self, dataset, output_column): + return 3 + + +class OrchestrationModel(BaseModel): + """Model that records the data it was trained with.""" + + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + + def __init__(self, **kwargs): + self.trained_with = None + + def save(self, filename): + joblib.dump({"trained_with": self.trained_with}, filename) + + def load(self, filename): + return joblib.load(filename) + + def predict(self, x): + return [0] * x.shape[0] + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.trained_with = { + "train": x_train.shape[0], + "validation": None if x_validation is None else x_validation.shape[0], + } + return self + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + +class TunableModelSchema(BaseSchema): + n_estimators: schema_field( + int_field(gt=0), + placeholder=2, + description="Number of estimators.", + ) # type: ignore + + +class TunableModel(BaseModel): + """Model with an optimizable parameter, to exercise the search branch.""" + + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + SCHEMA = TunableModelSchema + + def __init__(self, n_estimators=2, **kwargs): + self.n_estimators = n_estimators + + def save(self, filename): + joblib.dump({"n_estimators": self.n_estimators}, filename) + + def load(self, filename): + return joblib.load(filename) + + def predict(self, x): + return [0] * x.shape[0] + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + return self + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + +class OrchestrationMetric(BaseMetric): + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + MAXIMIZE = True + + @staticmethod + def score(true_labels, probs_pred_labels): + return 0.5 + + +class TunableMetric(BaseMetric): + """Metric whose score depends on the hyperparameter being searched.""" + + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + MAXIMIZE = True + + @staticmethod + def score(true_labels, probs_pred_labels): + return 0.25 + + +@pytest.fixture(scope="module", name="orchestration_registry", autouse=True) +def setup_orchestration_registry(client): + container = client.app.container + sentinel = object() + services = container._services + old = services.get("component_registry", sentinel) + + services["component_registry"] = ComponentRegistry( + initial_components=[ + OrchestrationTask, + OrchestrationModel, + TunableModel, + OrchestrationMetric, + TunableMetric, + CSVDataLoader, + ModelJob, + OptunaOptimizer, + HoldoutSplitter, + HoldoutEvaluationStrategy, + ] + ) + yield services["component_registry"] + if old is sentinel: + del services["component_registry"] + else: + services["component_registry"] = old + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session( + client: TestClient, dataset_1: Dataset, orchestration_registry +): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="OrchestrationSession", + task_name="OrchestrationTask", + input_columns=["SepalLengthCm", "SepalWidthCm"], + output_columns=["Species"], + train_metrics=["OrchestrationMetric"], + validation_metrics=["OrchestrationMetric"], + test_metrics=["OrchestrationMetric"], + evaluation_strategy="HoldoutEvaluationStrategy", + splits=json.dumps( + { + "splitter_name": "HoldoutSplitter", + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + "splitType": "random", + } + ), + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + yield model_session.id + + +def _create_run(client: TestClient, model_session_id: int, model_name: str) -> int: + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + model_name=model_name, + parameters={}, + optimizer_name="", + optimizer_parameters={}, + goal_metric="", + name="OrchestrationRun", + ) + db.add(run) + db.commit() + db.refresh(run) + return run.id + + +@pytest.fixture(name="finished_run") +def run_a_successful_job(client: TestClient, model_session_id: int) -> Run: + run_id = _create_run(client, model_session_id, "OrchestrationModel") + ModelJob(run_id=run_id).run() + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + return db.get(Run, run_id) + + +def test_successful_run_reaches_finished(finished_run: Run): + assert finished_run.status == RunStatus.FINISHED + + +def test_successful_run_stamps_its_timestamps(finished_run: Run): + assert finished_run.start_time is not None + assert finished_run.end_time is not None + assert finished_run.end_time >= finished_run.start_time + + +def test_successful_run_persists_the_split_indexes(finished_run: Run): + split_indexes = json.loads(finished_run.split_indexes) + + assert set(split_indexes) == {"train_indexes", "test_indexes", "val_indexes"} + assert len(split_indexes["train_indexes"]) > 0 + assert len(split_indexes["val_indexes"]) > 0 + assert len(split_indexes["test_indexes"]) > 0 + + +def test_successful_run_saves_the_model_artifact(finished_run: Run): + assert finished_run.run_path is not None + assert finished_run.run_path.endswith(str(finished_run.id)) + assert os.path.exists(finished_run.run_path) + + +def test_the_model_was_trained_with_the_train_and_validation_splits( + finished_run: Run, +): + saved = joblib.load(finished_run.run_path) + + assert saved["trained_with"]["train"] > 0 + assert saved["trained_with"]["validation"] > 0 + + +def test_successful_run_writes_a_last_metric_for_every_split( + client: TestClient, finished_run: Run +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + metrics = db.query(Metric).filter_by(run_id=finished_run.id).all() + + by_split = {metric.split: metric for metric in metrics} + + assert set(by_split) == {SplitEnum.TRAIN, SplitEnum.VALIDATION, SplitEnum.TEST} + for metric in metrics: + assert metric.level == LevelEnum.LAST + assert metric.name == "OrchestrationMetric" + assert metric.value == 0.5 + + +@pytest.fixture(scope="module", name="tuned_run") +def run_a_hyperparameter_search(client: TestClient, model_session_id: int) -> Run: + """Run the hyperparameter search branch end to end. + + This branch had no coverage at all: every other test creates runs with an + empty ``optimizer_name`` and no optimizable parameters, so the whole + optimize/plot path was never executed by the suite. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + model_name="TunableModel", + parameters={ + "n_estimators": { + "optimize": True, + "lower_bound": 1, + "upper_bound": 5, + "fixed_value": 2, + } + }, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": 3, + "sampler": "TPESampler", + "pruner": None, + }, + goal_metric="TunableMetric", + name="TunedRun", + ) + db.add(run) + db.commit() + db.refresh(run) + run_id = run.id + + ModelJob(run_id=run_id).run() + + with session_factory() as db: + return db.get(Run, run_id) + + +def test_hyperparameter_search_reaches_finished(tuned_run: Run): + assert tuned_run.status == RunStatus.FINISHED + + +def test_hyperparameter_search_persists_the_best_parameters(tuned_run: Run): + best = tuned_run.parameters["n_estimators"]["fixed_value"] + + assert 1 <= best <= 5 + + +def test_hyperparameter_search_saves_two_plots_for_a_single_parameter( + tuned_run: Run, +): + assert tuned_run.plot_history_path is not None + assert tuned_run.plot_slice_path is not None + assert os.path.exists(tuned_run.plot_history_path) + assert os.path.exists(tuned_run.plot_slice_path) + # create_plots only produces the contour and importance plots when more + # than one hyperparameter is being searched. + assert tuned_run.plot_contour_path is None + assert tuned_run.plot_importance_path is None + + +def test_the_tuned_model_still_logs_its_metrics(client: TestClient, tuned_run: Run): + """The optimizer must return the same model instance it received. + + ``ModelFactory`` hangs the run id, the splits and the metric classes off + the model. If an optimizer returned a fresh object instead, + ``calculate_metrics`` would silently return and the run would finish with + no metrics rather than fail. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + metrics = db.query(Metric).filter_by(run_id=tuned_run.id).all() + + last_splits = {m.split for m in metrics if m.level == LevelEnum.LAST} + assert last_splits == {SplitEnum.TRAIN, SplitEnum.VALIDATION, SplitEnum.TEST} + + # The optimizer logs a metric per trial while searching. + trial_metrics = [m for m in metrics if m.level == LevelEnum.TRIAL] + assert len(trial_metrics) > 0 + + +def test_a_missing_run_is_reported_instead_of_crashing(client: TestClient): + with pytest.raises(JobError, match="Run 987654 does not exist in DB."): + ModelJob(run_id=987654).run() + + +def test_an_unknown_model_leaves_the_run_in_error( + client: TestClient, model_session_id: int +): + run_id = _create_run(client, model_session_id, "ThereIsNoSuchModel") + + with pytest.raises(JobError, match="Unable to find Model with name"): + ModelJob(run_id=run_id).run() + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + assert db.get(Run, run_id).status == RunStatus.ERROR diff --git a/tests/back/api/test_predict_api.py b/tests/back/api/test_predict_api.py index 5c5249e97..e261ada46 100644 --- a/tests/back/api/test_predict_api.py +++ b/tests/back/api/test_predict_api.py @@ -450,10 +450,13 @@ def test_a_failing_prediction_reports_why( assert response.status_code == 200, response.text prediction_id = response.json()["id"] - def explode(**kwargs): + # Injected at the unit that predicts, which is where the three inline + # prediction steps this test used to patch now live. What is under test is + # not where the failure comes from but that its message survives the trip. + def explode(self, ctx): raise ValueError("ARIMA forecasts forward only") - monkeypatch.setattr("DashAI.back.job.predict_job._run_prediction_pipeline", explode) + monkeypatch.setattr("DashAI.back.units.predict_unit.PredictUnit.execute", explode) job = PredictJob(job_type="PredictJob", kwargs={"prediction_id": prediction_id}) with pytest.raises(JobError) as raised: diff --git a/tests/back/api/test_predict_job.py b/tests/back/api/test_predict_job.py new file mode 100644 index 000000000..cae6a9b61 --- /dev/null +++ b/tests/back/api/test_predict_job.py @@ -0,0 +1,521 @@ +"""End-to-end regression net for ``PredictJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so that the refactor has something to be measured +against. The assertions are deliberately explicit — exact status values, exact +columns on disk, exact error message fragments — instead of the looser +``status in ["finished", "error"]`` style used elsewhere in this suite, which +cannot tell a unit that silently stopped doing part of its work from one that +did it. + +``test_predict_api.py`` does not cover this: it enqueues the job and then only +exercises the CRUD endpoints, never checking the job's outcome nor the +``Prediction`` row. + +Tests named ``test_currently_*`` pin behaviour that is known to be wrong. They +exist so the refactor can be proven behaviour-preserving first; the fix lands +afterwards as its own change, which flips the assertion and renames the test. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import shutil +from pathlib import Path + +import pytest +from fastapi.exceptions import HTTPException +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import PredictionStatus +from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset +from DashAI.back.dependencies.database.models import ( + Dataset, + ModelSession, + Prediction, + Run, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.job.model_job import ModelJob +from DashAI.back.job.predict_job import PredictJob + +INPUT_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", +] +OUTPUT_COLUMN = "Species" +IRIS_ROWS = 150 + +SPLITS = json.dumps( + { + # The session names the splitter that produced its partitions, which is + # how every consumer of Run.split_indexes tells a holdout payload from + # a fold one. + "splitter_name": "HoldoutSplitter", + "splitType": "random", + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } +) + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session(client: TestClient, dataset_1: Dataset): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="PredictJobSession", + task_name="TabularClassificationTask", + input_columns=INPUT_COLUMNS, + output_columns=[OUTPUT_COLUMN], + train_metrics=[], + validation_metrics=[], + test_metrics=[], + evaluation_strategy="HoldoutEvaluationStrategy", + splits=SPLITS, + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + return model_session.id + + +@pytest.fixture(scope="module", name="trained_run_id") +def create_trained_run(client: TestClient, model_session_id: int): + """A genuinely trained run: the prediction path loads the model from disk.""" + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": 1, + "sampler": "TPESampler", + "pruner": "None", + }, + model_name="KNeighborsClassifier", + parameters={}, + name="PredictJobRun", + goal_metric="Accuracy", + ) + db.add(run) + db.commit() + db.refresh(run) + run_id = run.id + + ModelJob(run_id=run_id).run() + + with session_factory() as db: + run = db.get(Run, run_id) + assert run.run_path, "the run fixture did not produce a saved model" + return run_id + + +def _create_prediction(client, run_id, dataset_id=None): + response = client.post( + "/api/v1/predict/", + json={"run_id": run_id, "dataset_id": dataset_id}, + ) + assert response.status_code == 200, response.text + return response.json()["id"] + + +def _make_prediction_dataset(client, dataset_1: Dataset): + """A throwaway copy of the iris dataset, safe for a test to destroy.""" + import uuid + + config = client.app.container["config"] + session_factory = client.app.container["session_factory"] + + folder = Path(config["DATASETS_PATH"]) / f"predict-job-{uuid.uuid4()}" + shutil.copytree(Path(dataset_1.file_path), folder) + + with session_factory() as db: + row = Dataset(name=folder.name, file_path=str(folder)) + db.add(row) + db.commit() + db.refresh(row) + db.expunge(row) + return row + + +def _stored_prediction(client, prediction_id): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + prediction = db.get(Prediction, prediction_id) + return { + "status": prediction.status, + "results_path": prediction.results_path, + "start_time": prediction.start_time, + "end_time": prediction.end_time, + } + + +@pytest.fixture(name="restore_run") +def fixture_restore_run(client: TestClient, trained_run_id: int): + """Let a test corrupt the module-scoped Run row and put it back after. + + The run and the model session are module scoped because training is slow; + without this the error-branch tests would poison every test after them. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = db.get(Run, trained_run_id) + original = {"model_name": run.model_name, "run_path": run.run_path} + + yield + + with session_factory() as db: + run = db.get(Run, trained_run_id) + run.model_name = original["model_name"] + run.run_path = original["run_path"] + db.commit() + + +@pytest.fixture(name="restore_model_session") +def fixture_restore_model_session(client: TestClient, model_session_id: int): + """Same idea for the module-scoped ModelSession row.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + session_row = db.get(ModelSession, model_session_id) + original = { + "dataset_id": session_row.dataset_id, + "task_name": session_row.task_name, + "input_columns": list(session_row.input_columns), + "output_columns": list(session_row.output_columns), + } + + yield + + with session_factory() as db: + session_row = db.get(ModelSession, model_session_id) + for key, value in original.items(): + setattr(session_row, key, value) + db.commit() + + +def test_predict_job_writes_the_predictions_and_finishes( + client, trained_run_id, dataset_1 +): + """The happy path, end to end: status transitions and the dataset on disk. + + The saved dataset carries the input columns plus the predicted output + column, one row per row of the input. + """ + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + PredictJob(prediction_id=prediction_id).run() + + stored = _stored_prediction(client, prediction_id) + assert stored["status"] == PredictionStatus.FINISHED + assert stored["start_time"] is not None + assert stored["end_time"] is not None + assert stored["results_path"] is not None + + saved = load_dataset(str(Path(stored["results_path"]) / "dataset")) + assert saved.column_names == INPUT_COLUMNS + [OUTPUT_COLUMN] + assert len(saved) == IRIS_ROWS + + +def test_each_prediction_gets_its_own_results_folder(client, trained_run_id, dataset_1): + """The destination is a fresh uuid folder, so two runs never collide.""" + first = _create_prediction(client, trained_run_id, dataset_1.id) + second = _create_prediction(client, trained_run_id, dataset_1.id) + + PredictJob(prediction_id=first).run() + PredictJob(prediction_id=second).run() + + first_path = _stored_prediction(client, first)["results_path"] + second_path = _stored_prediction(client, second)["results_path"] + + assert first_path != second_path + assert Path(first_path).exists() + assert Path(second_path).exists() + + +def test_manual_input_predicts_without_a_dataset(client, trained_run_id): + """The manual branch builds the instances from typed values instead of disk.""" + prediction_id = _create_prediction(client, trained_run_id, dataset_id=None) + + PredictJob( + prediction_id=prediction_id, + manual_input_data=[ + { + "SepalLengthCm": 5.1, + "SepalWidthCm": 3.5, + "PetalLengthCm": 1.4, + "PetalWidthCm": 0.2, + } + ], + ).run() + + stored = _stored_prediction(client, prediction_id) + assert stored["status"] == PredictionStatus.FINISHED + + saved = load_dataset(str(Path(stored["results_path"]) / "dataset")) + assert saved.column_names == INPUT_COLUMNS + [OUTPUT_COLUMN] + assert len(saved) == 1 + + +def test_neither_a_dataset_nor_manual_input_is_rejected(client, trained_run_id): + prediction_id = _create_prediction(client, trained_run_id, dataset_id=None) + + with pytest.raises( + JobError, match="Either dataset_id or manual_input_data must be provided." + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_missing_prediction_row_is_a_404(client): + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=999999).run() + + assert excinfo.value.status_code == 404 + assert excinfo.value.detail == "Prediction not found for id 999999" + + +def test_an_unknown_model_name_reports_it_and_errors( + client, trained_run_id, dataset_1, restore_run +): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).model_name = "ThisModelDoesNotExist" + db.commit() + + with pytest.raises( + JobError, match="Model ThisModelDoesNotExist not found in the registry" + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_model_that_cannot_be_loaded_reports_the_path_and_errors( + client, trained_run_id, dataset_1, restore_run +): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).run_path = "nowhere/at/all" + db.commit() + + with pytest.raises( + JobError, + match="Failed to load model KNeighborsClassifier from path nowhere/at/all", + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_an_unknown_task_name_reports_it_and_errors( + client, trained_run_id, dataset_1, restore_model_session, model_session_id +): + """The task is resolved before the model, so this is the first error seen.""" + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "ThisTaskDoesNotExist" + db.commit() + + with pytest.raises( + JobError, match="Task ThisTaskDoesNotExist not found in the registry" + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_model_session_without_input_columns_is_a_422( + client, trained_run_id, dataset_1, restore_model_session, model_session_id +): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).input_columns = [] + db.commit() + + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert excinfo.value.status_code == 422 + assert excinfo.value.detail == "Model session has no input columns configured" + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_missing_training_dataset_row_leaves_the_row_in_error( + client, trained_run_id, dataset_1, restore_model_session, model_session_id +): + """The 404 must also mark the prediction as failed. + + This branch used to skip ``set_status_as_error``, unlike every one around + it, so the row stayed STARTED forever — nothing else marks it, because the + Huey error signal only writes to its own ``task_copy`` table and + ``_execute_base_job`` calls ``run()`` with no handler. + """ + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = 999999 + db.commit() + + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert excinfo.value.status_code == 404 + assert excinfo.value.detail == "Training dataset not found" + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_an_unreadable_training_dataset_leaves_the_row_in_error( + client, trained_run_id, dataset_1, tmp_path +): + """Same omission as above, on the branch that reads the training dataset. + + The message still has to be the specific one, not the generic prediction + wrapper: this load happens before the dataset to predict on is even + touched, and that ordering is what the message depends on. + """ + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + stored_folder = Path(dataset_1.file_path) / "dataset" + backup = tmp_path / "training-dataset-backup" + shutil.copytree(stored_folder, backup) + shutil.rmtree(stored_folder) + try: + with pytest.raises(JobError, match="Cannot load training dataset from"): + PredictJob(prediction_id=prediction_id).run() + + assert ( + _stored_prediction(client, prediction_id)["status"] + == PredictionStatus.ERROR + ) + finally: + shutil.copytree(backup, stored_folder) + + +def test_an_unreadable_prediction_dataset_is_reported_as_a_prediction_failure( + client, trained_run_id, dataset_1, tmp_path +): + """Loading the dataset to predict on shares the "prediction failed" wrapper. + + Unlike the *training* dataset, which has a message of its own, the + inference dataset is loaded inside the same ``try`` as the prediction, so a + read failure surfaces as the generic message with the real cause attached. + Pinned because it is the exact spot a unit boundary lands on, and an + improved message here would be a silent behaviour change. + """ + prediction_dataset = _make_prediction_dataset(client, dataset_1) + prediction_id = _create_prediction(client, trained_run_id, prediction_dataset.id) + + shutil.rmtree(Path(prediction_dataset.file_path) / "dataset") + + with pytest.raises(JobError, match="Model prediction failed"): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_prediction_that_blows_up_leaves_the_row_in_error( + client, trained_run_id, dataset_1, monkeypatch +): + """Any unexpected failure while predicting is reported as one message.""" + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + def _explode(self, x): + raise RuntimeError("the model itself blew up") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _explode) + + with pytest.raises(JobError, match="Model prediction failed") as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert "the model itself blew up" in str(excinfo.value.__cause__) + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_type_error_while_predicting_leaves_the_row_in_error( + client, trained_run_id, dataset_1, monkeypatch +): + """The ``TypeError`` branch is still a 400, but now it marks the row too. + + Its ``ValueError`` neighbour always did; this one did not, so a type + mismatch left the prediction STARTED forever. + """ + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + def _wrong_type(self, x): + raise TypeError("bad type somewhere in the input") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _wrong_type) + + # JobError rather than HTTPException: the exception leaves the worker + # through dill, which an HTTPException does not survive. + with pytest.raises(JobError) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert "Type validation failed" in str(excinfo.value) + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_value_error_while_predicting_leaves_the_row_in_error( + client, trained_run_id, dataset_1, monkeypatch +): + """The ``ValueError`` branch is reported as a JobError and marks the row.""" + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + def _bad_value(self, x): + raise ValueError("a value the model cannot use") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _bad_value) + + # JobError rather than HTTPException: the exception leaves the worker + # through dill, which an HTTPException does not survive. + with pytest.raises(JobError) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert "Invalid input data" in str(excinfo.value) + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_set_status_as_delivered_marks_the_row(client, trained_run_id, dataset_1): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + PredictJob(prediction_id=prediction_id).set_status_as_delivered() + + assert ( + _stored_prediction(client, prediction_id)["status"] + == PredictionStatus.DELIVERED + ) diff --git a/tests/back/api/test_predict_preview_api.py b/tests/back/api/test_predict_preview_api.py new file mode 100644 index 000000000..ce5f164ab --- /dev/null +++ b/tests/back/api/test_predict_preview_api.py @@ -0,0 +1,496 @@ +"""End-to-end regression net for ``POST /predict/preview``. + +The synchronous counterpart of ``PredictJob``: same prediction, no persistence, +and errors reported as HTTP responses instead of ``JobError``. Written before +``run_manual_prediction`` is decomposed into the same units the job uses, and +asserted against the pre-refactor implementation, so the refactor has something +to be measured against. + +Every failure mode gets its own test with its exact status code and detail +string, because that is the whole contract this endpoint has with the frontend's +manual-prediction form: eleven distinct responses, four status codes. A +migration that turned any of them into a generic 500 would be a silent UX +regression in the one feature the endpoint exists for. + +This endpoint had no tests at all before this file. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import shutil +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run +from DashAI.back.job.model_job import ModelJob + +INPUT_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", +] +OUTPUT_COLUMN = "Species" + +A_ROW = { + "SepalLengthCm": 5.1, + "SepalWidthCm": 3.5, + "PetalLengthCm": 1.4, + "PetalWidthCm": 0.2, +} + +SPLITS = json.dumps( + { + # The session names the splitter that produced its partitions, which is + # how every consumer of Run.split_indexes tells a holdout payload from + # a fold one. + "splitter_name": "HoldoutSplitter", + "splitType": "random", + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } +) + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session(client: TestClient, dataset_1: Dataset): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="PreviewSession", + task_name="TabularClassificationTask", + input_columns=INPUT_COLUMNS, + output_columns=[OUTPUT_COLUMN], + train_metrics=[], + validation_metrics=[], + test_metrics=[], + evaluation_strategy="HoldoutEvaluationStrategy", + splits=SPLITS, + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + return model_session.id + + +@pytest.fixture(scope="module", name="trained_run_id") +def create_trained_run(client: TestClient, model_session_id: int): + """A genuinely trained run: the preview path loads the model from disk.""" + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": 1, + "sampler": "TPESampler", + "pruner": "None", + }, + model_name="KNeighborsClassifier", + parameters={}, + name="PreviewRun", + goal_metric="Accuracy", + ) + db.add(run) + db.commit() + db.refresh(run) + run_id = run.id + + ModelJob(run_id=run_id).run() + + with session_factory() as db: + assert db.get(Run, run_id).run_path, "the fixture did not save a model" + return run_id + + +def _preview(client, run_id, rows=None): + return client.post( + "/api/v1/predict/preview", + data={ + "run_id": str(run_id), + "manual_input_data": json.dumps(rows if rows is not None else [A_ROW]), + }, + ) + + +@pytest.fixture(name="restore_run") +def fixture_restore_run(client: TestClient, trained_run_id: int): + """Let a test corrupt the module-scoped Run row and put it back after. + + The run is module scoped because training is slow; without this the + error-branch tests would poison every test after them. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = db.get(Run, trained_run_id) + original = { + "model_name": run.model_name, + "run_path": run.run_path, + "model_session_id": run.model_session_id, + } + + yield + + with session_factory() as db: + run = db.get(Run, trained_run_id) + for key, value in original.items(): + setattr(run, key, value) + db.commit() + + +@pytest.fixture(name="restore_model_session") +def fixture_restore_model_session(client: TestClient, model_session_id: int): + """Same idea for the module-scoped ModelSession row.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + row = db.get(ModelSession, model_session_id) + original = { + "dataset_id": row.dataset_id, + "task_name": row.task_name, + "input_columns": list(row.input_columns), + "output_columns": list(row.output_columns), + } + + yield + + with session_factory() as db: + row = db.get(ModelSession, model_session_id) + for key, value in original.items(): + setattr(row, key, value) + db.commit() + + +# --- the happy path ----------------------------------------------------- + + +def test_the_preview_returns_the_inputs_plus_the_prediction(client, trained_run_id): + response = _preview(client, trained_run_id) + + assert response.status_code == 200, response.text + body = response.json() + assert body["columns"] == INPUT_COLUMNS + [OUTPUT_COLUMN] + assert len(body["rows"]) == 1 + assert body["rows"][0][:4] == [5.1, 3.5, 1.4, 0.2] + # The label is decoded against the training dataset, not left as an index. + assert body["rows"][0][4] in {"Iris-setosa", "Iris-versicolor", "Iris-virginica"} + + +def test_the_preview_handles_several_rows_at_once(client, trained_run_id): + rows = [A_ROW, {**A_ROW, "PetalLengthCm": 5.9, "PetalWidthCm": 2.1}] + + response = _preview(client, trained_run_id, rows) + + assert response.status_code == 200, response.text + body = response.json() + assert len(body["rows"]) == 2 + assert body["rows"][1][2] == 5.9 + + +def test_the_preview_and_the_job_agree_on_the_same_input(client, trained_run_id): + """The point of sharing units: the two paths cannot answer differently. + + Same hand-typed row, one predicted synchronously for the preview and one + through ``PredictJob``. They used to run separate copies of the same three + steps, so nothing stopped them from drifting apart; this is what would fail + if the prediction were ever fixed in one place only. + """ + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + from DashAI.back.dependencies.database.models import Prediction + from DashAI.back.job.predict_job import PredictJob + + row = {**A_ROW, "PetalLengthCm": 4.7, "PetalWidthCm": 1.4} + + previewed = _preview(client, trained_run_id, [row]) + assert previewed.status_code == 200, previewed.text + preview_label = previewed.json()["rows"][0][4] + + created = client.post( + "/api/v1/predict/", json={"run_id": trained_run_id, "dataset_id": None} + ) + assert created.status_code == 200, created.text + prediction_id = created.json()["id"] + + PredictJob(prediction_id=prediction_id, manual_input_data=[row]).run() + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + results_path = db.get(Prediction, prediction_id).results_path + + saved = load_dataset(str(Path(results_path) / "dataset")) + assert saved[OUTPUT_COLUMN] == [preview_label] + + +def test_the_preview_persists_nothing(client, trained_run_id): + """It is a preview: no Prediction row, no results folder.""" + from DashAI.back.dependencies.database.models import Prediction + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + before = db.query(Prediction).count() + + assert _preview(client, trained_run_id).status_code == 200 + + with session_factory() as db: + assert db.query(Prediction).count() == before + + +# --- the four 404/422 checks the endpoint owns -------------------------- + + +def test_a_missing_run_is_a_404(client): + response = _preview(client, 999999) + + assert response.status_code == 404 + assert response.json()["detail"] == "Run not found for id 999999" + + +def test_a_missing_model_session_is_a_404(client, trained_run_id, restore_run): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).model_session_id = 999999 + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 404 + assert response.json()["detail"] == "Model session not found" + + +def test_a_missing_training_dataset_row_is_a_404( + client, trained_run_id, restore_model_session, model_session_id +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = 999999 + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 404 + assert response.json()["detail"] == "Training dataset not found" + + +def test_no_input_columns_is_a_422( + client, trained_run_id, restore_model_session, model_session_id +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).input_columns = [] + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 422 + assert response.json()["detail"] == "Model session has no input columns configured" + + +def test_no_output_columns_is_a_422( + client, trained_run_id, restore_model_session, model_session_id +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).output_columns = [] + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 422 + assert response.json()["detail"] == "Model session has no output columns configured" + + +# --- the registry and loading failures ---------------------------------- + + +def test_an_unknown_task_is_a_500_naming_the_task( + client, trained_run_id, restore_model_session, model_session_id +): + """The task is resolved before the model, so this wins when both are wrong. + + That ordering is behaviour: it decides which of the two the user is told + about, and a decomposition that resolves the model first would change it. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "NoSuchTask" + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Task NoSuchTask not found in the registry" + + +def test_an_unknown_model_is_a_500_naming_the_model( + client, trained_run_id, restore_run +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).model_name = "NoSuchModel" + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Model NoSuchModel not found in the registry" + + +def test_an_unreadable_model_is_a_500_naming_model_and_path( + client, trained_run_id, restore_run +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).run_path = "nowhere/at/all" + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == ( + "Failed to load model KNeighborsClassifier from path nowhere/at/all" + ) + + +def test_an_unreadable_training_dataset_is_a_500_without_the_path( + client, trained_run_id, dataset_1, tmp_path +): + """The detail is deliberately bare here — no path, unlike the model error. + + Pinned exactly because the unit that replaces this step reports a message of + its own that *does* carry the path; the endpoint has to keep saying this. + """ + stored = Path(dataset_1.file_path) / "dataset" + backup = tmp_path / "training-dataset-backup" + shutil.copytree(stored, backup) + shutil.rmtree(stored) + try: + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Cannot load training dataset" + finally: + shutil.copytree(backup, stored) + + +# --- the input and prediction failures ---------------------------------- + + +def test_an_unknown_input_column_is_a_400(client, trained_run_id): + response = _preview(client, trained_run_id, [{"NotAColumn": 1.0}]) + + assert response.status_code == 400 + detail = response.json()["detail"] + assert detail.startswith("Invalid input data: ") + assert "NotAColumn" in detail + + +def test_a_value_error_while_predicting_is_a_400(client, trained_run_id, monkeypatch): + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + def _bad_value(self, x): + raise ValueError("a value the model cannot use") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _bad_value) + + response = _preview(client, trained_run_id) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "Invalid input data: a value the model cannot use" + ) + + +def test_a_type_error_while_predicting_is_also_a_400_invalid_input( + client, trained_run_id, monkeypatch +): + """The sync path merges ``TypeError`` into the ``ValueError`` message. + + ``PredictJob`` keeps them apart ("Type validation failed" vs "Invalid input + data"). That divergence is pinned on both sides so neither drifts into the + other while they share units. + """ + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + def _wrong_type(self, x): + raise TypeError("bad type somewhere in the input") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _wrong_type) + + response = _preview(client, trained_run_id) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "Invalid input data: bad type somewhere in the input" + ) + + +def test_any_other_prediction_failure_is_a_500(client, trained_run_id, monkeypatch): + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + def _explode(self, x): + raise RuntimeError("the model itself blew up") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _explode) + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Model prediction failed" + + +# --- the request-shape checks, owned by the endpoint itself ------------- + + +@pytest.mark.parametrize( + ("payload", "detail"), + [ + ({}, "Missing run_id or manual_input_data"), + ( + {"run_id": "notanint", "manual_input_data": "[]"}, + "Invalid run_id: notanint", + ), + ( + {"run_id": "1", "manual_input_data": "[]"}, + "manual_input_data must be a non-empty JSON array of objects (list[dict]).", + ), + ( + {"run_id": "1", "manual_input_data": "[1, 2]"}, + "Each item in manual_input_data must be a JSON object (dict).", + ), + ], +) +def test_the_request_shape_is_validated_before_anything_is_loaded( + client, payload, detail +): + response = client.post("/api/v1/predict/preview", data=payload) + + assert response.status_code == 422 + assert response.json()["detail"] == detail + + +def test_malformed_manual_input_json_is_a_422(client): + response = client.post( + "/api/v1/predict/preview", + data={"run_id": "1", "manual_input_data": "{not json"}, + ) + + assert response.status_code == 422 + assert response.json()["detail"].startswith("Invalid manual_input_data JSON: ") diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py new file mode 100644 index 000000000..6cbb2bb43 --- /dev/null +++ b/tests/back/api/test_units_api.py @@ -0,0 +1,283 @@ +"""The atomic units are exposed as regular registry components.""" + +import pytest +from fastapi.testclient import TestClient + +EXPECTED_UNITS = { + "LoadDatasetUnit", + "PrepareAndSplitUnit", + "PrepareAndFoldUnit", + "BuildModelUnit", + "FitModelUnit", + "FitModelOverFoldsUnit", + "FitModelOverNestedFoldsUnit", + "EvaluateModelUnit", + "EvaluateModelToArtifactUnit", + "SaveModelUnit", + "ApplyConverterUnit", + "FitConverterUnit", + "TransformDatasetUnit", + "SaveDatasetUnit", + "RunExplorationUnit", + "SaveExplorationUnit", + "LoadTrainedModelUnit", + "LoadTrainingDatasetUnit", + "BuildManualInputUnit", + "PredictUnit", + "SavePredictionUnit", + "LoadRunModelUnit", + "BuildGlobalExplainerUnit", + "BuildLocalExplainerUnit", + "PrepareExplanationDataUnit", + "GenerateGlobalExplanationUnit", + "GenerateLocalExplanationUnit", + "LoadUploadedDatasetUnit", + "LoadDatafileDatasetUnit", + "InferDatasetTypesUnit", + "ApplyDatasetSchemaUnit", + "ComputeDatasetMetadataUnit", + "SaveDatasetToPathUnit", +} + + +@pytest.fixture(name="units", scope="module") +def get_units(client: TestClient): + response = client.get("/api/v1/component/?select_types=Unit") + assert response.status_code == 200, response.text + return {component["name"]: component for component in response.json()} + + +def test_every_unit_is_registered(units): + assert set(units) == EXPECTED_UNITS + + +def test_units_are_registered_under_the_unit_type(units): + for unit in units.values(): + assert unit["type"] == "Unit" + + +def test_units_expose_a_schema_the_front_can_render(units): + for name, unit in units.items(): + assert unit["configurable_object"] is True, name + assert "properties" in unit["schema"], name + + +def test_unit_schemas_describe_their_configuration(units): + assert set(units["LoadDatasetUnit"]["schema"]["properties"]) == { + "dataset_id", + "notebook_id", + } + # The two splitting units are the same form with a different family of + # splitters offered, which is the whole of what separates them. + for name in ("PrepareAndSplitUnit", "PrepareAndFoldUnit"): + assert set(units[name]["schema"]["properties"]) == { + "task_name", + "input_columns", + "output_columns", + "splitter", + }, name + assert ( + units["PrepareAndSplitUnit"]["schema"]["properties"]["splitter"]["parent"] + == "PartitionSplitter" + ) + assert ( + units["PrepareAndFoldUnit"]["schema"]["properties"]["splitter"]["parent"] + == "FoldSplitter" + ) + assert "model" in units["BuildModelUnit"]["schema"]["properties"] + assert "optimizer" in units["FitModelUnit"]["schema"]["properties"] + assert set(units["ApplyConverterUnit"]["schema"]["properties"]) == { + "converter", + "scope", + "target", + } + assert set(units["FitConverterUnit"]["schema"]["properties"]) == { + "converter", + "scope", + "target", + } + # No converter to pick: it arrives already fitted through the context. + assert set(units["TransformDatasetUnit"]["schema"]["properties"]) == { + "scope", + "target", + } + # SaveDatasetUnit is configuration-free: it saves where the load said. + assert units["SaveDatasetUnit"]["schema"]["properties"] == {} + assert set(units["RunExplorationUnit"]["schema"]["properties"]) == { + "explorer_id", + "explorer", + } + # SaveExplorationUnit only picks the destination; how the result is + # serialised belongs to the explorer that produced it. + assert set(units["SaveExplorationUnit"]["schema"]["properties"]) == {"explorer_id"} + assert set(units["LoadTrainedModelUnit"]["schema"]["properties"]) == {"run_id"} + assert set(units["PredictUnit"]["schema"]["properties"]) == { + "task_name", + "input_columns", + "output_columns", + } + assert set(units["SavePredictionUnit"]["schema"]["properties"]) == { + "input_columns", + "output_columns", + } + assert set(units["BuildGlobalExplainerUnit"]["schema"]["properties"]) == { + "explainer" + } + assert set(units["BuildLocalExplainerUnit"]["schema"]["properties"]) == { + "explainer" + } + assert set(units["GenerateGlobalExplanationUnit"]["schema"]["properties"]) == { + "explainer_id" + } + # temp_path is a runtime param: the endpoint makes the directory, so it + # never reaches the schema the front receives. + assert set(units["LoadUploadedDatasetUnit"]["schema"]["properties"]) == { + "dataloader", + "source", + "n_sample", + } + assert set(units["LoadDatafileDatasetUnit"]["schema"]["properties"]) == { + "dataloader", + "datafile_id", + "selected_file", + } + assert set(units["InferDatasetTypesUnit"]["schema"]["properties"]) == {"method"} + # The type declaration arrives through the context, not the configuration, + # so the only thing to configure here is the renaming. + assert set(units["ApplyDatasetSchemaUnit"]["schema"]["properties"]) == { + "column_renames" + } + assert set(units["ComputeDatasetMetadataUnit"]["schema"]["properties"]) == { + "compute_metadata", + } + # The sibling of SaveDatasetUnit: that one saves where the load said, this + # one is told where to save. The destination stays user-facing: a job + # happens to compute it, but on a canvas picking where to store a dataset + # is a decision, and a value a user can know while filling the form in. + assert set(units["SaveDatasetToPathUnit"]["schema"]["properties"]) == {"path"} + + +def test_component_fields_tell_the_front_which_components_to_offer(units): + """The recursive part of the schema system. + + A component field does not inline the chosen component's schema; it + carries a ``parent`` hint so the front can list the candidates and then + fetch that component's own schema to render the nested form. + """ + model = units["BuildModelUnit"]["schema"]["properties"]["model"] + optimizer = units["FitModelUnit"]["schema"]["properties"]["optimizer"] + converter = units["ApplyConverterUnit"]["schema"]["properties"]["converter"] + explorer = units["RunExplorationUnit"]["schema"]["properties"]["explorer"] + + uploaded = units["LoadUploadedDatasetUnit"]["schema"]["properties"]["dataloader"] + datafile = units["LoadDatafileDatasetUnit"]["schema"]["properties"]["dataloader"] + + assert model["parent"] == "BaseModel" + assert optimizer["parent"] == "BaseOptimizer" + assert converter["parent"] == "BaseConverter" + assert explorer["parent"] == "BaseExplorer" + # Both loading units offer the same readers; what differs is how each one + # finds the bytes to hand them. + assert uploaded["parent"] == "BaseDataLoader" + assert datafile["parent"] == "BaseDataLoader" + + # Global and local explainers are separate registries with separate base + # classes, and a component field carries a single parent hint. Hence two + # sibling units with one required field each: making a single field cover + # both scopes would need it to be optional, and an optional component field + # is emitted as an anyOf, which hides the hint from the front — that is what + # the assertions below would catch. + global_explainer = units["BuildGlobalExplainerUnit"]["schema"]["properties"] + local_explainer = units["BuildLocalExplainerUnit"]["schema"]["properties"] + assert global_explainer["explainer"]["parent"] == "BaseGlobalExplainer" + assert local_explainer["explainer"]["parent"] == "BaseLocalExplainer" + assert set(model["properties"]) == {"component", "params"} + assert set(converter["properties"]) == {"component", "params"} + + +def test_a_component_field_parent_resolves_to_real_components(client: TestClient): + response = client.get("/api/v1/component/?component_parent=BaseOptimizer") + assert response.status_code == 200, response.text + + names = {component["name"] for component in response.json()} + assert "OptunaOptimizer" in names + + +def test_units_do_not_leak_into_the_job_listing(client: TestClient): + response = client.get("/api/v1/component/?select_types=Job") + assert response.status_code == 200, response.text + + job_names = {component["name"] for component in response.json()} + assert not (job_names & EXPECTED_UNITS) + assert "ModelJob" in job_names + + +@pytest.fixture(name="unit_classes", scope="module") +def get_unit_classes(client: TestClient): + """The registered classes behind the response, from the same registry.""" + registry = client.app.container["component_registry"] + return {name: registry[name]["class"] for name in EXPECTED_UNITS} + + +#: Every runtime param in the palette. Configuration supplied by whatever runs +#: the unit -- a job, the DAG engine, an endpoint -- and never by a user, +#: because at the moment a form would be filled in the value is not knowable: +#: a path a job has yet to choose, a foreign key to a row nobody has seen, or +#: a name built from the id of a run that has not started. +EXPECTED_RUNTIME_PARAMS = { + ("LoadUploadedDatasetUnit", "temp_path"), + ("LoadTrainingDatasetUnit", "train_dataset_file_path"), + ("BuildManualInputUnit", "train_dataset_file_path"), + ("GenerateLocalExplanationUnit", "session_splits"), + ("ComputeDatasetMetadataUnit", "trust_inherited_metadata"), + ("BuildModelUnit", "run_id"), + ("EvaluateModelUnit", "run_id"), + ("FitModelUnit", "run_id"), + ("FitModelOverFoldsUnit", "run_id"), + ("FitModelOverFoldsUnit", "artifact_prefix"), + ("FitModelOverNestedFoldsUnit", "run_id"), + ("FitModelOverNestedFoldsUnit", "artifact_prefix"), + ("FitModelUnit", "artifact_prefix"), + ("SaveModelUnit", "artifact_prefix"), +} + + +def test_the_runtime_params_are_exactly_the_declared_ones(unit_classes): + """Adding or removing one has to be a deliberate edit, not a side effect.""" + declared = { + (name, param) + for name, cls in unit_classes.items() + for param in cls.RUNTIME_PARAMS + } + + assert declared == EXPECTED_RUNTIME_PARAMS + + +def test_no_runtime_param_reaches_the_schema_the_front_receives(units, unit_classes): + """The guarantee the whole arrangement exists for. + + A flag inside the schema would have left these on the wire, where every + renderer has to remember to skip them and a new one leaks by default. Being + a separate declaration means there is nothing to filter: the name is simply + not in what the front is given. + """ + for name, cls in unit_classes.items(): + exposed = set(units[name]["schema"]["properties"]) + leaked = exposed & set(cls.RUNTIME_PARAMS) + + assert not leaked, ( + f"{name} exposes {sorted(leaked)} in the schema the front renders, " + "even though it is declared as supplied by whatever runs the unit." + ) + + +def test_a_unit_can_have_no_user_configuration_at_all(units): + """Which is a real shape, not a degenerate one. + + Three units are told everything by their caller, and SaveDatasetUnit was + already configuration-free before any of this. An empty schema is valid and + renders as nothing. + """ + for name in ("SaveModelUnit", "LoadTrainingDatasetUnit", "SaveDatasetUnit"): + assert units[name]["schema"]["properties"] == {}, name + assert units[name]["configurable_object"] is True, name diff --git a/tests/back/conftest.py b/tests/back/conftest.py index e1a96dfb6..a9d525822 100644 --- a/tests/back/conftest.py +++ b/tests/back/conftest.py @@ -7,6 +7,7 @@ import pytest from DashAI.back.dependencies.job_queues.huey_job_queue import HueyJobQueue +from tests.back.scratch import clear_scratch TEST_PATH = pathlib.Path("tmp") TEST_DATASETS_PATH = pathlib.Path("./tests/back/test_datasets") @@ -18,6 +19,14 @@ def test_path(): return TEST_PATH +@pytest.fixture(scope="session", autouse=True) +def _clean_scratch_directories(): + """Drop the dataset caches the suite writes outside the repository.""" + clear_scratch() + yield + clear_scratch() + + @pytest.fixture(scope="session", autouse=True) def random_state(): return RANDOM_STATE diff --git a/tests/back/dag/__init__.py b/tests/back/dag/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/back/dag/test_engine.py b/tests/back/dag/test_engine.py new file mode 100644 index 000000000..6975b807b --- /dev/null +++ b/tests/back/dag/test_engine.py @@ -0,0 +1,474 @@ +"""The real engine, running real units as a graph, with no database. + +The graph below is the train/test converter flow, which is the smallest thing a +single shared context genuinely cannot run: two ``LoadDatasetUnit`` nodes both +write the fixed key ``dataset``, so in one context the second load erases the +first. + + load_train --dataset--+-------------------> tx_train --dataset--> save_train + | ^ ^ + +-> fit --converter--+---+ dataset_path | + | --+ + load_test --dataset------------------------|-> tx_test --dataset--> save_test + +Two fan-outs and two joins. The number that proves it worked is 2.0: +MinMaxScaler fitted on train [0, 5, 10] learns min=0 max=10, so a test value of +20 scales to 2.0. A refit on the test data would have produced 0.0. + +The spike that established this runs the same graph against a throwaway engine +and stays as a frozen record of the unit contract. This exercises the shipped +one, which differs in the two ways that matter: nodes name their unit instead +of carrying an instance, and there are no seeds. +""" + +from pathlib import Path + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.converters.scikit_learn.min_max_scaler import MinMaxScaler +from DashAI.back.dag.engine import NullSink, artifacts_of, run +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node, connect, sinks +from DashAI.back.dag.validate import instantiate, resolve_unit_class, validate +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.types.value_types import Float +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_model_unit import SaveModelUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit + +FULL_SCOPE = {"columns": [], "rows": []} + + +class _Row: + """Stand-in for a Dataset ORM row, as in tests/back/units.""" + + def __init__(self, file_path): + self.file_path = file_path + self.dataset_id = None + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """A class, not a lambda: kink resolves a registered lambda by calling it.""" + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +def _write(root, **columns): + frame = pd.DataFrame(columns) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + save_dataset(to_dashai_dataset(frame, types=types), str(root / "dataset")) + return root + + +def _values(path, column="a"): + return list(load_dataset(str(path)).to_pandas()[column]) + + +@pytest.fixture(name="stored_datasets") +def fixture_stored_datasets(tmp_path): + """Dataset 7 is the training data, dataset 8 the test data. + + The engine resolves a node's unit by name, so the registry has to hold the + unit classes as well as the converter. + """ + train = _write(tmp_path / "train", a=[0.0, 5.0, 10.0]) + test = _write(tmp_path / "test", a=[20.0]) + + di["session_factory"] = _FakeSessionFactory( + {"Dataset": {7: _Row(str(train)), 8: _Row(str(test))}} + ) + di["component_registry"] = { + "MinMaxScaler": {"class": MinMaxScaler}, + "LoadDatasetUnit": {"class": LoadDatasetUnit}, + "FitConverterUnit": {"class": FitConverterUnit}, + "TransformDatasetUnit": {"class": TransformDatasetUnit}, + "SaveDatasetUnit": {"class": SaveDatasetUnit}, + "SaveDatasetToPathUnit": {"class": SaveDatasetToPathUnit}, + "SaveModelUnit": {"class": SaveModelUnit}, + "NotAUnit": {"class": MinMaxScaler}, + } + yield train, test + del di["session_factory"] + del di["component_registry"] + + +def _graph(tmp_path, converter="MinMaxScaler"): + """The diamond, wired with the bundling rule wherever it applies.""" + load_train = Node("load_train", "LoadDatasetUnit", {"dataset_id": 7}) + load_test = Node("load_test", "LoadDatasetUnit", {"dataset_id": 8}) + fit = Node( + "fit", + "FitConverterUnit", + { + "converter": {"component": converter, "params": {}}, + "scope": FULL_SCOPE, + "target": None, + }, + ) + tx_train = Node( + "tx_train", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None} + ) + tx_test = Node( + "tx_test", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None} + ) + # Saves back over the training data, so it needs the ref `dataset_path` + # from the load as well as the live dataset from the transform: a join + # whose two inputs live in different halves of the context. + save_train = Node("save_train", "SaveDatasetUnit") + save_test = Node( + "save_test", + "SaveDatasetToPathUnit", + {"path": str(tmp_path / "out" / "dataset")}, + ) + + nodes = [load_train, load_test, fit, tx_train, tx_test, save_train, save_test] + edges = [ + *connect(load_train, fit), + Edge("load_train", "dataset", "tx_train", "dataset"), + Edge("load_train", "dataset_path", "save_train", "dataset_path"), + *connect(fit, tx_train), + *connect(fit, tx_test), + Edge("load_test", "dataset", "tx_test", "dataset"), + Edge("tx_train", "dataset", "save_train", "dataset"), + *connect(tx_test, save_test), + ] + return Graph(nodes, edges) + + +def test_the_diamond_runs_and_the_fitted_state_survives_the_branch( + stored_datasets, tmp_path +): + """One fit, two branches, no refit. + + 2.0 is only reachable if the converter fitted on the train branch reached + the test branch still holding the range it learned, after crossing two + context boundaries. + """ + train, _ = stored_datasets + + run(_graph(tmp_path)) + + assert _values(train / "dataset") == [0.0, 0.5, 1.0] + assert _values(tmp_path / "out" / "dataset") == [2.0] + + +def test_two_loads_of_the_same_unit_do_not_fight_over_the_key( + stored_datasets, tmp_path +): + """Both loads write the fixed key ``dataset`` and neither is lost. + + This is what a shared context cannot do, and the reason each node gets one + of its own. + """ + train, _ = stored_datasets + + run(_graph(tmp_path)) + + # A second load overwriting the first would have scaled the test value + # against the test data's own range, giving 0.0. + assert _values(tmp_path / "out" / "dataset") == [2.0] + assert _values(train / "dataset") == [0.0, 0.5, 1.0] + + +def test_a_node_names_its_unit_and_the_engine_builds_one_instance_each( + stored_datasets, tmp_path +): + """Each node owns its instance, which is what per-instance state needs. + + A unit memoizes work on itself -- ``FitConverterUnit`` its converter + class -- so two nodes sharing an instance would have the second silently + run with the first one's. + """ + a = instantiate( + Node("a", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + ) + b = instantiate( + Node("b", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + ) + + assert a is not b + assert type(a) is type(b) is TransformDatasetUnit + + +def test_the_execution_order_is_topological_not_the_order_given( + stored_datasets, tmp_path +): + """Order comes from the edges, not from the sequence the nodes arrived in. + + Its predecessor iterated the steps in whatever order the front sent them -- + creation order on the canvas -- and never read the edges it persisted. + """ + graph = _graph(tmp_path) + reversed_graph = Graph(list(reversed(graph.nodes)), graph.edges) + + order = validate(reversed_graph) + + assert order.index("load_train") < order.index("fit") + assert order.index("fit") < order.index("tx_test") + assert order.index("tx_train") < order.index("save_train") + + +def test_a_missing_input_is_caught_before_anything_runs(stored_datasets): + save = Node("save", "SaveDatasetUnit") + + with pytest.raises(GraphError, match="requires 'dataset'"): + validate(Graph([save], [])) + + +def test_two_edges_into_one_port_are_rejected(stored_datasets, tmp_path): + """The failure its predecessor had: a merge where the last writer wins.""" + graph = _graph(tmp_path) + doubled = Graph( + graph.nodes, + [*graph.edges, Edge("load_test", "dataset", "tx_train", "dataset")], + ) + + with pytest.raises(GraphError, match="from 2 edges"): + validate(doubled) + + +def test_a_cycle_is_reported(stored_datasets): + a = Node("a", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + b = Node("b", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + cyclic = Graph( + [a, b], + [ + Edge("a", "dataset", "b", "dataset"), + Edge("b", "dataset", "a", "dataset"), + ], + ) + + with pytest.raises(GraphError, match="cycle"): + validate(cyclic) + + +def test_an_edge_naming_a_key_the_unit_does_not_declare_is_rejected( + stored_datasets, tmp_path +): + graph = _graph(tmp_path) + bogus = Graph( + graph.nodes, + [*graph.edges, Edge("load_train", "nonsense", "save_train", "dataset_path")], + ) + + with pytest.raises(GraphError, match="does not provide 'nonsense'"): + validate(bogus) + + +def test_an_unknown_unit_name_is_rejected(stored_datasets): + with pytest.raises(GraphError, match="no unit named 'Nope'"): + validate(Graph([Node("x", "Nope")], [])) + + +def test_a_registered_component_that_is_not_a_unit_is_rejected(stored_datasets): + """The registry holds every kind of component, not only units.""" + with pytest.raises(GraphError, match="not a unit"): + validate(Graph([Node("x", "NotAUnit")], [])) + + +def test_duplicate_node_ids_are_rejected(stored_datasets, tmp_path): + twice = Node("same", "SaveDatasetToPathUnit", {"path": str(tmp_path / "d")}) + with pytest.raises(GraphError, match="Duplicate node ids"): + validate(Graph([twice, twice], [])) + + +def test_every_problem_is_reported_at_once(stored_datasets): + """A user fixing a graph wants the whole list, not one error per attempt.""" + graph = Graph( + [Node("save", "SaveDatasetUnit"), Node("save", "SaveDatasetUnit")], [] + ) + + with pytest.raises(GraphError) as caught: + validate(graph) + + message = str(caught.value) + assert "Duplicate node ids" in message + assert "requires 'dataset'" in message + + +def test_moving_a_value_across_an_edge_keeps_the_half_it_came_from( + stored_datasets, tmp_path +): + """Refs travel as refs, live objects as live objects. + + ``dataset_path`` is a reference and ``dataset`` is a cached object, and the + two have incompatible rules: a dataset handed to ``put_ref`` raises, and a + reference handed to ``put`` drops the copy-on-read guarantee it depended + on. + """ + contexts = run(_graph(tmp_path)) + + save_train = contexts["save_train"] + assert save_train.origin("dataset_path") == "ref" + assert save_train.origin("dataset") == "cache" + + +def test_a_context_is_released_once_nothing_downstream_reads_it( + stored_datasets, tmp_path +): + """Intermediate datasets do not stay alive for the length of the run.""" + graph = _graph(tmp_path) + + contexts = run(graph) + + assert set(contexts) == sinks(graph) + assert "load_train" not in contexts + + +def test_artifacts_are_the_serializable_outputs_only(stored_datasets, tmp_path): + """What gets recorded is the reference half of PROVIDES. + + The cache holds live datasets, models and tasks: not serializable, and + always derivable again. So ``LoadDatasetUnit`` records the id and the path + it published rather than the dataset itself, and ``TransformDatasetUnit``, + whose only output is the live dataset, records nothing -- a real case, not + a gap. + """ + train, _ = stored_datasets + sink = _RecordingSink() + + run(_graph(tmp_path), sink) + + recorded = {event[1]: event[2] for event in sink.events if event[0] == "finished"} + + assert set(recorded["load_train"]) == {"dataset_id", "dataset_path"} + assert recorded["load_train"]["dataset_id"] == 7 + # Compared as a path, not as text: the unit joins with a forward slash and + # Path renders a backslash on Windows. + assert Path(recorded["load_train"]["dataset_path"]) == train / "dataset" + assert recorded["tx_train"] == {} + + +def test_artifacts_of_reads_only_the_reference_half(stored_datasets): + """The same rule, stated directly against a hand-built context.""" + ctx = ExecutionContext() + ctx.put_ref("dataset_path", "/tmp/ds") + ctx.put_ref("dataset_id", 7) + ctx.put("dataset", object()) + + assert artifacts_of(resolve_unit_class("LoadDatasetUnit"), ctx) == { + "dataset_path": "/tmp/ds", + "dataset_id": 7, + } + + +def test_the_bundling_rule_picks_the_keys_two_nodes_agree_on(stored_datasets): + """One drawn edge between two nodes stands for a set of keys.""" + fit = Node("fit", "FitConverterUnit", {}) + tx = Node("tx", "TransformDatasetUnit", {}) + + assert connect(fit, tx) == ( + Edge("fit", "fitted_converter", "tx", "fitted_converter"), + ) + + +def test_a_node_belongs_to_itself_when_it_has_no_block(): + """Every node has a block, so the tracking never has a null to handle.""" + alone = Node("solo", "SaveDatasetUnit") + grouped = Node("inner", "SaveDatasetUnit", {}, block_id="train-1") + + assert alone.block_id == "solo" + assert grouped.block_id == "train-1" + + +class _RecordingSink(NullSink): + """Records the notifications the engine sends, in order.""" + + def __init__(self): + self.events = [] + + def run_started(self, order): + self.events.append(("run_started", tuple(order))) + + def node_started(self, node_id, payload): + self.events.append(("started", node_id, dict(payload))) + + def node_finished(self, node_id, artifacts, payload): + self.events.append(("finished", node_id, dict(artifacts))) + + def node_failed(self, node_id, message): + self.events.append(("failed", node_id)) + + def nodes_cancelled(self, node_ids): + self.events.append(("cancelled", tuple(node_ids))) + + def run_finished(self): + self.events.append(("run_finished",)) + + def run_failed(self, message): + self.events.append(("run_failed",)) + + +def test_the_sink_sees_every_node_start_and_finish(stored_datasets, tmp_path): + sink = _RecordingSink() + + run(_graph(tmp_path), sink) + + kinds = [event[0] for event in sink.events] + assert kinds[0] == "run_started" + assert kinds[-1] == "run_finished" + assert kinds.count("started") == 7 + assert kinds.count("finished") == 7 + assert "failed" not in kinds + + +def test_a_failing_node_cancels_everything_after_it(stored_datasets, tmp_path): + """A run that died halfway must not look like one still in flight.""" + graph = _graph(tmp_path, converter="NoSuchConverter") + sink = _RecordingSink() + + # The unit reports an unresolvable converter as a JobError, and the engine + # lets it through rather than wrapping it: the message the user reads is + # the unit's own. + with pytest.raises(JobError, match="NoSuchConverter"): + run(graph, sink) + + failed = [event for event in sink.events if event[0] == "failed"] + cancelled = [event for event in sink.events if event[0] == "cancelled"] + + assert [event[1] for event in failed] == ["fit"] + assert cancelled, "the nodes after the failure have to be reported" + assert "tx_train" in cancelled[0][1] + assert "save_test" in cancelled[0][1] + assert ("run_failed",) in sink.events + assert ("run_finished",) not in sink.events + + +def test_nothing_runs_when_the_graph_does_not_validate(stored_datasets, tmp_path): + """Validation comes first, so a broken graph costs no work.""" + sink = _RecordingSink() + save = Node("save", "SaveDatasetUnit") + + with pytest.raises(GraphError): + run(Graph([save], []), sink) + + assert sink.events == [] diff --git a/tests/back/dag/test_expand.py b/tests/back/dag/test_expand.py new file mode 100644 index 000000000..92bfbbc1e --- /dev/null +++ b/tests/back/dag/test_expand.py @@ -0,0 +1,285 @@ +"""Expanding what a canvas holds into the graph the engine runs.""" + +import pytest +from kink import di + +from DashAI.back.dag.expand import artifact_prefix, expand +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node +from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_model_unit import SaveModelUnit + + +@pytest.fixture(name="registry", autouse=True) +def fixture_registry(): + di["component_registry"] = { + "LoadDatasetUnit": {"class": LoadDatasetUnit}, + "SaveDatasetUnit": {"class": SaveDatasetUnit}, + "SaveModelUnit": {"class": SaveModelUnit}, + } + yield + del di["component_registry"] + + +def _block(block_id, unit, config=None, unit_id=None): + return { + "id": block_id, + "type": f"{unit}Block", + "units": [{"id": unit_id or block_id, "unit": unit, "config": config or {}}], + } + + +def test_one_unit_per_block_expands_to_the_graph_written_by_hand(): + steps = [ + _block("load-1", "LoadDatasetUnit", {"dataset_id": 7}), + _block("save-1", "SaveDatasetUnit"), + ] + edges = [{"source": "load-1", "target": "save-1"}] + + graph = expand(steps, edges, pipeline_run_id=3) + + assert graph == Graph( + [ + Node("load-1", "LoadDatasetUnit", {"dataset_id": 7}, block_id="load-1"), + Node("save-1", "SaveDatasetUnit", {}, block_id="save-1"), + ], + [ + Edge("load-1", "dataset", "save-1", "dataset"), + Edge("load-1", "dataset_path", "save-1", "dataset_path"), + ], + ) + + +def test_one_drawn_edge_carries_every_key_the_two_units_agree_on(): + """A canvas cannot draw one edge per key, so one stands for the set.""" + steps = [ + _block("load-1", "LoadDatasetUnit", {"dataset_id": 7}), + _block("save-1", "SaveDatasetUnit"), + ] + + graph = expand(steps, [{"source": "load-1", "target": "save-1"}], 3) + + assert {(edge.src_key, edge.dst_key) for edge in graph.edges} == { + ("dataset", "dataset"), + ("dataset_path", "dataset_path"), + } + + +def test_a_node_keeps_the_block_it_came_from(): + """The canvas colours a block from the status of its nodes.""" + steps = [_block("train-1", "SaveModelUnit", unit_id="train-1/save")] + + graph = expand(steps, [], pipeline_run_id=3) + + assert graph.nodes[0].id == "train-1/save" + assert graph.nodes[0].block_id == "train-1" + + +def test_the_engine_names_the_artifacts_of_a_unit_that_takes_a_prefix(): + steps = [_block("save-1", "SaveModelUnit")] + + graph = expand(steps, [], pipeline_run_id=3) + + assert graph.nodes[0].config["artifact_prefix"] == "pipeline-3-save-1" + + +def test_a_prefix_stored_in_the_graph_is_discarded(): + """No stored value can be right, so none is honoured. + + The prefix is built from the id of the run, and there is no run when a node + is configured -- so anything already there was chosen without the one thing + that decides it. Overriding is also what makes two nodes sharing a prefix + impossible rather than merely unlikely. + """ + steps = [_block("save-1", "SaveModelUnit", {"artifact_prefix": "mine"})] + + graph = expand(steps, [], pipeline_run_id=3) + + assert graph.nodes[0].config["artifact_prefix"] == "pipeline-3-save-1" + + +def test_two_saving_nodes_cannot_end_up_with_the_same_prefix(): + """Which is why the collision needs no validator to catch it.""" + steps = [ + _block("save-a", "SaveModelUnit", {"artifact_prefix": "same"}), + _block("save-b", "SaveModelUnit", {"artifact_prefix": "same"}), + ] + + prefixes = {node.config["artifact_prefix"] for node in expand(steps, [], 3).nodes} + + assert len(prefixes) == 2 + + +@pytest.mark.parametrize( + ("first", "second"), + [("save.a", "save_a"), ("train-1/build", "train-1_build"), ("a b", "a_b")], +) +def test_sanitizing_two_different_ids_cannot_produce_one_prefix(first, second): + """The escape has to be reversible, not merely safe. + + Replacing every unsafe character with a plain ``_`` is safe but collapses + distinct ids: ``save.a`` and ``save_a`` would land on the same directory and + the second saving node would overwrite the first one's model in silence. + That is why ``_`` is escaped along with everything else. + """ + assert artifact_prefix(3, first) != artifact_prefix(3, second) + + +def test_a_unit_that_takes_no_prefix_does_not_get_one(): + steps = [_block("load-1", "LoadDatasetUnit", {"dataset_id": 7})] + + graph = expand(steps, [], pipeline_run_id=3) + + assert "artifact_prefix" not in graph.nodes[0].config + + +def test_two_runs_of_one_pipeline_name_their_artifacts_apart(): + """Otherwise the second run would write over the first one's model.""" + steps = [_block("save-1", "SaveModelUnit")] + + first = expand(steps, [], pipeline_run_id=3).nodes[0].config["artifact_prefix"] + second = expand(steps, [], pipeline_run_id=4).nodes[0].config["artifact_prefix"] + + assert first != second + + +def test_the_prefix_can_never_be_read_as_a_path(): + """A node id a user chose has to be brought into the safe alphabet. + + ``os.path.join`` with a separator in the prefix produces a destination + outside the runs directory, and ``SaveModelUnit.validate`` refuses one -- + so an id with a slash in it has to be sanitized here rather than failing at + the node. + """ + # Each unsafe character becomes _, so the mapping is reversible. + assert artifact_prefix(3, "train-1/save") == "pipeline-3-train-1_2fsave" + # The dot is escaped too, which is what leaves a parent reference inert + # rather than merely separator-free. + assert artifact_prefix(3, "../escape") == "pipeline-3-_2e_2e_2fescape" + + steps = [_block("save-1", "SaveModelUnit", unit_id="train-1/save")] + prefix = expand(steps, [], 3).nodes[0].config["artifact_prefix"] + + # And what comes out passes the guard at the unit. + SaveModelUnit(artifact_prefix=prefix).validate(None) + + +def test_a_pipeline_run_id_alone_would_collide_with_a_real_run(): + """The whole reason the prefix is not just the id. + + RUNS_PATH is shared with every real Run, and pipeline_run.id and run.id are + independent sequences that both start at 1. + """ + assert artifact_prefix(3, "save-1") != "3" + assert artifact_prefix(3, "save-1").startswith("pipeline-") + + +def test_a_block_with_no_units_is_refused_by_name(): + """Rows written by the previous subsystem land here. + + Their steps name a node type and carry a single config, with no unit to + resolve, so they are refused with a message that says so rather than + failing somewhere deeper. + """ + old_style = [{"id": "DataSelector-1", "type": "DataSelector", "config": {}}] + + with pytest.raises(GraphError, match="previous pipeline subsystem"): + expand(old_style, [], 3) + + +def test_more_than_one_unit_per_block_is_refused_rather_than_guessed(): + """Wiring across a sequence boundary has more than one defensible answer.""" + steps = [ + { + "id": "train-1", + "units": [ + {"id": "a", "unit": "LoadDatasetUnit", "config": {"dataset_id": 7}}, + {"id": "b", "unit": "SaveDatasetUnit", "config": {}}, + ], + } + ] + + with pytest.raises(GraphError, match="one unit per block"): + expand(steps, [], 3) + + +def test_a_block_without_an_id_is_refused(): + with pytest.raises(GraphError, match="no id"): + expand([{"units": []}], [], 3) + + +def test_two_blocks_with_the_same_id_are_refused(): + steps = [ + _block("same", "LoadDatasetUnit", {"dataset_id": 7}), + _block("same", "SaveDatasetUnit"), + ] + + with pytest.raises(GraphError, match="more than one block"): + expand(steps, [], 3) + + +def test_an_edge_to_a_block_that_is_not_there_is_refused(): + steps = [_block("load-1", "LoadDatasetUnit", {"dataset_id": 7})] + + with pytest.raises(GraphError, match="not a block"): + expand(steps, [{"source": "load-1", "target": "ghost"}], 3) + + +def test_an_edge_that_would_carry_nothing_is_refused(): + """Two units with no key in common cannot be usefully connected.""" + steps = [ + _block("save-1", "SaveModelUnit"), + _block("load-1", "LoadDatasetUnit", {"dataset_id": 7}), + ] + + with pytest.raises(GraphError, match="would carry nothing"): + expand(steps, [{"source": "save-1", "target": "load-1"}], 3) + + +def test_a_node_that_takes_a_run_id_is_told_there_is_no_run(): + """A pipeline has no Run row, and that is a value rather than an omission. + + It is how a node says the model it builds belongs to no run, which is what + keeps that model from trying to log metrics against a foreign key that + points at nothing. + """ + di["component_registry"]["FitModelUnit"] = {"class": FitModelUnit} + steps = [ + _block( + "fit-1", + "FitModelUnit", + { + "optimizer": {"component": "OptunaOptimizer", "params": {}}, + "goal_metric": "Accuracy", + }, + ) + ] + + config = expand(steps, [], pipeline_run_id=3).nodes[0].config + + assert config["run_id"] is None + assert config["artifact_prefix"] == "pipeline-3-fit-1" + + +def test_a_run_id_stored_in_the_graph_is_discarded(): + """A pipeline has no run, so a stored one is overridden rather than trusted. + + This is what makes the sandbox mechanical instead of conventional: a + hand-edited row cannot point a training node at a real run and start + writing Metric rows into it. + """ + di["component_registry"]["FitModelUnit"] = {"class": FitModelUnit} + steps = [ + _block( + "fit-1", + "FitModelUnit", + { + "optimizer": {"component": "OptunaOptimizer", "params": {}}, + "goal_metric": "Accuracy", + "run_id": 17, + }, + ) + ] + + assert expand(steps, [], 3).nodes[0].config["run_id"] is None diff --git a/tests/back/dag/test_pipeline_job.py b/tests/back/dag/test_pipeline_job.py new file mode 100644 index 000000000..8c1327c1a --- /dev/null +++ b/tests/back/dag/test_pipeline_job.py @@ -0,0 +1,339 @@ +"""PipelineJob end to end, against a real database. + +Its predecessor implemented two of BaseJob's four abstract methods, so it could +not be instantiated at all: asking the job endpoint for a PipelineJob raised +TypeError. These check the job runs, and that the four methods are there. +""" + +import pytest +from kink import di +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import StaticPool + +from DashAI.back.core.enums.status import NodeRunStatus, PipelineRunStatus +from DashAI.back.dependencies.database.models import ( + Base, + NodeArtifact, + NodeRun, + Pipeline, + PipelineRun, +) +from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.job.pipeline_job import PipelineJob +from DashAI.back.units.base_unit import BaseUnit + + +class _EmitUnit(BaseUnit): + REQUIRES = () + PROVIDES = ("dataset_path", "dataset") + + def execute(self, ctx): + ctx.put_ref("dataset_path", "/tmp/ds") + ctx.put("dataset", object()) + + +class _ConsumeUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = ("results_path",) + + def execute(self, ctx): + ctx.require("dataset") + ctx.put_ref("results_path", "/tmp/results.json") + + +class _FailUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = () + + def execute(self, ctx): + raise JobError("this node was always going to fail") + + +@pytest.fixture(name="database") +def fixture_database(): + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + Base.metadata.create_all(engine) + factory = sessionmaker(bind=engine) + + class _Factory: + def __call__(self): + return factory() + + di["session_factory"] = _Factory() + di["component_registry"] = { + "EmitUnit": {"class": _EmitUnit}, + "ConsumeUnit": {"class": _ConsumeUnit}, + "FailUnit": {"class": _FailUnit}, + } + yield factory + del di["session_factory"] + del di["component_registry"] + engine.dispose() + + +def _pipeline(factory, steps, edges, name="a pipeline"): + with factory() as db: + pipeline = Pipeline(name=name, steps=steps, edges=edges) + db.add(pipeline) + db.commit() + return pipeline.id + + +def _block(block_id, unit): + return { + "id": block_id, + "units": [{"id": block_id, "unit": unit, "config": {}}], + } + + +def test_the_job_implements_every_abstract_method_of_base_job(): + """The reason its predecessor could not be instantiated at all.""" + missing = { + name + for name in getattr(BaseJob, "__abstractmethods__", set()) + if getattr(PipelineJob, name) is getattr(BaseJob, name) + } + assert not missing + assert PipelineJob(pipeline_id=1) is not None + + +def test_a_pipeline_runs_and_records_what_each_node_produced(database): + pipeline_id = _pipeline( + database, + [_block("emit-1", "EmitUnit"), _block("consume-1", "ConsumeUnit")], + [{"source": "emit-1", "target": "consume-1"}], + ) + + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.FINISHED + assert {row.node_id for row in pipeline_run.node_runs} == { + "emit-1", + "consume-1", + } + assert {row.status for row in pipeline_run.node_runs} == { + NodeRunStatus.FINISHED + } + artifacts = { + (row.node_run.node_id, row.key) for row in db.query(NodeArtifact).all() + } + + assert artifacts == { + ("emit-1", "dataset_path"), + ("consume-1", "results_path"), + } + + +def test_the_run_freezes_the_expanded_graph(database): + """A run records units and keys, not the blocks a canvas drew.""" + pipeline_id = _pipeline( + database, + [_block("emit-1", "EmitUnit"), _block("consume-1", "ConsumeUnit")], + [{"source": "emit-1", "target": "consume-1"}], + ) + + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + + assert pipeline_run.edges == [ + { + "src": "emit-1", + "src_key": "dataset", + "dst": "consume-1", + "dst_key": "dataset", + } + ] + assert [step["unit"] for step in pipeline_run.steps] == [ + "EmitUnit", + "ConsumeUnit", + ] + + +def test_a_graph_that_cannot_work_says_why_and_marks_the_run(database): + """An unfed input is a mistake in the graph, reported as one.""" + pipeline_id = _pipeline(database, [_block("consume-1", "ConsumeUnit")], []) + + with pytest.raises(JobError, match="requires 'dataset'"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "requires 'dataset'" in pipeline_run.error_message + # Nothing ran, so there is nothing to show per node. + assert db.query(NodeRun).count() == 0 + + +def test_steps_from_the_previous_subsystem_are_refused_by_name(database): + """Rows the old canvas saved name a node type and have no unit.""" + pipeline_id = _pipeline( + database, + [{"id": "DataSelector-1", "type": "DataSelector", "config": {}}], + [], + ) + + with pytest.raises(JobError, match="previous pipeline subsystem"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).one().status == PipelineRunStatus.ERROR + + +def test_a_pipeline_with_no_steps_is_refused_before_a_run_is_created(database): + pipeline_id = _pipeline(database, [], []) + + with pytest.raises(JobError, match="no steps"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).count() == 0 + + +def test_a_missing_pipeline_is_refused(database): + with pytest.raises(JobError, match="does not exist"): + PipelineJob(pipeline_id=9999).run() + + +def test_two_runs_of_one_pipeline_name_their_artifacts_apart(database): + """The reason the run row is created before the graph is expanded.""" + pipeline_id = _pipeline(database, [_block("emit-1", "EmitUnit")], []) + + PipelineJob(pipeline_id=pipeline_id).run() + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + runs = db.query(PipelineRun).order_by(PipelineRun.id).all() + assert len(runs) == 2 + assert runs[0].id != runs[1].id + # Both runs recorded their own history rather than overwriting one. + assert {row.status for run in runs for row in run.node_runs} == { + NodeRunStatus.FINISHED + } + + +def test_the_job_name_comes_from_the_pipeline(database): + pipeline_id = _pipeline(database, [], [], name="my flow") + + assert PipelineJob(pipeline_id=pipeline_id).get_job_name() == "Pipeline: my flow" + assert PipelineJob(pipeline_id=9999).get_job_name() == "Pipeline (9999)" + + +def test_the_job_kwargs_are_only_plain_data(database): + """The whole job is serialized with dill to reach the worker process.""" + import dill + + job = PipelineJob(pipeline_id=7) + + assert job.kwargs == {"pipeline_id": 7} + assert dill.loads(dill.dumps(job)).kwargs == {"pipeline_id": 7} + + +def test_the_id_the_front_actually_sends_is_accepted(database): + """Regression: the wire contract predates this job. + + ``api/job.ts`` posts ``kwargs: {id: pipelineId}``, and ``jobs.py`` passes + kwargs through verbatim. Reading only ``pipeline_id`` meant every enqueued + pipeline died in the worker with a raw KeyError, before any run row existed. + """ + pipeline_id = _pipeline(database, [_block("emit-1", "EmitUnit")], []) + + PipelineJob(id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).one().status == PipelineRunStatus.FINISHED + + +def test_either_name_reaches_the_same_pipeline(database): + pipeline_id = _pipeline(database, [], [], name="named") + + assert PipelineJob(id=pipeline_id).get_job_name() == "Pipeline: named" + assert PipelineJob(pipeline_id=pipeline_id).get_job_name() == "Pipeline: named" + + +def test_a_job_with_no_id_at_all_says_so(database): + with pytest.raises(JobError, match="No pipeline id"): + PipelineJob().run() + + assert PipelineJob().get_job_name() == "Pipeline" + + +def test_a_run_is_never_left_looking_like_it_is_still_going(database, monkeypatch): + """The backstop for the paths where the sink recorded nothing. + + The sink's own calls sit outside the engine's try, so one of them failing + -- a locked database, a run row deleted underneath -- raises without any + status having been written. A run stuck in STARTED with no error message is + indistinguishable from one still in flight. + """ + from DashAI.back.dag import tracking + + def explode(self, node_id, payload): + raise RuntimeError("the sink could not write") + + monkeypatch.setattr(tracking.DatabaseSink, "node_started", explode) + + pipeline_id = _pipeline(database, [_block("emit-1", "EmitUnit")], []) + + with pytest.raises(JobError, match="could not write"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "could not write" in pipeline_run.error_message + + +def test_the_backstop_does_not_overwrite_what_the_sink_recorded(database): + """A node's own message is the better one, so it has to survive.""" + pipeline_id = _pipeline( + database, + [_block("emit-1", "EmitUnit"), _block("boom-1", "FailUnit")], + [{"source": "emit-1", "target": "boom-1"}], + ) + + with pytest.raises(JobError, match="always going to fail"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "always going to fail" in pipeline_run.error_message + + +class _NeedsARuntimeParam(BaseUnit): + """Stands in for a unit only a particular job knows how to configure.""" + + REQUIRES = () + PROVIDES = ("dataset_path",) + RUNTIME_PARAMS = ("temp_path",) + + def execute(self, ctx): + ctx.put_ref("dataset_path", self.config["temp_path"]) + + +def test_a_runtime_param_nobody_supplies_fails_before_anything_runs(database): + """Not halfway through, after earlier nodes already wrote to disk. + + Only two runtime params are about a pipeline run, so the engine can answer + those. A unit whose runtime params only a particular job knows how to fill + is not usable as a node yet, and this is where a user is told so. + """ + di["component_registry"]["NeedsARuntimeParam"] = {"class": _NeedsARuntimeParam} + pipeline_id = _pipeline(database, [_block("odd-1", "NeedsARuntimeParam")], []) + + with pytest.raises(JobError, match=r"needs \['temp_path'\]"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).one().status == PipelineRunStatus.ERROR + # Nothing ran, so there is nothing to show per node. + assert db.query(NodeRun).count() == 0 diff --git a/tests/back/dag/test_tracking.py b/tests/back/dag/test_tracking.py new file mode 100644 index 000000000..ec1c6b885 --- /dev/null +++ b/tests/back/dag/test_tracking.py @@ -0,0 +1,266 @@ +"""What a run writes down, against a real database. + +The engine is exercised with a sink that records nothing in test_engine.py; +this checks the sink that persists, including the two things "no row" cannot +express: a node still waiting its turn, and a node whose run died before +reaching it. +""" + +import pytest +from kink import di +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import StaticPool + +from DashAI.back.core.enums.status import NodeRunStatus, PipelineRunStatus +from DashAI.back.dag.engine import run +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node, dump, load +from DashAI.back.dag.tracking import DatabaseSink +from DashAI.back.dependencies.database.models import ( + Base, + NodeArtifact, + NodeRun, + Pipeline, + PipelineRun, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit + + +class _EmitUnit(BaseUnit): + """Publishes a reference and a live object, so both halves are covered.""" + + REQUIRES = () + PROVIDES = ("dataset_path", "dataset") + + def execute(self, ctx): + ctx.put_ref("dataset_path", self.config["path"]) + ctx.put("dataset", object()) + + +class _ConsumeUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = ("results_path",) + + def execute(self, ctx): + ctx.require("dataset") + ctx.put_ref("results_path", "/tmp/results.json") + + +class _FailUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = () + + def execute(self, ctx): + raise JobError("this node was always going to fail") + + +@pytest.fixture(name="database") +def fixture_database(): + """An in-memory database shared across sessions. + + StaticPool because every notification opens its own session: without it + each connection would get a fresh empty database. + """ + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + Base.metadata.create_all(engine) + factory = sessionmaker(bind=engine) + + class _Factory: + """Not a lambda: kink calls a registered lambda when resolving it.""" + + def __call__(self): + return factory() + + di["session_factory"] = _Factory() + di["component_registry"] = { + "EmitUnit": {"class": _EmitUnit}, + "ConsumeUnit": {"class": _ConsumeUnit}, + "FailUnit": {"class": _FailUnit}, + } + + with factory() as db: + pipeline = Pipeline(name="a pipeline", steps=[], edges=[]) + db.add(pipeline) + db.commit() + # The run row is the caller's to create: naming a node's artifacts + # needs its id, and that comes before the graph exists. + pipeline_run = PipelineRun(pipeline_id=pipeline.id) + db.add(pipeline_run) + db.commit() + pipeline_run_id = pipeline_run.id + + yield factory, pipeline_run_id + + del di["session_factory"] + del di["component_registry"] + engine.dispose() + + +def _pair_graph(): + emit = Node("emit", "EmitUnit", {"path": "/tmp/ds"}) + consume = Node("consume", "ConsumeUnit") + return Graph([emit, consume], [Edge("emit", "dataset", "consume", "dataset")]) + + +def test_a_finished_run_records_every_node_and_its_artifacts(database): + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, sink.pipeline_run_id) + assert pipeline_run.status == PipelineRunStatus.FINISHED + assert pipeline_run.start_time is not None + assert pipeline_run.end_time is not None + + node_runs = {row.node_id: row for row in pipeline_run.node_runs} + assert set(node_runs) == {"emit", "consume"} + for row in node_runs.values(): + assert row.status == NodeRunStatus.FINISHED + assert row.start_time is not None + assert row.end_time is not None + + artifacts = { + (row.node_run.node_id, row.key): row.value + for row in db.query(NodeArtifact).all() + } + + # The live dataset is not recorded: the cache half is not serializable and + # is always derivable again. The path it published is. + assert artifacts == { + ("emit", "dataset_path"): "/tmp/ds", + ("consume", "results_path"): "/tmp/results.json", + } + + +def test_the_run_freezes_the_graph_it_executed(database): + """A past run stays readable after the pipeline it came from is edited.""" + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, sink.pipeline_run_id) + frozen_steps, frozen_edges = pipeline_run.steps, pipeline_run.edges + + steps, edges = dump(graph) + assert frozen_steps == steps + assert frozen_edges == edges + # And the frozen form is enough to rebuild the graph, without re-running + # whatever produced it. + assert load(frozen_steps, frozen_edges) == graph + + +def test_a_failure_marks_the_node_and_cancels_the_rest(database): + """CANCELLED is what NOT_STARTED cannot say. + + Without it, a run that died before reaching a node is indistinguishable + from one still in flight. + """ + factory, pipeline_run_id = database + emit = Node("emit", "EmitUnit", {"path": "/tmp/ds"}) + boom = Node("boom", "FailUnit") + consume = Node("consume", "ConsumeUnit") + graph = Graph( + [emit, boom, consume], + [ + Edge("emit", "dataset", "boom", "dataset"), + Edge("emit", "dataset", "consume", "dataset"), + ], + ) + sink = DatabaseSink(pipeline_run_id, graph) + + with pytest.raises(JobError, match="always going to fail"): + run(graph, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, sink.pipeline_run_id) + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "always going to fail" in pipeline_run.error_message + + rows = {row.node_id: row for row in pipeline_run.node_runs} + + assert rows["emit"].status == NodeRunStatus.FINISHED + assert rows["boom"].status == NodeRunStatus.ERROR + assert "always going to fail" in rows["boom"].error_message + assert rows["consume"].status == NodeRunStatus.CANCELLED + + +def test_every_node_has_a_row_before_the_first_one_runs(database): + """Rows exist up front, so a node waiting its turn is visible as waiting.""" + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + sink.run_started(["emit", "consume"]) + + with factory() as db: + rows = db.query(NodeRun).all() + assert {row.node_id for row in rows} == {"emit", "consume"} + assert {row.status for row in rows} == {NodeRunStatus.NOT_STARTED} + + +def test_a_node_belongs_to_a_block_even_when_it_is_alone(database): + """block_id is never null, so grouping later needs no migration.""" + factory, pipeline_run_id = database + graph = Graph( + [Node("inner", "EmitUnit", {"path": "/tmp/ds"}, block_id="train-1")], [] + ) + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + row = db.query(NodeRun).one() + assert row.node_id == "inner" + assert row.block_id == "train-1" + assert row.node_type == "EmitUnit" + + +def test_a_graph_that_does_not_validate_leaves_the_run_untouched(database): + """Validation happens before the engine reports anything. + + So the run stays exactly as its caller made it -- no node rows, no start + time, still NOT_STARTED -- and saying why is the caller's to do. Marking it + started and then failed would claim a run began when nothing did. + """ + factory, pipeline_run_id = database + orphan = Graph([Node("consume", "ConsumeUnit")], []) + sink = DatabaseSink(pipeline_run_id, orphan) + + with pytest.raises(GraphError, match="requires 'dataset'"): + run(orphan, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, pipeline_run_id) + assert pipeline_run.status == PipelineRunStatus.NOT_STARTED + assert pipeline_run.start_time is None + assert db.query(NodeRun).count() == 0 + + +def test_the_node_input_and_output_are_the_serializable_halves(database): + """What crossed into a node and what it left behind, as plain data.""" + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + rows = {row.node_id: row for row in db.query(NodeRun).all()} + + # `dataset` is a live object, so it appears in neither: this is what + # replaced skipping the key by name when serializing. + assert rows["emit"].input == {} + assert rows["emit"].output == {"dataset_path": "/tmp/ds"} + assert rows["consume"].input == {} + assert rows["consume"].output == {"results_path": "/tmp/results.json"} diff --git a/tests/back/dataloaders/base_tabular_dataloader_tests.py b/tests/back/dataloaders/base_tabular_dataloader_tests.py index 7a6a43c23..4736b66bd 100644 --- a/tests/back/dataloaders/base_tabular_dataloader_tests.py +++ b/tests/back/dataloaders/base_tabular_dataloader_tests.py @@ -10,6 +10,7 @@ from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset, split_dataset from DashAI.back.dataloaders.classes.dataloader import BaseDataLoader +from tests.back.scratch import scratch_dir # TODO: Test no header, empty file, bad split folder structure. @@ -70,7 +71,7 @@ def _test_load_data_from_file( # load data dataset = dataloder_instance.load_data( filepath_or_buffer=file, - temp_path="tests/back/dataloaders", + temp_path=scratch_dir("dataloaders"), params=params, ) @@ -116,7 +117,7 @@ def _test_load_data_from_zip( dataset = dataloder_instance.load_data( filepath_or_buffer=str(dataset_path), - temp_path="tests/back/dataloaders/iris", + temp_path=scratch_dir("dataloaders", "iris"), params=params, ) dataset = split_dataset(dataset) @@ -171,7 +172,7 @@ def _test_dataloader_with_missing_required_params( ): dataloder_instance.load_data( filepath_or_buffer=file, - temp_path="tests/back/dataloaders/iris", + temp_path=scratch_dir("dataloaders", "iris"), params=params, ) @@ -205,6 +206,6 @@ def _test_dataloader_try_to_load_a_invalid_datasets( ): dataloder_instance.load_data( filepath_or_buffer=file, - temp_path="tests/back/dataloaders/iris", + temp_path=scratch_dir("dataloaders", "iris"), params=params, ) diff --git a/tests/back/dataloaders/test_dashai_dataset.py b/tests/back/dataloaders/test_dashai_dataset.py index 0aec2c50c..5dbd572b9 100644 --- a/tests/back/dataloaders/test_dashai_dataset.py +++ b/tests/back/dataloaders/test_dashai_dataset.py @@ -23,6 +23,7 @@ update_dataset_splits, validate_inputs_outputs, ) +from tests.back.scratch import scratch_dir from tests.back.test_datasets_generator import CSVTestDatasetGenerator @@ -61,7 +62,7 @@ def load_test_datasetdict(test_datasets_path: pathlib.Path) -> DatasetDict: test_datasetdict = CSVDataLoader().load_data( filepath_or_buffer=file, - temp_path=str(test_datasets_path), + temp_path=scratch_dir("dashai_dataset"), params={"separator": ";"}, ) diff --git a/tests/back/evaluation/test_forecasting_strategies.py b/tests/back/evaluation/test_forecasting_strategies.py index a2f8b3c17..abee9c8f7 100644 --- a/tests/back/evaluation/test_forecasting_strategies.py +++ b/tests/back/evaluation/test_forecasting_strategies.py @@ -127,17 +127,40 @@ def test_the_ordinary_holdout_strategy_still_scores_all_three(): ) +def _fit_unit(**overrides): + """A FitModelUnit configured for an ordinary run. + + The strategies declare how a run is evaluated; this unit is what carries it + out. The tests below live next to the declarations because that is the pair + that has to stay consistent -- a strategy that says it does not score the + training partition, and a fit that then does not. + """ + from DashAI.back.units.fit_model_unit import FitModelUnit + + config = { + "optimizer": {"component": "", "params": {}}, + "goal_metric": "", + "run_id": None, + "artifact_prefix": None, + } + config.update(overrides) + return FitModelUnit(**config) + + # --- the final fit ----------------------------------------------------------- def test_the_kept_model_stops_at_the_end_of_training(): + """Handing the validation partition to the fit is not fitting on it. + + A model is given it to watch and stop early. A forecaster must not be + advanced through it: the kept model has to end where the training rows end, + or the validation metrics describe a model that already saw them. + """ xs, ys, _ = _split() - strategy = ForecastingHoldoutEvaluationStrategy.__new__( - ForecastingHoldoutEvaluationStrategy - ) model = ExponentialSmoothing(seasonal="add", season_length=12) - strategy._fit_final_model(model, xs, ys) + _fit_unit()._fit_kept_model(model, xs, ys) last_train_date = pd.to_datetime(xs["train"].to_pandas().iloc[:, 0]).max() assert model._last_train_date == last_train_date @@ -147,15 +170,12 @@ def test_the_kept_model_is_the_one_the_validation_metrics_describe(): from DashAI.back.metrics.regression.mae import MAE xs, ys, _ = _split() - strategy = ForecastingHoldoutEvaluationStrategy.__new__( - ForecastingHoldoutEvaluationStrategy - ) scored = ExponentialSmoothing(seasonal="add", season_length=12) scored.train(xs["train"], ys["train"]) kept = ExponentialSmoothing(seasonal="add", season_length=12) - strategy._fit_final_model(kept, xs, ys) + _fit_unit()._fit_kept_model(kept, xs, ys) assert MAE.score(ys["validation"], kept.predict(xs["validation"])) == MAE.score( ys["validation"], scored.predict(xs["validation"]) @@ -163,16 +183,18 @@ def test_the_kept_model_is_the_one_the_validation_metrics_describe(): def test_a_session_without_validation_rows_still_fits(): + """Present and empty, which is not the same as absent. + + The fit is handed a validation partition with no rows in it rather than + none at all, and has to cope: a session may reserve nothing for it. + """ xs, ys, _ = _split() empty = xs["validation"].select(range(0)) xs = {**xs, "validation": empty} ys = {**ys, "validation": ys["validation"].select(range(0))} - strategy = ForecastingHoldoutEvaluationStrategy.__new__( - ForecastingHoldoutEvaluationStrategy - ) model = ExponentialSmoothing(seasonal="add", season_length=12) - strategy._fit_final_model(model, xs, ys) + _fit_unit()._fit_kept_model(model, xs, ys) last_train_date = pd.to_datetime(xs["train"].to_pandas().iloc[:, 0]).max() assert model._last_train_date == last_train_date @@ -202,25 +224,37 @@ def calculate_metrics(self, split, level=None, **kwargs): self.scored.append(split) -def _evaluate_with(strategy_class): +def _trial_of(strategy_class): + """Run one trial configured the way this strategy's runs are configured. + + Which partitions a trial records is no longer the strategy's own code: the + job reads SCORED_SPLITS off the class and hands it to the fitting unit, + minus the test partition, which a trial may never score. This does the same + thing so the declaration stays connected to what it produces. + """ from DashAI.back.metrics.regression.mae import MAE xs, ys, _ = _split() - strategy = strategy_class.__new__(strategy_class) model = _RecordingModel() - strategy.evaluate(model, xs, ys, MAE) + trial_splits = [ + split.name + for split in strategy_class.SCORED_SPLITS + if split is not SplitEnum.TEST + ] + + _fit_unit(trial_splits=trial_splits)._score_one_trial(model, xs, ys, MAE) return model.scored def test_a_forecasting_trial_never_scores_the_training_partition(): - scored = _evaluate_with(ForecastingHoldoutEvaluationStrategy) + scored = _trial_of(ForecastingHoldoutEvaluationStrategy) assert SplitEnum.TRAIN not in scored assert SplitEnum.VALIDATION in scored def test_an_ordinary_trial_still_scores_both(): - scored = _evaluate_with(HoldoutEvaluationStrategy) + scored = _trial_of(HoldoutEvaluationStrategy) assert SplitEnum.TRAIN in scored assert SplitEnum.VALIDATION in scored diff --git a/tests/back/explainers/test_explainers.py b/tests/back/explainers/test_explainers.py index 2c854683a..14a3b8648 100644 --- a/tests/back/explainers/test_explainers.py +++ b/tests/back/explainers/test_explainers.py @@ -24,6 +24,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir INPUT_COLUMNS = [ "SepalLengthCm", @@ -46,7 +47,7 @@ def tabular_model_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/explainers/test_lib_explainers.py b/tests/back/explainers/test_lib_explainers.py index dbd3790df..4c16c6d4b 100644 --- a/tests/back/explainers/test_lib_explainers.py +++ b/tests/back/explainers/test_lib_explainers.py @@ -22,6 +22,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir INPUT_COLUMNS = [ "SepalLengthCm", @@ -44,7 +45,7 @@ def tabular_dataset_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/explainers/test_new_explainers.py b/tests/back/explainers/test_new_explainers.py index 63418b633..a21324557 100644 --- a/tests/back/explainers/test_new_explainers.py +++ b/tests/back/explainers/test_new_explainers.py @@ -21,6 +21,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir INPUT_COLUMNS = [ "SepalLengthCm", @@ -43,7 +44,7 @@ def tabular_model_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/explainers/test_shap_predictor_handover.py b/tests/back/explainers/test_shap_predictor_handover.py new file mode 100644 index 000000000..dcce6df25 --- /dev/null +++ b/tests/back/explainers/test_shap_predictor_handover.py @@ -0,0 +1,88 @@ +"""SHAP must not be handed a bound method of the model. + +``shap.utils._legacy.convert_to_model`` suppresses scikit-learn's "X does not +have valid feature names" warning by blanking ``feature_names_in_`` on the +object the callable is bound to, reached through ``__self__``. It assumes that +attribute is writable, which does not hold for every estimator: those that +expose ``feature_names_in_`` as a read-only ``property`` raise instead, and the +explanation dies before it starts. + +``as_shap_predictor`` hands over a plain closure, which has no ``__self__``, so +SHAP skips that step entirely. These tests pin the mechanism. +""" + +import pytest + +from DashAI.back.explainability.model_input import as_shap_predictor + + +def test_the_wrapped_predictor_forwards_to_predict_prepared_positionally(): + """SHAP calls the model with one positional argument; that must not change.""" + seen = {} + + class Model: + def predict_prepared(self, x): + seen["arg"] = x + return [0] + + predictor = as_shap_predictor(Model()) + assert predictor("the frame") == [0] + assert seen["arg"] == "the frame" + + +def test_the_wrapped_predictor_never_routes_through_predict(): + """The callers hand SHAP a background already in the model's feature space. + + Routing to ``predict`` would prepare an already prepared matrix a second + time, and SHAP perturbs it into plain arrays that the preparation cannot + consume at all — which surfaces far away, as + ``'numpy.ndarray' object has no attribute 'types'``. + """ + + class Model: + def predict(self, x): + raise AssertionError("predict would prepare an already prepared matrix") + + def predict_prepared(self, x): + return [1] + + assert as_shap_predictor(Model())("the frame") == [1] + + +def test_the_wrapped_predictor_hides_the_model_from_shap(): + """The whole mechanism: no ``__self__`` means SHAP never reaches the model.""" + + class Model: + def predict_prepared(self, x): + return [0] + + model = Model() + assert getattr(model.predict_prepared, "__self__", None) is model + assert getattr(as_shap_predictor(model), "__self__", None) is None + + +def test_the_wrapped_predictor_survives_shaps_model_conversion(): + """``convert_to_model`` must accept the closure without touching a model.""" + from shap.utils._legacy import convert_to_model + + class ReadOnlyFeatureNames: + """Stands in for an estimator whose ``feature_names_in_`` has no setter.""" + + @property + def feature_names_in_(self): + return ["a", "b"] + + def predict_prepared(self, x): + return [0] + + model = ReadOnlyFeatureNames() + + # The failure this exists to prevent: SHAP reaches the model through + # ``__self__`` and tries to blank the attribute. + with pytest.raises(AttributeError, match="feature_names_in_"): + convert_to_model(model.predict_prepared) + + converted = convert_to_model(as_shap_predictor(model)) + + assert converted.f is not None + assert list(model.feature_names_in_) == ["a", "b"] diff --git a/tests/back/explainers/test_task_explainers.py b/tests/back/explainers/test_task_explainers.py index 65549c706..d05efe4c6 100644 --- a/tests/back/explainers/test_task_explainers.py +++ b/tests/back/explainers/test_task_explainers.py @@ -26,6 +26,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir REGRESSION_INPUT_COLUMNS = [ "SepalLengthCm", @@ -42,7 +43,7 @@ def regression_dataset_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/models/test_bow_text_class_model.py b/tests/back/models/test_bow_text_class_model.py index 1e3dff685..0828c5985 100644 --- a/tests/back/models/test_bow_text_class_model.py +++ b/tests/back/models/test_bow_text_class_model.py @@ -26,6 +26,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(autouse=True, name="test_registry") @@ -58,7 +59,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_deberta_v3_transformer.py b/tests/back/models/test_deberta_v3_transformer.py index a92f4fd5e..fb35b1a21 100644 --- a/tests/back/models/test_deberta_v3_transformer.py +++ b/tests/back/models/test_deberta_v3_transformer.py @@ -13,6 +13,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="splited_dataset") @@ -22,7 +23,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_distilbert_transformer.py b/tests/back/models/test_distilbert_transformer.py index ddf233ef3..60cbf0aa7 100644 --- a/tests/back/models/test_distilbert_transformer.py +++ b/tests/back/models/test_distilbert_transformer.py @@ -13,6 +13,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="splited_dataset") @@ -22,7 +23,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_modernbert_transformer.py b/tests/back/models/test_modernbert_transformer.py index 0801d1ce4..1080a0a50 100644 --- a/tests/back/models/test_modernbert_transformer.py +++ b/tests/back/models/test_modernbert_transformer.py @@ -13,6 +13,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="splited_dataset") @@ -22,7 +23,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_tabular_class_models.py b/tests/back/models/test_tabular_class_models.py index 3d99798d0..c5bdf537e 100644 --- a/tests/back/models/test_tabular_class_models.py +++ b/tests/back/models/test_tabular_class_models.py @@ -34,6 +34,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="divided_dataset") @@ -43,7 +44,7 @@ def tabular_model_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "separator": ",", "schema": { diff --git a/tests/back/optimizers/test_optuna_pruning_integration.py b/tests/back/optimizers/test_optuna_pruning_integration.py index 0f08be0d7..9078cfb21 100644 --- a/tests/back/optimizers/test_optuna_pruning_integration.py +++ b/tests/back/optimizers/test_optuna_pruning_integration.py @@ -29,18 +29,32 @@ from DashAI.back.evaluation.holdout import HoldoutEvaluationStrategy from DashAI.back.models.base_model import BaseModel from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.units.fit_model_unit import FitModelUnit def _holdout_evaluate(model, input_dataset, output_dataset, metric): - """The real holdout evaluation path, on a strategy with no factory. - - `evaluate` reads which partitions its strategy scores, so it needs a real - instance rather than None for self. Building one through __init__ would - need a `ModelFactory` this test does not have, and does not need: the only - thing read off the instance is a class attribute. + """The real objective of a holdout search: one fit, then score validation. + + It used to be ``HoldoutEvaluationStrategy.evaluate``. The strategies now + only declare how a run is evaluated, and the unit that fits carries it out, + so the objective the optimizer measures comes from there. Which partitions + a trial records is still the strategy's declaration -- the job reads + SCORED_SPLITS off it and hands it over, minus the test partition, which a + trial may never score -- so it is read off the class here too. """ - strategy = HoldoutEvaluationStrategy.__new__(HoldoutEvaluationStrategy) - return strategy.evaluate(model, input_dataset, output_dataset, metric) + trial_splits = [ + split.name + for split in HoldoutEvaluationStrategy.SCORED_SPLITS + if split.name != "TEST" + ] + unit = FitModelUnit( + optimizer={"component": "", "params": {}}, + goal_metric="", + run_id=None, + artifact_prefix=None, + trial_splits=trial_splits, + ) + return unit._score_one_trial(model, input_dataset, output_dataset, metric) EPOCHS = 12 diff --git a/tests/back/optimizers/test_optuna_real_model.py b/tests/back/optimizers/test_optuna_real_model.py index 9eeb7dba7..0a49431d4 100644 --- a/tests/back/optimizers/test_optuna_real_model.py +++ b/tests/back/optimizers/test_optuna_real_model.py @@ -40,18 +40,32 @@ from DashAI.back.metrics.classification.accuracy import Accuracy from DashAI.back.models.mlp_image_classifier import MLPImageClassifier from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.units.fit_model_unit import FitModelUnit def _holdout_evaluate(model, input_dataset, output_dataset, metric): - """The real holdout evaluation path, on a strategy with no factory. - - `evaluate` reads which partitions its strategy scores, so it needs a real - instance rather than None for self. Building one through __init__ would - need a `ModelFactory` this test does not have, and does not need: the only - thing read off the instance is a class attribute. + """The real objective of a holdout search: one fit, then score validation. + + It used to be ``HoldoutEvaluationStrategy.evaluate``. The strategies now + only declare how a run is evaluated, and the unit that fits carries it out, + so the objective the optimizer measures comes from there. Which partitions + a trial records is still the strategy's declaration -- the job reads + SCORED_SPLITS off it and hands it over, minus the test partition, which a + trial may never score -- so it is read off the class here too. """ - strategy = HoldoutEvaluationStrategy.__new__(HoldoutEvaluationStrategy) - return strategy.evaluate(model, input_dataset, output_dataset, metric) + trial_splits = [ + split.name + for split in HoldoutEvaluationStrategy.SCORED_SPLITS + if split.name != "TEST" + ] + unit = FitModelUnit( + optimizer={"component": "", "params": {}}, + goal_metric="", + run_id=None, + artifact_prefix=None, + trial_splits=trial_splits, + ) + return unit._score_one_trial(model, input_dataset, output_dataset, metric) EPOCHS = 3 diff --git a/tests/back/pipeline_spike/__init__.py b/tests/back/pipeline_spike/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/back/pipeline_spike/dag_engine.py b/tests/back/pipeline_spike/dag_engine.py new file mode 100644 index 000000000..b446ea085 --- /dev/null +++ b/tests/back/pipeline_spike/dag_engine.py @@ -0,0 +1,259 @@ +"""Minimal sequential DAG engine over the existing atomic units. + +A spike, not a deliverable. It exists to find out which parts of the unit +contract (``BaseUnit`` + ``ExecutionContext``) survive being driven by a graph +instead of by a job, before more jobs get atomized. It deliberately has no UI, +no database, no persistence, no API and no parallelism. + +The one design decision it puts to the test is the one proposed in section 9 of +ATOMIZING_JOBS.md: a graph engine does not need a shared context. It gives every +node an ``ExecutionContext`` of its own, preloaded with exactly the keys that +node declares in ``REQUIRES``, and the renaming lives on the edge rather than in +the unit. That is what lets two ``LoadDatasetUnit`` instances -- both of which +write the fixed key ``dataset`` -- coexist in one graph. + +Nothing here imports from ``DashAI/back/pipeline/``: that subsystem does not run +and is being redesigned. Nothing here modifies a unit either. Whatever the +engine cannot express with the contract as it stands is reported, not patched. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Sequence, Set, Tuple + +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + + +class GraphError(Exception): + """Raised when a graph is not executable as declared.""" + + +@dataclass(frozen=True) +class Node: + """A unit instance placed in a graph. + + Parameters + ---------- + id : str + Identifier of the node inside the graph. + unit : BaseUnit + The unit instance to run. Already configured. + seeds : Mapping[str, Any] + Constants injected straight into this node's context, for keys no + upstream unit publishes. ``run_id`` is the real case: four units + require it and none provides it, because in a job it arrived through + ``self.kwargs``. See the open question in the spike's report. + """ + + id: str + unit: BaseUnit + seeds: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Edge: + """A single key travelling from one node's output to another's input. + + The pair of key names is what makes fixed ``REQUIRES``/``PROVIDES`` strings + behave as port names: ``src_key`` is drawn from the source's ``PROVIDES`` + and ``dst_key`` from the target's ``REQUIRES``, and the edge maps one onto + the other. + """ + + src: str + src_key: str + dst: str + dst_key: str + + +def connect(src: Node, dst: Node) -> Tuple[Edge, ...]: + """Wire every key the two nodes agree on: ``PROVIDES`` against ``REQUIRES``. + + A visual canvas cannot draw one edge per key: ``FitModelUnit`` alone + requires eight. This is the bundling rule a real engine would use -- one + drawn edge between two nodes stands for this whole set. + """ + shared = sorted(set(src.unit.PROVIDES) & set(dst.unit.REQUIRES)) + return tuple(Edge(src.id, key, dst.id, key) for key in shared) + + +@dataclass(frozen=True) +class Graph: + """A set of nodes and the edges between them.""" + + nodes: Sequence[Node] + edges: Sequence[Edge] + + def node(self, node_id: str) -> Node: + for node in self.nodes: + if node.id == node_id: + return node + raise GraphError(f"Unknown node id: {node_id}") + + +def validate(graph: Graph) -> List[str]: + """Check the graph statically and return a topological execution order. + + This is the static DAG validator ATOMIZING_JOBS.md lists as missing. It is + almost free: ``REQUIRES``/``PROVIDES`` already carry everything it needs, + and no unit has to run. + + Raises + ------ + GraphError + If any check fails, with every problem found listed at once. + """ + problems: List[str] = [] + ids = [node.id for node in graph.nodes] + if len(set(ids)) != len(ids): + problems.append("Duplicate node ids in the graph.") + + by_id = {node.id: node for node in graph.nodes} + inbound: Dict[str, Dict[str, List[Edge]]] = {node_id: {} for node_id in by_id} + + for edge in graph.edges: + if edge.src not in by_id or edge.dst not in by_id: + problems.append(f"Edge {edge} points at a node that is not in the graph.") + continue + if edge.src_key not in by_id[edge.src].unit.PROVIDES: + problems.append( + f"{edge.src} does not provide '{edge.src_key}' " + f"(provides: {list(by_id[edge.src].unit.PROVIDES)})." + ) + if edge.dst_key not in by_id[edge.dst].unit.REQUIRES: + problems.append( + f"{edge.dst} does not require '{edge.dst_key}' " + f"(requires: {list(by_id[edge.dst].unit.REQUIRES)})." + ) + inbound[edge.dst].setdefault(edge.dst_key, []).append(edge) + + for node in graph.nodes: + for key in node.seeds: + if key not in node.unit.REQUIRES: + problems.append( + f"{node.id} is seeded with '{key}', which it never uses." + ) + for key in node.unit.REQUIRES: + sources = len(inbound[node.id].get(key, [])) + (key in node.seeds) + if sources == 0: + problems.append(f"{node.id} requires '{key}' and nothing supplies it.") + elif sources > 1: + problems.append( + f"{node.id} gets '{key}' from {sources} sources; a port takes one." + ) + + order = _topological_order(graph, by_id, problems) + if problems: + raise GraphError("\n".join(problems)) + return order + + +def _topological_order( + graph: Graph, by_id: Mapping[str, Node], problems: List[str] +) -> List[str]: + """Kahn's algorithm. Leftover nodes mean a cycle.""" + pending = dict.fromkeys(by_id, 0) + successors: Dict[str, List[str]] = {node_id: [] for node_id in by_id} + for edge in graph.edges: + if edge.src in by_id and edge.dst in by_id: + pending[edge.dst] += 1 + successors[edge.src].append(edge.dst) + + ready = [node_id for node_id, count in pending.items() if count == 0] + order: List[str] = [] + while ready: + node_id = ready.pop(0) + order.append(node_id) + for successor in successors[node_id]: + pending[successor] -= 1 + if pending[successor] == 0: + ready.append(successor) + + if len(order) != len(by_id): + problems.append( + "The graph has a cycle: " + ", ".join(sorted(set(by_id) - set(order))) + ) + return order + + +#: How many times the engine had to deep-copy the whole reference half of a +#: context just to ask which half a single key lives in. See the report. +REF_PROBES = [0] + + +def _transport(src_ctx: ExecutionContext, key: str) -> Tuple[str, Any]: + """Read a value out of a context, keeping the half it came from. + + The engine has to move a value from one context to another, and the two + halves have incompatible rules: ``put_ref`` validates with ``json.dumps`` + and deep-copies, ``put`` stores the live object by reference. Handing a + ``DashAIDataset`` to ``put_ref`` raises; handing a dict meant as a reference + to ``put`` silently drops the copy-on-write guarantee. + + ``ExecutionContext`` has no public way to ask which half a key is in: + ``get``/``has`` merge them and ``refs``/``to_dict`` deep-copy the entire + reference half. So this probe is correct but expensive, and the cost is + counted rather than hidden. + """ + REF_PROBES[0] += 1 + if key in src_ctx.refs: + return "ref", src_ctx.get(key) + return "cache", src_ctx.require(key) + + +def run(graph: Graph) -> Dict[str, ExecutionContext]: + """Execute the graph sequentially and return each node's own context. + + Parameters + ---------- + graph : Graph + A graph that passes :func:`validate`. + + Returns + ------- + Dict[str, ExecutionContext] + The context of every node whose values were still needed at the end, + keyed by node id. Contexts nobody reads any more are dropped as the run + proceeds so intermediate datasets do not pile up. + """ + order = validate(graph) + by_id = {node.id: node for node in graph.nodes} + + consumers: Dict[str, int] = dict.fromkeys(by_id, 0) + for edge in graph.edges: + consumers[edge.src] += 1 + + contexts: Dict[str, ExecutionContext] = {} + for node_id in order: + node = by_id[node_id] + ctx = ExecutionContext() + + for key, value in node.seeds.items(): + ctx.put_ref(key, value) + + for edge in graph.edges: + if edge.dst != node_id: + continue + half, value = _transport(contexts[edge.src], edge.src_key) + if half == "ref": + ctx.put_ref(edge.dst_key, value) + else: + ctx.put(edge.dst_key, value) + + node.unit(ctx) + contexts[node_id] = ctx + + for edge in graph.edges: + if edge.dst != node_id: + continue + consumers[edge.src] -= 1 + if consumers[edge.src] == 0: + contexts.pop(edge.src).clear_cache() + + return contexts + + +def sinks(graph: Graph) -> Set[str]: + """Node ids nothing consumes. They are what a run exists to produce.""" + consumed = {edge.src for edge in graph.edges} + return {node.id for node in graph.nodes} - consumed diff --git a/tests/back/pipeline_spike/test_dag_engine_spike.py b/tests/back/pipeline_spike/test_dag_engine_spike.py new file mode 100644 index 000000000..19600ec55 --- /dev/null +++ b/tests/back/pipeline_spike/test_dag_engine_spike.py @@ -0,0 +1,383 @@ +"""Runs real units as a DAG, to find out what the unit contract cannot express. + +The graph below is the train/test converter flow, which is the smallest thing a +single shared context genuinely cannot run: two ``LoadDatasetUnit`` instances +both write the fixed key ``dataset``, so in one context the second load erases +the first. Section 9 of ATOMIZING_JOBS.md proposes giving each node a context of +its own and putting the renaming on the edge; these tests check whether that +actually works with the units exactly as they are. + + load_train --dataset--+-------------------> tx_train --dataset--> save_train + | ^ ^ + +-> fit --converter--+---+ dataset_path | + | --+ + load_test --dataset------------------------|-> tx_test --dataset--> save_test + +Two fan-outs (``load_train`` feeds two nodes, ``fit`` feeds two nodes) and two +joins (each ``tx`` node takes a dataset from one branch and a fitted converter +from another). The number that proves it worked is 2.0: MinMaxScaler fitted on +train [0, 5, 10] learns min=0 max=10, so a test value of 20 scales to 2.0. A +refit on the test data would have produced 0.0 instead. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.converters.scikit_learn.min_max_scaler import MinMaxScaler +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + to_dashai_dataset, +) +from DashAI.back.types.value_types import Float +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_model_unit import SaveModelUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit +from tests.back.pipeline_spike.dag_engine import ( + Edge, + Graph, + GraphError, + Node, + connect, + run, + sinks, + validate, +) + +FULL_SCOPE = {"columns": [], "rows": []} +_MIN_MAX = {"component": "MinMaxScaler", "params": {}} + + +class _Row: + """Stand-in for a Dataset ORM row, as in tests/back/units.""" + + def __init__(self, file_path): + self.file_path = file_path + self.dataset_id = None + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """A class, not a lambda: kink resolves a registered lambda by calling it.""" + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +def _write(root, **columns): + frame = pd.DataFrame(columns) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + save_dataset(to_dashai_dataset(frame, types=types), str(root / "dataset")) + return root + + +def _values(path, column="a"): + return list(load_dataset(str(path)).to_pandas()[column]) + + +@pytest.fixture(name="stored_datasets") +def fixture_stored_datasets(tmp_path): + """Dataset 7 is the training data, dataset 8 the test data.""" + train = _write(tmp_path / "train", a=[0.0, 5.0, 10.0]) + test = _write(tmp_path / "test", a=[20.0]) + + di["session_factory"] = _FakeSessionFactory( + {"Dataset": {7: _Row(str(train)), 8: _Row(str(test))}} + ) + di["component_registry"] = {"MinMaxScaler": {"class": MinMaxScaler}} + yield train, test + del di["session_factory"] + del di["component_registry"] + + +def _graph(tmp_path, converter="MinMaxScaler"): + """The diamond, wired with the bundling rule wherever it applies.""" + load_train = Node("load_train", LoadDatasetUnit(dataset_id=7)) + load_test = Node("load_test", LoadDatasetUnit(dataset_id=8)) + fit = Node( + "fit", + FitConverterUnit( + converter={"component": converter, "params": {}}, + scope=FULL_SCOPE, + target=None, + ), + ) + tx_train = Node("tx_train", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + tx_test = Node("tx_test", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + # Saves back over the training data, so it needs the ref `dataset_path` from + # the load as well as the live dataset from the transform: a join whose two + # inputs live in different halves of the context. + save_train = Node("save_train", SaveDatasetUnit()) + save_test = Node( + "save_test", SaveDatasetToPathUnit(path=str(tmp_path / "out" / "dataset")) + ) + + nodes = [load_train, load_test, fit, tx_train, tx_test, save_train, save_test] + edges = [ + *connect(load_train, fit), + Edge("load_train", "dataset", "tx_train", "dataset"), + Edge("load_train", "dataset_path", "save_train", "dataset_path"), + *connect(fit, tx_train), + *connect(fit, tx_test), + Edge("load_test", "dataset", "tx_test", "dataset"), + Edge("tx_train", "dataset", "save_train", "dataset"), + *connect(tx_test, save_test), + ] + return Graph(nodes, edges) + + +def test_the_diamond_runs_and_the_fitted_state_survives_the_branch( + stored_datasets, tmp_path +): + """The headline case: one fit, two branches, no refit. + + 2.0 on the test branch is only reachable if the converter fitted in the + ``fit`` node arrived at ``tx_test`` still carrying the range it learned from + the training data, having crossed two context boundaries on the way. + """ + train, _test = stored_datasets + + run(_graph(tmp_path)) + + assert _values(train / "dataset") == [0.0, 0.5, 1.0] + assert _values(tmp_path / "out" / "dataset") == [2.0] + + +def test_two_loads_of_the_same_unit_do_not_fight_over_the_key( + stored_datasets, tmp_path +): + """Section 9's proposal, checked directly. + + Both loads publish the fixed key ``dataset``. In one shared context the + second would overwrite the first and the test branch would silently scale + the training data. + """ + graph = _graph(tmp_path) + contexts = run(graph) + + assert sinks(graph) == {"save_train", "save_test"} + # Every context still alive at the end belongs to a sink; the loads and the + # fit were dropped as soon as nothing needed them. + assert set(contexts) == {"save_train", "save_test"} + + +def test_the_engine_never_needed_a_shared_key_to_be_renamed_by_a_unit(tmp_path): + """The edge does the renaming; the units keep their fixed key names.""" + graph = _graph(tmp_path) + + assert set(validate(graph)[:2]) == {"load_train", "load_test"} + assert all(edge.src_key == edge.dst_key for edge in graph.edges) + + +def test_the_bundling_rule_picks_the_keys_two_nodes_agree_on(): + """One drawn edge between two nodes stands for a set of keys.""" + load = Node("load", LoadDatasetUnit(dataset_id=1)) + save = Node("save", SaveDatasetUnit()) + + assert {edge.src_key for edge in connect(load, save)} == {"dataset", "dataset_path"} + + +def test_a_missing_input_is_caught_before_anything_runs(): + tx = Node("tx", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + load = Node("load", LoadDatasetUnit(dataset_id=1)) + + with pytest.raises(GraphError, match="tx requires 'fitted_converter'"): + validate(Graph([load, tx], list(connect(load, tx)))) + + +def test_two_edges_into_one_port_are_rejected(): + """The failure the old engine had: a merge where the last writer wins.""" + load_a = Node("a", LoadDatasetUnit(dataset_id=1)) + load_b = Node("b", LoadDatasetUnit(dataset_id=2)) + fit = Node( + "fit", + FitConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + graph = Graph([load_a, load_b, fit], [*connect(load_a, fit), *connect(load_b, fit)]) + + with pytest.raises(GraphError, match="fit gets 'dataset' from 2 sources"): + validate(graph) + + +def test_a_cycle_is_reported(): + fit = Node( + "fit", + FitConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + tx = Node("tx", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + graph = Graph([fit, tx], [*connect(fit, tx), *connect(tx, fit)]) + + with pytest.raises(GraphError, match="cycle"): + validate(graph) + + +def test_run_id_is_no_longer_an_orphan_input(): + """``run_id`` used to be a context key no unit published. + + Four units required it and none provided it, so a graph had to inject it + from outside and the validator could not tell that injection apart from a + wire the user forgot to draw. It was never data flowing through the graph: + in a job it arrived through ``self.kwargs``. + + It is configuration now, which is what closed that hole. Every key left in + ``REQUIRES`` has a unit that publishes it, so "nothing supplies this" means + a missing edge and nothing else. The seeds mechanism this spike used for it + has no remaining caller. + """ + save = Node("save", SaveModelUnit(artifact_prefix="pipeline-1-save")) + + with pytest.raises(GraphError, match="save requires 'model'"): + validate(Graph([save], [])) + + assert "run_id" not in SaveModelUnit.REQUIRES + + +def test_a_seed_for_a_key_the_unit_never_uses_is_rejected(): + save = Node( + "save", + SaveModelUnit(artifact_prefix="pipeline-1-save"), + seeds={"nonsense": 2}, + ) + + with pytest.raises(GraphError, match="seeded with 'nonsense'"): + validate(Graph([save], [])) + + +def test_an_optional_output_cannot_be_wired_at_all(): + """``PROVIDES`` is the whole vocabulary an edge can name. + + ``FitModelUnit`` produces ``best_parameters`` only on the hyperparameter + search branch, so it cannot declare it (``__call__`` checks ``PROVIDES`` + unconditionally). ``ModelJob`` copes by reading it with ``ctx.has(...)``. + A graph cannot: a key outside ``PROVIDES`` is not addressable by an edge, + so the limitation goes from "not verifiable" to "not connectable". + """ + fit = Node( + "fit", + FitConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + tx = Node("tx", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + graph = Graph([fit, tx], [Edge("fit", "best_parameters", "tx", "dataset")]) + + with pytest.raises(GraphError, match="fit does not provide 'best_parameters'"): + validate(graph) + + +def test_a_key_a_middle_node_does_not_republish_needs_an_edge_around_it(): + """A converter in the middle does not carry the ids the load published. + + ``ApplyConverterUnit`` provides only ``dataset`` and ``fitted_converter``, + and it is right not to republish ``dataset_id`` -- section 4.3 forbids + anything derived from the object being transformed from crossing the + boundary. But ``PrepareAndSplitUnit`` downstream needs both, so ``dataset`` + comes through the converter and ``dataset_id`` has to jump over it. + + The engine accepts the bypass. What it costs is on the canvas: the user has + to draw an arrow from the loader past the converters to the split. + """ + load = Node("load", LoadDatasetUnit(dataset_id=1)) + apply_ = Node( + "apply", + ApplyConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + split = Node("split", PrepareAndSplitUnit(splitter={})) + + bundled = Graph( + [load, apply_, split], [*connect(load, apply_), *connect(apply_, split)] + ) + with pytest.raises(GraphError, match="split requires 'dataset_id'"): + validate(bundled) + + bypassed = Graph( + [load, apply_, split], + [ + *connect(load, apply_), + *connect(apply_, split), + Edge("load", "dataset_id", "split", "dataset_id"), + ], + ) + assert validate(bypassed) == ["load", "apply", "split"] + + +class _CountingConverter: + """Records how many times it was asked to transform, and who built it.""" + + instances = [] + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + _CountingConverter.instances.append(self) + self.transform_calls = 0 + + def fit(self, x, y=None): + return self + + def transform(self, x, y=None): + self.transform_calls += 1 + return x + + +def test_a_fan_out_hands_both_branches_the_same_live_object(stored_datasets, tmp_path): + """Identity survives the fan-out, which it has to (ATOMIZING_JOBS 5.7). + + Some converters cache against the object they were given, so a copy would + silently recompute. It also means two consumers of one fitted converter + share mutable state -- harmless while the engine is sequential, and a + reason to keep it that way. + """ + _CountingConverter.instances.clear() + di["component_registry"]["Counting"] = {"class": _CountingConverter} + + graph = _graph(tmp_path, converter="Counting") + run(graph) + + assert len(_CountingConverter.instances) == 1 + assert _CountingConverter.instances[0].transform_calls == 2 + + +def test_moving_a_value_across_an_edge_keeps_the_half_it_came_from( + stored_datasets, tmp_path +): + """Refs travel as refs, live objects as live objects. + + ``dataset_path`` is a ref and ``dataset`` is a cached object, and the two + have incompatible rules: ``put_ref`` validates with ``json.dumps``, so a + dataset sent that way raises, and ``put`` would drop the copy-on-write + guarantee a ref depends on. The engine has to ask which half a key is in -- + and the only public way to ask deep-copies the entire reference half, once + per edge. + """ + from tests.back.pipeline_spike import dag_engine + + graph = _graph(tmp_path) + dag_engine.REF_PROBES[0] = 0 + contexts = run(graph) + + assert dag_engine.REF_PROBES[0] == len(graph.edges) + save_train = contexts["save_train"] + assert "dataset_path" in save_train.refs + assert "dataset" not in save_train.refs + assert save_train.has("dataset") diff --git a/tests/back/scratch.py b/tests/back/scratch.py new file mode 100644 index 000000000..747564b43 --- /dev/null +++ b/tests/back/scratch.py @@ -0,0 +1,50 @@ +"""Scratch directories for tests that load datasets. + +Dataloaders take a ``temp_path`` and hand it straight to HuggingFace as +``cache_dir`` (see ``CSVDataLoader.load_data``), so every test that loads a +dataset leaves an arrow cache and a ``dataset_info.json`` behind. Pointing that +path at the system temp directory keeps those artifacts out of the repository +tree instead of scattering them across ``tests/back/``. +""" + +import os +import pathlib +import shutil +import tempfile + +# Scoped by pid so two pytest processes running at once (two terminals, a CI +# matrix on one machine) never share a root: without this, one session's +# start-of-run ``clear_scratch()`` would rmtree the arrow cache another, +# still-running session has memory mapped. +SCRATCH_ROOT = ( + pathlib.Path(tempfile.gettempdir()) / f"dashai-test-scratch-{os.getpid()}" +) + + +def scratch_dir(*parts: str) -> str: + """Return a scratch directory, creating it if it does not exist. + + Parameters + ---------- + *parts : str + Path segments below the scratch root. Use them to keep unrelated test + modules from sharing a cache, the way the old in-repo paths did. + + Returns + ------- + str + Absolute path to the directory. + """ + path = SCRATCH_ROOT.joinpath(*parts) + path.mkdir(parents=True, exist_ok=True) + return str(path) + + +def clear_scratch() -> None: + """Remove the scratch root, ignoring files still held open. + + Arrow caches stay memory mapped on Windows, so a best effort delete is the + most that can be promised here; the point is that whatever survives lives + outside the repository. + """ + shutil.rmtree(SCRATCH_ROOT, ignore_errors=True) diff --git a/tests/back/tasks/test_tasks.py b/tests/back/tasks/test_tasks.py index 12377696d..08d4d690a 100644 --- a/tests/back/tasks/test_tasks.py +++ b/tests/back/tasks/test_tasks.py @@ -20,6 +20,7 @@ from DashAI.back.tasks.text_to_image_generation_task import TextToImageGenerationTask from DashAI.back.tasks.text_to_text_generation_task import TextToTextGenerationTask from DashAI.back.tasks.translation_task import TranslationTask +from tests.back.scratch import scratch_dir def load_csv_into_datasetdict_iris(file_name): @@ -28,7 +29,7 @@ def load_csv_into_datasetdict_iris(file_name): datasetdict = csv_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={"separator": ","}, ) schema = { @@ -48,7 +49,7 @@ def load_csv_into_datasetdict_iris_extra(file_name): datasetdict = csv_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={"separator": ","}, ) schema = { @@ -152,7 +153,7 @@ def text_classification_dataset_fixture(): dataset = json_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={ "data_key": "data", }, @@ -215,7 +216,7 @@ def translation_dataset_fixture(): dataset = json_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={"data_key": "data"}, ) @@ -279,8 +280,7 @@ def sample_image_fixture(): @pytest.fixture(scope="module", name="temp_path") def temp_path_fixture(): - temp_path = pathlib.Path("tests") / "back" / "tasks" / "temp" - os.makedirs(temp_path, exist_ok=True) + temp_path = pathlib.Path(scratch_dir("tasks", "temp")) yield temp_path # Cleanup after all tests in the module using this fixture have finished if temp_path.exists() and temp_path.is_dir(): diff --git a/tests/back/units/__init__.py b/tests/back/units/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/back/units/test_apply_converter_unit.py b/tests/back/units/test_apply_converter_unit.py new file mode 100644 index 000000000..cc960504a --- /dev/null +++ b/tests/back/units/test_apply_converter_unit.py @@ -0,0 +1,388 @@ +"""Contract tests for ApplyConverterUnit, isolated from ConverterJob. + +The job only ever runs one of these per invocation, so an end-to-end run cannot +show whether the unit is safe to chain. These tests build the context by hand +and run several units against it, which is what the future DAG will do. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.types.value_types import Integer +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError + + +def _dataset(**columns): + frame = pd.DataFrame(columns) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +class _DropScopedColumns: + """Converter that removes whatever is in scope, like ColumnRemover.""" + + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + self.params = params + self.columns = [] + + def fit(self, x, y=None): + self.columns = x.column_names + return self + + def transform(self, x, y=None): + return x.remove_columns(self.columns) + + +class _ReplaceWithScope: + """Converter that replaces the dataset with the scoped columns only.""" + + CHANGES_ROW_COUNT = True + + def __init__(self, **params): + self.params = params + + def fit(self, x, y=None): + return self + + def transform(self, x, y=None): + return x + + +class _RecordingConverter: + """Converter that records what fit and transform were handed.""" + + CHANGES_ROW_COUNT = True + + def __init__(self, **params): + self.params = params + self.fit_x = None + self.fit_y = None + self.transform_x = None + self.transform_y = None + + def fit(self, x, y=None): + self.fit_x = x + self.fit_y = y + return self + + def transform(self, x, y=None): + self.transform_x = x + self.transform_y = y + return x + + +class _FailingFit: + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + pass + + def fit(self, x, y=None): + raise ValueError("bad input") + + def transform(self, x, y=None): # pragma: no cover - never reached + return x + + +class _FailingTransform: + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + pass + + def fit(self, x, y=None): + return self + + def transform(self, x, y=None): + raise RuntimeError("boom") + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "DropScopedColumns": {"class": _DropScopedColumns}, + "ReplaceWithScope": {"class": _ReplaceWithScope}, + "RecordingConverter": {"class": _RecordingConverter}, + "FailingFit": {"class": _FailingFit}, + "FailingTransform": {"class": _FailingTransform}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _unit(name, scope=None, target=None, params=None): + return ApplyConverterUnit( + converter={"component": name, "params": params or {}}, + scope=scope, + target=target, + ) + + +def _ctx(dataset): + ctx = ExecutionContext() + ctx.put("dataset", dataset) + return ctx + + +def test_the_unit_refuses_to_run_without_a_dataset(registry): + with pytest.raises(UnitContractError, match="'dataset' is not available"): + _unit("DropScopedColumns")(ExecutionContext()) + + +def test_a_scoped_column_is_resolved_by_one_based_index(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + _unit("DropScopedColumns", scope={"columns": [{"idx": 2}], "rows": []})(ctx) + + assert ctx.require("dataset").column_names == ["a", "c"] + + +def test_an_empty_column_scope_selects_every_column(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + + _unit("ReplaceWithScope", scope={"columns": [], "rows": []})(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + + +def test_a_null_scope_is_treated_as_an_empty_one(registry): + """The API schema writes the ``scope`` key but allows it to be null. + + ``dict.get("scope", default)`` returns ``None`` rather than the default when + the key is present and null, so the unit has to coalesce it explicitly. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + + _unit("ReplaceWithScope", scope=None)(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + + +def test_chained_units_resolve_their_indexes_against_the_current_dataset(registry): + """The reason no column identity may cross the context boundary. + + The first unit drops column ``a``, so index 1 means ``b`` by the time the + second unit runs. A column list resolved once and published to the context + would make the second unit drop the wrong column. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + assert ctx.require("dataset").column_names == ["b", "c"] + + _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + assert ctx.require("dataset").column_names == ["c"] + + +def test_the_unit_publishes_no_column_state_into_the_context(registry): + """Nothing resolved from the dataset may outlive the unit that resolved it.""" + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + + _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + + for leaked in ( + "column_names", + "scope_column_names", + "scope_column_indexes", + "target_column_name", + "converter", + ): + assert not ctx.has(leaked), leaked + + +def test_two_units_in_one_context_do_not_share_their_resolved_converter(registry): + """Registry lookups are memoized on the instance, never in the context. + + A class cached under a context key would make the second unit silently run + the first one's converter. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + first = _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []}) + second = _unit("ReplaceWithScope", scope={"columns": [{"idx": 1}], "rows": []}) + + first(ctx) + second(ctx) + + assert first._converter_class is _DropScopedColumns + assert second._converter_class is _ReplaceWithScope + # ReplaceWithScope keeps only what is in scope, which is now "b". + assert ctx.require("dataset").column_names == ["b"] + + +def test_a_row_scope_narrows_fit_but_not_transform(registry): + ctx = _ctx(_dataset(a=[1, 2, 3, 4], b=[5, 6, 7, 8])) + recorded = {} + + class _Spy(_RecordingConverter): + def fit(self, x, y=None): + recorded["fit_rows"] = len(x) + return super().fit(x, y) + + def transform(self, x, y=None): + recorded["transform_rows"] = len(x) + recorded["same_object"] = x is self.fit_x + return super().transform(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit("Spy", scope={"columns": [{"idx": 1}], "rows": [1, 2]})(ctx) + + assert recorded["fit_rows"] == 2 + assert recorded["transform_rows"] == 4 + assert recorded["same_object"] is False + + +def test_no_row_scope_hands_transform_the_same_object_as_fit(registry): + """Object identity, not equality, is part of the contract here. + + ``TypeCastConverter`` caches converted columns during ``fit`` and reuses + them in ``transform`` only when handed the same dataset object. Losing that + identity does not fail anything — it just silently recomputes. + """ + ctx = _ctx(_dataset(a=[1, 2, 3, 4], b=[5, 6, 7, 8])) + recorded = {} + + class _Spy(_RecordingConverter): + def transform(self, x, y=None): + recorded["same_object"] = x is self.fit_x + return super().transform(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit("Spy", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + + assert recorded["same_object"] is True + + +def test_a_target_column_is_resolved_and_handed_over_as_y(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], label=[5, 6])) + recorded = {} + + class _Spy(_RecordingConverter): + def fit(self, x, y=None): + recorded["fit_y"] = None if y is None else y.column_names + return super().fit(x, y) + + def transform(self, x, y=None): + recorded["transform_y"] = None if y is None else y.column_names + return super().transform(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit( + "Spy", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 3}, + )(ctx) + + assert recorded["fit_y"] == ["label"] + assert recorded["transform_y"] == ["label"] + + +def test_no_target_means_no_y(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + recorded = {} + + class _Spy(_RecordingConverter): + def fit(self, x, y=None): + recorded["fit_y"] = y + return super().fit(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit("Spy", scope={"columns": [{"idx": 1}], "rows": []}, target=None)(ctx) + + assert recorded["fit_y"] is None + + +def test_validate_rejects_an_out_of_bounds_target_without_running(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + unit = _unit("DropScopedColumns", target={"idx": 99}) + + with pytest.raises(JobError, match="Target column index 99 is out of bounds"): + unit.validate(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + + +def test_the_target_bound_is_rechecked_against_the_dataset_of_the_moment(registry): + """A target that was valid before the previous converter ran may not be now. + + ``validate`` runs against whatever dataset is in the context when it is + called; chaining means ``execute`` can face a narrower one. Re-checking + turns what would be a bare IndexError into the same JobError. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + unit = _unit( + "DropScopedColumns", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 3}, + ) + + unit.validate(ctx) # 3 columns: idx 3 is fine right now. + + # Another converter narrows the dataset behind its back. + _unit("ReplaceWithScope", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + + with pytest.raises(JobError, match="Target column index 3 is out of bounds"): + unit(ctx) + + +def test_an_unknown_converter_names_the_culprit(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + with pytest.raises(JobError, match="Error importing converter NotRegistered"): + _unit("NotRegistered")(ctx) + + +def test_a_value_error_during_fit_is_reported_as_a_validation_error(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + with pytest.raises(JobError, match="Validation error fitting FailingFit"): + _unit("FailingFit")(ctx) + + +def test_any_other_error_during_fit_is_reported_as_a_fit_error(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + class _Boom(_FailingFit): + def fit(self, x, y=None): + raise RuntimeError("nope") + + registry["Boom"] = {"class": _Boom} + + with pytest.raises(JobError, match="Error fitting converter Boom"): + _unit("Boom")(ctx) + + +def test_an_error_during_transform_is_reported_as_a_transform_error(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + with pytest.raises(JobError, match="Error transforming data with FailingTransform"): + _unit("FailingTransform")(ctx) + + +def test_converter_params_reach_the_constructor(registry): + ctx = _ctx(_dataset(a=[1, 2])) + unit = _unit("ReplaceWithScope", params={"threshold": 3}) + + unit(ctx) + + assert unit._converter_class is _ReplaceWithScope + + +def test_a_changes_row_count_converter_replaces_the_whole_dataset(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + _unit("ReplaceWithScope", scope={"columns": [{"idx": 2}], "rows": []})(ctx) + + assert ctx.require("dataset").column_names == ["b"] diff --git a/tests/back/units/test_base_unit.py b/tests/back/units/test_base_unit.py new file mode 100644 index 000000000..bba78d84d --- /dev/null +++ b/tests/back/units/test_base_unit.py @@ -0,0 +1,132 @@ +"""Tests for the atomic unit base class and its registration.""" + +import pytest + +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError + + +class DummyUnit(BaseUnit): + """Unit that copies a required key into a provided one.""" + + REQUIRES = ("dataset",) + PROVIDES = ("x",) + + def execute(self, ctx: ExecutionContext) -> None: + ctx.put("x", f"prepared {ctx.require('dataset')}") + + +class LyingUnit(BaseUnit): + """Unit that declares an output it never writes.""" + + PROVIDES = ("x",) + + def execute(self, ctx: ExecutionContext) -> None: + pass + + +class ValidatingUnit(BaseUnit): + """Unit that rejects its configuration before running.""" + + def __init__(self, **config) -> None: + super().__init__(**config) + self.validate_calls = 0 + + def validate(self, ctx: ExecutionContext) -> None: + self.validate_calls += 1 + if not self.config.get("allowed", True): + raise ValueError("not allowed") + + def execute(self, ctx: ExecutionContext) -> None: + ctx.put("ran", True) + + +def test_unit_stores_its_configuration(): + unit = DummyUnit(dataset_id=3, columns=["a"]) + + assert unit.config == {"dataset_id": 3, "columns": ["a"]} + + +def test_unit_runs_and_writes_its_output(): + ctx = ExecutionContext() + ctx.put("dataset", "a dataset") + + DummyUnit()(ctx) + + assert ctx.require("x") == "prepared a dataset" + + +def test_unit_refuses_to_run_without_its_required_keys(): + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="'dataset' is not available"): + DummyUnit()(ctx) + + +def test_unit_fails_when_it_does_not_deliver_what_it_promised(): + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="'x' is not available"): + LyingUnit()(ctx) + + +def test_validate_is_a_noop_by_default(): + DummyUnit().validate(ExecutionContext()) + + +def test_validate_can_reject_a_configuration_without_executing(): + ctx = ExecutionContext() + unit = ValidatingUnit(allowed=False) + + with pytest.raises(ValueError, match="not allowed"): + unit.validate(ctx) + + assert not ctx.has("ran") + + +def test_call_invokes_validate_automatically_before_execute(): + """Regression: ``unit(ctx)`` alone must run validate(), not just execute(). + + Before this fix a unit's ``validate`` (e.g. ``BuildModelUnit``'s download + gate) only ran if the caller remembered to invoke it separately — an easy + step to forget for any future caller that just does ``unit(ctx)``, the one + call the base class marks as the sanctioned entry point. + """ + ctx = ExecutionContext() + unit = ValidatingUnit(allowed=True) + + unit(ctx) + + assert unit.validate_calls == 1 + assert ctx.require("ran") is True + + +def test_call_stops_at_validate_and_never_reaches_execute(): + ctx = ExecutionContext() + unit = ValidatingUnit(allowed=False) + + with pytest.raises(ValueError, match="not allowed"): + unit(ctx) + + assert not ctx.has("ran") + + +def test_units_register_under_their_own_registry_type(): + """The registry derives TYPE by walking the MRO for a single "Base" + ancestor declaring it. BaseUnit must be that single candidate.""" + registry = ComponentRegistry(initial_components=[DummyUnit]) + + assert registry["DummyUnit"]["type"] == "Unit" + assert [c["name"] for c in registry.get_components_by_types(select="Unit")] == [ + "DummyUnit" + ] + + +def test_registered_units_are_configurable_objects_with_a_schema(): + registry = ComponentRegistry(initial_components=[DummyUnit]) + + component = registry["DummyUnit"] + + assert component["configurable_object"] is True + assert component["schema"] is not None diff --git a/tests/back/units/test_build_model_unit.py b/tests/back/units/test_build_model_unit.py new file mode 100644 index 000000000..6e17fbb4a --- /dev/null +++ b/tests/back/units/test_build_model_unit.py @@ -0,0 +1,112 @@ +"""Tests for BuildModelUnit's model-class resolution and its caching.""" + +import pytest +from kink import di + +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.context import ExecutionContext + + +class _ModelA: + def __init__(self, **kwargs): + pass + + +class _ModelB: + def __init__(self, **kwargs): + pass + + +@pytest.fixture +def fake_registry(): + """A minimal dict-shaped registry: enough for name -> class lookups.""" + registry = { + "ModelA": {"class": _ModelA}, + "ModelB": {"class": _ModelB}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _build_unit(model_name, run_id=1): + return BuildModelUnit( + model={"component": model_name, "params": {}}, + train_metrics=[], + validation_metrics=[], + test_metrics=[], + run_id=run_id, + ) + + +def test_two_build_model_units_in_one_context_resolve_independently(fake_registry): + """Regression: the model-class cache must not be keyed by the context. + + Two ``BuildModelUnit`` instances configured for different models, run + against the same context, must each build their own model. Before this + fix the class was memoized under a context-global ``"model_class"`` key, + so the second unit would silently reuse the first one's resolved class — + and its download-gate check would validate the wrong model too. + """ + ctx = ExecutionContext() + ctx.put("x", {"train": None, "validation": None}) + ctx.put("y", {"train": None, "validation": None}) + ctx.put("n_labels", None) + ctx.put_ref("task_name", "ATask") + + a = _build_unit("ModelA") + a(ctx) + model_a = ctx.get("model") + + b = _build_unit("ModelB") + b(ctx) + model_b = ctx.get("model") + + assert isinstance(model_a, _ModelA) + assert isinstance(model_b, _ModelB) + + +def test_resolve_model_class_is_memoized_per_instance_not_shared(fake_registry): + a = _build_unit("ModelA") + b = _build_unit("ModelB") + + assert a._resolve_model_class() is _ModelA + assert b._resolve_model_class() is _ModelB + # Calling again must return the same, still-correct class from the cache. + assert a._resolve_model_class() is _ModelA + + +def test_a_run_id_of_none_leaves_the_model_detached_from_any_run(fake_registry): + """No run id means no run, and a model with no run logs no metrics. + + ``ModelFactory`` hangs the run id on the model instance, and that + attribute is the only thing ``BaseModel.calculate_metrics`` consults + before deciding whether to write anything: ``if not metrics or not + self.run_id`` returns early. So a caller with no run -- a pipeline -- gets + a model that computes nothing into the ``Metric`` table, during training + or after it, without the caller having to intercept anything. + + This is the sole reason ``run_id`` is nullable rather than required. + """ + ctx = ExecutionContext() + ctx.put("x", {"train": None, "validation": None}) + ctx.put("y", {"train": None, "validation": None}) + ctx.put("n_labels", None) + ctx.put_ref("task_name", "ATask") + + _build_unit("ModelA", run_id=None)(ctx) + + assert ctx.get("model").run_id is None + + +def test_a_real_run_id_is_attached_to_the_model(fake_registry): + """The mirror of the above: a run's model has to be able to log.""" + ctx = ExecutionContext() + ctx.put("x", {"train": None, "validation": None}) + ctx.put("y", {"train": None, "validation": None}) + ctx.put("n_labels", None) + ctx.put_ref("task_name", "ATask") + + _build_unit("ModelA", run_id=17)(ctx) + + assert ctx.get("model").run_id == 17 diff --git a/tests/back/units/test_context.py b/tests/back/units/test_context.py new file mode 100644 index 000000000..45bce97a9 --- /dev/null +++ b/tests/back/units/test_context.py @@ -0,0 +1,219 @@ +"""Tests for the execution context shared between atomic units.""" + +import pytest + +from DashAI.back.units.context import ExecutionContext, UnitContractError + + +def test_put_ref_stores_a_serializable_value(): + ctx = ExecutionContext() + ctx.put_ref("run_id", 5) + + assert ctx.get("run_id") == 5 + assert ctx.refs == {"run_id": 5} + + +def test_put_ref_rejects_a_non_serializable_value(): + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="not JSON serializable"): + ctx.put_ref("model", object()) + + +def test_constructor_validates_the_initial_refs(): + with pytest.raises(UnitContractError, match="not JSON serializable"): + ExecutionContext(refs={"model": object()}) + + +def test_put_accepts_a_live_object(): + ctx = ExecutionContext() + sentinel = object() + ctx.put("model", sentinel) + + assert ctx.get("model") is sentinel + assert ctx.refs == {} + + +def test_get_prefers_the_cache_over_the_refs(): + ctx = ExecutionContext(refs={"model_path": "/on/disk"}) + ctx.put("model_path", "/in/memory") + + assert ctx.get("model_path") == "/in/memory" + + +def test_get_falls_back_to_the_refs_and_then_the_default(): + ctx = ExecutionContext(refs={"dataset_id": 3}) + + assert ctx.get("dataset_id") == 3 + assert ctx.get("missing") is None + assert ctx.get("missing", "fallback") == "fallback" + + +def test_require_returns_the_value_when_present(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + assert ctx.require("run_id") == 7 + assert ctx.require("model") == "a model" + + +def test_require_raises_and_lists_the_present_keys(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + with pytest.raises(UnitContractError) as exc_info: + ctx.require("x") + + message = str(exc_info.value) + assert "'x' is not available" in message + assert "'model'" in message + assert "'run_id'" in message + + +def test_has_checks_both_halves(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + assert ctx.has("run_id") + assert ctx.has("model") + assert not ctx.has("x") + + +def test_clear_cache_drops_live_objects_and_keeps_refs(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + ctx.clear_cache() + + assert not ctx.has("model") + assert ctx.require("run_id") == 7 + + +def test_refs_property_returns_a_copy(): + ctx = ExecutionContext(refs={"run_id": 7}) + + ctx.refs["run_id"] = 99 + + assert ctx.require("run_id") == 7 + + +def test_round_trip_through_to_dict_keeps_refs_and_drops_the_cache(): + ctx = ExecutionContext( + refs={"run_id": 7, "split_indexes": {"train_indexes": [0, 1]}} + ) + ctx.put("model", object()) + + restored = ExecutionContext.from_dict(ctx.to_dict()) + + assert restored.require("run_id") == 7 + assert restored.require("split_indexes") == {"train_indexes": [0, 1]} + assert not restored.has("model") + + +def test_to_dict_is_detached_from_the_context(): + ctx = ExecutionContext(refs={"run_id": 7}) + + serialized = ctx.to_dict() + serialized["run_id"] = 99 + + assert ctx.require("run_id") == 7 + + +def test_put_ref_isolates_from_later_mutation_of_the_original_object(): + """Regression: ``put_ref`` must not alias a caller's mutable dict. + + ``BuildModelUnit`` passes ``run.parameters`` — the dict SQLAlchemy attaches + to the ``Run`` row — straight into ``put_ref``. Without a defensive copy, + a later in-place edit reachable through the context (e.g. + ``ModelFactory.update_parameters`` rewriting a nested ``fixed_value`` + during hyperparameter search) would silently write through into the ORM + object, corrupting ``run.parameters`` before it was ever meant to change. + """ + original = {"n_estimators": {"fixed_value": 2}} + ctx = ExecutionContext() + ctx.put_ref("model_parameters", original) + + original["n_estimators"]["fixed_value"] = 999 + + assert ctx.require("model_parameters")["n_estimators"]["fixed_value"] == 2 + + +def test_require_returns_an_isolated_copy_of_a_reference(): + ctx = ExecutionContext( + refs={"model_parameters": {"n_estimators": {"fixed_value": 2}}} + ) + + fetched = ctx.require("model_parameters") + fetched["n_estimators"]["fixed_value"] = 999 + + assert ctx.require("model_parameters")["n_estimators"]["fixed_value"] == 2 + + +def test_get_returns_an_isolated_copy_of_a_reference(): + ctx = ExecutionContext( + refs={"model_parameters": {"n_estimators": {"fixed_value": 2}}} + ) + + fetched = ctx.get("model_parameters") + fetched["n_estimators"]["fixed_value"] = 999 + + assert ctx.get("model_parameters")["n_estimators"]["fixed_value"] == 2 + + +def test_cached_live_objects_are_still_returned_by_reference(): + """The copy-on-read fix must only apply to refs, never to the cache. + + A model, an in-memory dataset, an open session — these have to be the + same live object on every ``get``/``require``. Deep-copying them would + silently break training (the model you fit is not the model you saved) + while looking, from the outside, like nothing went wrong. + """ + + class LiveThing: + pass + + live_object = LiveThing() + ctx = ExecutionContext() + ctx.put("model", live_object) + + assert ctx.get("model") is live_object + assert ctx.require("model") is live_object + + +def test_origin_says_which_half_a_key_is_in(): + ctx = ExecutionContext(refs={"dataset_path": "/tmp/ds"}) + ctx.put("dataset", object()) + + assert ctx.origin("dataset_path") == "ref" + assert ctx.origin("dataset") == "cache" + assert ctx.origin("nothing_here") is None + + +def test_origin_prefers_the_cache_the_same_way_get_does(): + """One key in both halves has to have one answer, not two. + + ``get`` and ``require`` read the cache first, so anything deciding how to + move a value has to agree with them or the value would be read from one + half and written as if it came from the other. + """ + live = object() + ctx = ExecutionContext(refs={"x": 1}) + ctx.put("x", live) + + assert ctx.origin("x") == "cache" + assert ctx.get("x") is live + assert ctx.require("x") is live + + +def test_origin_does_not_copy_the_reference_half(): + """The reason the method exists: asking must not cost a deep copy. + + ``key in ctx.refs`` answers the same question, but ``refs`` returns a deep + copy of every reference the context holds, so asking once per edge of a + graph would copy the whole reference half once per edge. + """ + ctx = ExecutionContext(refs={"nested": {"deep": [1, 2, 3]}}) + inner = ctx._refs["nested"] + + assert ctx.origin("nested") == "ref" + assert ctx._refs["nested"] is inner diff --git a/tests/back/units/test_converter_fit_transform_split.py b/tests/back/units/test_converter_fit_transform_split.py new file mode 100644 index 000000000..00cba74a1 --- /dev/null +++ b/tests/back/units/test_converter_fit_transform_split.py @@ -0,0 +1,278 @@ +"""Fitting a converter on one dataset and applying it to another. + +The point of splitting ``ApplyConverterUnit`` into ``FitConverterUnit`` + +``TransformDatasetUnit``: a scaler, encoder or imputer must learn its statistics +from the training data only and then be applied unchanged to the test data. +Refitting on test would leak the test distribution into the evaluation, which is +exactly what the fused unit forced. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.converters.scikit_learn.min_max_scaler import MinMaxScaler +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.types.value_types import Float +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit + +FULL_SCOPE = {"columns": [], "rows": []} + + +def _dataset(**columns): + frame = pd.DataFrame(columns) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +def _values(dataset, column="a"): + return list(dataset.to_pandas()[column]) + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = {"MinMaxScaler": {"class": MinMaxScaler}} + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def test_a_converter_fitted_on_train_is_applied_to_test_without_refitting(registry): + """The headline case. + + MinMaxScaler fitted on [0, 5, 10] learns min=0, max=10. Applied to a test + value of 20 it must yield 2.0 — outside [0, 1] precisely because the range + came from train. A refit on the test data would have produced 0.0 instead, + so the number is what proves the fitted state survived. + """ + ctx = ExecutionContext() + + # Fit on train. + ctx.put("dataset", _dataset(a=[0.0, 5.0, 10.0])) + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + # Transform train with it. + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + assert _values(ctx.require("dataset")) == [0.0, 0.5, 1.0] + + # Swap in the test dataset and transform with the *same* fitted converter. + ctx.put("dataset", _dataset(a=[20.0])) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + assert _values(ctx.require("dataset")) == [2.0] + + +def test_fitting_leaves_the_dataset_untouched(registry): + """FitConverterUnit produces a converter, not data.""" + ctx = ExecutionContext() + original = _dataset(a=[0.0, 5.0, 10.0]) + ctx.put("dataset", original) + + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + assert ctx.require("dataset") is original + assert _values(ctx.require("dataset")) == [0.0, 5.0, 10.0] + + +def test_the_fitted_converter_is_published_live_not_copied(registry): + """It has to be the same object, or the learned state would be lost.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0])) + + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + fitted = ctx.require("fitted_converter") + assert isinstance(fitted, MinMaxScaler) + assert ctx.require("fitted_converter") is fitted + # The learned statistics are what makes it worth reusing. + assert list(fitted.data_min_) == [0.0] + assert list(fitted.data_max_) == [10.0] + + +def test_transforming_without_a_fitted_converter_is_a_contract_error(): + """A missing converter is a wiring mistake, not "nothing to apply".""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1.0])) + + with pytest.raises(UnitContractError, match="'fitted_converter' is not available"): + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + +def test_transforming_without_a_dataset_is_a_contract_error(registry): + ctx = ExecutionContext() + ctx.put("fitted_converter", MinMaxScaler()) + + with pytest.raises(UnitContractError, match="'dataset' is not available"): + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + +def test_the_fused_unit_also_publishes_its_fitted_converter(registry): + """ApplyConverterUnit stays usable as the source of a reusable converter. + + So the single-dataset path and the train/test path are the same mechanism: + whoever fitted the converter publishes it, and any number of + TransformDatasetUnits can then apply it elsewhere. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 5.0, 10.0])) + + ApplyConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + assert _values(ctx.require("dataset")) == [0.0, 0.5, 1.0] + + ctx.put("dataset", _dataset(a=[20.0])) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + assert _values(ctx.require("dataset")) == [2.0] + + +def test_the_split_pair_matches_the_fused_unit_on_a_single_dataset(registry): + """The two paths must not drift: same input, same output.""" + fused_ctx = ExecutionContext() + fused_ctx.put("dataset", _dataset(a=[1.0, 2.0, 3.0, 4.0])) + ApplyConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(fused_ctx) + + split_ctx = ExecutionContext() + split_ctx.put("dataset", _dataset(a=[1.0, 2.0, 3.0, 4.0])) + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(split_ctx) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(split_ctx) + + assert _values(fused_ctx.require("dataset")) == _values( + split_ctx.require("dataset") + ) + + +def test_a_row_scope_narrows_the_fit_but_the_transform_still_sees_every_row(registry): + """Row scope belongs to fitting; transform always covers the dataset. + + Fitting on rows 1-2 of [0, 10, 100] learns min=0, max=10, so the third row + scales to 10.0 rather than 1.0. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0, 100.0])) + + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope={"columns": [], "rows": [1, 2]}, + target=None, + )(ctx) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + assert _values(ctx.require("dataset")) == [0.0, 1.0, 10.0] + + +def test_fit_names_the_culprit_when_the_converter_is_not_registered(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1.0])) + + with pytest.raises(JobError, match="Error importing converter NotRegistered"): + FitConverterUnit( + converter={"component": "NotRegistered", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + +def test_fit_rejects_an_out_of_bounds_target_before_running(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1.0], b=[2.0])) + unit = FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target={"idx": 99}, + ) + + with pytest.raises(JobError, match="Target column index 99 is out of bounds"): + unit.validate(ctx) + + assert not ctx.has("fitted_converter") + + +def test_two_fit_units_in_one_context_resolve_their_converters_independently(registry): + """Registry lookups are memoized on the instance, never in the context.""" + + class _Other(MinMaxScaler): + pass + + registry["Other"] = {"class": _Other} + + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0])) + + first = FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + ) + second = FitConverterUnit( + converter={"component": "Other", "params": {}}, + scope=FULL_SCOPE, + target=None, + ) + first(ctx) + second(ctx) + + assert first._converter_class is MinMaxScaler + assert second._converter_class is _Other + + +def test_the_units_register_under_the_unit_type(): + """Guards the mixin's name. + + ``ConverterScopeMixin`` must not be called ``Base*``: the registry rejects a + component whose MRO has more than one "Base" ancestor declaring a TYPE, so + an intermediate ``BaseConverterUnit`` would break registration for all three + converter units at once. + """ + from DashAI.back.dependencies.registry import ComponentRegistry + + units = [ApplyConverterUnit, FitConverterUnit, TransformDatasetUnit] + component_registry = ComponentRegistry(initial_components=units) + + for unit in units: + assert component_registry[unit.__name__]["type"] == "Unit" + + +def test_one_fitted_converter_feeds_several_transforms(registry): + """Nothing about a transform consumes or invalidates the fitted converter.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0])) + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + for value, expected in ((5.0, 0.5), (20.0, 2.0), (-10.0, -1.0)): + ctx.put("dataset", _dataset(a=[value])) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + assert _values(ctx.require("dataset")) == [expected] diff --git a/tests/back/units/test_dataset_ingest_units.py b/tests/back/units/test_dataset_ingest_units.py new file mode 100644 index 000000000..bb2f86398 --- /dev/null +++ b/tests/back/units/test_dataset_ingest_units.py @@ -0,0 +1,522 @@ +"""Contract tests for the units DatasetJob is built from. + +The context is built by hand rather than through the job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +import pandas as pd +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.units.apply_dataset_schema_unit import ApplyDatasetSchemaUnit +from DashAI.back.units.compute_dataset_metadata_unit import ( + EXTENDED_METADATA_KEYS, + ComputeDatasetMetadataUnit, +) +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.infer_dataset_types_unit import InferDatasetTypesUnit +from DashAI.back.units.load_datafile_dataset_unit import LoadDatafileDatasetUnit +from DashAI.back.units.load_uploaded_dataset_unit import LoadUploadedDatasetUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit + +SCHEMA = { + "n": {"type": "Float", "dtype": "float64"}, + "label": {"type": "Categorical", "dtype": "string"}, +} + + +@pytest.fixture(name="dataset") +def fixture_dataset(): + frame = pd.DataFrame({"n": [1.0, 2.0, 3.0], "label": ["a", "b", "a"]}) + return to_dashai_dataset(frame) + + +@pytest.fixture(name="ctx") +def fixture_ctx(dataset): + ctx = ExecutionContext() + ctx.put("dataset", dataset) + return ctx + + +# --------------------------------------------------------------------------- # +# ComputeDatasetMetadataUnit +# --------------------------------------------------------------------------- # + + +def test_metadata_unit_computes_the_extended_keys_by_default(ctx): + ComputeDatasetMetadataUnit(compute_metadata=True)(ctx) + + splits = ctx.require("dataset").splits + assert splits["total_rows"] == 3 + for key in EXTENDED_METADATA_KEYS: + assert key in splits + + +def test_metadata_unit_strips_extended_keys_it_did_not_ask_for(ctx): + """A dataset can arrive carrying metadata from wherever it was copied from. + + Asking for base only has to mean base only, whatever the file happened to + bring — otherwise the result depends on the dataset's history. + """ + dataset = ctx.require("dataset") + dataset.splits["correlations"] = {"stale": "value"} + dataset.splits["general_info"] = {"stale": "value"} + + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + splits = ctx.require("dataset").splits + assert splits["total_rows"] == 3 + for key in EXTENDED_METADATA_KEYS: + assert key not in splits + + +def test_metadata_unit_reuses_trusted_extended_metadata(ctx): + """Trusting what is there is the whole point of the flag: the marker value + below survives only because nothing was recomputed.""" + dataset = ctx.require("dataset") + for key in EXTENDED_METADATA_KEYS: + dataset.splits[key] = {"marker": key} + + ComputeDatasetMetadataUnit(compute_metadata=True, trust_inherited_metadata=True)( + ctx + ) + + splits = ctx.require("dataset").splits + assert splits["correlations"] == {"marker": "correlations"} + + +def test_metadata_unit_computes_when_there_is_nothing_to_trust(ctx): + """Trusting metadata that is not there would silently store none at all.""" + ComputeDatasetMetadataUnit(compute_metadata=True, trust_inherited_metadata=True)( + ctx + ) + + splits = ctx.require("dataset").splits + assert splits["total_rows"] == 3 + for key in EXTENDED_METADATA_KEYS: + assert key in splits + + +def test_metadata_unit_keeps_the_same_dataset_object(ctx): + """It fills the dataset in place, so downstream units that already hold a + reference see the metadata too. Asserted with ``is``: an equal copy would + pass ``==`` and still break that.""" + before = ctx.require("dataset") + + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + assert ctx.require("dataset") is before + + +def test_metadata_unit_refuses_to_run_without_a_dataset(): + with pytest.raises(UnitContractError): + ComputeDatasetMetadataUnit(compute_metadata=True)(ExecutionContext()) + + +# --------------------------------------------------------------------------- # +# InferDatasetTypesUnit +# --------------------------------------------------------------------------- # + + +def test_infer_types_unit_publishes_a_type_per_column(ctx): + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + + inferred = ctx.require("inferred_types") + assert set(inferred) == {"n", "label"} + + +def test_infer_types_unit_prefers_the_types_the_dataset_already_carries(ctx): + """A reader that got types from the source knows better than inference over + the values, so those have to win.""" + typed = ApplyDatasetSchemaUnit(column_renames=None) + ctx.put_ref("inferred_types", SCHEMA) + typed(ctx) + + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + + inferred = ctx.require("inferred_types") + assert inferred["label"]["type"] == "Categorical" + + +def test_infer_types_unit_publishes_something_json_serializable(ctx): + """It is a ref, not a live object: it has to survive ``put_ref``'s check.""" + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + + # Round-trips through the serializable half without raising. + assert "inferred_types" in ctx.to_dict() + + +# --------------------------------------------------------------------------- # +# ApplyDatasetSchemaUnit +# --------------------------------------------------------------------------- # + + +def test_apply_schema_unit_casts_the_declared_types(ctx): + ctx.put_ref("inferred_types", SCHEMA) + + ApplyDatasetSchemaUnit(column_renames=None)(ctx) + + types = ctx.require("dataset").types + assert type(types["label"]).__name__ == "Categorical" + + +def test_apply_schema_unit_carries_the_types_through_a_rename(ctx): + """The declared type has to follow the column, not the old name.""" + ctx.put_ref("inferred_types", SCHEMA) + + ApplyDatasetSchemaUnit(column_renames={"label": "variety"})(ctx) + + dataset = ctx.require("dataset") + assert dataset.column_names == ["n", "variety"] + assert type(dataset.types["variety"]).__name__ == "Categorical" + + +def test_apply_schema_unit_rejects_renames_that_collide(ctx): + ctx.put_ref("inferred_types", SCHEMA) + + with pytest.raises(JobError) as excinfo: + ApplyDatasetSchemaUnit(column_renames={"label": "n"})(ctx) + + assert "contain duplicates: ['n']" in str(excinfo.value) + + +def test_apply_schema_unit_rejects_a_declaration_for_columns_that_are_gone(ctx): + """The underlying transform passes unknown columns through untouched, so a + stale declaration would silently leave columns with the wrong types. This is + the check that turns that into an error. + """ + ctx.put_ref("inferred_types", dict(SCHEMA, removed_column={"type": "Float"})) + + with pytest.raises(JobError) as excinfo: + ApplyDatasetSchemaUnit(column_renames=None)(ctx) + + assert "removed_column" in str(excinfo.value) + + +def test_apply_schema_unit_validate_runs_before_any_work(ctx): + """``validate`` is a precondition check, so it must not need the transform to + have run — and ``__call__`` runs it on its own.""" + ctx.put_ref("inferred_types", {"not_a_column": {"type": "Float"}}) + + with pytest.raises(JobError): + ApplyDatasetSchemaUnit(column_renames=None).validate(ctx) + + +def test_apply_schema_unit_refuses_to_run_without_a_declaration(ctx): + """Declared in REQUIRES, so a missing value is a wiring mistake and has to + read as one instead of as "no types to apply".""" + with pytest.raises(UnitContractError): + ApplyDatasetSchemaUnit(column_renames=None)(ctx) + + +def test_infer_then_apply_compose_over_the_same_key(ctx): + """The pair is meant to chain: one publishes what the other consumes.""" + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + ApplyDatasetSchemaUnit(column_renames={"n": "number"})(ctx) + + assert ctx.require("dataset").column_names == ["number", "label"] + + +# --------------------------------------------------------------------------- # +# SaveDatasetToPathUnit +# --------------------------------------------------------------------------- # + + +def test_save_to_path_unit_writes_where_it_is_told(ctx, tmp_path): + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + destination = tmp_path / "somewhere" / "dataset" + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + SaveDatasetToPathUnit(path=str(destination))(ctx) + + assert load_dataset(str(destination)).column_names == ["n", "label"] + + +def test_save_to_path_unit_ignores_dataset_path_entirely(ctx, tmp_path): + """Its whole reason to exist: a context left over from a load must not be + able to redirect the save back onto the source.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + source = tmp_path / "source" + destination = tmp_path / "destination" + ctx.put_ref("dataset_path", str(source)) + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + SaveDatasetToPathUnit(path=str(destination))(ctx) + + assert not source.exists() + assert load_dataset(str(destination)).column_names == ["n", "label"] + + +def test_save_to_path_unit_keeps_the_underlying_error_text(ctx, tmp_path): + """A save failure is an infrastructure failure, and only the outermost + message reaches the user — so the original text has to stay in it.""" + destination = tmp_path / "dataset" + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_save = dashai_dataset.save_dataset + + def fail(dataset, path): + raise OSError("no space left on device") + + dashai_dataset.save_dataset = fail + try: + with pytest.raises(JobError) as excinfo: + SaveDatasetToPathUnit(path=str(destination))(ctx) + finally: + dashai_dataset.save_dataset = real_save + + message = str(excinfo.value) + assert str(destination) in message + assert "no space left on device" in message + + +# --------------------------------------------------------------------------- # +# the loading units +# --------------------------------------------------------------------------- # + + +@pytest.fixture(name="csv_file") +def fixture_csv_file(tmp_path): + path = tmp_path / "data.csv" + path.write_text("n,label\n1.0,a\n2.0,b\n", encoding="utf-8") + return path + + +@pytest.fixture(name="registry") +def fixture_registry(): + """A registry holding only the CSV reader, injected without an app.""" + from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader + from DashAI.back.dependencies.registry import ComponentRegistry + + registry = ComponentRegistry(initial_components=[CSVDataLoader]) + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def test_uploaded_unit_reads_the_file_and_publishes_the_dataset( + registry, csv_file, tmp_path +): + ctx = ExecutionContext() + + LoadUploadedDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + source=str(csv_file), + temp_path=str(tmp_path), + n_sample=None, + )(ctx) + + assert ctx.require("dataset").column_names == ["n", "label"] + + +def test_uploaded_unit_publishes_no_path_and_no_id(registry, csv_file, tmp_path): + """An uploaded file is not a stored dataset yet: there is no id to correlate + against and nowhere it belongs on disk. Publishing either would hand a later + unit a value that describes something else.""" + ctx = ExecutionContext() + + LoadUploadedDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + source=str(csv_file), + temp_path=str(tmp_path), + n_sample=None, + )(ctx) + + assert not ctx.has("dataset_path") + assert not ctx.has("dataset_id") + + +def test_uploaded_unit_reports_an_unknown_reader_with_the_registry_wording( + registry, csv_file, tmp_path +): + ctx = ExecutionContext() + + with pytest.raises(KeyError) as excinfo: + LoadUploadedDatasetUnit( + dataloader={"component": "NoSuchLoader", "params": {}}, + source=str(csv_file), + temp_path=str(tmp_path), + n_sample=None, + )(ctx) + + assert "does not exists in the registry" in str(excinfo.value) + + +class _DatafileRow: + def __init__(self, local_path, status): + self.local_path = local_path + self.status = status + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """A class, not a lambda: kink calls any registered lambda with the container + to resolve it, so a lambda would be invoked as a service factory instead of + handed to the unit as one.""" + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +@pytest.fixture(name="datafile_rows") +def fixture_datafile_rows(): + rows = {} + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +def test_datafile_unit_reads_the_selected_file( + registry, datafile_rows, csv_file, tmp_path +): + from DashAI.back.core.enums.status import DatafileStatus + + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + ctx = ExecutionContext() + + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + datafile_id=7, + selected_file="data.csv", + )(ctx) + + assert ctx.require("dataset").column_names == ["n", "label"] + + +def test_datafile_unit_skips_dotted_paths_when_choosing( + registry, datafile_rows, csv_file, tmp_path +): + """Download tools leave metadata directories behind that sort before the real + data, so without the filter the wrong file would win.""" + from DashAI.back.core.enums.status import DatafileStatus + + hidden = tmp_path / ".cache" + hidden.mkdir() + (hidden / "aaa.csv").write_text("junk\n", encoding="utf-8") + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + ctx = ExecutionContext() + + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + datafile_id=7, + selected_file=None, + )(ctx) + + assert ctx.require("dataset").column_names == ["n", "label"] + + +def test_datafile_unit_rejects_a_download_that_is_not_finished( + registry, datafile_rows, tmp_path +): + from DashAI.back.core.enums.status import DatafileStatus + + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.DOWNLOADING) + + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {}}, + datafile_id=7, + selected_file=None, + )(ExecutionContext()) + + assert str(excinfo.value) == "Datafile 7 is not ready." + + +def test_datafile_unit_treats_a_missing_row_the_same_way(registry, datafile_rows): + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {}}, + datafile_id=99, + selected_file=None, + )(ExecutionContext()) + + assert str(excinfo.value) == "Datafile 99 is not ready." + + +def test_datafile_unit_reports_an_empty_download(registry, datafile_rows, tmp_path): + from DashAI.back.core.enums.status import DatafileStatus + + empty = tmp_path / "empty" + empty.mkdir() + datafile_rows[7] = _DatafileRow(str(empty), DatafileStatus.READY) + + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {}}, + datafile_id=7, + selected_file=None, + )(ExecutionContext()) + + assert str(excinfo.value) == "Hub download directory is empty." + + +def test_datafile_unit_has_its_own_wording_for_an_unknown_reader( + registry, datafile_rows, csv_file, tmp_path +): + """Deliberately different from the uploaded unit's, which surfaces the + registry's own ``KeyError``. Both texts reach users today.""" + from DashAI.back.core.enums.status import DatafileStatus + + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "NoSuchLoader", "params": {}}, + datafile_id=7, + selected_file="data.csv", + )(ExecutionContext()) + + assert str(excinfo.value) == "DataLoader 'NoSuchLoader' not found in registry." + + +def test_datafile_unit_rejects_a_component_that_is_not_a_reader( + datafile_rows, csv_file, tmp_path +): + """A registered component of some other kind is still not a reader. + + The registry's ``registry[name]`` indexer searches every type at once, so + looking the name up that way would find, say, a metric and call it as if it + could parse a file. The lookup is deliberately scoped to the readers. + """ + from DashAI.back.core.enums.status import DatafileStatus + from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader + from DashAI.back.dependencies.registry import ComponentRegistry + from DashAI.back.metrics.classification.accuracy import Accuracy + + di["component_registry"] = ComponentRegistry( + initial_components=[CSVDataLoader, Accuracy] + ) + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + try: + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "Accuracy", "params": {}}, + datafile_id=7, + selected_file="data.csv", + )(ExecutionContext()) + finally: + del di["component_registry"] + + assert str(excinfo.value) == "DataLoader 'Accuracy' not found in registry." diff --git a/tests/back/units/test_evaluate_model_to_artifact_unit.py b/tests/back/units/test_evaluate_model_to_artifact_unit.py new file mode 100644 index 000000000..53c8344d2 --- /dev/null +++ b/tests/back/units/test_evaluate_model_to_artifact_unit.py @@ -0,0 +1,199 @@ +"""Tests for EvaluateModelToArtifactUnit, and its agreement with its sibling. + +The agreement test is the load-bearing one. Two ways to score the same model +that both pass their own tests and return different numbers is the exact +failure this repository has already had once, when a prediction endpoint kept +its own copy of the three steps a prediction takes. +""" + +import pytest + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.evaluate_model_to_artifact_unit import ( + EvaluateModelToArtifactUnit, +) + + +class _Accuracy: + @staticmethod + def score(y_true, y_pred): + return 0.75 + + +class _AlwaysNaN: + @staticmethod + def score(y_true, y_pred): + return float("nan") + + +class _Model: + """A model with the runtime state ModelFactory would have attached.""" + + def __init__(self, run_id=None, metrics=(_Accuracy,), with_data=True): + self.run_id = run_id + self.train_metrics = list(metrics) + self.validation_metrics = list(metrics) + self.test_metrics = list(metrics) + splits = {"train": "x-train", "validation": "x-val", "test": "x-test"} + self.x_data = splits if with_data else None + self.y_data = splits if with_data else None + self.saved = [] + + def predict(self, x_data): + return [0, 1] + + def prepare_output(self, y_data, is_fit=False): + return [0, 1] + + # The two methods under test come from BaseModel; this stands in for it + # closely enough to exercise the unit without training anything. + compute_metrics = None # replaced below + + # Declared because ``calculate_metrics`` reads it directly: it is a class + # attribute of ``BaseModel``, so every real model has it, but this double + # borrows the method instead of inheriting it. + _epoch_reporter = None + + +def _model(**kwargs): + """A ``_Model`` borrowing BaseModel's real compute_metrics/_save_metrics.""" + from DashAI.back.models.base_model import BaseModel + + model = _Model(**kwargs) + model.compute_metrics = BaseModel.compute_metrics.__get__(model, _Model) + return model + + +def test_the_metrics_are_published_per_split(): + ctx = ExecutionContext() + ctx.put("model", _model()) + + EvaluateModelToArtifactUnit(splits=["TRAIN", "TEST"])(ctx) + + assert ctx.require("metrics") == { + "train": {"_Accuracy": 0.75}, + "test": {"_Accuracy": 0.75}, + } + + +def test_the_metrics_travel_as_plain_data(): + """A ref, not a cached object: what a caller records is plain numbers.""" + ctx = ExecutionContext() + ctx.put("model", _model()) + + EvaluateModelToArtifactUnit(splits=["TRAIN"])(ctx) + + assert ctx.origin("metrics") == "ref" + + +def test_a_model_with_no_run_is_still_evaluated(): + """The whole point: no Run row, and the numbers still come out. + + ``EvaluateModelUnit`` cannot do this. Its metric rows are keyed by a + foreign key to ``run.id``, so with no run it has nowhere to write, and + ``calculate_metrics`` returns without scoring anything. + """ + ctx = ExecutionContext() + ctx.put("model", _model(run_id=None)) + + EvaluateModelToArtifactUnit(splits=["VALIDATION"])(ctx) + + assert ctx.require("metrics") == {"validation": {"_Accuracy": 0.75}} + + +def test_a_split_with_nothing_to_score_is_left_out(): + """Absent is not the same as scored zero metrics. + + An empty entry would claim the split was evaluated and produced no + metrics, which is a different statement from there being no data for it. + """ + ctx = ExecutionContext() + ctx.put("model", _model(with_data=False)) + + EvaluateModelToArtifactUnit(splits=["TRAIN", "TEST"])(ctx) + + assert ctx.require("metrics") == {} + + +def test_a_non_finite_score_is_dropped_but_the_split_is_still_reported(): + """A split whose every metric was non-finite was still evaluated.""" + ctx = ExecutionContext() + ctx.put("model", _model(metrics=(_AlwaysNaN,))) + + EvaluateModelToArtifactUnit(splits=["TRAIN"])(ctx) + + assert ctx.require("metrics") == {"train": {}} + + +def test_the_unit_needs_a_model(): + with pytest.raises(UnitContractError, match="'model'"): + EvaluateModelToArtifactUnit(splits=["TRAIN"])(ExecutionContext()) + + +def test_the_two_evaluation_paths_agree_on_the_same_model_and_split(monkeypatch): + """The numbers the artifact carries are the numbers a run would have logged. + + Both paths score through ``BaseModel.compute_metrics``, so this fixes that + they cannot drift into two answers for the same model. + """ + from DashAI.back.models.base_model import BaseModel + + logged = {} + + # ``**kwargs`` rather than the exact signature: the real ``_save_metrics`` + # grew ``fold_index`` and ``inner_fold_index`` with cross-validation, and a + # double that has to be edited every time a persistence-only argument is + # added fails for a reason that has nothing to do with what it checks. + def _capture(self, split, level, results, log_index=None, **kwargs): + logged[split.value] = results + + monkeypatch.setattr(BaseModel, "_save_metrics", _capture) + + # The run path: a model with a run id, going through calculate_metrics. + with_run = _model(run_id=7) + with_run._save_metrics = BaseModel._save_metrics.__get__(with_run, _Model) + BaseModel.calculate_metrics.__get__(with_run, _Model)( + split=SplitEnum.TEST, level=LevelEnum.LAST + ) + + # The artifact path: the same model, no run, through the unit. + ctx = ExecutionContext() + ctx.put("model", _model(run_id=None)) + EvaluateModelToArtifactUnit(splits=["TEST"])(ctx) + + assert logged["test"] == ctx.require("metrics")["test"] + + +def test_base_model_defines_compute_metrics_exactly_once(): + """A second definition would silently win, and nothing else would notice. + + Cross-validation arrived on ``develop`` with its own copy of the scoring + body, added to ``BaseModel`` under this same name while this branch was + extracting the original one. The two do not overlap textually, so git + merges them without a conflict and Python keeps whichever comes last -- + changing what every caller of ``calculate_metrics`` computes, including the + filter that drops non-finite scores. + + Parsed rather than introspected: ``getattr`` sees one attribute either way, + which is exactly why the duplicate went unnoticed. + """ + import ast + import inspect + + from DashAI.back.models import base_model + + tree = ast.parse(inspect.getsource(base_model)) + model_class = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "BaseModel" + ) + definitions = [ + node.name + for node in model_class.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == "compute_metrics" + ] + + assert definitions == ["compute_metrics"] diff --git a/tests/back/units/test_evaluate_model_unit.py b/tests/back/units/test_evaluate_model_unit.py new file mode 100644 index 000000000..6b4a95b6c --- /dev/null +++ b/tests/back/units/test_evaluate_model_unit.py @@ -0,0 +1,39 @@ +"""Tests for EvaluateModelUnit's contract, independent of a real database.""" + +import pytest + +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit + + +def test_call_refuses_to_run_without_a_run_id(): + """Regression: a missing run id must fail loudly, not silently no-op. + + The idempotency check filters an existing-metric query by ``run_id``; a + ``None`` value would match no row regardless of what was already logged, + and -- if the model were also detached from its run -- + ``BaseModel.calculate_metrics`` no-ops on a falsy ``run_id`` too. Either + way the unit could "succeed" having written zero metrics instead of + surfacing the missing wiring. + + The run id is configuration rather than a context key (no unit publishes + it), so the guard lives in ``validate`` instead of in ``REQUIRES``. + """ + ctx = ExecutionContext() + ctx.put("model", object()) + + with pytest.raises(JobError, match="no run id"): + EvaluateModelUnit(run_id=None)(ctx) + + +def test_validate_refuses_a_missing_run_id_before_any_metric_is_computed(): + """``__call__`` validates first, so nothing is written on the way out.""" + with pytest.raises(JobError, match="no run id"): + EvaluateModelUnit(run_id=None).validate(ExecutionContext()) + + +def test_the_unit_still_needs_a_model(): + """``run_id`` left REQUIRES; ``model`` did not.""" + with pytest.raises(UnitContractError, match="'model'"): + EvaluateModelUnit(run_id=1)(ExecutionContext()) diff --git a/tests/back/units/test_explanation_units.py b/tests/back/units/test_explanation_units.py new file mode 100644 index 000000000..8d99124e1 --- /dev/null +++ b/tests/back/units/test_explanation_units.py @@ -0,0 +1,663 @@ +"""Contract tests for the explanation units, isolated from any orchestrating job. + +The context is built by hand rather than through a job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +import json +import pickle +from pathlib import Path + +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.build_global_explainer_unit import BuildGlobalExplainerUnit +from DashAI.back.units.build_local_explainer_unit import BuildLocalExplainerUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.generate_global_explanation_unit import ( + GenerateGlobalExplanationUnit, +) +from DashAI.back.units.generate_local_explanation_unit import ( + GenerateLocalExplanationUnit, +) +from DashAI.back.units.load_run_model_unit import LoadRunModelUnit +from DashAI.back.units.prepare_explanation_data_unit import PrepareExplanationDataUnit + +SPLITS = { + "train_indexes": [0, 1, 2], + "test_indexes": [3, 4], + "val_indexes": [5], +} + + +class _RunRow: + def __init__( + self, model_name="RecordingModel", run_path="somewhere", parameters=None + ): + self.id = 5 + self.model_name = model_name + self.run_path = run_path + self.parameters = parameters if parameters is not None else {"depth": 3} + + +class _DatasetRow: + def __init__(self, file_path): + self.file_path = file_path + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """Stand-in for a ``sessionmaker``. + + A class rather than a lambda on purpose: kink invokes any registered lambda + with the container to resolve it. + """ + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +class RecordingModel: + """Model that records how it was built and what it was asked to encode.""" + + built_with = None + + def __init__(self, **kwargs): + RecordingModel.built_with = kwargs + + @staticmethod + def load(filename): + # Bypasses __init__ on purpose: this is what every real model does — + # joblib or a checkpoint rebuilds the object, and the instance the unit + # constructed beforehand is thrown away. Going through __init__ here + # would overwrite the record of how that instance was built. + model = object.__new__(RecordingModel) + model.loaded_from = filename + return model + + def prepare_output(self, dataset, is_fit=False): + return dataset + + +class UninstantiableModel(RecordingModel): + def __init__(self, **kwargs): + raise RuntimeError("this model refuses to be built") + + +class UnloadableModel(RecordingModel): + @staticmethod + def load(filename): + raise OSError("the artifact is not there") + + +class RecordingTask: + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + def process_manual_input(self, rows, dataset_path): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame(rows) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +class RecordingGlobalExplainer: + def __init__(self, model, **kwargs): + self.model = model + self.kwargs = kwargs + self.seen = None + + def explain(self, dataset): + self.seen = dataset + return {"importance": [1, 2]} + + def plot(self, explanation): + return "a plot" + + +class ExplodingGlobalExplainer(RecordingGlobalExplainer): + def explain(self, dataset): + raise RuntimeError("the explanation itself blew up") + + +class RecordingLocalExplainer: + fitted_with = None + + def __init__(self, model, **kwargs): + self.model = model + self.explained_columns = None + + def fit(self, dataset, **kwargs): + RecordingLocalExplainer.fitted_with = kwargs + return self + + def explain_instance(self, instances): + columns = instances.column_names + if isinstance(columns, dict): + columns = [c for split in columns.values() for c in split] + self.explained_columns = columns + return {"local": True} + + def plot(self, explanation): + return "a plot" + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "RecordingModel": {"class": RecordingModel}, + "UninstantiableModel": {"class": UninstantiableModel}, + "UnloadableModel": {"class": UnloadableModel}, + "RecordingTask": {"class": RecordingTask}, + "RecordingGlobalExplainer": {"class": RecordingGlobalExplainer}, + "ExplodingGlobalExplainer": {"class": ExplodingGlobalExplainer}, + "RecordingLocalExplainer": {"class": RecordingLocalExplainer}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _dataset(rows=6): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame( + {"a": list(range(rows)), "b": list(range(rows)), "target": [0, 1] * (rows // 2)} + ) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +@pytest.fixture(name="stored_instances") +def fixture_stored_instances(tmp_path): + root = tmp_path / "instances" + save_dataset(_dataset(), str(root / "dataset")) + return root + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(stored_instances): + rows = { + "Run": {5: _RunRow()}, + "Dataset": {9: _DatasetRow(str(stored_instances))}, + } + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +@pytest.fixture(name="explanations_path") +def fixture_explanations_path(tmp_path): + target = tmp_path / "explanations" + target.mkdir() + di["config"] = {"EXPLANATIONS_PATH": target} + yield target + del di["config"] + + +# --- LoadRunModelUnit --------------------------------------------------- + + +def test_the_run_model_is_built_with_its_parameters_before_loading(registry, fake_db): + """The extra construction step is the difference from ``LoadTrainedModelUnit``. + + It has no effect for models whose ``load`` is a static or class method — all + of the real ones — but it is the inherited behaviour of this flow, and this + test is what would notice if the two units were quietly merged. + """ + ctx = ExecutionContext() + RecordingModel.built_with = None + + LoadRunModelUnit(run_id=5)(ctx) + + assert RecordingModel.built_with == {"depth": 3} + assert ctx.require("model").loaded_from == "somewhere" + + +def test_a_missing_run_is_reported_by_id(registry, fake_db): + with pytest.raises(JobError, match="Run 99 does not exist in DB."): + LoadRunModelUnit(run_id=99)(ExecutionContext()) + + +def test_an_unknown_model_name_uses_the_explanation_wording(registry, fake_db): + """Word for word different from ``LoadTrainedModelUnit``'s message, which is + why the two units are not merged.""" + fake_db["Run"][5].model_name = "NoSuchModel" + + with pytest.raises( + JobError, match="Unable to find Model with name NoSuchModel in registry." + ): + LoadRunModelUnit(run_id=5)(ExecutionContext()) + + +def test_a_model_that_cannot_be_built_is_reported_separately_from_loading( + registry, fake_db +): + fake_db["Run"][5].model_name = "UninstantiableModel" + + with pytest.raises(JobError, match="Unable to instantiate model") as excinfo: + LoadRunModelUnit(run_id=5)(ExecutionContext()) + + assert "refuses to be built" in str(excinfo.value.__cause__) + + +def test_a_model_that_cannot_be_loaded_names_the_path(registry, fake_db): + fake_db["Run"][5].model_name = "UnloadableModel" + fake_db["Run"][5].run_path = "gone" + + with pytest.raises(JobError, match="Can not load model from path gone"): + LoadRunModelUnit(run_id=5)(ExecutionContext()) + + +# --- the two build units ------------------------------------------------ + + +def test_the_global_build_unit_binds_the_model_from_the_context(registry): + ctx = ExecutionContext() + model = RecordingModel() + ctx.put("model", model) + + BuildGlobalExplainerUnit( + explainer={"component": "RecordingGlobalExplainer", "params": {"n": 5}} + )(ctx) + + explainer = ctx.require("explainer") + assert explainer.model is model + assert explainer.kwargs == {"n": 5} + + +def test_the_local_build_unit_produces_the_same_context_key(registry): + """Both scopes publish ``explainer``, so whatever generates the explanation + afterwards does not have to know which one ran.""" + ctx = ExecutionContext() + ctx.put("model", RecordingModel()) + + BuildLocalExplainerUnit( + explainer={"component": "RecordingLocalExplainer", "params": {}} + )(ctx) + + assert isinstance(ctx.require("explainer"), RecordingLocalExplainer) + + +def test_building_without_a_model_is_rejected_before_it_starts(registry): + with pytest.raises(UnitContractError, match="Context key 'model'"): + BuildGlobalExplainerUnit( + explainer={"component": "RecordingGlobalExplainer", "params": {}} + )(ExecutionContext()) + + +def test_each_build_unit_names_its_own_scope_in_its_errors(registry): + """The messages are worded per scope and are user-visible.""" + ctx = ExecutionContext() + ctx.put("model", RecordingModel()) + + with pytest.raises(JobError, match="Unable to find the global explainer with name"): + BuildGlobalExplainerUnit( + explainer={"component": "NoSuchExplainer", "params": {}} + )(ctx) + + with pytest.raises(JobError, match="Unable to find the local explainer with name"): + BuildLocalExplainerUnit( + explainer={"component": "NoSuchExplainer", "params": {}} + )(ctx) + + +# --- PrepareExplanationDataUnit ----------------------------------------- + + +def _prepared_context(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put("model", RecordingModel()) + ctx.put_ref("dataset_id", 9) + ctx.put_ref("split_indexes", SPLITS) + return ctx + + +def _prepare_unit(**overrides): + config = { + "task_name": "RecordingTask", + "input_columns": ["a", "b"], + "output_columns": ["target"], + } + config.update(overrides) + return PrepareExplanationDataUnit(**config) + + +def test_prepare_replays_the_recorded_split(registry): + """The indexes come from the run, not from a ratio: the explanation has to + be about the rows the model actually saw.""" + ctx = _prepared_context(registry) + + _prepare_unit()(ctx) + + data_x = ctx.require("data_x") + assert sorted(data_x.keys()) == ["test", "train", "validation"] + assert len(data_x["train"]) == 3 + assert len(data_x["test"]) == 2 + assert len(data_x["validation"]) == 1 + assert data_x["train"].column_names == ["a", "b"] + assert ctx.require("data_y")["train"].column_names == ["target"] + + +def test_prepare_publishes_the_task_for_the_local_path(registry): + ctx = _prepared_context(registry) + + _prepare_unit()(ctx) + + assert isinstance(ctx.require("task"), RecordingTask) + + +def test_prepare_validates_the_task_before_it_runs(registry): + """``validate`` is called by the orchestrator outside the block that wraps + preparation failures, so a missing task stays a registry error.""" + with pytest.raises( + JobError, match="Unable to find Task with name NoSuchTask in registry" + ): + _prepare_unit(task_name="NoSuchTask").validate(ExecutionContext()) + + +def test_prepare_without_split_indexes_is_rejected_before_it_starts(registry): + """A missing key means "nothing published them", not "there is no split".""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put("model", RecordingModel()) + ctx.put_ref("dataset_id", 9) + + with pytest.raises(UnitContractError, match="Context key 'split_indexes'"): + _prepare_unit()(ctx) + + +def test_prepare_without_a_model_is_rejected_before_it_starts(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put_ref("dataset_id", 9) + ctx.put_ref("split_indexes", SPLITS) + + with pytest.raises(UnitContractError, match="Context key 'model'"): + _prepare_unit()(ctx) + + +def test_prepare_composes_after_a_loader_that_publishes_no_dataset_id(registry): + """A dataset with no id attached is enough to run. + + ``REQUIRES`` is demanded unconditionally, so listing a key the unit never + reads would silently restrict what it can follow. ``BuildManualInputUnit`` + publishes ``dataset`` alone, and this unit has to work after it. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put("model", RecordingModel()) + ctx.put_ref("split_indexes", SPLITS) + + _prepare_unit()(ctx) + + assert not ctx.has("dataset_id") + assert len(ctx.require("data_x")["train"]) == 3 + + +# --- GenerateGlobalExplanationUnit -------------------------------------- + + +def test_the_global_explanation_pickles_both_artifacts(registry, explanations_path): + ctx = ExecutionContext() + explainer = RecordingGlobalExplainer(RecordingModel()) + ctx.put("explainer", explainer) + ctx.put("data_x", {"train": 1}) + ctx.put("data_y", {"train": 2}) + + GenerateGlobalExplanationUnit(explainer_id=7)(ctx) + + explanation_path = Path(ctx.require("explanation_path")) + plot_path = Path(ctx.require("plot_path")) + assert explanation_path.name == "global_explanation_7.pickle" + assert plot_path.name == "global_explanation_plot_7.pickle" + + with open(explanation_path, "rb") as handle: + assert pickle.load(handle) == {"importance": [1, 2]} + # The explainer receives the two halves as a pair, in order. + assert explainer.seen == ({"train": 1}, {"train": 2}) + + +def test_the_global_unit_never_writes_the_row(registry, explanations_path): + """It publishes where it wrote; the row belongs to the job. + + The unit takes only an id, so there is nothing for it to write a row with — + which is the point. + """ + assert set(GenerateGlobalExplanationUnit.SCHEMA.model_fields) == {"explainer_id"} + assert GenerateGlobalExplanationUnit.PROVIDES == ( + "explanation_path", + "plot_path", + ) + + +def test_a_failing_global_explanation_is_wrapped(registry, explanations_path): + ctx = ExecutionContext() + ctx.put("explainer", ExplodingGlobalExplainer(RecordingModel())) + ctx.put("data_x", {}) + ctx.put("data_y", {}) + + with pytest.raises(JobError, match="Failed to generate the explanation") as excinfo: + GenerateGlobalExplanationUnit(explainer_id=7)(ctx) + + assert "the explanation itself blew up" in str(excinfo.value.__cause__) + + +def test_generating_without_an_explainer_is_rejected_before_it_starts( + registry, explanations_path +): + ctx = ExecutionContext() + ctx.put("data_x", {}) + ctx.put("data_y", {}) + + with pytest.raises(UnitContractError, match="Context key 'explainer'"): + GenerateGlobalExplanationUnit(explainer_id=7)(ctx) + + +# --- GenerateLocalExplanationUnit --------------------------------------- + + +def _local_context(registry): + ctx = ExecutionContext() + ctx.put("explainer", RecordingLocalExplainer(RecordingModel())) + ctx.put("task", RecordingTask()) + ctx.put("data_x", {"train": 1}) + ctx.put("data_y", {"train": 2}) + ctx.put_ref("split_indexes", SPLITS) + return ctx + + +def _local_unit(**overrides): + config = { + "explainer_id": 7, + "instance_dataset_id": 9, + "scope": {"split": "test", "percentage": 100}, + "fit_parameters": {"nsamples": 10}, + "input_columns": ["a", "b"], + "output_columns": ["target"], + "manual_input_data": None, + "same_dataset": True, + "session_splits": None, + } + config.update(overrides) + return GenerateLocalExplanationUnit(**config) + + +def test_the_local_explanation_writes_three_artifacts( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + _local_unit()(ctx) + + assert Path(ctx.require("explanation_path")).name == "local_explanation_7.pickle" + assert Path(ctx.require("plots_path")).name == "local_explanation_plots_7.pickle" + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + assert saved.column_names == ["a", "b"] + assert len(saved) == 2 + + +def test_the_local_explanation_forwards_its_fit_parameters( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + RecordingLocalExplainer.fitted_with = None + + _local_unit()(ctx) + + assert RecordingLocalExplainer.fitted_with == {"nsamples": 10} + + +def test_instances_from_another_dataset_recompute_the_split( + registry, fake_db, explanations_path +): + """The run's row indexes are meaningless over a different dataset. + + When the instances do not come from the dataset the model was trained on, + replaying ``split_indexes`` would address rows that do not correspond, so + the split is recomputed from the session's ratios over the dataset in hand. + That derived state is resolved inside ``execute`` and never published — this + is the branch that proves the recompute actually happens. + """ + ctx = _local_context(registry) + + _local_unit( + same_dataset=False, + session_splits=json.dumps( + { + "train": 0.5, + "test": 0.5, + "validation": 0.0, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } + ), + scope={"split": "test", "percentage": 100}, + )(ctx) + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + # Half of the six stored rows, not the two the run's test_indexes name. + assert len(saved) == 3 + # The recomputed split never leaks back into the context. + assert ctx.require("split_indexes") == SPLITS + + +def test_a_rows_scope_selects_exactly_the_valid_indexes( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + _local_unit(scope={"mode": "rows", "row_indexes": [0, 3, 5]})(ctx) + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + assert len(saved) == 3 + + +def test_a_manual_scope_builds_the_instances_from_the_given_rows( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + _local_unit( + scope={"mode": "manual"}, + manual_input_data=[{"a": 1, "b": 2}, {"a": 3, "b": 4}], + )(ctx) + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + assert saved.column_names == ["a", "b"] + assert len(saved) == 2 + + +def test_the_three_selection_complaints_are_swallowed_by_one_wrapper( + registry, fake_db, explanations_path +): + """All three modes report through the same message, keeping their own + complaint only as ``__cause__``. Pinned because it is an easy detail to + "fix" by accident.""" + cases = [ + ({"split": "notasplit", "percentage": 100}, None, "not a valid split"), + ( + {"mode": "rows", "row_indexes": [10**9]}, + None, + "No valid row indexes provided", + ), + ({"mode": "manual"}, None, "No manual input data provided"), + ] + + for scope, manual, cause in cases: + ctx = _local_context(registry) + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + _local_unit(scope=scope, manual_input_data=manual)(ctx) + assert cause in str(excinfo.value.__cause__), scope + + +def test_a_missing_instance_dataset_is_reported_by_id( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + with pytest.raises( + JobError, match="Dataset 99 to be explained does not exist in DB." + ): + _local_unit(instance_dataset_id=99)(ctx) + + +def test_the_local_unit_needs_the_task_and_says_so( + registry, fake_db, explanations_path +): + """The manual mode calls into the task, so it is a declared requirement even + though the other two modes barely touch it.""" + ctx = ExecutionContext() + ctx.put("explainer", RecordingLocalExplainer(RecordingModel())) + ctx.put("data_x", {"train": 1}) + ctx.put("data_y", {"train": 2}) + ctx.put_ref("split_indexes", SPLITS) + + with pytest.raises(UnitContractError, match="Context key 'task'"): + _local_unit()(ctx) + + +def test_the_published_paths_are_plain_strings(registry, fake_db, explanations_path): + """All three travel as refs, so they have to be JSON data.""" + ctx = _local_context(registry) + + _local_unit()(ctx) + + refs = ctx.to_dict() + for key in ("explanation_path", "plots_path", "input_dataset_path"): + assert isinstance(refs[key], str), key diff --git a/tests/back/units/test_exploration_units.py b/tests/back/units/test_exploration_units.py new file mode 100644 index 000000000..0a104fcb0 --- /dev/null +++ b/tests/back/units/test_exploration_units.py @@ -0,0 +1,326 @@ +"""Contract tests for the exploration units, isolated from any orchestrating job. + +The context is built by hand rather than through a job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +import pathlib + +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.exploration.base_explorer import BaseExplorer +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.run_exploration_unit import RunExplorationUnit +from DashAI.back.units.save_exploration_unit import SaveExplorationUnit + + +class _ExplorerRow: + """Stand-in for an Explorer ORM row.""" + + def __init__( + self, notebook_id=3, columns=None, exploration_type="RecordingExplorer" + ): + self.id = 11 + self.notebook_id = notebook_id + self.columns = columns if columns is not None else [{"columnName": "a"}] + self.exploration_type = exploration_type + self.name = "an exploration" + + +class _NotebookRow: + """Stand-in for a Notebook ORM row.""" + + def __init__(self, notebook_id=3): + self.id = notebook_id + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """Stand-in for a ``sessionmaker``. + + A class rather than a lambda on purpose: kink invokes any registered lambda + with the container to resolve it, so a lambda here would be called as a + service factory instead of being handed to the unit as one. + """ + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +class RecordingExplorer(BaseExplorer): + """Explorer that records what the units hand it, and when.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.marker = kwargs.get("marker") + self.seen_columns = None + self.seen_row_name = None + + def prepare_dataset(self, loaded_dataset, columns): + self.seen_columns = [column["columnName"] for column in columns] + return loaded_dataset.select_columns(self.seen_columns) + + def launch_exploration(self, dataset, explorer_info): + self.seen_row_name = explorer_info.name + return {"rows": len(dataset), "columns": dataset.column_names} + + def save_notebook(self, notebook_info, explorer_info, save_path, result): + # Reads instance state set at construction time, the way CorrMatrix + # reads self.plot, and returns a str the way DescribeExplorer does. + target = pathlib.Path(save_path) / f"{explorer_info.id}-{self.marker}.txt" + target.write_text(str(result), encoding="utf-8") + return target.as_posix() + + def get_results(self, exploration_path, options): + return [] + + +class BadPathExplorer(RecordingExplorer): + """Explorer whose save returns something that is not a path.""" + + def save_notebook(self, notebook_info, explorer_info, save_path, result): + return 42 + + +class ExplodingExplorer(RecordingExplorer): + def launch_exploration(self, dataset, explorer_info): + raise RuntimeError("the exploration itself blew up") + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "RecordingExplorer": {"class": RecordingExplorer}, + "BadPathExplorer": {"class": BadPathExplorer}, + "ExplodingExplorer": {"class": ExplodingExplorer}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(): + rows = { + "Explorer": {11: _ExplorerRow()}, + "Notebook": {3: _NotebookRow()}, + } + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +@pytest.fixture(name="notebook_path") +def fixture_notebook_path(tmp_path): + config = {"NOTEBOOK_PATH": tmp_path / "notebooks"} + di["config"] = config + yield config["NOTEBOOK_PATH"] + del di["config"] + + +@pytest.fixture(name="ctx") +def fixture_ctx(): + """A context holding a three-column dataset, as a loader would leave it.""" + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + types = {name: Integer(arrow_type=pa.int64()) for name in ("a", "b", "c")} + + context = ExecutionContext() + context.put("dataset", to_dashai_dataset(frame, types=types)) + return context + + +def _explorer(component="RecordingExplorer", **params): + return {"component": component, "params": {"marker": "x", **params}} + + +# --- RunExplorationUnit ------------------------------------------------- + + +def test_run_exploration_publishes_the_result_and_the_explorer(ctx, registry, fake_db): + """Both outputs are declared, so both must be there. + + The explorer instance is an output and not a private detail: saving is a + method on it, and the save unit has to receive the object that ran. + """ + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + assert ctx.require("exploration_result") == {"rows": 3, "columns": ["a"]} + assert isinstance(ctx.require("explorer"), RecordingExplorer) + + +def test_run_exploration_narrows_the_dataset_to_the_rows_selected_columns( + ctx, registry, fake_db +): + """The column list comes from the row, resolved against the dataset the + context holds right now — never against a list captured earlier.""" + fake_db["Explorer"][11].columns = [{"columnName": "b"}, {"columnName": "c"}] + + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + assert ctx.require("explorer").seen_columns == ["b", "c"] + + +def test_run_exploration_hands_the_row_to_the_component(ctx, registry, fake_db): + """``launch_exploration`` takes the ORM row; the unit re-reads it itself.""" + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + assert ctx.require("explorer").seen_row_name == "an exploration" + + +def test_run_exploration_without_a_dataset_is_rejected_before_it_starts( + registry, fake_db +): + """REQUIRES is enforced by ``__call__``, so a wiring mistake is not + mistaken for an empty dataset.""" + with pytest.raises(UnitContractError, match="Context key 'dataset'"): + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ExecutionContext()) + + +def test_an_unknown_explorer_is_reported_by_name(ctx, registry, fake_db): + with pytest.raises(JobError, match="Explorer NoSuchExplorer not found in the reg"): + RunExplorationUnit( + explorer_id=11, explorer=_explorer(component="NoSuchExplorer") + )(ctx) + + +def test_a_missing_explorer_row_is_reported_by_id(ctx, registry, fake_db): + with pytest.raises(JobError, match="Explorer with id 99 not found."): + RunExplorationUnit(explorer_id=99, explorer=_explorer())(ctx) + + +def test_a_column_absent_from_the_dataset_becomes_a_preparation_error( + ctx, registry, fake_db +): + fake_db["Explorer"][11].columns = [{"columnName": "nope"}] + + with pytest.raises( + JobError, match="Error preparing the dataset for the exploration" + ): + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + +def test_a_failing_exploration_is_wrapped_with_the_component_name( + ctx, registry, fake_db +): + with pytest.raises( + JobError, match="Error launching the exploration ExplodingExplorer." + ) as excinfo: + RunExplorationUnit( + explorer_id=11, explorer=_explorer(component="ExplodingExplorer") + )(ctx) + + assert "the exploration itself blew up" in str(excinfo.value.__cause__) + + +def test_two_exploration_units_do_not_share_a_resolved_class(ctx, registry, fake_db): + """The registry lookup is memoized on the instance, not in the context. + + Two exploration nodes in one context must each resolve their own component; + a context-global cache key would make the second silently reuse the first. + """ + first = RunExplorationUnit(explorer_id=11, explorer=_explorer()) + second = RunExplorationUnit( + explorer_id=11, explorer=_explorer(component="BadPathExplorer") + ) + + first(ctx) + second(ctx) + + assert first._explorer_class is RecordingExplorer + assert second._explorer_class is BadPathExplorer + + +# --- SaveExplorationUnit ------------------------------------------------ + + +def test_save_exploration_writes_under_the_notebook_folder( + ctx, registry, fake_db, notebook_path +): + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + SaveExplorationUnit(explorer_id=11)(ctx) + + written = pathlib.Path(ctx.require("exploration_path")) + assert written.exists() + assert written.parent == notebook_path / "3" + assert written.name == "11-x.txt" + + +def test_save_exploration_uses_the_explorer_that_ran( + ctx, registry, fake_db, notebook_path +): + """Identity, not equality: the saved file name carries state the instance + was built with, so rebuilding a second explorer from the same config would + pass this by accident. Asserting ``is`` is what makes it real.""" + RunExplorationUnit(explorer_id=11, explorer=_explorer(marker="carried"))(ctx) + ran = ctx.require("explorer") + + SaveExplorationUnit(explorer_id=11)(ctx) + + assert ctx.require("explorer") is ran + assert pathlib.Path(ctx.require("exploration_path")).name == "11-carried.txt" + + +def test_save_exploration_creates_the_notebook_folder_when_absent( + ctx, registry, fake_db, notebook_path +): + assert not notebook_path.exists() + + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + SaveExplorationUnit(explorer_id=11)(ctx) + + assert (notebook_path / "3").is_dir() + + +def test_save_exploration_without_a_result_is_rejected( + registry, fake_db, notebook_path +): + with pytest.raises(UnitContractError, match="Context key 'exploration_result'"): + SaveExplorationUnit(explorer_id=11)(ExecutionContext()) + + +def test_a_save_that_does_not_return_a_path_is_reported( + ctx, registry, fake_db, notebook_path +): + RunExplorationUnit(explorer_id=11, explorer=_explorer(component="BadPathExplorer"))( + ctx + ) + + with pytest.raises(JobError, match="save path is not a pathlib.Path"): + SaveExplorationUnit(explorer_id=11)(ctx) + + +def test_the_published_path_is_a_plain_string(ctx, registry, fake_db, notebook_path): + """``exploration_path`` travels as a ref, so it has to be JSON data. + + A ``pathlib.Path`` would raise on ``put_ref``; this pins that the unit + converts before publishing. + """ + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + SaveExplorationUnit(explorer_id=11)(ctx) + + assert isinstance(ctx.to_dict()["exploration_path"], str) diff --git a/tests/back/units/test_fit_model_over_folds_unit.py b/tests/back/units/test_fit_model_over_folds_unit.py new file mode 100644 index 000000000..cb9587ff2 --- /dev/null +++ b/tests/back/units/test_fit_model_over_folds_unit.py @@ -0,0 +1,310 @@ +"""Contract tests for the unit that fits a model across folds. + +Built on a hand-made ``ExecutionContext`` rather than through a job. The +end-to-end net already runs this unit inside a real run; what it cannot show is +what the unit reads, promises and refuses on its own, which is where the +composability mistakes live. +""" + +import pytest +from kink import di + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit + +#: Three folds and the trailing entry, which is not one: it holds every row the +#: folds could use and the rows reserved from them. +FOLDS = [ + {"train": "train-0", "validation": "val-0"}, + {"train": "train-1", "validation": "val-1"}, + {"train": "train-2", "validation": "val-2"}, + {"train": "pool", "test": "reserved"}, +] + + +class _RecordingModel: + """Records what was asked of it, and scores by naming what it saw.""" + + def __init__(self, run_id=7, metrics=("Accuracy",)): + self.run_id = run_id + self._metrics = list(metrics) + self.fits = [] + self.logged = [] + self.saved = [] + self.x_data = None + self.y_data = None + self._score = 0.0 + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.fits.append({"train": x_train, "validation": x_validation}) + + def predict(self, x_data): + return f"predictions-for-{x_data}" + + def prepare_output(self, y_data, is_fit=False): + return f"expected-from-{y_data}" + + def compute_metrics(self, split, x_data=None, y_data=None): + if not self._metrics: + return None + # A different number every time, so a mean over folds is not the same + # as any one of them. + self._score += 1.0 + return dict.fromkeys(self._metrics, self._score) + + def calculate_metrics(self, split, level, fold_index=None, **kwargs): + if not self.run_id: + return None + results = self.compute_metrics(split) + if results is None: + return None + self.logged.append((split, level, fold_index)) + return results + + def _save_metrics(self, results, split, level, **kwargs): + self.saved.append((split, level, results)) + + +class _Metric: + """Scores by naming what it was given.""" + + __name__ = "Accuracy" + + @staticmethod + def score(expected, predictions): + return 1.0 + + +class _Factory: + @staticmethod + def update_parameters(old, best): + return dict(old) + + +def _unit(**overrides): + config = { + "optimizer": {"component": "AnOptimizer", "params": {}}, + "goal_metric": "Accuracy", + "run_id": 7, + "artifact_prefix": "7", + } + config.update(overrides) + return FitModelOverFoldsUnit(**config) + + +def _context(model, folds=None): + ctx = ExecutionContext() + ctx.put("model", model) + ctx.put("x_folds", folds if folds is not None else FOLDS) + ctx.put("y_folds", folds if folds is not None else FOLDS) + ctx.put("optimizable_parameters", []) + ctx.put("factory", _Factory) + ctx.put_ref("model_parameters", {}) + return ctx + + +# --------------------------------------------------------------------------- # +# The fold loop +# --------------------------------------------------------------------------- # + + +def test_every_fold_is_fitted_and_the_kept_model_once_more(): + """Three folds and a refit, and the refit is on the pooled rows.""" + model = _RecordingModel() + + _unit()(_context(model)) + + assert [fit["train"] for fit in model.fits] == [ + "train-0", + "train-1", + "train-2", + "pool", + ] + + +def test_the_trailing_entry_is_never_scored_as_a_fold(): + """It is not a fold: it fits the model that gets kept. + + Scoring it as one would put the pooled rows -- which every fold trained on + -- into the mean that is meant to describe how the model does on rows it + has not seen. + """ + model = _RecordingModel() + + _unit()(_context(model)) + + fold_indexes = sorted( + {index for _, level, index in model.logged if index is not None} + ) + assert fold_indexes == [0, 1, 2] + + +def test_no_fold_fit_receives_validation_data(): + """A fold is scored on the rows it held back, so the fit may not see them. + + Nothing raises when it does; the score simply comes out better than the + model deserves. Unlike the holdout sibling this is not configurable, + because there is no reading of a fold under which it would be right. + """ + model = _RecordingModel() + + _unit()(_context(model)) + + assert all(fit["validation"] is None for fit in model.fits) + + +def test_the_fold_rows_are_written_at_the_fold_level_and_indexed_from_zero(): + """Contiguous from zero, which the repeated-cross-validation charts assume. + + They bucket folds into repetitions by integer division, so a gap moves a + fold into the wrong repetition rather than raising. + """ + model = _RecordingModel() + + _unit()(_context(model)) + + by_split = {} + for split, level, index in model.logged: + assert level is LevelEnum.FOLD + by_split.setdefault(split, []).append(index) + + assert set(by_split) == {SplitEnum.TRAIN, SplitEnum.VALIDATION} + for split, indexes in by_split.items(): + assert indexes == [0, 1, 2], split + + +def test_the_per_fold_scores_are_published_for_whoever_aggregates_them(): + """The unit does not summarise them itself. + + A summary row carries a standard deviation, and the one sanctioned write in + the domain layer has nowhere to put one -- so the numbers are handed over + and the caller persists the summary. + """ + model = _RecordingModel() + ctx = _context(model) + + _unit()(ctx) + + fold_metrics = ctx.require("fold_metrics") + assert set(fold_metrics) == {"TRAIN", "VALIDATION"} + for split_name, by_metric in fold_metrics.items(): + assert set(by_metric) == {"Accuracy"}, split_name + assert len(by_metric["Accuracy"]) == 3, split_name + + +def test_the_model_ends_pointing_at_the_partitions_it_was_last_fitted_on(): + """Whatever scores it afterwards reads the data off the instance.""" + model = _RecordingModel() + + _unit()(_context(model)) + + assert model.x_data is FOLDS[-1] + assert model.y_data is FOLDS[-1] + + +def test_a_split_with_no_metrics_configured_is_absent_rather_than_empty(): + """Present and empty would claim it was scored and produced nothing.""" + model = _RecordingModel(metrics=()) + ctx = _context(model) + + _unit()(ctx) + + assert ctx.require("fold_metrics") == {} + + +# --------------------------------------------------------------------------- # +# The objective the search measures +# --------------------------------------------------------------------------- # + + +def test_one_trial_fits_every_fold_and_returns_their_mean(): + """A trial costs k fits. That is what cross-validation buys, and its price.""" + model = _RecordingModel() + + score = _unit()._score_one_trial(model, FOLDS, FOLDS, _Metric) + + assert [fit["train"] for fit in model.fits] == ["train-0", "train-1", "train-2"] + assert score == 1.0 + + +def test_one_trial_records_one_row_per_split_holding_the_mean_of_its_folds(): + """Not one row per fold: those measure a setting, not the kept model.""" + model = _RecordingModel() + + _unit()._score_one_trial(model, FOLDS, FOLDS, _Metric) + + assert [(split, level) for split, level, _ in model.saved] == [ + (SplitEnum.TRAIN, LevelEnum.TRIAL), + (SplitEnum.VALIDATION, LevelEnum.TRIAL), + ] + # Three folds scoring 1, 3, 5 on TRAIN (the double counts every call) means + # the row is their mean and not any one of them. + train_row = model.saved[0][2] + assert train_row["Accuracy"] == 3.0 + + +def test_a_trial_writes_no_rows_when_there_is_no_run_to_write_against(): + """``_save_metrics`` does not guard this for itself. + + A caller with no run -- a pipeline -- would write rows against a foreign + key pointing at nothing, and they insert without complaint because nothing + enforces it. That is the failure a metrics unit already exists to refuse. + """ + model = _RecordingModel(run_id=None) + + _unit()._score_one_trial(model, FOLDS, FOLDS, _Metric) + + assert model.saved == [] + assert model.fits, "the folds are still fitted; only the recording is refused" + + +# --------------------------------------------------------------------------- # +# The contract itself +# --------------------------------------------------------------------------- # + + +def test_the_unit_refuses_to_run_before_the_model_was_built(): + """``__call__`` checks REQUIRES, so the failure names the missing key.""" + with pytest.raises(UnitContractError): + _unit()(ExecutionContext()) + + +def test_validate_raises_when_called_before_the_model_was_built(): + """A missing key is a call-order mistake, not "nothing to optimize". + + The two are different: the key present and empty is a model that declares + no optimizable parameters, and that one legitimately skips the checks. + """ + with pytest.raises(UnitContractError, match="'optimizable_parameters'"): + _unit().validate(ExecutionContext()) + + +def test_two_units_in_one_context_do_not_share_their_optimizer(): + """Instance state stays on the instance, which is what lets a graph hold two. + + A context-global cache would give the second unit the first one's + optimizer, and the second search would silently run the first one's. + """ + + class _Optimizer: + def __init__(self, **params): + pass + + registry = { + "AnOptimizer": {"class": _Optimizer}, + "Accuracy": {"class": _Metric, "metadata": {"maximize": True}}, + } + di["component_registry"] = registry + try: + ctx = ExecutionContext() + ctx.put("optimizable_parameters", ["lr"]) + + first, second = _unit(), _unit() + first.validate(ctx) + second.validate(ctx) + + assert first._optimizer is not second._optimizer + assert not ctx.has("optimizer") + finally: + del di["component_registry"] diff --git a/tests/back/units/test_fit_model_unit.py b/tests/back/units/test_fit_model_unit.py new file mode 100644 index 000000000..814835607 --- /dev/null +++ b/tests/back/units/test_fit_model_unit.py @@ -0,0 +1,404 @@ +"""Tests for FitModelUnit's validation, independent of an actual training run.""" + +import pytest +from kink import di + +from DashAI.back.job.base_job import JobError +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.fit_model_unit import FitModelUnit + +# Two trials over one hyperparameter: the shape create_plots expects, +# small enough that only the filenames are under test here. +_TRIALS = [ + {"params": {"C": 0.1}, "value": 0.5}, + {"params": {"C": 1.0}, "value": 0.9}, +] +_GOAL_METRIC = {"name": "Accuracy", "metadata": {"maximize": True}} + + +def _unit( + optimizer_name="OptunaOptimizer", + goal_metric="Accuracy", + run_id=None, + artifact_prefix=None, +): + # run_id and artifact_prefix are read straight from the config, with no + # default: omitting one is a KeyError rather than a silently disabled + # runtime-state assertion or a plot filename nobody chose. + return FitModelUnit( + optimizer={"component": optimizer_name, "params": {}}, + goal_metric=goal_metric, + run_id=run_id, + artifact_prefix=artifact_prefix, + ) + + +def test_validate_raises_when_called_before_build_model_unit_has_run(): + """Regression: a missing key must not read as "nothing to optimize". + + ``optimizable_parameters`` is only absent from the context when + ``BuildModelUnit`` hasn't run yet — a call-order mistake, not a model with + no optimizable parameters (that case has the key present but empty). + Before this fix, ``validate`` used ``ctx.get`` and treated both the same, + silently skipping the optimizer/goal-metric checks it exists to run. + """ + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="'optimizable_parameters'"): + _unit().validate(ctx) + + +def test_validate_is_a_noop_when_there_are_genuinely_no_optimizable_parameters(): + ctx = ExecutionContext() + ctx.put("optimizable_parameters", []) + + # Should not raise, and should not need the optimizer/goal_metric to + # resolve in the registry. + unit = _unit(optimizer_name="DoesNotExist", goal_metric="DoesNotExist") + unit.validate(ctx) + + assert unit._optimizer is None + assert unit._goal_metric is None + + +def test_the_optimizer_is_kept_on_the_unit_not_in_the_shared_context(): + """Regression: the optimizer is this unit's own state, not an output. + + It used to be written to the context by ``validate`` and read back by + ``execute``, using the shared context as a scratchpad between one unit's + own two phases. Two FitModelUnits in one context would overwrite each + other's optimizer, and the second would silently run the first one's. + """ + + class _Optimizer: + def __init__(self, **params): + pass + + registry = { + "AnOptimizer": {"class": _Optimizer}, + "Accuracy": {"class": object, "metadata": {"maximize": True}}, + } + di["component_registry"] = registry + try: + ctx = ExecutionContext() + ctx.put("optimizable_parameters", ["lr"]) + + unit = _unit(optimizer_name="AnOptimizer") + unit.validate(ctx) + + assert isinstance(unit._optimizer, _Optimizer) + assert not ctx.has("optimizer") + assert not ctx.has("goal_metric") + finally: + del di["component_registry"] + + +def test_the_plot_filenames_come_from_the_artifact_prefix_when_there_is_one(): + """A caller that is not a run names the plots itself. + + Two pipeline executions have no run id to tell them apart, so without a + prefix both would write ``history_objective_plot_None.pickle`` into the + shared runs directory and the second would overwrite the first. + """ + optimizer = OptunaOptimizer() + + filenames, _ = optimizer.create_plots( + _TRIALS, + None, + n_params=1, + goal_metric=_GOAL_METRIC, + artifact_prefix="pipeline-3-fit", + ) + + assert filenames == [ + "history_objective_plot_pipeline-3-fit.pickle", + "slice_plot_pipeline-3-fit.pickle", + ] + + +def test_without_an_artifact_prefix_the_filenames_still_come_from_the_run_id(): + """The default has to leave a real run's filenames exactly as they were.""" + optimizer = OptunaOptimizer() + + named_by_default, _ = optimizer.create_plots( + _TRIALS, 42, n_params=1, goal_metric=_GOAL_METRIC + ) + named_explicitly, _ = optimizer.create_plots( + _TRIALS, 42, n_params=1, goal_metric=_GOAL_METRIC, artifact_prefix=None + ) + + assert named_by_default == [ + "history_objective_plot_42.pickle", + "slice_plot_42.pickle", + ] + assert named_explicitly == named_by_default + + +class _Detached: + """What an optimizer must never return: a model with no data attached.""" + + x_data = None + run_id = None + + +class _Attached: + """What ModelFactory produces: the splits hang off the instance.""" + + x_data = {"train": "x"} + run_id = None + + +def test_a_model_detached_from_its_data_is_refused_even_with_no_run(): + """The guard used to key on run_id, which made it a no-op for a pipeline. + + A pipeline always has run_id None, so returning early on that skipped the + check for the one caller with no other signal: it would finish with an + empty metrics artifact instead of an error. What both callers need is that + the model still carries what scoring reads, so that is what is checked. + """ + with pytest.raises(JobError, match="detached from its data"): + FitModelUnit._assert_model_keeps_its_runtime_state(_Detached()) + + +def test_a_model_that_kept_its_data_passes_without_a_run(): + FitModelUnit._assert_model_keeps_its_runtime_state(_Attached()) + + +# --------------------------------------------------------------------------- # +# What the fit is allowed to look at +# --------------------------------------------------------------------------- # + + +class _RecordingModel: + """Records what was asked of it, and does nothing else.""" + + def __init__(self): + self.fits = [] + self.logged = [] + self.x_data = None + self.y_data = None + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.fits.append( + {"train": x_train, "validation": x_validation}, + ) + + def predict(self, x_data): + return f"predictions-for-{x_data}" + + def prepare_output(self, y_data, is_fit=False): + return f"expected-from-{y_data}" + + def calculate_metrics(self, split, level, **kwargs): + self.logged.append((split, level)) + + +def _fit_context(model, x, y): + ctx = ExecutionContext() + ctx.put("model", model) + ctx.put("x", x) + ctx.put("y", y) + ctx.put("optimizable_parameters", []) + ctx.put("factory", object()) + ctx.put_ref("model_parameters", {}) + ctx.put("task", object()) + return ctx + + +_HOLDOUT = {"train": "x-train", "validation": "x-val", "test": "x-test"} +#: What a fold splitter's trailing entry looks like: the pooled rows and the +#: reserved ones, and no validation partition at all. +_POOLED = {"train": "x-pool", "test": "x-reserved"} + + +def test_an_ordinary_fit_hands_the_validation_partition_over(): + """Models use it to watch the fit and stop early, which holdout wants.""" + model = _RecordingModel() + + _unit()(_fit_context(model, _HOLDOUT, _HOLDOUT)) + + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + + +def test_a_fit_that_will_be_scored_on_validation_does_not_look_at_it(): + """A fold is scored on the rows it held back from its own training. + + Handing them to the fit would measure it on data it was allowed to watch, + and nothing about that failure raises -- the score simply comes out better + than the model deserves. + """ + model = _RecordingModel() + unit = _unit() + unit.config["validation_during_fit"] = False + + unit(_fit_context(model, _HOLDOUT, _HOLDOUT)) + + assert model.fits == [{"train": "x-train", "validation": None}] + + +def test_a_partition_set_without_a_validation_split_still_fits(): + """The trailing entry of a fold splitter has nothing to hand over. + + Reading ``x["validation"]`` unconditionally, as this unit used to, is a + KeyError on the very partition set that fits the model which gets kept. + """ + model = _RecordingModel() + + _unit()(_fit_context(model, _POOLED, _POOLED)) + + assert model.fits == [{"train": "x-pool", "validation": None}] + + +def test_the_fit_points_the_model_at_the_data_it_is_being_fitted_on(): + """Built once, fitted many times: the data comes with the fit, not the build. + + The metric methods read these attributes off the instance to decide what + they are scoring, so over folds they have to follow the iteration. + """ + model = _RecordingModel() + # Same partition names on both sides, as a splitter always produces, and + # distinguishable values so the two cannot be confused for one another. + y = {name: value.replace("x-", "y-") for name, value in _HOLDOUT.items()} + + _unit()(_fit_context(model, _HOLDOUT, y)) + + assert model.x_data is _HOLDOUT + assert model.y_data is y + + +# --------------------------------------------------------------------------- # +# The objective the search measures +# --------------------------------------------------------------------------- # + + +class _NamedMetric: + """Scores by naming what it was given, so the arguments can be checked.""" + + @staticmethod + def score(expected, predictions): + return f"{expected}|{predictions}" + + +def test_one_trial_is_a_fit_and_a_score_of_the_validation_partition(): + """What the optimizer measures, and where the numbers come from. + + The objective used to be the optimizer's own business: it fitted and scored + inline, and the sixth argument of ``optimize`` was the task. Making it a + callable the unit supplies is what lets the same search be reused over + anything that can be fitted and scored. + """ + model = _RecordingModel() + + score = _unit()._score_one_trial(model, _HOLDOUT, _HOLDOUT, _NamedMetric) + + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + assert score == "expected-from-x-val|predictions-for-x-val" + + +def test_one_trial_logs_the_metrics_of_that_trial(): + """Written by the objective, not by the optimizer. + + What counts as a scored partition is a property of the thing being fitted + rather than of the search, so it is decided here. A partition with no + metrics configured writes nothing, because ``calculate_metrics`` finds + nothing to score and returns. + """ + from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum + + model = _RecordingModel() + + _unit()._score_one_trial(model, _HOLDOUT, _HOLDOUT, _NamedMetric) + + assert model.logged == [ + (SplitEnum.TRAIN, LevelEnum.TRIAL), + (SplitEnum.VALIDATION, LevelEnum.TRIAL), + ] + + +class _RecordingOptimizer: + """Stands in for a real optimizer to watch what it is handed.""" + + last_call = None + + def optimize(self, model, x, y, parameters, metric, strategy): + type(self).last_call = { + "model": model, + "parameters": parameters, + "metric": metric, + "strategy": strategy, + } + # A real optimizer leaves the model fitted at the best point it found. + strategy(model, x, y, _NamedMetric) + + def get_model(self): + return type(self).last_call["model"] + + def get_best_params(self): + return {} + + def get_trials_values(self): + return [] + + def create_plots(self, trials, run_id, n_params, goal_metric, artifact_prefix): + return [], [] + + +class _Factory: + @staticmethod + def update_parameters(old, best): + return dict(old) + + +def test_the_search_is_handed_the_units_own_objective(tmp_path): + """The wiring, pinned separately from what the objective computes. + + Passing the wrong sixth argument is silent until the optimizer calls it: + the task used to sit in that position, and a task is not callable, so the + mistake surfaced from inside a trial rather than from the call. + """ + registry = { + "RecordingOptimizer": {"class": _RecordingOptimizer}, + "Accuracy": {"class": _NamedMetric, "metadata": {"maximize": True}}, + } + di["component_registry"] = registry + # A search names the plots it produces after the run, so the runs directory + # is a real dependency of this path even when the double produces none. + di["config"] = {"RUNS_PATH": str(tmp_path)} + try: + model = _RecordingModel() + ctx = _fit_context(model, _HOLDOUT, _HOLDOUT) + ctx.put("optimizable_parameters", [("obj", "C", (0, 1), "number")]) + ctx.put("factory", _Factory) + + unit = _unit(optimizer_name="RecordingOptimizer") + unit(ctx) + + assert _RecordingOptimizer.last_call["strategy"] == unit._score_one_trial + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + finally: + del di["component_registry"] + del di["config"] + + +def test_a_run_with_no_optimizer_fits_once_even_if_a_parameter_is_optimizable(): + """Both halves are needed to call something a search: a tuner and a target. + + ``Run.optimizer_name`` is a plain string and the wizard leaves it empty when + the user does not ask for a search, while the model may still declare a + parameter as optimizable. That combination has always meant "fit it once + with the values given". Checking only the parameters turns it into a lookup + of the empty string in the registry, which surfaces as a complaint about + the metric being incompatible with the task -- a message with nothing to do + with what happened. + """ + model = _RecordingModel() + ctx = _fit_context(model, _HOLDOUT, _HOLDOUT) + ctx.put("optimizable_parameters", [("obj", "C", (0, 1), "number")]) + + unit = _unit(optimizer_name="", goal_metric="") + unit.validate(ctx) + unit(ctx) + + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + assert not ctx.has("best_parameters") diff --git a/tests/back/units/test_load_dataset_unit.py b/tests/back/units/test_load_dataset_unit.py new file mode 100644 index 000000000..1c5bbc994 --- /dev/null +++ b/tests/back/units/test_load_dataset_unit.py @@ -0,0 +1,151 @@ +"""Contract tests for LoadDatasetUnit, isolated from any orchestrating job. + +The context is built by hand rather than through a job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + save_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit + + +class _Row: + """Stand-in for a Dataset or Notebook ORM row.""" + + def __init__(self, file_path=None, dataset_id=None): + self.file_path = file_path + self.dataset_id = dataset_id + + +class _FakeSession: + """Session that answers ``get`` from a table -> {id: row} mapping.""" + + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """Stand-in for a ``sessionmaker``. + + A class rather than a lambda on purpose: kink invokes any registered lambda + with the container to resolve it, so a lambda here would be called as a + service factory instead of being handed to the unit as one. + """ + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +@pytest.fixture(name="stored_dataset") +def fixture_stored_dataset(tmp_path): + """A real two-column dataset written to ``/store/dataset``.""" + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + types = { + "a": Integer(arrow_type=pa.int64()), + "b": Integer(arrow_type=pa.int64()), + } + root = tmp_path / "store" + save_dataset(to_dashai_dataset(frame, types=types), str(root / "dataset")) + return root + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(stored_dataset): + """Point both a Dataset row and a Notebook row at the stored dataset.""" + rows = { + "Dataset": {7: _Row(file_path=str(stored_dataset))}, + "Notebook": {3: _Row(file_path=str(stored_dataset), dataset_id=7)}, + } + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +def test_loading_by_dataset_id_publishes_the_whole_contract(fake_db): + ctx = ExecutionContext() + + LoadDatasetUnit(dataset_id=7)(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + assert ctx.require("dataset_id") == 7 + assert ctx.require("dataset_path").endswith("dataset") + + +def test_loading_by_notebook_id_resolves_the_source_dataset_id(fake_db): + """The notebook branch must still publish a dataset id. + + Downstream error messages identify the work by dataset id, so a notebook + load that left the key unset would report ``None`` instead of the dataset. + """ + ctx = ExecutionContext() + + LoadDatasetUnit(notebook_id=3)(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + assert ctx.require("dataset_id") == 7 + + +def test_the_two_starting_points_are_mutually_exclusive(fake_db): + with pytest.raises(JobError, match="exactly one of dataset_id or notebook_id"): + LoadDatasetUnit(dataset_id=7, notebook_id=3)(ExecutionContext()) + + +def test_no_starting_point_at_all_is_rejected(fake_db): + with pytest.raises(JobError, match="exactly one of dataset_id or notebook_id"): + LoadDatasetUnit()(ExecutionContext()) + + +def test_a_missing_dataset_row_is_reported_by_id(fake_db): + with pytest.raises(JobError, match="Dataset 99 does not exist in DB."): + LoadDatasetUnit(dataset_id=99)(ExecutionContext()) + + +def test_a_missing_notebook_row_is_reported_by_id(fake_db): + with pytest.raises(JobError, match="Notebook 99 does not exist in DB."): + LoadDatasetUnit(notebook_id=99)(ExecutionContext()) + + +def test_an_unreadable_path_becomes_a_job_error(fake_db, tmp_path): + fake_db["Notebook"][4] = _Row(file_path=str(tmp_path / "nowhere"), dataset_id=7) + + with pytest.raises(JobError, match="Can not load dataset from path"): + LoadDatasetUnit(notebook_id=4)(ExecutionContext()) + + +def test_the_dataset_is_cached_live_not_copied(fake_db): + """The dataset must come back as the same object, not a copy. + + ``ctx.get`` deep-copies the refs half and returns the cache half by + reference; a dataset that came back copied would mean every unit downstream + transformed a different object than the one that gets saved. + """ + ctx = ExecutionContext() + + LoadDatasetUnit(dataset_id=7)(ctx) + + assert ctx.require("dataset") is ctx.require("dataset") diff --git a/tests/back/units/test_prediction_units.py b/tests/back/units/test_prediction_units.py new file mode 100644 index 000000000..25a937b0f --- /dev/null +++ b/tests/back/units/test_prediction_units.py @@ -0,0 +1,423 @@ +"""Contract tests for the prediction units, isolated from any orchestrating job. + +The context is built by hand rather than through a job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +from pathlib import Path + +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.load_trained_model_unit import LoadTrainedModelUnit +from DashAI.back.units.load_training_dataset_unit import LoadTrainingDatasetUnit +from DashAI.back.units.predict_unit import PredictUnit +from DashAI.back.units.save_prediction_unit import SavePredictionUnit + + +class _RunRow: + """Stand-in for a Run ORM row.""" + + def __init__(self, model_name="RecordingModel", run_path="somewhere"): + self.model_name = model_name + self.run_path = run_path + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """Stand-in for a ``sessionmaker``. + + A class rather than a lambda on purpose: kink invokes any registered lambda + with the container to resolve it, so a lambda here would be called as a + service factory instead of being handed to the unit as one. + """ + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +class RecordingModel: + """Model whose ``load`` is a staticmethod, the way every real one is.""" + + loaded_from = None + + def __init__(self): + self.seen_columns = None + + @staticmethod + def load(filename): + model = RecordingModel() + RecordingModel.loaded_from = filename + return model + + def predict(self, x): + self.seen_columns = x.column_names + return [0] * len(x) + + +class UnloadableModel(RecordingModel): + @staticmethod + def load(filename): + raise OSError("the artifact is not there") + + +class RecordingTask: + """Task that records the training dataset it was given for decoding.""" + + seen_train_columns = None + + def process_predictions(self, train_dataset, y_pred_proba, output_column): + RecordingTask.seen_train_columns = train_dataset.column_names + return [f"label-{int(value)}" for value in y_pred_proba] + + def process_manual_input(self, rows, dataset_path): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame(rows) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "RecordingModel": {"class": RecordingModel}, + "UnloadableModel": {"class": UnloadableModel}, + "RecordingTask": {"class": RecordingTask}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(): + rows = {"Run": {5: _RunRow()}} + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +def _dataset(**columns): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame(columns) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +@pytest.fixture(name="stored_training_dataset") +def fixture_stored_training_dataset(tmp_path): + """A training dataset on disk, laid out the way a Dataset row points at it.""" + root = tmp_path / "training" + save_dataset(_dataset(a=[1, 2, 3], b=[4, 5, 6]), str(root / "dataset")) + return root + + +@pytest.fixture(name="datasets_path") +def fixture_datasets_path(tmp_path): + config = {"DATASETS_PATH": tmp_path / "datasets"} + di["config"] = config + yield config["DATASETS_PATH"] + del di["config"] + + +# --- LoadTrainedModelUnit ----------------------------------------------- + + +def test_the_model_comes_from_the_path_the_run_recorded(registry, fake_db): + """Neither the component nor the path is configuration: both are read off + the run, so the model restored is always the one that run saved.""" + ctx = ExecutionContext() + fake_db["Run"][5].run_path = "the/recorded/path" + + LoadTrainedModelUnit(run_id=5)(ctx) + + assert isinstance(ctx.require("model"), RecordingModel) + assert RecordingModel.loaded_from == "the/recorded/path" + + +def test_a_missing_run_is_reported_by_id(registry, fake_db): + with pytest.raises(JobError, match="Run 99 does not exist in DB."): + LoadTrainedModelUnit(run_id=99)(ExecutionContext()) + + +def test_an_unknown_model_name_is_reported_by_name(registry, fake_db): + fake_db["Run"][5].model_name = "NoSuchModel" + + with pytest.raises(JobError, match="Model NoSuchModel not found in the registry"): + LoadTrainedModelUnit(run_id=5)(ExecutionContext()) + + +def test_an_artifact_that_cannot_be_read_names_the_model_and_the_path( + registry, fake_db +): + fake_db["Run"][5].model_name = "UnloadableModel" + fake_db["Run"][5].run_path = "gone" + + with pytest.raises( + JobError, match="Failed to load model UnloadableModel from path gone" + ) as excinfo: + LoadTrainedModelUnit(run_id=5)(ExecutionContext()) + + assert "the artifact is not there" in str(excinfo.value.__cause__) + + +def test_two_model_units_do_not_share_a_resolved_class(registry, fake_db): + """The registry lookup is memoized on the instance, not in the context.""" + fake_db["Run"][6] = _RunRow(model_name="UnloadableModel", run_path="gone") + + first = LoadTrainedModelUnit(run_id=5) + second = LoadTrainedModelUnit(run_id=6) + + first(ExecutionContext()) + with pytest.raises(JobError): + second(ExecutionContext()) + + assert first._model_class is RecordingModel + assert second._model_class is UnloadableModel + + +# --- LoadTrainingDatasetUnit -------------------------------------------- + + +def test_the_training_dataset_lands_under_its_own_key(stored_training_dataset): + """Not ``dataset``: this one is a reference for decoding and typing, and + would otherwise collide with the dataset actually being predicted on.""" + ctx = ExecutionContext() + + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + assert ctx.require("train_dataset").column_names == ["a", "b"] + assert not ctx.has("dataset") + + +def test_the_training_dataset_can_coexist_with_the_one_being_predicted_on( + stored_training_dataset, +): + """The whole reason for the separate key: both datasets are live at once.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[9], b=[9])) + + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + assert ctx.require("dataset")["a"] == [9] + assert ctx.require("train_dataset")["a"] == [1, 2, 3] + + +def test_the_declared_types_travel_as_plain_data(stored_training_dataset): + """``train_dataset_types`` is a ref so the saving step never reopens the + file. ``put_ref`` rejects anything that is not JSON data, which is what + stops a live type object from being smuggled across the boundary — note + ``to_string`` returns a dict despite its name. + """ + import json + + ctx = ExecutionContext() + + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + types = ctx.to_dict()["train_dataset_types"] + assert types == { + "a": {"type": "Integer", "dtype": "int64"}, + "b": {"type": "Integer", "dtype": "int64"}, + } + json.dumps(types) + + +def test_an_unreadable_training_dataset_names_the_folder(tmp_path): + with pytest.raises(JobError, match="Cannot load training dataset from"): + LoadTrainingDatasetUnit(train_dataset_file_path=str(tmp_path / "nowhere"))( + ExecutionContext() + ) + + +# --- PredictUnit -------------------------------------------------------- + + +def _ready_context(stored_training_dataset, dataset=None): + ctx = ExecutionContext() + ctx.put("dataset", dataset if dataset is not None else _dataset(a=[1, 2], b=[3, 4])) + ctx.put("model", RecordingModel()) + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + return ctx + + +def _predict_unit(**overrides): + config = { + "task_name": "RecordingTask", + "input_columns": ["a"], + "output_columns": ["target"], + } + config.update(overrides) + return PredictUnit(**config) + + +def test_predict_publishes_decoded_labels(registry, stored_training_dataset): + ctx = _ready_context(stored_training_dataset) + + _predict_unit()(ctx) + + assert ctx.require("y_pred") == ["label-0", "label-0"] + + +def test_predict_hands_the_model_only_the_input_columns( + registry, stored_training_dataset +): + """Selected against the dataset in the context right now, so whatever + produced it — a load or hand-typed rows — is free to differ in shape.""" + ctx = _ready_context(stored_training_dataset) + + _predict_unit()(ctx) + + assert ctx.require("model").seen_columns == ["a"] + + +def test_predict_decodes_against_the_training_dataset( + registry, stored_training_dataset +): + ctx = _ready_context(stored_training_dataset) + + _predict_unit()(ctx) + + assert RecordingTask.seen_train_columns == ["a", "b"] + + +def test_predict_validates_the_task_before_it_runs(registry, stored_training_dataset): + """``validate`` is what the orchestrator calls early, so a missing task is + reported as a task problem rather than being overtaken by a later failure.""" + with pytest.raises(JobError, match="Task NoSuchTask not found in the registry"): + _predict_unit(task_name="NoSuchTask").validate(ExecutionContext()) + + +def test_predict_without_a_model_is_rejected_before_it_starts( + registry, stored_training_dataset +): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + with pytest.raises(UnitContractError, match="Context key 'model'"): + _predict_unit()(ctx) + + +def test_predict_without_a_training_dataset_is_rejected_before_it_starts(registry): + """A missing key means "the loader did not run", not "there is nothing to + decode against" — so it has to fail loudly instead of predicting anyway.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + ctx.put("model", RecordingModel()) + + with pytest.raises(UnitContractError, match="Context key 'train_dataset'"): + _predict_unit()(ctx) + + +# --- SavePredictionUnit ------------------------------------------------- + + +def _save_unit(**overrides): + config = {"input_columns": ["a"], "output_columns": ["target"]} + config.update(overrides) + return SavePredictionUnit(**config) + + +def test_save_writes_the_inputs_plus_the_predicted_column( + registry, stored_training_dataset, datasets_path +): + ctx = _ready_context(stored_training_dataset) + _predict_unit()(ctx) + + _save_unit()(ctx) + + saved = load_dataset(str(Path(ctx.require("results_path")) / "dataset")) + assert saved.column_names == ["a", "b", "target"] + assert saved["target"] == ["label-0", "label-0"] + + +def test_save_resolves_the_columns_against_the_dataset_it_is_handed( + registry, stored_training_dataset, datasets_path +): + """The column list is read at the top of execute, never published earlier. + + Here the dataset already carries a column named like the output one; it has + to be replaced, not duplicated — which only works if the names are resolved + from the dataset in hand. + """ + ctx = _ready_context( + stored_training_dataset, dataset=_dataset(a=[1, 2], target=[7, 8]) + ) + _predict_unit()(ctx) + + _save_unit()(ctx) + + saved = load_dataset(str(Path(ctx.require("results_path")) / "dataset")) + assert saved.column_names == ["a", "target"] + assert saved["target"] == ["label-0", "label-0"] + + +def test_two_saves_never_collide(registry, stored_training_dataset, datasets_path): + """A prediction has no natural key to overwrite, so each run gets a folder.""" + ctx = _ready_context(stored_training_dataset) + _predict_unit()(ctx) + + _save_unit()(ctx) + first = ctx.require("results_path") + _save_unit()(ctx) + second = ctx.require("results_path") + + assert first != second + assert Path(first).exists() + assert Path(second).exists() + + +def test_save_without_a_prediction_is_rejected_before_it_starts( + registry, stored_training_dataset, datasets_path +): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + ctx.put_ref("train_dataset_types", {}) + + with pytest.raises(UnitContractError, match="Context key 'y_pred'"): + _save_unit()(ctx) + + +def test_the_published_results_path_is_a_plain_string( + registry, stored_training_dataset, datasets_path +): + """``results_path`` travels as a ref, so it has to be JSON data.""" + ctx = _ready_context(stored_training_dataset) + _predict_unit()(ctx) + + _save_unit()(ctx) + + assert isinstance(ctx.to_dict()["results_path"], str) diff --git a/tests/back/units/test_save_dataset_unit.py b/tests/back/units/test_save_dataset_unit.py new file mode 100644 index 000000000..a1e534c08 --- /dev/null +++ b/tests/back/units/test_save_dataset_unit.py @@ -0,0 +1,64 @@ +"""Contract tests for SaveDatasetUnit.""" + +import pandas as pd +import pyarrow as pa +import pytest + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.types.value_types import Integer +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit + + +def _dataset(**columns): + frame = pd.DataFrame(columns) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +def test_the_dataset_is_written_where_the_path_says(tmp_path): + destination = str(tmp_path / "notebook" / "dataset") + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1, 2], b=[3, 4])) + ctx.put_ref("dataset_path", destination) + + SaveDatasetUnit()(ctx) + + assert load_dataset(destination).column_names == ["a", "b"] + + +def test_saving_without_a_dataset_is_a_contract_error(tmp_path): + ctx = ExecutionContext() + ctx.put_ref("dataset_path", str(tmp_path / "dataset")) + + with pytest.raises(UnitContractError, match="'dataset' is not available"): + SaveDatasetUnit()(ctx) + + +def test_saving_without_a_path_is_a_contract_error(): + """A missing path is a wiring mistake, not a "nowhere to save" decision. + + The unit has no fallback destination on purpose: silently picking one would + write the dataset somewhere nobody asked for. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + + with pytest.raises(UnitContractError, match="'dataset_path' is not available"): + SaveDatasetUnit()(ctx) + + +def test_an_unwritable_destination_becomes_a_job_error(tmp_path): + blocker = tmp_path / "blocker" + blocker.write_text("not a directory", encoding="utf-8") + + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + ctx.put_ref("dataset_path", str(blocker / "dataset")) + + with pytest.raises(JobError, match="Can not save dataset to path"): + SaveDatasetUnit()(ctx) diff --git a/tests/back/units/test_save_model_unit.py b/tests/back/units/test_save_model_unit.py new file mode 100644 index 000000000..ff98c22a0 --- /dev/null +++ b/tests/back/units/test_save_model_unit.py @@ -0,0 +1,135 @@ +"""Tests for SaveModelUnit's contract, independent of a real model or disk.""" + +import os + +import pytest +from kink import di + +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.save_model_unit import SaveModelUnit + + +class _RecordingModel: + """Stands in for a trained model, recording where it was asked to go. + + It writes a real file at that path, because the save is now made atomic: + the model is handed a temporary path and what it leaves there is moved + into place. A double that recorded without writing would leave nothing to + move, which is the same failure a model that silently saved nothing would + produce -- and is exactly what should be reported rather than hidden. + """ + + def __init__(self) -> None: + self.saved_to = None + + def save(self, path) -> None: + self.saved_to = path + with open(path, "w", encoding="utf-8") as file: + file.write("a saved model") + + +@pytest.fixture(name="runs_path") +def fixture_runs_path(tmp_path): + di["config"] = {"RUNS_PATH": str(tmp_path)} + yield str(tmp_path) + del di["config"] + + +def test_the_model_lands_in_a_directory_named_by_the_prefix(runs_path): + ctx = ExecutionContext() + model = _RecordingModel() + ctx.put("model", model) + + SaveModelUnit(artifact_prefix="pipeline-3-save")(ctx) + + # Where the artifact ended up, not where the model was told to write. The + # save is atomic, so the model writes to a temporary sibling and what it + # left there is moved into place; the destination is what the caller and + # the row care about. + expected = os.path.join(runs_path, "pipeline-3-save") + assert os.path.exists(expected) + assert ctx.require("model_path") == expected + + +def test_a_numeric_run_id_is_a_valid_prefix(runs_path): + """The path a real run writes to has to keep working unchanged. + + ``ModelJob`` passes ``str(run.id)``, and + ``test_model_job_orchestration.py`` asserts the directory is named exactly + that. A guard that rejected it would break every training run. + """ + ctx = ExecutionContext() + model = _RecordingModel() + ctx.put("model", model) + + SaveModelUnit(artifact_prefix="17")(ctx) + + assert os.path.exists(os.path.join(runs_path, "17")) + assert ctx.require("model_path") == os.path.join(runs_path, "17") + + +@pytest.mark.parametrize( + "prefix", + ["../escape", "a/b", "a\b", "", "C:nope", ".", "with space"], +) +def test_a_prefix_that_could_be_read_as_a_path_is_refused(runs_path, prefix): + """A prefix is a directory name, so it must not be able to leave RUNS_PATH. + + ``os.path.join`` with a separator or a parent reference in the prefix + happily produces a destination outside the runs directory, and the model + would be written there without any error. + """ + ctx = ExecutionContext() + model = _RecordingModel() + ctx.put("model", model) + + with pytest.raises(JobError, match="artifact prefix"): + SaveModelUnit(artifact_prefix=prefix)(ctx) + + assert model.saved_to is None + + +def test_validate_runs_before_anything_is_written(runs_path): + """``__call__`` validates first, so a bad prefix never reaches ``save``.""" + unit = SaveModelUnit(artifact_prefix="../escape") + ctx = ExecutionContext() + ctx.put("model", _RecordingModel()) + + with pytest.raises(JobError): + unit.validate(ctx) + + +def test_the_unit_still_needs_a_model(runs_path): + """``run_id`` left REQUIRES; ``model`` did not.""" + with pytest.raises(UnitContractError, match="'model'"): + SaveModelUnit(artifact_prefix="1")(ExecutionContext()) + + +def test_a_save_that_dies_halfway_leaves_the_previous_artifact_alone(runs_path): + """The reason the save is atomic. + + A model that raises partway through writing used to leave a truncated + artifact at the destination, which the row then pointed at as if it were a + model. Now the destination is only replaced once the new one is complete. + """ + destination = os.path.join(runs_path, "42") + with open(destination, "w", encoding="utf-8") as file: + file.write("the previous model") + + class _DyingModel: + def save(self, path): + with open(path, "w", encoding="utf-8") as file: + file.write("half a model") + raise RuntimeError("out of disk") + + ctx = ExecutionContext() + ctx.put("model", _DyingModel()) + + with pytest.raises(JobError, match="Model saving failed"): + SaveModelUnit(artifact_prefix="42")(ctx) + + with open(destination, encoding="utf-8") as file: + assert file.read() == "the previous model" + # And nothing half-written was left lying next to it. + assert os.listdir(runs_path) == ["42"] diff --git a/tests/back/units/test_splitting_units.py b/tests/back/units/test_splitting_units.py new file mode 100644 index 000000000..e403db6fd --- /dev/null +++ b/tests/back/units/test_splitting_units.py @@ -0,0 +1,265 @@ +"""Contract tests for the two units that prepare a dataset and partition it. + +Built on a hand-made ``ExecutionContext`` rather than through a job: the point +is what each unit reads, promises and refuses on its own, which a job that +always wires the context correctly cannot show. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.splitters.holdout import HoldoutSplitter +from DashAI.back.splitters.k_fold import KFoldSplitter +from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.types.value_types import Float +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit + +ROWS = 20 + + +class SplitTask(BaseTask): + name: str = "SplitTask" + metadata: dict = { + "inputs_types": [], + "outputs_types": [], + "inputs_cardinality": "n", + "outputs_cardinality": 1, + } + + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + def num_labels(self, dataset, output_column): + return 2 + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "SplitTask": {"class": SplitTask}, + "HoldoutSplitter": {"class": HoldoutSplitter}, + "KFoldSplitter": {"class": KFoldSplitter}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _dataset(): + frame = pd.DataFrame( + { + "a": [float(i) for i in range(ROWS)], + "b": [float(i % 3) for i in range(ROWS)], + "y": [float(i % 2) for i in range(ROWS)], + } + ) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +def _context(): + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put_ref("dataset_id", 7) + return ctx + + +def _split_unit(**splitter_params): + params = { + "train": 0.5, + "test": 0.25, + "validation": 0.25, + "shuffle": True, + "stratify": False, + "random_state": 42, + } + params.update(splitter_params) + return PrepareAndSplitUnit( + task_name="SplitTask", + input_columns=["a", "b"], + output_columns=["y"], + splitter={"component": "HoldoutSplitter", "params": params}, + ) + + +def _fold_unit(**splitter_params): + params = { + "n_splits": 4, + "test_size": 0.25, + "shuffle": True, + "random_state": 42, + } + params.update(splitter_params) + return PrepareAndFoldUnit( + task_name="SplitTask", + input_columns=["a", "b"], + output_columns=["y"], + splitter={"component": "KFoldSplitter", "params": params}, + ) + + +# --------------------------------------------------------------------------- # +# What each unit promises +# --------------------------------------------------------------------------- # + + +def test_the_holdout_unit_publishes_one_dataset_dict_per_side(registry): + ctx = _context() + + _split_unit()(ctx) + + assert set(ctx.require("x")) == {"train", "test", "validation"} + assert set(ctx.require("y")) == {"train", "test", "validation"} + assert set(ctx.require("split_indexes")) == { + "train_indexes", + "test_indexes", + "val_indexes", + } + + +def test_the_fold_unit_publishes_a_list_and_not_the_holdout_keys(registry): + """The reason the two are separate units and not one with a flag. + + Reusing ``x`` for a list would give one key two possible types, which a + contract that compares key names cannot express. + """ + ctx = _context() + + _fold_unit()(ctx) + + x_folds = ctx.require("x_folds") + assert isinstance(x_folds, list) + # Four folds plus the trailing entry, which is not a fold. + assert len(x_folds) == 5 + assert not ctx.has("x") + assert not ctx.has("y") + + +def test_every_fold_holds_a_train_and_a_validation_partition(registry): + ctx = _context() + + _fold_unit()(ctx) + + x_folds = ctx.require("x_folds") + for fold in x_folds[:-1]: + assert set(fold) == {"train", "validation"} + # The trailing entry fits the kept model and is scored on the reserved rows. + assert set(x_folds[-1]) == {"train", "test"} + + +def test_the_fold_unit_keeps_the_reserved_rows_out_of_every_fold(registry): + """Checked on the indexes the unit itself published.""" + ctx = _context() + + _fold_unit()(ctx) + + split_indexes = ctx.require("split_indexes") + reserved = set(split_indexes["full_dataset"]["test_indexes"]) + assert reserved + + for name, partitions in split_indexes.items(): + if name == "full_dataset": + continue + seen = set(partitions["train_indexes"]) | set(partitions["validation_indexes"]) + assert reserved.isdisjoint(seen), name + + +def test_a_session_that_reserves_nothing_still_produces_folds(registry): + ctx = _context() + + _fold_unit(test_size=0)(ctx) + + assert ctx.require("split_indexes")["full_dataset"]["test_indexes"] == [] + assert len(ctx.require("x_folds")) == 5 + + +# --------------------------------------------------------------------------- # +# What the two share +# --------------------------------------------------------------------------- # + + +def test_the_two_units_prepare_the_dataset_the_same_way(registry): + """Both run the same body, so the parts before the split cannot diverge. + + The one thing that differs is what the splitter returns, and everything + upstream of it -- the task, the label count, the columns selected -- has to + match or the two would be scoring different datasets. + """ + holdout_ctx, fold_ctx = _context(), _context() + + _split_unit()(holdout_ctx) + _fold_unit()(fold_ctx) + + assert holdout_ctx.require("n_labels") == fold_ctx.require("n_labels") + assert holdout_ctx.require("task_name") == fold_ctx.require("task_name") + assert type(holdout_ctx.require("task")) is type(fold_ctx.require("task")) + + holdout_columns = holdout_ctx.require("x")["train"].column_names + fold_columns = fold_ctx.require("x_folds")[0]["train"].column_names + assert holdout_columns == fold_columns == ["a", "b"] + + +@pytest.mark.parametrize("build", [_split_unit, _fold_unit]) +def test_a_unit_refuses_to_run_before_a_dataset_was_loaded(registry, build): + """``__call__`` checks REQUIRES, so the failure names the missing key.""" + with pytest.raises(UnitContractError): + build()(ExecutionContext()) + + +@pytest.mark.parametrize("build", [_split_unit, _fold_unit]) +def test_an_unknown_task_is_reported_by_name(registry, build): + unit = build() + unit.config["task_name"] = "ThereIsNoSuchTask" + + with pytest.raises(JobError, match="Unable to find Task with name"): + unit(_context()) + + +@pytest.mark.parametrize("build", [_split_unit, _fold_unit]) +def test_an_unknown_splitter_is_reported_by_name(registry, build): + unit = build() + unit.config["splitter"] = {"component": "ThereIsNoSuchSplitter", "params": {}} + + with pytest.raises(JobError, match="Unable to find Splitter with name"): + unit(_context()) + + +def test_the_splitters_own_refusal_is_passed_through_undecorated(registry): + """More folds than rows is the splitter's diagnosis, not the unit's. + + It already names the numbers that explain the refusal, and a caller that + wants to say which run this was frames it from outside. Decorating it here + would push the numbers into the middle of someone else's sentence. + """ + with pytest.raises(JobError) as error: + _fold_unit(n_splits=ROWS + 1, test_size=0)(_context()) + + message = " ".join(str(error.value).split()) + assert message == ( + f"Number of splits (n_splits={ROWS + 1}) cannot be greater " + f"than the number of samples ({ROWS})." + ) + + +def test_two_units_in_one_context_do_not_share_their_resolved_task(registry): + """Instance state stays on the instance, which is what lets a graph hold two. + + A context-global cache would give the second unit the first one's task, + and the failure would be silent: the wrong task prepares the dataset + without complaining. + """ + first, second = _split_unit(), _fold_unit() + ctx = _context() + + first(ctx) + second(ctx) + + assert first._task is not None + assert second._task is not None + assert first._task is not second._task diff --git a/tests/back/units/test_unit_contracts.py b/tests/back/units/test_unit_contracts.py new file mode 100644 index 000000000..439ce19a3 --- /dev/null +++ b/tests/back/units/test_unit_contracts.py @@ -0,0 +1,484 @@ +"""A contract audit over every registered unit, enforced as a test. + +The individual unit tests check behaviour; this one checks that the *declared* +contract matches the code. Undeclared context reads are the recurring mistake in +this design: they never break the job that happens to wire the context by hand, +so they survive every end-to-end test and only surface when something reuses the +unit — which is the whole point of having units. +""" + +import ast +import pathlib +import re + +import pytest + +UNITS_DIR = pathlib.Path(__file__).resolve().parents[3] / "DashAI" / "back" / "units" + +#: Keys a unit reads that its own execution produces, so they need no declaration. +SELF_PRODUCED = {"dataset"} + + +def _unit_modules(): + for path in sorted(UNITS_DIR.glob("*.py")): + if path.name in {"__init__.py", "base_unit.py", "context.py"}: + continue + yield path + + +def _string_literals(node): + return { + element.value + for element in ast.walk(node) + if isinstance(element, ast.Constant) and isinstance(element.value, str) + } + + +def _all_classes(): + """Every class defined under ``units/``, by name. + + Built first so a unit that inherits from another unit can be resolved: the + audit reads source rather than importing, so a base class is just a name + until something maps it back to a definition. + """ + classes = {} + for path in _unit_modules(): + tree = ast.parse(path.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + classes[node.name] = node + return classes + + +_CLASSES = _all_classes() + + +def _is_unit(node): + """A class deriving from ``BaseUnit``, directly or through another unit.""" + for base in node.bases: + if not isinstance(base, ast.Name): + continue + if base.id == "BaseUnit": + return True + parent = _CLASSES.get(base.id) + if parent is not None and parent is not node and _is_unit(parent): + return True + return False + + +def _unit_class(tree): + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and _is_unit(node): + return node + return None + + +def _lineage(cls): + """The class and the units it inherits from, nearest first. + + Everything below reads the whole lineage rather than one class body. A unit + that extends another one inherits both its declarations and the code that + honours them, so auditing only its own body would report that it promises + keys it never writes -- which is the same blindness a shared helper causes, + arriving by a different road. + """ + chain = [cls] + for base in cls.bases: + if not isinstance(base, ast.Name): + continue + parent = _CLASSES.get(base.id) + if parent is not None and parent is not cls and _is_unit(parent): + chain.extend(_lineage(parent)) + return chain + + +def _declared(cls, name): + # Nearest declaration wins: a subclass that redeclares PROVIDES replaces + # what it inherited rather than adding to it, the way Python resolves it. + for ancestor in _lineage(cls): + for node in ancestor.body: + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == name for t in node.targets + ): + return _string_literals(node.value) + return set() + + +def _context_calls(cls, methods): + """Every ``ctx.("key")`` literal in the class and what it extends.""" + keys = set() + for ancestor in _lineage(cls): + keys |= _context_calls_in(ancestor, methods) + return keys + + +def _context_calls_in(cls, methods): + keys = set() + for node in ast.walk(cls): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in methods + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "ctx" + and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + keys.add(node.args[0].value) + return keys + + +def _session_names(cls): + """Names bound by ``with session_factory() as :`` inside the class.""" + names = set() + for node in ast.walk(cls): + if not isinstance(node, (ast.With, ast.AsyncWith)): + continue + for item in node.items: + call = item.context_expr + if ( + isinstance(call, ast.Call) + and isinstance(call.func, ast.Name) + and "session" in call.func.id.lower() + and isinstance(item.optional_vars, ast.Name) + ): + names.add(item.optional_vars.id) + return names + + +def _parsed_units(): + units = [] + for path in _unit_modules(): + tree = ast.parse(path.read_text(encoding="utf-8")) + cls = _unit_class(tree) + if cls is not None: + units.append((path.name, cls)) + return units + + +UNITS = _parsed_units() + + +def test_the_audit_actually_found_the_units(): + """Guards the audit itself: a broken parser would make it vacuously pass.""" + assert {name for name, _ in UNITS} >= { + "load_dataset_unit.py", + "apply_converter_unit.py", + "save_dataset_unit.py", + } + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_every_context_key_a_unit_reads_is_declared_in_requires(name, cls): + read = _context_calls(cls, {"require", "get", "has"}) + declared = _declared(cls, "REQUIRES") | _declared(cls, "PROVIDES") | SELF_PRODUCED + + undeclared = read - declared + assert not undeclared, ( + f"{name} reads {sorted(undeclared)} from the context without declaring " + "them in REQUIRES. A caller inspecting the contract cannot know they " + "are needed, and a missing value reads as 'not applicable' instead of " + "'wiring mistake'." + ) + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_a_unit_does_not_require_a_key_it_never_reads(name, cls): + """The mirror of the test above, and just as load-bearing. + + ``__call__`` demands every key in ``REQUIRES`` unconditionally, so a key + listed but never read is not harmless documentation: it rejects any upstream + that does not happen to publish it. ``PrepareExplanationDataUnit`` used to + require ``dataset_id`` — left over from an error message that moved to the + job — which would have made it impossible to compose after + ``BuildManualInputUnit``, whose ``PROVIDES`` is just ``("dataset",)``. + + It passed every end-to-end test because the one job wiring it happened to + run a loader that publishes the id first. That is exactly the class of + mistake this file exists to catch. + """ + required = _declared(cls, "REQUIRES") + read = _context_calls(cls, {"require", "get", "has"}) + + unread = required - read + assert not unread, ( + f"{name} declares {sorted(unread)} in REQUIRES but never reads them. " + "Every declared key is demanded before the unit runs, so an unused one " + "only narrows what the unit can be composed after." + ) + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_every_key_a_unit_promises_is_actually_written(name, cls): + written = _context_calls(cls, {"put", "put_ref"}) + promised = _declared(cls, "PROVIDES") + + unwritten = promised - written + assert not unwritten, ( + f"{name} promises {sorted(unwritten)} in PROVIDES but never writes it." + ) + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_a_unit_does_not_use_the_context_as_its_own_scratchpad(name, cls): + """A key written and read by the same unit, and promised to nobody. + + That is instance state wearing a context key's clothes: two units of the + same class in one context would overwrite each other. It belongs on + ``self``, memoized, the way the registry lookups are. + """ + written = _context_calls(cls, {"put", "put_ref"}) + read = _context_calls(cls, {"require", "get", "has"}) + promised = _declared(cls, "PROVIDES") + + scratch = (written & read) - promised - SELF_PRODUCED + assert not scratch, ( + f"{name} writes and reads {sorted(scratch)} without promising it. " + "Keep per-instance state on the unit, not in the shared context." + ) + + +def test_no_unit_requires_a_key_that_no_unit_provides(): + """Every declared input has to have a declared producer. + + This is what makes a static graph validator decidable. As long as some key + has no producer anywhere in the palette, "nothing supplies this key" is + ambiguous: it could be an edge the user forgot to draw, or a constant the + engine is expected to inject. A validator cannot tell those apart, and the + consequences are not symmetric -- a wrongly injected id is read as a + legitimate value and fails silently. + + ``run_id`` was the only such key: four units required it and none published + it, because in a job it arrived through ``self.kwargs``. It lives in unit + configuration now, where the schema validates it. The rule that keeps it + that way: a key in ``REQUIRES`` is something an upstream unit produces; a + constant the caller chooses is configuration. + """ + required = {} + provided = set() + for name, cls in UNITS: + provided |= _declared(cls, "PROVIDES") + for key in _declared(cls, "REQUIRES"): + required.setdefault(key, []).append(name) + + orphans = {key: names for key, names in required.items() if key not in provided} + assert not orphans, ( + "These keys are required but no unit provides them: " + f"{ {key: sorted(names) for key, names in sorted(orphans.items())} }. " + "A key with no producer makes 'nothing supplies this' ambiguous for the " + "graph validator. Put caller-chosen constants in the unit's config, not " + "in the context." + ) + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_a_unit_does_not_write_domain_rows(name, cls): + """Units read the database; jobs and endpoints write it. + + A unit that persists an application entity cannot be reused by a caller + that has no such entity to persist against -- and it is the caller, not the + unit, that owns transaction boundaries and status transitions. Several units + do open a session, always read-only (``db.get`` / ``db.query``); the ones + whose verb is "save" write to disk and publish the path. + + The one place a unit reaches a write is indirect and named: + ``EvaluateModelUnit`` calls ``BaseModel.calculate_metrics``, which persists + through a session of its own. That is why that unit needs a real ``Run`` + row and refuses to run without one. + """ + # Whatever the unit bound its session to, rather than a fixed list of + # names: a unit that writes through ``as s:`` would otherwise slip past. + sessions = {"db", "session", "sess"} | _session_names(cls) + + writes = set() + for node in ast.walk(cls): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in {"add", "add_all", "delete", "commit", "flush"} + and isinstance(node.func.value, ast.Name) + and node.func.value.id in sessions + ): + writes.add(f"{node.func.value.id}.{node.func.attr}") + + assert not writes, ( + f"{name} writes to the database ({sorted(writes)}). Transaction " + "boundaries and status transitions belong to the caller, and a unit " + "that persists an application entity cannot be reused by a caller that " + "has none." + ) + + +# --- The configuration side of the contract --------------------------------- +# +# Until now this file audited only the context. Configuration went unchecked, +# and it could not be checked here: telling a user-facing field from an +# engine-supplied one meant reading a flag off the emitted JSON schema, which +# needs the imported class. Splitting the two declarations apart put both within +# reach of the AST. + +#: Names that usually belong to plumbing: an id, a path, a prefix. Not a +#: definition, a net -- a new field shaped like this has to be classified rather +#: than land in a form by default. +_PLUMBING_SHAPED = re.compile(r"(_id|_path|_prefix)$|^path$") + +#: Fields a user does choose, but that a plain text input cannot render: they +#: need to select an entity, or to type a row. They stay in the schema, because +#: needing a purpose-built widget is a rendering problem and hiding them would +#: answer the wrong question. The value records what each one needs, so the day +#: a node form is built it can tell "do not show this" from "show this +#: differently" instead of rediscovering the distinction. +NEEDS_A_WIDGET = { + ("save_dataset_to_path_unit.py", "path"): "directory picker", + ("load_dataset_unit.py", "dataset_id"): "dataset selector", + ("load_dataset_unit.py", "notebook_id"): "notebook selector", + ("load_datafile_dataset_unit.py", "datafile_id"): "datafile selector", + ("load_run_model_unit.py", "run_id"): "run selector", + ("load_trained_model_unit.py", "run_id"): "run selector", + ("run_exploration_unit.py", "explorer_id"): "exploration selector", + ("save_exploration_unit.py", "explorer_id"): "exploration selector", + ("generate_global_explanation_unit.py", "explainer_id"): "explainer selector", + ("generate_local_explanation_unit.py", "explainer_id"): "explainer selector", + ("generate_local_explanation_unit.py", "instance_dataset_id"): "dataset selector", + ("build_manual_input_unit.py", "manual_input_data"): "typed row editor", + ("generate_local_explanation_unit.py", "manual_input_data"): "typed row editor", +} + + +def _schema_fields(tree): + """Names annotated in the module's ``*Schema`` class, from the AST.""" + names = set() + for node in ast.walk(tree): + if not (isinstance(node, ast.ClassDef) and node.name.endswith("Schema")): + continue + for statement in node.body: + if isinstance(statement, ast.AnnAssign) and isinstance( + statement.target, ast.Name + ): + names.add(statement.target.id) + return names + + +def _config_reads(cls): + reads = set() + for ancestor in _lineage(cls): + reads |= _config_reads_in(ancestor) + return reads + + +def _config_reads_in(cls): + """Every ``self.config["key"]`` and ``self.config.get("key")`` in the class.""" + keys = set() + for node in ast.walk(cls): + if ( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Attribute) + and node.value.attr == "config" + and isinstance(node.slice, ast.Constant) + and isinstance(node.slice.value, str) + ): + keys.add(node.slice.value) + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "get" + and isinstance(node.func.value, ast.Attribute) + and node.func.value.attr == "config" + and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + keys.add(node.args[0].value) + return keys + + +#: (filename, unit class node, module tree) for every unit. +UNITS_WITH_TREES = [] +for _path in _unit_modules(): + _tree = ast.parse(_path.read_text(encoding="utf-8")) + _cls = _unit_class(_tree) + if _cls is not None: + UNITS_WITH_TREES.append((_path.name, _cls, _tree)) + + +_IDS = [name for name, _, _ in UNITS_WITH_TREES] + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_every_config_key_a_unit_reads_is_declared(name, cls, tree): + """In the schema if a user fills it in, in RUNTIME_PARAMS if not. + + An undeclared key is a raw ``KeyError`` waiting for the first caller that + builds the unit from its declarations instead of copying an existing call. + """ + read = _config_reads(cls) + declared = _schema_fields(tree) | _declared(cls, "RUNTIME_PARAMS") + + undeclared = read - declared + assert not undeclared, ( + f"{name} reads {sorted(undeclared)} from its configuration without " + "declaring it. Put it in the schema if a user chooses it, or in " + "RUNTIME_PARAMS if whatever runs the unit supplies it." + ) + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_a_runtime_param_is_not_also_in_the_schema(name, cls, tree): + """The two answers are mutually exclusive. + + A name in both would be exposed to the front and overridden by the caller + at once -- a field the user is invited to fill in and whose value is then + discarded. + """ + overlap = _declared(cls, "RUNTIME_PARAMS") & _schema_fields(tree) + assert not overlap, f"{name} declares {sorted(overlap)} twice." + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_a_runtime_param_is_actually_read(name, cls, tree): + """A declared name nothing reads is a promise to nobody.""" + unread = _declared(cls, "RUNTIME_PARAMS") - _config_reads(cls) + assert not unread, ( + f"{name} declares {sorted(unread)} in RUNTIME_PARAMS but never reads it." + ) + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_a_plumbing_shaped_field_is_classified_one_way_or_the_other(name, cls, tree): + """A new field named like plumbing cannot default into a form. + + The net has a known hole: ``session_splits`` and + ``trust_inherited_metadata`` are runtime params and match nothing, which is + why the exact set in ``tests/back/api/test_units_api.py`` exists as well. + What this catches is the common case -- a new ``*_id`` or ``*_path`` added + without anyone deciding what it is. + """ + runtime = _declared(cls, "RUNTIME_PARAMS") + + unclassified = [ + field + for field in _schema_fields(tree) + if _PLUMBING_SHAPED.search(field) + and field not in runtime + and (name, field) not in NEEDS_A_WIDGET + ] + + assert not unclassified, ( + f"{name} has schema fields named like plumbing that nobody classified: " + f"{sorted(unclassified)}. Move them to RUNTIME_PARAMS if whatever runs " + "the unit supplies them, or add them to NEEDS_A_WIDGET with the widget " + "they need." + ) + + +def test_every_entry_in_the_widget_registry_still_exists(): + """Guards the registry itself: a stale entry would make it decorative.""" + declared = { + (name, field) + for name, _, tree in UNITS_WITH_TREES + for field in _schema_fields(tree) + } + + missing = sorted(set(NEEDS_A_WIDGET) - declared) + assert not missing, missing