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, )