diff --git a/docs/hamilton-ui/ui.rst b/docs/hamilton-ui/ui.rst index 3e1811eba..6517026f7 100644 --- a/docs/hamilton-ui/ui.rst +++ b/docs/hamilton-ui/ui.rst @@ -248,6 +248,43 @@ View a history of runs, telemetry on runs/comparison, and data for specific runs .. image:: ../_static/run_data.png :alt: Run Data +The result builder node +----------------------- + +The combined result your driver returns is assembled by a +:doc:`result builder ` once the last node has finished. That +happens outside the dataflow, so there was no node to attach a data summary to, and what a run +actually produced did not show up in the UI. + +The tracker synthesizes a node named `_result_builder` to stand in for it. You don't write this +node and it is not part of your dataflow -- the tracker adds it to the DAG it registers, and +logs a run for it on every successful tracked run, with the same data observability as any +other node. Nothing is needed to enable it; just attach a `HamiltonTracker` as usual. + +Which nodes fed the result varies from run to run, so the node declares no dependencies and +records the outputs it combined on each run instead. In the DAG view it therefore appears off to +one side, rather than downstream of the nodes it summarizes. + +Note: the node is part of what identifies a DAG version, so the first tracked run after +upgrading registers a new version of each of your dataflows. Your existing versions are left +as they are, with their runs intact; new runs attach to the new version. + +A few things to be aware of: + +1. The result is profiled a second time. Each output node is already summarized, and the + combined result contains those same objects, so expect roughly twice the profiling time and + payload. The options under `Changing behavior of what is captured`_ apply here too. +2. Only `Driver.execute()` sees a real built result. `raw_execute()` and `materialize()` never + call a result builder, and for async drivers the result is built after the tracker has + logged. The node is still emitted in those cases, summarizing the raw dictionary of outputs. + No lifecycle hook tells the tracker which result builder ran, or whether one ran at all. +3. Failed runs emit nothing, so the node renders as not executed, like any node the run never + reached. A result builder that returns `None` still counts as having run. +4. Underscore-prefixed function names never become nodes, so `_result_builder` is out of reach + of your transforms. If a dataflow defines a node with that name some other way -- an external + input, or a decorator-generated name -- the tracker steps aside entirely for that dataflow: + it logs a warning, registers no node of its own, and records nothing against yours. + ------------------ SDK Configuration diff --git a/ui/backend/server/trackingserver_template/migrations/0003_alter_nodetemplate_classifications.py b/ui/backend/server/trackingserver_template/migrations/0003_alter_nodetemplate_classifications.py new file mode 100644 index 000000000..71ab2601b --- /dev/null +++ b/ui/backend/server/trackingserver_template/migrations/0003_alter_nodetemplate_classifications.py @@ -0,0 +1,47 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Generated by Django 5.2.15 on 2026-08-02 05:16 + +import django.contrib.postgres.fields +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("trackingserver_template", "0002_alter_dagtemplate_unique_together"), + ] + + operations = [ + migrations.AlterField( + model_name="nodetemplate", + name="classifications", + field=django.contrib.postgres.fields.ArrayField( + base_field=models.CharField( + choices=[ + ("transform", "Transform"), + ("data_saver", "DataSaver"), + ("data_loader", "DataLoader"), + ("input", "Input"), + ("placeholder", "Placeholder"), + ("result_builder", "ResultBuilder"), + ] + ), + size=None, + ), + ), + ] diff --git a/ui/backend/server/trackingserver_template/migrations_sqlite/0001_initial.py b/ui/backend/server/trackingserver_template/migrations_sqlite/0001_initial.py index c4320ce47..1cc38dd1a 100644 --- a/ui/backend/server/trackingserver_template/migrations_sqlite/0001_initial.py +++ b/ui/backend/server/trackingserver_template/migrations_sqlite/0001_initial.py @@ -170,6 +170,7 @@ class Migration(migrations.Migration): ("data_loader", "DataLoader"), ("input", "Input"), ("placeholder", "Placeholder"), + ("result_builder", "ResultBuilder"), ] ) ), diff --git a/ui/backend/server/trackingserver_template/models.py b/ui/backend/server/trackingserver_template/models.py index 352265103..ab16d480b 100644 --- a/ui/backend/server/trackingserver_template/models.py +++ b/ui/backend/server/trackingserver_template/models.py @@ -154,6 +154,8 @@ class NodeType(models.TextChoices): data_loader = "data_loader", _("DataLoader") input = "input", _("Input") # input, not actually run placeholder = "placeholder", _("Placeholder") + # Synthesized by the SDK tracker -- represents the driver's built result, not a DAG node + result_builder = "result_builder", _("ResultBuilder") # Nodes have a unique name (up to 511 chars) # In Hamilton's case, this includes a .-separated namespace diff --git a/ui/frontend/src/state/api/friendlyApi.ts b/ui/frontend/src/state/api/friendlyApi.ts index 0632b0197..04df254bf 100644 --- a/ui/frontend/src/state/api/friendlyApi.ts +++ b/ui/frontend/src/state/api/friendlyApi.ts @@ -137,7 +137,8 @@ export type Classification = | "artifact" | "data_loader" | "data_saver" - | "input"; + | "input" + | "result_builder"; const codeVersionTypeMap = { CodeVersionGit1: { version: 1, type: "git" }, diff --git a/ui/sdk/src/hamilton_sdk/adapters.py b/ui/sdk/src/hamilton_sdk/adapters.py index f63248ee7..9f462b289 100644 --- a/ui/sdk/src/hamilton_sdk/adapters.py +++ b/ui/sdk/src/hamilton_sdk/adapters.py @@ -31,7 +31,7 @@ from types import ModuleType from typing import Any, Optional -from collections.abc import Callable +from collections.abc import Callable, Mapping from hamilton import graph as h_graph from hamilton import node @@ -56,6 +56,105 @@ def get_node_name(node_: node.Node, task_id: Optional[str]) -> str: LONG_SCALE = float(0xFFFFFFFFFFFFFFF) +def _result_attribute(node_name: str, name: str, observation: dict) -> dict: + """Shapes one observation into the attribute dict the tracking API expects.""" + return dict( + node_name=node_name, + name=name, + type=observation["observability_type"], + # 0.0.3 -> 3 + schema_version=int(observation["observability_schema_version"].split(".")[-1]), + value=observation["observability_value"], + attribute_role="result_summary", + ) + + +def _result_attributes( + node_name: str, result_summary: dict, schema: Optional[dict], additional: list[dict] +) -> list[dict]: + """Builds the attribute list for a successful task run. + + `result_summary` is first because the order influences UI display order. + """ + others = ([schema] if schema is not None else []) + additional + return [_result_attribute(node_name, "result_summary", result_summary)] + [ + # retrieve name if specified + _result_attribute(node_name, other.get("name", f"Attribute {i + 1}"), other) + for i, other in enumerate(others) + ] + + +def _observability_failure_summary() -> dict: + """The result summary used when profiling the result did not produce one.""" + return { + "observability_type": "observability_failure", + "observability_schema_version": "0.0.3", + "observability_value": { + "type": str(str), + "value": "Failed to process result.", + }, + } + + +def _result_builder_dependencies(results: Any, final_vars: list[str]) -> list[str]: + """The requested outputs that actually reached the result being reported. + + Narrowing is only sound when the result is keyed by node name, as ``DictResult`` and the raw + dict ``materialize`` hands over both are. A result that cannot say what went into it -- a + dataframe, a custom builder's own dict -- keeps the full list rather than be credited with + nothing. + + :param results: What the driver passed to ``post_graph_execute``. + :param final_vars: The outputs requested at ``pre_graph_execute``. + :return: The node names to record as this run's dependencies, always a subset of + ``final_vars``. + """ + if isinstance(results, Mapping) and set(results).issubset(final_vars): + return [var for var in final_vars if var in results] + return final_vars + + +def _result_builder_payload( + results: Any, timestamp: datetime.datetime, final_vars: list[str] +) -> tuple[TaskRun, list[dict], dict]: + """Builds everything the synthetic ``_result_builder`` task run needs to be sent. + + Free of I/O, so the sync and async trackers can share it and each send it their own way. + + Not guarded on ``results`` being non-None: a result builder with a side effect and no return + value still ran. See "The result builder node" in docs/hamilton-ui/ui.rst for which execution + paths hand this a built result and which hand it the raw output dict. + + :param results: The combined result the driver produced. + :param timestamp: Time to stamp the task run with -- the builder runs between the last + node and ``post_graph_execute``, and its duration is not observable. + :param final_vars: The outputs this run asked for, used where the result cannot say. + :return: The task run, its attributes, and the task update to send. + """ + node_name = driver.RESULT_BUILDER_NODE_NAME + # process_result only reads `.name` and `.tags`; there is no real node to pass. + stand_in = node.Node(node_name, Any, callabl=lambda: None) + task_run = TaskRun(node_name=node_name, is_in_sample=True) + task_run.status = Status.SUCCESS + task_run.start_time = timestamp + task_run.end_time = timestamp + task_run.result_type = type(results) + result_summary, schema, additional_attributes = runs.process_result(results, stand_in) + if result_summary is None: + result_summary = _observability_failure_summary() + task_run.result_summary = result_summary + attributes = _result_attributes(node_name, result_summary, schema, additional_attributes) + task_update = dict( + node_template_name=node_name, + node_name=node_name, + realized_dependencies=_result_builder_dependencies(results, final_vars), + status=task_run.status, + start_time=task_run.start_time, + end_time=task_run.end_time, + ) + return task_run, attributes, task_update + + class HamiltonTracker( base.BasePostGraphConstruct, base.BasePreGraphExecute, @@ -120,6 +219,8 @@ def __init__( self.tracking_states = {} self.dw_run_ids = {} self.task_runs = {} + # requested outputs per run -- the result-builder node's per-run dependencies + self.final_vars = {} super().__init__() # set this to a float to sample blocks. 0.1 means 10% of blocks will be sampled. # set this to an int to sample blocks by modulo. @@ -145,14 +246,16 @@ def post_graph_construct( return module_hash = driver._get_modules_hash(modules) vcs_info = driver._derive_version_control_info(module_hash) - dag_hash = driver.hash_dag(graph) + dag_hash = driver.hash_dag(graph, include_result_builder=True) code_hash = driver.hash_dag_modules(graph, modules) dag_template_id = self.client.register_dag_template_if_not_exists( project_id=self.project_id, dag_hash=dag_hash, code_hash=code_hash, name=self.dag_name, - nodes=driver._extract_node_templates_from_function_graph(graph), + nodes=driver._extract_node_templates_from_function_graph( + graph, include_result_builder=True + ), code_artifacts=driver.extract_code_artifacts_from_function_graph( graph, vcs_info, vcs_info.local_repo_base_path ), @@ -189,6 +292,7 @@ def pre_graph_execute( ) self.dw_run_ids[run_id] = dw_run_id self.task_runs[run_id] = {} + self.final_vars[run_id] = final_vars logger.warning( f"\nCapturing execution run. Results can be found at " f"{self.hamilton_ui_url}/dashboard/project/{self.project_id}/runs/{dw_run_id}\n" @@ -363,6 +467,29 @@ def post_node_execute( in_samples=[task_run.is_in_sample for _ in attributes], ) + def _emit_result_builder_task_run( + self, run_id: str, results: Any, timestamp: datetime.datetime + ): + """Emits the task run for the synthetic ``_result_builder`` node. + + Failures are logged and swallowed: this runs before ``log_dag_run_end``, and an + otherwise-successful run should not be left rendering as still-running because + profiling or sending the combined result went wrong. + """ + try: + task_run, attributes, task_update = _result_builder_payload( + results, timestamp, self.final_vars.get(run_id, []) + ) + self.tracking_states[run_id].update_task(task_run.node_name, task_run) + self.client.update_tasks( + self.dw_run_ids[run_id], + attributes=attributes, + task_updates=[task_update for _ in attributes], + in_samples=[True for _ in attributes], + ) + except Exception: + logger.exception("Failed to emit the %s task run.", driver.RESULT_BUILDER_NODE_NAME) + def post_graph_execute( self, run_id: str, @@ -393,6 +520,8 @@ def post_graph_execute( task_run.error = ["Run was likely aborted."] if task_run.end_time is None and task_run.status == Status.SUCCESS: task_run.end_time = finally_block_time + elif driver._should_register_result_builder(graph): + self._emit_result_builder_task_run(run_id, results, finally_block_time) self.client.log_dag_run_end( dag_run_id=dw_run_id, @@ -439,6 +568,8 @@ def __init__( self.tracking_states = {} self.dw_run_ids = {} self.task_runs = {} + # requested outputs per run -- the result-builder node's per-run dependencies + self.final_vars = {} self.initialized = False super().__init__() @@ -479,14 +610,16 @@ async def post_graph_construct( return module_hash = driver._get_modules_hash(modules) vcs_info = driver._derive_version_control_info(module_hash) - dag_hash = driver.hash_dag(graph) + dag_hash = driver.hash_dag(graph, include_result_builder=True) code_hash = driver.hash_dag_modules(graph, modules) dag_template_id = await self.client.register_dag_template_if_not_exists( project_id=self.project_id, dag_hash=dag_hash, code_hash=code_hash, name=self.dag_name, - nodes=driver._extract_node_templates_from_function_graph(graph), + nodes=driver._extract_node_templates_from_function_graph( + graph, include_result_builder=True + ), code_artifacts=driver.extract_code_artifacts_from_function_graph( graph, vcs_info, vcs_info.local_repo_base_path ), @@ -523,6 +656,7 @@ async def pre_graph_execute( ) self.dw_run_ids[run_id] = dw_run_id self.task_runs[run_id] = {} + self.final_vars[run_id] = final_vars async def pre_node_execute( self, run_id: str, node_: node.Node, kwargs: dict[str, Any], task_id: Optional[str] = None @@ -640,6 +774,31 @@ async def post_node_execute( in_samples=[task_run.is_in_sample for _ in attributes], ) + async def _emit_result_builder_task_run( + self, run_id: str, results: Any, timestamp: datetime.datetime + ): + """Emits the task run for the synthetic ``_result_builder`` node. + + ``results`` here is always the raw output dict, never a built result: + ``async_driver.execute()`` awaits ``raw_execute()`` -- whose ``finally`` fires this hook + -- and only then calls ``do_build_result``, so the tracker cannot observe the builder. + + Failures are logged and swallowed, as in the sync tracker. + """ + try: + task_run, attributes, task_update = _result_builder_payload( + results, timestamp, self.final_vars.get(run_id, []) + ) + self.tracking_states[run_id].update_task(task_run.node_name, task_run) + await self.client.update_tasks( + self.dw_run_ids[run_id], + attributes=attributes, + task_updates=[task_update for _ in attributes], + in_samples=[True for _ in attributes], + ) + except Exception: + logger.exception("Failed to emit the %s task run.", driver.RESULT_BUILDER_NODE_NAME) + async def post_graph_execute( self, run_id: str, @@ -668,6 +827,8 @@ async def post_graph_execute( task_run.error = ["Run was likely aborted."] if task_run.end_time is None and task_run.status == Status.SUCCESS: task_run.end_time = finally_block_time + elif driver._should_register_result_builder(graph): + await self._emit_result_builder_task_run(run_id, results, finally_block_time) # TODO: only update things that have changed? # self.client.update_tasks( diff --git a/ui/sdk/src/hamilton_sdk/driver.py b/ui/sdk/src/hamilton_sdk/driver.py index 6f63f58a4..9514c54b2 100644 --- a/ui/sdk/src/hamilton_sdk/driver.py +++ b/ui/sdk/src/hamilton_sdk/driver.py @@ -431,13 +431,20 @@ def _get_fully_qualified_function_path(fn: Callable) -> str: return fn_name -def hash_dag(dag: graph.FunctionGraph) -> str: +def hash_dag(dag: graph.FunctionGraph, include_result_builder: bool = False) -> str: """Hashes a DAG. :param dag: DAG to hash + :param include_result_builder: Whether the synthetic result-builder node will be part of the + registered template. ``register_dag_template_if_not_exists`` matches on this hash alone, + so without folding the node in, a template registered before it existed would be reused + and the run would log task runs against a node that template does not have. Off by + default, so the legacy ``Driver`` below keeps hashing the way it always has. :return: Hash of the DAG """ digest = hashlib.sha256() + if include_result_builder and _should_register_result_builder(dag): + digest.update(RESULT_BUILDER_NODE_NAME.encode()) hashing_node_fields = { "name": str, @@ -523,10 +530,65 @@ def _convert_classifications(node_: Node) -> list[str]: return out -def _extract_node_templates_from_function_graph(fn_graph: graph.FunctionGraph) -> list[dict]: +#: Name of the synthetic node the tracker adds to represent the result builder's output. +#: The leading underscore keeps it clear of user functions -- ``graph_utils.py`` excludes +#: ``_``-prefixed functions from becoming nodes -- but names that do not come from a function +#: bypass that filter, so a collision is still possible. See +#: ``tests/resources/dag_with_reserved_node_name.py``. +RESULT_BUILDER_NODE_NAME = "_result_builder" + + +def _should_register_result_builder(fn_graph: graph.FunctionGraph) -> bool: + """Whether the synthetic result-builder node can be added to this graph's template. + + ``NodeTemplate`` is unique on ``("name", "dag_template")``, so a second template with this + name would fail DAG registration outright. Node templates, the DAG hash and the tracker's + task run all ask here, so they cannot disagree about whether the node exists. + + :param fn_graph: The function graph being registered. + :return: False if the dataflow already has a node by that name, True otherwise. + """ + return RESULT_BUILDER_NODE_NAME not in fn_graph.nodes + + +def _result_builder_node_template() -> dict: + """Builds the synthetic node template for the result builder. + + The output type is ``typing.Any`` because it varies per run and per result builder; the + real type shows up in the per-run result summary. + + :return: A node template dict, shaped like the ones built from real nodes. + """ + return dict( + name=RESULT_BUILDER_NODE_NAME, + output={"type_name": str(Any)}, + output_type="python_type", + output_schema_version=1, + documentation=( + "The combined result of the run, as assembled by the driver's result builder. " + "This node is synthesized by the Hamilton tracker -- it does not exist in the " + "DAG itself. Its dependencies vary per run: they are the requested outputs." + ), + tags={}, + classifications=["result_builder"], + code_artifact_pointers=[], # originates from no user function + # Dependencies are per-run (the requested outputs), so the template has none. + dependencies=[], + dependency_specs=[], + dependency_specs_type="python_type", + dependency_specs_schema_version=1, + ) + + +def _extract_node_templates_from_function_graph( + fn_graph: graph.FunctionGraph, include_result_builder: bool = False +) -> list[dict]: """Converts a function graph to a list of nodes that the DAGWorks graph can understand. @param fn: Function graph to convert + @param include_result_builder: Whether to append the synthetic result-builder node. Off by + default, since a caller that registers it without also emitting a task run would render + it as never-executed on every run -- as the legacy ``Driver`` in this module would. @return: A list of node objects """ node_templates = [] @@ -555,6 +617,17 @@ def _extract_node_templates_from_function_graph(fn_graph: graph.FunctionGraph) - **_convert_node_dependencies(node_), ) ) + if include_result_builder: + if _should_register_result_builder(fn_graph): + node_templates.append(_result_builder_node_template()) + else: + # Skipping costs this run the node; registering both would cost the user their run. + logger.warning( + "Not registering the synthetic %s node: this dataflow already has a node by " + "that name, and registering both would fail. The combined result of the run " + "will not be shown in the Hamilton UI.", + RESULT_BUILDER_NODE_NAME, + ) return node_templates diff --git a/ui/sdk/tests/resources/dag_with_reserved_node_name.py b/ui/sdk/tests/resources/dag_with_reserved_node_name.py new file mode 100644 index 000000000..21043ee0b --- /dev/null +++ b/ui/sdk/tests/resources/dag_with_reserved_node_name.py @@ -0,0 +1,28 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""A dataflow that already contains a node named ``_result_builder``. + +``graph_utils.find_functions`` filters out ``_``-prefixed *functions*, which is why the tracker's +synthetic node is named this way. An external input takes its name from a parameter, and +parameters are never filtered, so the name is reachable after all. Decorator-generated names +(``@extract_columns("_result_builder")``) bypass the filter the same way. +""" + + +def uses_the_reserved_name(_result_builder: int) -> int: + return _result_builder + 1 diff --git a/ui/sdk/tests/test_adapters.py b/ui/sdk/tests/test_adapters.py index 754cec68d..9a033410d 100644 --- a/ui/sdk/tests/test_adapters.py +++ b/ui/sdk/tests/test_adapters.py @@ -15,14 +15,18 @@ # specific language governing permissions and limitations # under the License. +import asyncio import os.path +from types import SimpleNamespace import pytest from hamilton_sdk import adapters -from hamilton import driver +from hamilton import driver, lifecycle +from hamilton.io.materialization import to import tests.resources.basic_dag_with_config +import tests.resources.dag_with_reserved_node_name import tests.resources.parallel_dag import tests.resources.parallel_dag_error from tests import test_tracking @@ -52,12 +56,236 @@ def test_adapters(): assert result == {"a": 1, "b": 3, "c": 6} -# def test_async(): -# # TODO: complete Async -# kwargs = adapter_kwargs | dict( -# dag_name="async_test_dag", -# ) -# [adapters.AsyncHamiltonAdapter(**kwargs)] +class RecordingClient(test_tracking.MockHamiltonClient): + """Mock client that keeps every task update and attribute, not just the latest.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.all_task_updates = [] + self.all_attributes = [] + + def update_tasks(self, dag_run_id, attributes, task_updates, in_samples=None): + self.all_task_updates.extend(task_updates) + self.all_attributes.extend(attributes) + + +def _run_and_record(final_vars, inputs, client, *extra_adapters): + tracker = adapters.HamiltonTracker( + **( + adapter_kwargs + | dict(dag_name="result_builder_dag", client_factory=lambda *a, **kw: client) + ) + ) + dr = ( + driver.Builder() + .with_modules(tests.resources.basic_dag_with_config) + .with_config({"foo": "baz"}) + .with_adapters(tracker, *extra_adapters) + .build() + ) + return dr.execute(final_vars=final_vars, inputs=inputs) + + +def _emitted_result_builder_tasks(client): + return [u for u in client.all_task_updates if u["node_name"] == "_result_builder"] + + +def test_result_builder_task_run_is_emitted(): + """A successful run emits a _result_builder task whose deps are the requested outputs.""" + client = RecordingClient() + _run_and_record(["a", "b", "c"], {"a": 1}, client) + + emitted = _emitted_result_builder_tasks(client) + assert len(emitted) == 1 + assert emitted[0]["realized_dependencies"] == ["a", "b", "c"] + assert emitted[0]["node_template_name"] == "_result_builder" + assert emitted[0]["status"] == adapters.Status.SUCCESS + + +def test_no_result_builder_task_run_when_the_name_is_taken(): + """The template is skipped on a collision, so the task run is skipped with it. + + Emitting anyway would post a task update keyed on the user's own node of that name. + """ + client = RecordingClient() + tracker = adapters.HamiltonTracker( + **( + adapter_kwargs + | dict(dag_name="reserved_name_dag", client_factory=lambda *a, **kw: client) + ) + ) + dr = ( + driver.Builder() + .with_modules(tests.resources.dag_with_reserved_node_name) + .with_adapters(tracker) + .build() + ) + assert dr.execute(final_vars=["uses_the_reserved_name"], inputs={"_result_builder": 1}) == { + "uses_the_reserved_name": 2 + } + + assert not _emitted_result_builder_tasks(client) + + +def test_result_builder_task_run_carries_a_result_summary(): + """The combined result is profiled, which is the point of the node. + + The other tests assert on task updates; this one asserts on the attributes, since a + node that renders with no data observability would satisfy all of them. + """ + client = RecordingClient() + _run_and_record(["a", "b", "c"], {"a": 1}, client) + + summaries = [ + a + # pre_node_execute sends a literal [None] attribute, so entries can be empty + for a in client.all_attributes + if a and a["node_name"] == "_result_builder" and a["name"] == "result_summary" + ] + assert len(summaries) == 1 + assert summaries[0]["attribute_role"] == "result_summary" + # process_result profiled the built dict rather than falling back to a failure summary. + assert summaries[0]["type"] == "dict" + assert summaries[0]["value"]["value"].keys() == {"a", "b", "c"} + + +def test_result_builder_task_run_not_emitted_on_failure(): + """A failed run gets no result-builder task, so the UI shows not-executed.""" + client = RecordingClient() + with pytest.raises(Exception): + _run_and_record(["c"], {"a": 1, "should_fail": True}, client) + + assert not _emitted_result_builder_tasks(client) + + +def test_result_builder_failure_does_not_break_the_run(caplog): + """It runs just before `log_dag_run_end`, so an exception escaping here would leave an + otherwise-successful run rendering as still-running forever. + """ + + class BrokenClient(RecordingClient): + def update_tasks(self, dag_run_id, attributes, task_updates, in_samples=None): + if any(u["node_name"] == "_result_builder" for u in task_updates): + raise RuntimeError("tracking server is down") + super().update_tasks(dag_run_id, attributes, task_updates, in_samples) + + client = BrokenClient() + assert _run_and_record(["a", "b", "c"], {"a": 1}, client) == {"a": 1, "b": 3, "c": 6} + assert client.log_dag_run_end_latest_kwargs["status"] == adapters.Status.SUCCESS.value + assert "Failed to emit the _result_builder task run." in caplog.text + + +class SideEffectResultBuilder(lifecycle.ResultBuilder): + """A result builder that has a side effect and returns nothing -- e.g. saves to disk.""" + + def build_result(self, **outputs): + return None + + +def test_result_builder_task_run_is_emitted_when_the_builder_returns_none(): + """Emitted on every successful run, even when the built result is None.""" + client = RecordingClient() + assert _run_and_record(["a", "b", "c"], {"a": 1}, client, SideEffectResultBuilder()) is None + + emitted = _emitted_result_builder_tasks(client) + assert len(emitted) == 1 + assert emitted[0]["status"] == adapters.Status.SUCCESS + assert emitted[0]["realized_dependencies"] == ["a", "b", "c"] + + +def test_result_builder_dependencies_under_materialize(tmp_path): + """materialize() asks pre_graph_execute for final_vars + materializer_vars, but hands + post_graph_execute only the final_vars slice -- so the recorded deps must not name the + materializer. + """ + client = RecordingClient() + tracker = adapters.HamiltonTracker( + **( + adapter_kwargs + | dict(dag_name="materialize_dag", client_factory=lambda *a, **kw: client) + ) + ) + dr = ( + driver.Builder() + .with_modules(tests.resources.basic_dag_with_config) + .with_config({"foo": "baz"}) + .with_adapters(tracker) + .build() + ) + dr.materialize( + to.pickle(id="save_c", dependencies=["c"], path=str(tmp_path / "c.pkl")), + additional_vars=["a", "b"], + inputs={"a": 1}, + ) + + emitted = _emitted_result_builder_tasks(client) + assert len(emitted) == 1 + assert emitted[0]["realized_dependencies"] == ["a", "b"] + + +def test_result_builder_dependencies_keeps_the_requested_list_for_a_builder_of_its_own_keys(): + """A custom builder's keys are its own, so there is nothing safe to narrow away.""" + assert adapters._result_builder_dependencies({"renamed": 3}, ["a", "b"]) == ["a", "b"] + assert adapters._result_builder_dependencies({"a": 1, "b": 2}, ["a", "b"]) == ["a", "b"] + assert adapters._result_builder_dependencies({"a": 1}, ["a", "b"]) == ["a"] + + +class RecordingAsyncClient(RecordingClient): + """Async twin of RecordingClient -- the async tracker awaits every client call.""" + + async def create_and_start_dag_run(self, **kwargs): + return 1 + + async def update_tasks(self, dag_run_id, attributes, task_updates, in_samples=None): + self.all_task_updates.extend(task_updates) + self.all_attributes.extend(attributes) + + async def log_dag_run_end(self, dag_run_id, status): + pass + + +def _run_async_graph_hooks(final_vars, results, success): + """Drives the async tracker's graph hooks directly. + + There is no async DAG fixture here and standing one up would need pytest-asyncio; the + graph hooks are what carry the result-builder behaviour, so they are what gets exercised. + The stand-in graph needs an identity and an empty `nodes`; the seeded template cache stands + in for the `post_graph_construct` this skips. + """ + client = RecordingAsyncClient() + tracker = adapters.AsyncHamiltonTracker( + **( + adapter_kwargs + | dict(dag_name="async_result_builder_dag", client_factory=lambda *a, **kw: client) + ) + ) + graph = SimpleNamespace(nodes={}) + tracker.dag_template_id_cache[id(graph)] = 1 + + async def run(): + await tracker.pre_graph_execute("run-1", graph, final_vars, {}, {}) + await tracker.post_graph_execute("run-1", graph, success, None, results) + + asyncio.run(run()) + return client + + +def test_async_result_builder_task_run_is_emitted(): + """The async tracker emits the same task run as the sync one.""" + client = _run_async_graph_hooks(["a", "b", "c"], {"a": 1, "b": 3, "c": 6}, success=True) + + emitted = _emitted_result_builder_tasks(client) + assert len(emitted) == 1 + assert emitted[0]["realized_dependencies"] == ["a", "b", "c"] + assert emitted[0]["node_template_name"] == "_result_builder" + assert emitted[0]["status"] == adapters.Status.SUCCESS + + +def test_async_result_builder_task_run_not_emitted_on_failure(): + """A failed run gets no result-builder task -- async side.""" + client = _run_async_graph_hooks(["c"], None, success=False) + + assert not _emitted_result_builder_tasks(client) def test_parallel_ray(): diff --git a/ui/sdk/tests/test_driver.py b/ui/sdk/tests/test_driver.py index b59ff5e07..ee39934a1 100644 --- a/ui/sdk/tests/test_driver.py +++ b/ui/sdk/tests/test_driver.py @@ -20,7 +20,12 @@ from types import ModuleType from unittest.mock import mock_open, patch -from hamilton_sdk.driver import _hash_module +from hamilton_sdk.driver import ( + RESULT_BUILDER_NODE_NAME, + _extract_node_templates_from_function_graph, + _hash_module, + hash_dag, +) @patch("builtins.open", new_callable=mock_open, read_data=b"print('hello world')\n") @@ -112,3 +117,74 @@ def test_hash_module_file_is_none(caplog): assert "Skipping hash" in caplog.text assert result.hexdigest() == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + + +def _basic_function_graph(): + from hamilton import graph + + from tests.resources import basic_dag_with_config + + return graph.FunctionGraph.from_modules(basic_dag_with_config, config={"foo": "bar"}) + + +def _function_graph_with_a_colliding_node(): + """A graph containing a real node named ``_result_builder``, via an external input.""" + from hamilton import graph + + from tests.resources import dag_with_reserved_node_name + + return graph.FunctionGraph.from_modules(dag_with_reserved_node_name, config={}) + + +def test_extract_node_templates_appends_result_builder(): + """The synthetic result builder node is appended when the caller asks for it.""" + fg = _basic_function_graph() + templates = _extract_node_templates_from_function_graph(fg, include_result_builder=True) + + assert len(templates) == len(fg.nodes) + 1 + result_builder = templates[-1] + assert result_builder["name"] == RESULT_BUILDER_NODE_NAME + assert result_builder["classifications"] == ["result_builder"] + # Deps are per-run (the requested outputs), so the template carries none. + assert result_builder["dependencies"] == [] + assert result_builder["code_artifact_pointers"] == [] + assert result_builder["output"] == {"type_name": "typing.Any"} + # It must not shadow a real node. + assert RESULT_BUILDER_NODE_NAME not in fg.nodes + + +def test_extract_node_templates_omits_result_builder_by_default(): + """Only callers that also emit a task run should register the node. + + The legacy ``hamilton_sdk.driver.Driver`` shares this function but emits no task run, so + registering the node there would leave every run rendering a node that never executes. + """ + fg = _basic_function_graph() + templates = _extract_node_templates_from_function_graph(fg) + + assert len(templates) == len(fg.nodes) + assert RESULT_BUILDER_NODE_NAME not in {t["name"] for t in templates} + + +def test_result_builder_node_is_skipped_when_the_name_is_taken(caplog): + """A second template of the same name would fail registration, so the node is dropped.""" + fg = _function_graph_with_a_colliding_node() + assert RESULT_BUILDER_NODE_NAME in fg.nodes, "test graph does not reproduce the collision" + + templates = _extract_node_templates_from_function_graph(fg, include_result_builder=True) + + assert len(templates) == len(fg.nodes) + assert len([t for t in templates if t["name"] == RESULT_BUILDER_NODE_NAME]) == 1 + assert "Not registering the synthetic _result_builder node" in caplog.text + + +def test_result_builder_node_changes_the_dag_hash(): + """The hash identifies the template, so it has to move when the node set does.""" + fg = _basic_function_graph() + assert hash_dag(fg, include_result_builder=True) != hash_dag(fg) + + +def test_dag_hash_is_unchanged_when_the_name_is_taken(): + """No node registered means no new template -- the three uses have to agree.""" + fg = _function_graph_with_a_colliding_node() + assert hash_dag(fg, include_result_builder=True) == hash_dag(fg)