def create_test_pipeline_execution_context(logger_defs=None): loggers = check.opt_dict_param(logger_defs, "logger_defs", key_type=str, value_type=LoggerDefinition) mode_def = ModeDefinition(logger_defs=loggers) pipeline_def = PipelineDefinition(name="test_legacy_context", solid_defs=[], mode_defs=[mode_def]) run_config = {"loggers": {key: {} for key in loggers}} pipeline_run = PipelineRun(pipeline_name="test_legacy_context", run_config=run_config) instance = DagsterInstance.ephemeral() execution_plan = create_execution_plan(pipeline=pipeline_def, run_config=run_config) creation_data = create_context_creation_data( InMemoryPipeline(pipeline_def), execution_plan, run_config, pipeline_run, instance) log_manager = create_log_manager(creation_data) scoped_resources_builder = ScopedResourcesBuilder() executor = create_executor(creation_data) return PlanExecutionContext( plan_data=create_plan_data(creation_data, True, executor.retries), execution_data=create_execution_data( context_creation_data=creation_data, scoped_resources_builder=scoped_resources_builder, ), log_manager=log_manager, output_capture=None, )
def test_clean_event_generator_exit(): """Testing for generator cleanup (see https://amir.rachum.com/blog/2017/03/03/generator-cleanup/) """ from dagster.core.execution.context.init import InitResourceContext from dagster.core.definitions.resource_definition import ScopedResourcesBuilder pipeline_def = gen_basic_resource_pipeline() instance = DagsterInstance.ephemeral() execution_plan = create_execution_plan(pipeline_def) pipeline_run = instance.create_run_for_pipeline( pipeline_def=pipeline_def, execution_plan=execution_plan) log_manager = DagsterLogManager.create(loggers=[], pipeline_run=pipeline_run) resolved_run_config = ResolvedRunConfig.build(pipeline_def) execution_plan = create_execution_plan(pipeline_def) resource_name, resource_def = next( iter(pipeline_def.get_default_mode().resource_defs.items())) resource_context = InitResourceContext( resource_def=resource_def, resources=ScopedResourcesBuilder().build(None), resource_config=None, pipeline_run=pipeline_run, instance=instance, ) generator = single_resource_event_generator(resource_context, resource_name, resource_def) next(generator) generator.close() resource_defs = pipeline_def.get_mode_definition(resolved_run_config.mode) generator = resource_initialization_event_generator( resource_defs=resource_defs, resource_configs=resolved_run_config.resources, log_manager=log_manager, execution_plan=execution_plan, pipeline_run=pipeline_run, resource_keys_to_init={"a"}, instance=instance, emit_persistent_events=True, pipeline_def_for_backwards_compat=pipeline_def, ) next(generator) generator.close() generator = PlanExecutionContextManager( # pylint: disable=protected-access pipeline=InMemoryPipeline(pipeline_def), execution_plan=execution_plan, run_config={}, pipeline_run=pipeline_run, instance=instance, retry_mode=RetryMode.DISABLED, scoped_resources_builder_cm=resource_initialization_manager, ).get_generator() next(generator) generator.close()
def __init__( self, run_id: str, log_manager: DagsterLogManager, scoped_resources_builder: ScopedResourcesBuilder, dagster_type: DagsterType, ): self._run_id = run_id self._log = log_manager self._resources = scoped_resources_builder.build(dagster_type.required_resource_keys)
def for_type(self, dagster_type: DagsterType) -> TypeCheckContext: resources = cast(NamedTuple, self.resources) return TypeCheckContext(self.run_id, self.log, ScopedResourcesBuilder(resources._asdict()), dagster_type)
def _core_resource_initialization_event_generator( resource_defs: Dict[str, ResourceDefinition], resource_configs: Dict[str, ResourceConfig], resource_log_manager: DagsterLogManager, resource_managers: Deque[EventGenerationManager], execution_plan: Optional[ExecutionPlan], pipeline_run: Optional[PipelineRun], resource_keys_to_init: Optional[AbstractSet[str]], instance: Optional[DagsterInstance], emit_persistent_events: Optional[bool], pipeline_def_for_backwards_compat: Optional[PipelineDefinition], ): pipeline_name = None contains_generator = False if emit_persistent_events: check.invariant( pipeline_run and execution_plan, "If emit_persistent_events is enabled, then pipeline_run and execution_plan must be provided", ) pipeline_name = cast(PipelineRun, pipeline_run).pipeline_name resource_keys_to_init = check.opt_set_param(resource_keys_to_init, "resource_keys_to_init") resource_instances: Dict[str, "InitializedResource"] = {} resource_init_times = {} try: if emit_persistent_events and resource_keys_to_init: yield DagsterEvent.resource_init_start( cast(str, pipeline_name), cast(ExecutionPlan, execution_plan), resource_log_manager, resource_keys_to_init, ) resource_dependencies = resolve_resource_dependencies(resource_defs) for level in toposort(resource_dependencies): for resource_name in level: resource_def = resource_defs[resource_name] if not resource_name in resource_keys_to_init: continue resource_fn = cast(Callable[[InitResourceContext], Any], resource_def.resource_fn) resources = ScopedResourcesBuilder(resource_instances).build( resource_def.required_resource_keys) resource_context = InitResourceContext( resource_def=resource_def, resource_config=resource_configs[resource_name].config, pipeline_run=pipeline_run, # Add tags with information about the resource log_manager=resource_log_manager.with_tags( resource_name=resource_name, resource_fn_name=str(resource_fn.__name__), ), resources=resources, instance=instance, pipeline_def_for_backwards_compat= pipeline_def_for_backwards_compat, ) manager = single_resource_generation_manager( resource_context, resource_name, resource_def) for event in manager.generate_setup_events(): if event: yield event initialized_resource = check.inst(manager.get_object(), InitializedResource) resource_instances[ resource_name] = initialized_resource.resource resource_init_times[ resource_name] = initialized_resource.duration contains_generator = contains_generator or initialized_resource.is_generator resource_managers.append(manager) if emit_persistent_events and resource_keys_to_init: yield DagsterEvent.resource_init_success( cast(str, pipeline_name), cast(ExecutionPlan, execution_plan), resource_log_manager, resource_instances, resource_init_times, ) yield ScopedResourcesBuilder(resource_instances, contains_generator) except DagsterUserCodeExecutionError as dagster_user_error: # Can only end up in this state if we attempt to initialize a resource, so # resource_keys_to_init cannot be empty if emit_persistent_events: yield DagsterEvent.resource_init_failure( cast(str, pipeline_name), cast(ExecutionPlan, execution_plan), resource_log_manager, resource_keys_to_init, serializable_error_info_from_exc_info( dagster_user_error.original_exc_info), ) raise dagster_user_error