def load_input_object(self, step_context):
source_handle = self.step_output_handle
if self.input_def.manager_key:
loader = getattr(step_context.resources,
self.input_def.manager_key)
return loader.load_input(self.get_load_context(step_context))
object_manager = step_context.get_output_manager(source_handle)
check.invariant(
isinstance(object_manager, InputManager),
f'Input "{self.input_def.name}" for step "{step_context.step.key}" is depending on '
f'the manager of upstream output "{source_handle.output_name}" from step '
f'"{source_handle.step_key}" to load it, but that manager is not an InputManager. '
f"Please ensure that the resource returned for resource key "
f'"{step_context.execution_plan.get_manager_key(source_handle)}" is an InputManager.',
)
obj = object_manager.load_input(self.get_load_context(step_context))
output_def = step_context.execution_plan.get_step_output(
source_handle).output_def
# TODO yuhan retire ObjectStoreOperation https://github.com/dagster-io/dagster/issues/3043
if isinstance(obj, ObjectStoreOperation):
return obj
else:
from dagster.core.storage.asset_store import AssetStoreHandle
return AssetStoreOperation(
AssetStoreOperationType.GET_ASSET,
source_handle,
AssetStoreHandle(output_def.manager_key, output_def.metadata),
obj=obj,
)
def _set_objects(step_context, step_output, step_output_handle, output):
from dagster.core.storage.asset_store import AssetStoreHandle
output_def = step_output.output_def
output_manager = step_context.get_output_manager(step_output_handle)
output_context = step_context.get_output_context(step_output_handle)
materializations = output_manager.handle_output(output_context,
output.value)
# TODO yuhan retire ObjectStoreOperation https://github.com/dagster-io/dagster/issues/3043
if isinstance(materializations, ObjectStoreOperation):
yield DagsterEvent.object_store_operation(
step_context,
ObjectStoreOperation.serializable(
materializations, value_name=step_output_handle.output_name),
)
else:
for evt in _materializations_to_events(step_context,
step_output_handle,
materializations):
yield evt
# SET_ASSET operation by AssetStore
yield DagsterEvent.asset_store_operation(
step_context,
AssetStoreOperation(
AssetStoreOperationType.SET_ASSET,
step_output_handle,
AssetStoreHandle(output_def.manager_key, output_def.metadata),
),
)
def test_asset_store_context_from_execution_plan():
pipeline_def = define_asset_pipeline(fs_asset_store, {})
execution_plan = create_execution_plan(pipeline_def)
asset_store_context = execution_plan.construct_asset_store_context(
step_output_handle=StepOutputHandle("solid_a.compute", "result"),
asset_store_handle=AssetStoreHandle("asset_store"),
)
assert asset_store_context == AssetStoreContext(
step_key="solid_a.compute",
output_name="result",
asset_metadata={},
pipeline_name=pipeline_def.name,
solid_def=pipeline_def.solid_def_named("solid_a"),
)
def __init__(
self,
dagster_type=None,
name=None,
description=None,
is_required=None,
asset_store_key=None,
asset_metadata=None,
):
from dagster.core.storage.asset_store import AssetStoreHandle
self._name = check_valid_name(check.opt_str_param(name, "name", DEFAULT_OUTPUT))
self._dagster_type = resolve_dagster_type(dagster_type)
self._description = check.opt_str_param(description, "description")
self._is_required = check.opt_bool_param(is_required, "is_required", default=True)
self._asset_store_handle = (
AssetStoreHandle(asset_store_key, asset_metadata) if asset_store_key else None
)
def _set_objects(step_context, step_output, step_output_handle, output):
from dagster.core.storage.asset_store import AssetStoreHandle
output_def = step_output.output_def
output_manager = step_context.get_output_manager(step_output_handle)
output_context = step_context.get_output_context(step_output_handle)
with user_code_error_boundary(
DagsterExecutionHandleOutputError,
control_flow_exceptions=[Failure, RetryRequested],
msg_fn=lambda:
(f"Error occurred during the the handling of step output:"
f' step key: "{step_context.step.key}"'
f' output name: "{output_context.name}"'),
step_key=step_context.step.key,
output_name=output_context.name,
):
materializations = output_manager.handle_output(
output_context, output.value)
# TODO yuhan retire ObjectStoreOperation https://github.com/dagster-io/dagster/issues/3043
if isinstance(materializations, ObjectStoreOperation):
yield DagsterEvent.object_store_operation(
step_context,
ObjectStoreOperation.serializable(
materializations, value_name=step_output_handle.output_name),
)
else:
for evt in _materializations_to_events(step_context,
step_output_handle,
materializations):
yield evt
# SET_ASSET operation by AssetStore
yield DagsterEvent.asset_store_operation(
step_context,
AssetStoreOperation.serializable(
AssetStoreOperation(
AssetStoreOperationType.SET_ASSET,
step_output_handle,
AssetStoreHandle(output_def.manager_key,
output_def.metadata),
)),
)
def create_compute_step(pipeline_name, environment_config, solid, step_inputs,
handle):
from dagster.core.storage.asset_store import AssetStoreHandle
check.str_param(pipeline_name, "pipeline_name")
check.inst_param(solid, "solid", Solid)
check.list_param(step_inputs, "step_inputs", of_type=StepInput)
check.opt_inst_param(handle, "handle", SolidHandle)
# the environment config has the solid output name configured
config_output_names = set()
current_handle = handle
while current_handle:
solid_config = environment_config.solids.get(
current_handle.to_string())
current_handle = current_handle.parent
for output_spec in solid_config.outputs:
config_output_names = config_output_names.union(output_spec.keys())
return ExecutionStep(
pipeline_name=pipeline_name,
key_suffix="compute",
step_inputs=step_inputs,
step_outputs=[
StepOutput(
name=name,
dagster_type=output_def.dagster_type,
optional=output_def.optional,
should_materialize=name in config_output_names,
asset_store_handle=AssetStoreHandle(output_def.asset_store_key,
output_def.asset_metadata),
) for name, output_def in solid.definition.output_dict.items()
],
compute_fn=lambda step_context, inputs: _execute_core_compute(
step_context.for_compute(), inputs, solid.definition.compute_fn),
kind=StepKind.COMPUTE,
solid_handle=handle,
solid=solid,
)