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