def __call__(self, *args, **kwargs) -> Any: from .composition import is_in_composition from dagster.core.execution.context.invocation import DirectSolidExecutionContext if is_in_composition(): return super(SolidDefinition, self).__call__(*args, **kwargs) else: experimental_functionality_warning( "Solid invocation out of composition scope") if self._context_arg_provided: if len(args) == 0: raise DagsterInvalidInvocationError( f"Compute function of solid '{self.name}' has context argument, but no context " "was provided when invoking.") elif args[0] is not None and not isinstance( args[0], DirectSolidExecutionContext): raise DagsterInvalidInvocationError( f"Compute function of solid '{self.name}' has context argument, but no context " "was provided when invoking.") context = args[0] return solid_invocation_result(self, context, *args[1:], **kwargs) else: if len(args) > 0 and isinstance(args[0], DirectSolidExecutionContext): raise DagsterInvalidInvocationError( f"Compute function of solid '{self.name}' has no context argument, but " "context was provided when invoking.") return solid_invocation_result(self, None, *args, **kwargs)
def _store_output( step_context: StepExecutionContext, step_output_handle: StepOutputHandle, output: Union[Output, DynamicOutput], input_lineage: List[AssetLineageInfo], ) -> Iterator[DagsterEvent]: output_def = step_context.solid_def.output_def_named( step_output_handle.output_name) output_manager = step_context.get_io_manager(step_output_handle) output_context = step_context.get_output_context(step_output_handle) with solid_execution_error_boundary( DagsterExecutionHandleOutputError, msg_fn=lambda: (f'Error occurred while handling output "{output_context.name}" of ' f'step "{step_context.step.key}":'), step_context=step_context, step_key=step_context.step.key, output_name=output_context.name, ): handle_output_res = output_manager.handle_output( output_context, output.value) manager_materializations = [] manager_metadata_entries = [] if handle_output_res is not None: for elt in ensure_gen(handle_output_res): if isinstance(elt, AssetMaterialization): manager_materializations.append(elt) elif isinstance(elt, (EventMetadataEntry, PartitionMetadataEntry)): experimental_functionality_warning( "Yielding metadata from an IOManager's handle_output() function" ) manager_metadata_entries.append(elt) else: raise DagsterInvariantViolationError( f"IO manager on output {output_def.name} has returned " f"value {elt} of type {type(elt).__name__}. The return type can only be " "one of AssetMaterialization, EventMetadataEntry, PartitionMetadataEntry." ) # do not alter explicitly created AssetMaterializations for materialization in manager_materializations: yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage) asset_key, partitions = _asset_key_and_partitions_for_output( output_context, output_def, output_manager) if asset_key: for materialization in _get_output_asset_materializations( asset_key, partitions, output, output_def, manager_metadata_entries, ): yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage) yield DagsterEvent.handled_output( step_context, output_name=step_output_handle.output_name, manager_key=output_def.io_manager_key, message_override= f'Handled input "{step_output_handle.output_name}" using intermediate storage' if isinstance(output_manager, IntermediateStorageAdapter) else None, metadata_entries=[ entry for entry in manager_metadata_entries if isinstance(entry, EventMetadataEntry) ], )
def _store_output( step_context: StepExecutionContext, step_output_handle: StepOutputHandle, output: Union[Output, DynamicOutput], input_lineage: List[AssetLineageInfo], ) -> Iterator[DagsterEvent]: output_def = step_context.solid_def.output_def_named(step_output_handle.output_name) output_manager = step_context.get_io_manager(step_output_handle) output_context = step_context.get_output_context(step_output_handle) manager_materializations = [] manager_metadata_entries: List[Union[PartitionMetadataEntry, MetadataEntry]] = [] # output_manager.handle_output is either a generator function, or a normal function with or # without a return value. In the case that handle_output is a normal function, we need to # catch errors should they be raised before a return value. We can do this by wrapping # handle_output in a generator so that errors will be caught within iterate_with_context. if not inspect.isgeneratorfunction(output_manager.handle_output): def _gen_fn(): gen_output = output_manager.handle_output(output_context, output.value) for event in output_context.consume_events(): yield event if gen_output: yield gen_output handle_output_gen = _gen_fn() else: handle_output_gen = output_manager.handle_output(output_context, output.value) for elt in iterate_with_context( lambda: solid_execution_error_boundary( DagsterExecutionHandleOutputError, msg_fn=lambda: ( f'Error occurred while handling output "{output_context.name}" of ' f'step "{step_context.step.key}":' ), step_context=step_context, step_key=step_context.step.key, output_name=output_context.name, ), handle_output_gen, ): for event in output_context.consume_events(): yield event manager_metadata_entries.extend(output_context.consume_logged_metadata_entries()) if isinstance(elt, DagsterEvent): yield elt elif isinstance(elt, AssetMaterialization): manager_materializations.append(elt) elif isinstance(elt, (MetadataEntry, PartitionMetadataEntry)): experimental_functionality_warning( "Yielding metadata from an IOManager's handle_output() function" ) manager_metadata_entries.append(elt) else: raise DagsterInvariantViolationError( f"IO manager on output {output_def.name} has returned " f"value {elt} of type {type(elt).__name__}. The return type can only be " "one of AssetMaterialization, MetadataEntry, PartitionMetadataEntry." ) for event in output_context.consume_events(): yield event manager_metadata_entries.extend(output_context.consume_logged_metadata_entries()) # do not alter explicitly created AssetMaterializations for materialization in manager_materializations: if materialization.metadata_entries and manager_metadata_entries: raise DagsterInvariantViolationError( f"When handling output '{output_context.name}' of {output_context.solid_def.node_type_str} '{output_context.solid_def.name}', received a materialization with metadata, while context.add_output_metadata was used within the same call to handle_output. Due to potential conflicts, this is not allowed. Please specify metadata in one place within the `handle_output` function." ) if manager_metadata_entries: materialization = AssetMaterialization( asset_key=materialization.asset_key, description=materialization.description, metadata_entries=manager_metadata_entries, partition=materialization.partition, tags=materialization.tags, metadata=None, ) yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage) asset_key, partitions = _asset_key_and_partitions_for_output( output_context, output_def, output_manager ) if asset_key: for materialization in _get_output_asset_materializations( asset_key, partitions, output, output_def, manager_metadata_entries, ): yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage) yield DagsterEvent.handled_output( step_context, output_name=step_output_handle.output_name, manager_key=output_def.io_manager_key, metadata_entries=[ entry for entry in manager_metadata_entries if isinstance(entry, MetadataEntry) ], )
def _store_output( step_context: StepExecutionContext, step_output_handle: StepOutputHandle, output: Union[Output, DynamicOutput], input_lineage: List[AssetLineageInfo], ) -> Iterator[DagsterEvent]: output_def = step_context.solid_def.output_def_named(step_output_handle.output_name) output_manager = step_context.get_io_manager(step_output_handle) output_context = step_context.get_output_context(step_output_handle) manager_materializations = [] manager_metadata_entries = [] # output_manager.handle_output is either a generator function, or a normal function with or # without a return value. In the case that handle_output is a normal function, we need to # catch errors should they be raised before a return value. We can do this by wrapping # handle_output in a generator so that errors will be caught within iterate_with_context. if not inspect.isgeneratorfunction(output_manager.handle_output): def _gen_fn(): gen_output = output_manager.handle_output(output_context, output.value) if gen_output: yield gen_output handle_output_gen = _gen_fn() else: handle_output_gen = output_manager.handle_output(output_context, output.value) for elt in iterate_with_context( lambda: solid_execution_error_boundary( DagsterExecutionHandleOutputError, msg_fn=lambda: ( f'Error occurred while handling output "{output_context.name}" of ' f'step "{step_context.step.key}":' ), step_context=step_context, step_key=step_context.step.key, output_name=output_context.name, ), handle_output_gen, ): if isinstance(elt, AssetMaterialization): manager_materializations.append(elt) elif isinstance(elt, (EventMetadataEntry, PartitionMetadataEntry)): experimental_functionality_warning( "Yielding metadata from an IOManager's handle_output() function" ) manager_metadata_entries.append(elt) else: raise DagsterInvariantViolationError( f"IO manager on output {output_def.name} has returned " f"value {elt} of type {type(elt).__name__}. The return type can only be " "one of AssetMaterialization, EventMetadataEntry, PartitionMetadataEntry." ) # do not alter explicitly created AssetMaterializations for materialization in manager_materializations: yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage) asset_key, partitions = _asset_key_and_partitions_for_output( output_context, output_def, output_manager ) if asset_key: for materialization in _get_output_asset_materializations( asset_key, partitions, output, output_def, manager_metadata_entries, ): yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage) yield DagsterEvent.handled_output( step_context, output_name=step_output_handle.output_name, manager_key=output_def.io_manager_key, metadata_entries=[ entry for entry in manager_metadata_entries if isinstance(entry, EventMetadataEntry) ], )