Allow HookContext to access Output Metadata

[JYoussouf]

What's the use case?

Upon hook execution, one can access output values via the HookContext (HookContext.op_output_values, but there is currently no way to access output metadata in the same way. Enabling this would allow users to pass information more dynamically to their hooks, flexible with respect to each individual ops Output metadata.

One application of this is in dynamically applying tags to runs based on the metadata of the last failed/successful run.

Ideas of implementation

If a hook is present, we currently initialize a step_output_capture dict here:

dagster/python_modules/dagster/dagster/_core/execution/context/system.py

Lines 565 to 569 in 29f6aa7

	# Enable step output capture if there are any hooks which will receive them.
	# Expect in the future that hooks may control whether or not they get outputs,
	# but for now presence of any will cause output capture.
	if self.job_def.get_all_hooks_for_handle(self.node_handle):
	self._step_output_capture = {}

then, we capture the step output value here:

dagster/python_modules/dagster/dagster/_core/execution/plan/execute_step.py

Lines 573 to 574 in 29f6aa7

	if step_context.step_output_capture is not None:
	step_context.step_output_capture[step_output_handle] = output.value

and finally, the hook accesses this the step_output_capture dict here:

dagster/python_modules/dagster/dagster/_core/execution/context/hook.py

Lines 160 to 186 in 29f6aa7

	@property
	def solid_output_values(self) -> Mapping[str, Union[Any, Mapping[str, Any]]]:
	"""The computed output values.
	Returns a dictionary where keys are output names and the values are:
	* the output values in the normal case
	* a dictionary from mapping key to corresponding value in the mapped case
	"""
	results: Dict[str, Union[Any, Dict[str, Any]]] = {}
	captured = self._step_execution_context.step_output_capture
	if captured is None:
	check.failed("Outputs were unexpectedly not captured for hook")
	# make the returned values more user-friendly
	for step_output_handle, value in captured.items():
	if step_output_handle.mapping_key:
	if results.get(step_output_handle.output_name) is None:
	results[step_output_handle.output_name] = {
	step_output_handle.mapping_key: value
	}
	else:
	results[step_output_handle.output_name][step_output_handle.mapping_key] = value
	else:
	results[step_output_handle.output_name] = value
	return results

Implementation:

Option 1: One proposed solution would be to change what gets stored in step_context.step_output_capture, to include both {"value": output.value, "metadata": output.metadata}. This is a relatively simple patch, but it will break nearly all applications that currently call HookContext.op_output_values, since we would be returning a dict now, instead of a single output value.

    if step_context.step_output_capture is not None:
        step_context.step_output_capture[step_output_handle] = {"value": output.value, "metadata": output.metadata}

Option 2: My preferred solution is to replicate the step_output_capture behaviour with the _output_metadata property currently available here:

dagster/python_modules/dagster/dagster/_core/execution/context/system.py

Line 571 in 29f6aa7

self._output_metadata: Dict[str, Any] = {}

We could then add the following line:
step_context.output_metadata[step_output_handle] = output.metadata to execute_step.py
and a solid_output_metadata/op_output_metadata method to hooks.py

Additional information

N/A

Message from the maintainers

Impacted by this issue? Give it a 👍! We factor engagement into prioritization.

[JYoussouf]

I've created a PR for this here: #20027