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_addressable_asset(context, step_output_handle, asset_store_handle,
value):
check.inst_param(asset_store_handle, "asset_store_handle",
AssetStoreHandle)
asset_store = context.get_asset_store(asset_store_handle.asset_store_key)
materializations = asset_store.set_asset(context, step_output_handle,
value,
asset_store_handle.asset_metadata)
# Allow zero, one, or multiple AssetMaterialization yielded by set_asset
if materializations is not None:
for materialization in ensure_gen(materializations):
if not isinstance(materialization, AssetMaterialization):
raise DagsterInvariantViolationError((
"asset_store on output {output_name} has returned "
"value {value} of type {python_type}. The return type can only be "
"AssetMaterialization.").format(
output_name=step_output_handle.output_name,
value=repr(materialization),
python_type=type(materialization).__name__,
))
yield materialization
# SET_ASSET operation by AssetStore
yield AssetStoreOperation(AssetStoreOperationType.SET_ASSET,
step_output_handle, asset_store_handle)
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 _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 _set_addressable_asset(context, step_output_handle, asset_store_handle,
value):
check.inst_param(asset_store_handle, "asset_store_handle",
AssetStoreHandle)
asset_store = context.get_asset_store(asset_store_handle.asset_store_key)
asset_store.set_asset(context, step_output_handle, value,
asset_store_handle.asset_metadata)
return AssetStoreOperation(AssetStoreOperationType.SET_ASSET,
step_output_handle, asset_store_handle)
def _get_addressable_asset(context, step_output_handle):
asset_store_handle = context.execution_plan.get_asset_store_handle(
step_output_handle)
asset_store = context.get_asset_store(asset_store_handle.asset_store_key)
obj = asset_store.get_asset(context, step_output_handle,
asset_store_handle.asset_metadata)
return AssetStoreOperation(
AssetStoreOperationType.GET_ASSET,
step_output_handle,
asset_store_handle,
obj=obj,
)
def core_dagster_event_sequence_for_step(step_context, prior_attempt_count):
"""
Execute the step within the step_context argument given the in-memory
events. This function yields a sequence of DagsterEvents, but without
catching any exceptions that have bubbled up during the computation
of the step.
"""
check.inst_param(step_context, "step_context", SystemStepExecutionContext)
check.int_param(prior_attempt_count, "prior_attempt_count")
if prior_attempt_count > 0:
yield DagsterEvent.step_restarted_event(step_context,
prior_attempt_count)
else:
yield DagsterEvent.step_start_event(step_context)
inputs = {}
for input_name, input_value in _load_input_values(step_context):
# TODO yuhan retire ObjectStoreOperation https://github.com/dagster-io/dagster/issues/3043
if isinstance(input_value, ObjectStoreOperation):
yield DagsterEvent.object_store_operation(
step_context,
ObjectStoreOperation.serializable(input_value,
value_name=input_name))
inputs[input_name] = input_value.obj
elif isinstance(input_value, FanInStepInputValuesWrapper):
final_values = []
for inner_value in input_value:
# inner value is either a store interaction
# TODO yuhan retire ObjectStoreOperation https://github.com/dagster-io/dagster/issues/3043
if isinstance(inner_value, ObjectStoreOperation):
yield DagsterEvent.object_store_operation(
step_context,
ObjectStoreOperation.serializable(
inner_value, value_name=input_name),
)
final_values.append(inner_value.obj)
elif isinstance(inner_value, AssetStoreOperation):
yield DagsterEvent.asset_store_operation(
step_context,
AssetStoreOperation.serializable(inner_value))
final_values.append(inner_value.obj)
# or the value directly
else:
final_values.append(inner_value)
inputs[input_name] = final_values
elif isinstance(input_value, AssetStoreOperation):
yield DagsterEvent.asset_store_operation(
step_context, AssetStoreOperation.serializable(input_value))
inputs[input_name] = input_value.obj
else:
inputs[input_name] = input_value
for input_name, input_value in inputs.items():
for evt in check.generator(
_type_checked_event_sequence_for_input(step_context,
input_name,
input_value)):
yield evt
with time_execution_scope() as timer_result:
user_event_sequence = check.generator(
_user_event_sequence_for_step_compute_fn(step_context, inputs))
# It is important for this loop to be indented within the
# timer block above in order for time to be recorded accurately.
for user_event in check.generator(
_step_output_error_checked_user_event_sequence(
step_context, user_event_sequence)):
if isinstance(user_event, (Output, DynamicOutput)):
for evt in _create_step_events_for_output(
step_context, user_event):
yield evt
elif isinstance(user_event,
(AssetMaterialization, Materialization)):
yield DagsterEvent.step_materialization(
step_context, user_event)
elif isinstance(user_event, ExpectationResult):
yield DagsterEvent.step_expectation_result(
step_context, user_event)
else:
check.failed(
"Unexpected event {event}, should have been caught earlier"
.format(event=user_event))
yield DagsterEvent.step_success_event(
step_context, StepSuccessData(duration_ms=timer_result.millis))