def __new__(
cls,
resource_config,
pipeline_def,
resource_def,
run_id,
log_manager=None,
resource_instance_dict=None,
required_resource_keys=None,
):
check.opt_dict_param(resource_instance_dict, "resource_instance_dict")
required_resource_keys = check.opt_set_param(required_resource_keys,
"required_resource_keys")
scoped_resources_builder = ScopedResourcesBuilder(
resource_instance_dict)
return super(InitResourceContext, cls).__new__(
cls,
resource_config,
check.inst_param(pipeline_def, "pipeline_def", PipelineDefinition),
check.inst_param(resource_def, "resource_def", ResourceDefinition),
check.str_param(run_id, "run_id"),
check.opt_inst_param(log_manager, "log_manager",
DagsterLogManager),
resources=scoped_resources_builder.build(required_resource_keys),
)
def __new__(
cls,
resource_config: Any,
resource_def: ResourceDefinition,
pipeline_run: Optional[PipelineRun] = None,
log_manager: Optional[DagsterLogManager] = None,
resource_instance_dict: Optional[Dict[str, Any]] = None,
required_resource_keys: Optional[AbstractSet[str]] = None,
instance: Optional[DagsterInstance] = None,
pipeline_def_for_backwards_compat: Optional[PipelineDefinition] = None,
):
check.opt_dict_param(resource_instance_dict, "resource_instance_dict")
required_resource_keys = check.opt_set_param(required_resource_keys,
"required_resource_keys")
scoped_resources_builder = ScopedResourcesBuilder(
resource_instance_dict)
return super(InitResourceContext, cls).__new__(
cls,
resource_config,
check.inst_param(resource_def, "resource_def", ResourceDefinition),
check.opt_inst_param(log_manager, "log_manager",
DagsterLogManager),
resources=scoped_resources_builder.build(required_resource_keys),
instance=check.opt_inst_param(instance, "instance",
DagsterInstance),
pipeline_def_for_backwards_compat=check.opt_inst_param(
pipeline_def_for_backwards_compat,
"pipeline_def_for_backwards_compat",
PipelineDefinition,
),
pipeline_run=check.opt_inst_param(pipeline_run, "pipeline_run",
PipelineRun),
)
def construct_pipeline_execution_context(
context_creation_data,
scoped_resources_builder,
system_storage_data,
log_manager,
executor,
raise_on_error,
):
check.inst_param(context_creation_data, 'context_creation_data', ContextCreationData)
scoped_resources_builder = check.inst_param(
scoped_resources_builder if scoped_resources_builder else ScopedResourcesBuilder(),
'scoped_resources_builder',
ScopedResourcesBuilder,
)
check.inst_param(system_storage_data, 'system_storage_data', SystemStorageData)
check.inst_param(log_manager, 'log_manager', DagsterLogManager)
check.inst_param(executor, 'executor', Executor)
return SystemPipelineExecutionContext(
SystemExecutionContextData(
pipeline=context_creation_data.pipeline,
mode_def=context_creation_data.mode_def,
system_storage_def=context_creation_data.system_storage_def,
pipeline_run=context_creation_data.pipeline_run,
scoped_resources_builder=scoped_resources_builder,
environment_config=context_creation_data.environment_config,
instance=context_creation_data.instance,
intermediates_manager=system_storage_data.intermediates_manager,
file_manager=system_storage_data.file_manager,
raise_on_error=raise_on_error,
retries=executor.retries,
),
executor=executor,
log_manager=log_manager,
)
def create(self):
for resource_name, resource_def in sorted(
self.mode_definition.resource_defs.items()):
if not resource_name in self.resource_keys_to_init:
continue
user_fn = create_resource_fn_lambda(
self.pipeline_def,
resource_def,
self.environment_config.resources.get(resource_name,
{}).get('config'),
self.pipeline_run.run_id,
self.log_manager,
)
def _create_msg_fn(rn):
return lambda: 'Error executing resource_fn on ResourceDefinition {name}'.format(
name=rn)
resource_obj = self.stack.enter_context(
user_code_context_manager(user_fn,
DagsterResourceFunctionError,
_create_msg_fn(resource_name)))
self.resource_instances[resource_name] = resource_obj
return ScopedResourcesBuilder(self.resource_instances)
def create_test_pipeline_execution_context(logger_defs=None):
from dagster.core.storage.intermediate_storage import build_in_mem_intermediates_storage
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(execution_plan, run_config, pipeline_run, instance)
log_manager = create_log_manager(creation_data)
scoped_resources_builder = ScopedResourcesBuilder()
executor = create_executor(creation_data)
return SystemPipelineExecutionContext(
construct_execution_context_data(
context_creation_data=creation_data,
scoped_resources_builder=scoped_resources_builder,
intermediate_storage=build_in_mem_intermediates_storage(pipeline_run.run_id),
log_manager=log_manager,
retries=executor.retries,
raise_on_error=True,
),
executor=executor,
log_manager=log_manager,
)
def construct_execution_context_data(
context_creation_data,
scoped_resources_builder,
system_storage_data,
intermediate_storage,
log_manager,
retries,
raise_on_error,
):
check.inst_param(context_creation_data, 'context_creation_data', ContextCreationData)
scoped_resources_builder = check.inst_param(
scoped_resources_builder if scoped_resources_builder else ScopedResourcesBuilder(),
'scoped_resources_builder',
ScopedResourcesBuilder,
)
check.inst_param(system_storage_data, 'system_storage_data', SystemStorageData)
check.inst_param(intermediate_storage, 'intermediate_storage', IntermediateStorage)
check.inst_param(log_manager, 'log_manager', DagsterLogManager)
check.inst_param(retries, 'retries', Retries)
return SystemExecutionContextData(
pipeline=context_creation_data.pipeline,
mode_def=context_creation_data.mode_def,
system_storage_def=context_creation_data.system_storage_def,
intermediate_storage_def=context_creation_data.intermediate_storage_def,
pipeline_run=context_creation_data.pipeline_run,
scoped_resources_builder=scoped_resources_builder,
environment_config=context_creation_data.environment_config,
instance=context_creation_data.instance,
intermediate_storage=intermediate_storage,
file_manager=system_storage_data.file_manager,
raise_on_error=raise_on_error,
retries=retries,
)
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]
)
environment_dict = {'loggers': {key: {} for key in loggers}}
pipeline_run = PipelineRun(
pipeline_name='test_legacy_context', environment_dict=environment_dict
)
instance = DagsterInstance.ephemeral()
execution_plan = create_execution_plan(pipeline=pipeline_def, environment_dict=environment_dict)
creation_data = create_context_creation_data(
pipeline_def, environment_dict, pipeline_run, instance, execution_plan
)
log_manager = create_log_manager(creation_data)
scoped_resources_builder = ScopedResourcesBuilder()
executor_config = create_executor_config(creation_data)
return construct_pipeline_execution_context(
context_creation_data=creation_data,
scoped_resources_builder=scoped_resources_builder,
system_storage_data=SystemStorageData(
intermediates_manager=InMemoryIntermediatesManager(),
file_manager=LocalFileManager.for_instance(instance, pipeline_run.run_id),
),
log_manager=log_manager,
executor_config=executor_config,
raise_on_error=True,
)
def create_test_pipeline_execution_context(
logger_defs=None, scoped_resources_builder=None, tags=None, run_config_loggers=None
):
run_id = str(uuid.uuid4())
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 = check.opt_list_param(
run_config_loggers, 'run_config_loggers', of_type=logging.Logger
)
run_config = RunConfig(run_id, tags=tags, loggers=run_config_loggers)
environment_dict = {'loggers': {key: {} for key in loggers}}
creation_data = create_context_creation_data(pipeline_def, environment_dict, run_config)
log_manager = create_log_manager(creation_data)
scoped_resources_builder = check.opt_inst_param(
scoped_resources_builder,
'scoped_resources_builder',
ScopedResourcesBuilder,
default=ScopedResourcesBuilder(),
)
return construct_pipeline_execution_context(
context_creation_data=creation_data,
scoped_resources_builder=scoped_resources_builder,
system_storage_data=SystemStorageData(
run_storage=InMemoryRunStorage(),
intermediates_manager=InMemoryIntermediatesManager(),
file_manager=LocalFileManager.for_run_id(run_id),
),
log_manager=log_manager,
)
def construct_execution_context_data(
context_creation_data,
scoped_resources_builder,
intermediate_storage,
log_manager,
retries,
raise_on_error,
):
check.inst_param(context_creation_data, "context_creation_data",
ContextCreationData)
scoped_resources_builder = check.inst_param(
scoped_resources_builder
if scoped_resources_builder else ScopedResourcesBuilder(),
"scoped_resources_builder",
ScopedResourcesBuilder,
)
check.inst_param(intermediate_storage, "intermediate_storage",
IntermediateStorage)
check.inst_param(log_manager, "log_manager", DagsterLogManager)
check.inst_param(retries, "retries", Retries)
return SystemExecutionContextData(
pipeline=context_creation_data.pipeline,
mode_def=context_creation_data.mode_def,
intermediate_storage_def=context_creation_data.
intermediate_storage_def,
pipeline_run=context_creation_data.pipeline_run,
scoped_resources_builder=scoped_resources_builder,
environment_config=context_creation_data.environment_config,
instance=context_creation_data.instance,
intermediate_storage=intermediate_storage,
raise_on_error=raise_on_error,
retries=retries,
execution_plan=context_creation_data.execution_plan,
)
def create_intermediate_storage(
context_creation_data: ContextCreationData,
intermediate_storage_data: Optional[IntermediateStorage],
scoped_resources_builder: ScopedResourcesBuilder,
) -> IntermediateStorage:
environment_config, pipeline_def, intermediate_storage_def, pipeline_run = (
context_creation_data.environment_config,
context_creation_data.pipeline_def,
context_creation_data.intermediate_storage_def,
context_creation_data.pipeline_run,
)
intermediate_storage_data = (
intermediate_storage_data
if intermediate_storage_data else construct_intermediate_storage_data(
InitIntermediateStorageContext(
pipeline_def=pipeline_def,
mode_def=context_creation_data.mode_def,
intermediate_storage_def=intermediate_storage_def,
intermediate_storage_config=environment_config.
intermediate_storage.intermediate_storage_config,
pipeline_run=pipeline_run,
instance=context_creation_data.instance,
environment_config=environment_config,
type_storage_plugin_registry=
construct_type_storage_plugin_registry(
pipeline_def, intermediate_storage_def),
resources=scoped_resources_builder.build(
context_creation_data.intermediate_storage_def.
required_resource_keys, ),
)))
return intermediate_storage_data
def construct_pipeline_execution_context(context_creation_data,
scoped_resources_builder,
system_storage_data, log_manager):
check.inst_param(context_creation_data, 'context_creation_data',
ContextCreationData)
scoped_resources_builder = check.inst_param(
scoped_resources_builder
if scoped_resources_builder else ScopedResourcesBuilder(),
'scoped_resources_builder',
ScopedResourcesBuilder,
)
check.inst_param(system_storage_data, 'system_storage_data',
SystemStorageData)
check.inst_param(log_manager, 'log_manager', DagsterLogManager)
return SystemPipelineExecutionContext(
SystemPipelineExecutionContextData(
pipeline_def=context_creation_data.pipeline_def,
mode_def=context_creation_data.mode_def,
system_storage_def=context_creation_data.system_storage_def,
run_config=context_creation_data.run_config,
scoped_resources_builder=scoped_resources_builder,
environment_config=context_creation_data.environment_config,
run_storage=system_storage_data.run_storage,
intermediates_manager=system_storage_data.intermediates_manager,
file_manager=system_storage_data.file_manager,
execution_target_handle=context_creation_data.
execution_target_handle,
),
log_manager=log_manager,
)
def for_type(self, dagster_type: DagsterType) -> TypeCheckContext:
return TypeCheckContext(
self.run_id,
self.log,
ScopedResourcesBuilder(cast(NamedTuple, self.resources)._asdict()),
dagster_type,
)
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 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(run_id=pipeline_run.run_id,
logging_tags={},
loggers=[])
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 resource_initialization_event_generator(pipeline_def, environment_config,
pipeline_run, log_manager,
resource_keys_to_init):
check.inst_param(pipeline_def, 'pipeline_def', PipelineDefinition)
check.inst_param(environment_config, 'environment_config',
EnvironmentConfig)
check.inst_param(pipeline_run, 'pipeline_run', PipelineRun)
check.inst_param(log_manager, 'log_manager', DagsterLogManager)
check.set_param(resource_keys_to_init,
'resource_keys_to_init',
of_type=str)
resource_instances = {}
mode_definition = pipeline_def.get_mode_definition(pipeline_run.mode)
resource_managers = deque()
generator_closed = False
try:
for resource_name, resource_def in sorted(
mode_definition.resource_defs.items()):
if not resource_name in resource_keys_to_init:
continue
resource_context = InitResourceContext(
pipeline_def=pipeline_def,
resource_def=resource_def,
resource_config=environment_config.resources.get(
resource_name, {}).get('config'),
run_id=pipeline_run.run_id,
log_manager=log_manager,
)
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_managers.append(manager)
yield ScopedResourcesBuilder(resource_instances)
except GeneratorExit:
# Shouldn't happen, but avoid runtime-exception in case this generator gets GC-ed
# (see https://amir.rachum.com/blog/2017/03/03/generator-cleanup/).
generator_closed = True
raise
finally:
if not generator_closed:
error = None
while len(resource_managers) > 0:
manager = resource_managers.pop()
try:
for event in manager.generate_teardown_events():
yield event
except DagsterUserCodeExecutionError as dagster_user_error:
error = dagster_user_error
if error:
raise error
def _core_resource_initialization_event_generator(
execution_plan,
environment_config,
pipeline_run,
resource_keys_to_init,
resource_log_manager,
resource_managers,
):
pipeline_def = execution_plan.pipeline_def
resource_instances = {}
mode_definition = pipeline_def.get_mode_definition(pipeline_run.mode)
resource_init_times = {}
try:
if resource_keys_to_init:
yield DagsterEvent.resource_init_start(
execution_plan, resource_log_manager, resource_keys_to_init,
)
for resource_name, resource_def in sorted(mode_definition.resource_defs.items()):
if not resource_name in resource_keys_to_init:
continue
resource_context = InitResourceContext(
pipeline_def=pipeline_def,
resource_def=resource_def,
resource_config=environment_config.resources.get(resource_name, {}).get("config"),
run_id=pipeline_run.run_id,
# Add tags with information about the resource
log_manager=resource_log_manager.with_tags(
resource_name=resource_name,
resource_fn_name=str(resource_def.resource_fn.__name__),
),
)
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
resource_managers.append(manager)
if resource_keys_to_init:
yield DagsterEvent.resource_init_success(
execution_plan, resource_log_manager, resource_instances, resource_init_times
)
yield ScopedResourcesBuilder(resource_instances)
except DagsterUserCodeExecutionError as dagster_user_error:
yield DagsterEvent.resource_init_failure(
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
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 __new__(
cls,
resource_config,
resource_def,
pipeline_run,
log_manager=None,
resource_instance_dict=None,
required_resource_keys=None,
instance_for_backwards_compat=None,
pipeline_def_for_backwards_compat=None,
):
check.opt_dict_param(resource_instance_dict, "resource_instance_dict")
required_resource_keys = check.opt_set_param(required_resource_keys,
"required_resource_keys")
scoped_resources_builder = ScopedResourcesBuilder(
resource_instance_dict)
return super(InitResourceContext, cls).__new__(
cls,
resource_config,
check.inst_param(resource_def, "resource_def", ResourceDefinition),
check.inst_param(pipeline_run, "pipeline_run", PipelineRun),
check.opt_inst_param(log_manager, "log_manager",
DagsterLogManager),
resources=scoped_resources_builder.build(required_resource_keys),
# The following are used internally for adapting intermediate storage defs to resources
instance_for_backwards_compat=check.opt_inst_param(
instance_for_backwards_compat, "instance_for_backwards_compat",
DagsterInstance),
pipeline_def_for_backwards_compat=check.opt_inst_param(
pipeline_def_for_backwards_compat,
"pipeline_def_for_backwards_compat",
PipelineDefinition,
),
)
def create_test_pipeline_execution_context(logger_defs=None,
scoped_resources_builder=None,
tags=None):
run_id = str(uuid.uuid4())
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 = RunConfig(run_id, tags=tags)
environment_dict = {'loggers': {key: {} for key in loggers}}
instance = DagsterInstance.ephemeral()
creation_data = create_context_creation_data(pipeline_def,
environment_dict, run_config,
instance)
log_manager = create_log_manager(creation_data)
scoped_resources_builder = check.opt_inst_param(
scoped_resources_builder,
'scoped_resources_builder',
ScopedResourcesBuilder,
default=ScopedResourcesBuilder(),
)
executor_config = create_executor_config(creation_data)
return construct_pipeline_execution_context(
context_creation_data=creation_data,
scoped_resources_builder=scoped_resources_builder,
system_storage_data=SystemStorageData(
intermediates_manager=InMemoryIntermediatesManager(),
file_manager=LocalFileManager.for_instance(instance, run_id),
),
log_manager=log_manager,
executor_config=executor_config,
raise_on_error=True,
)
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
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[Set[str]],
instance: Optional[DagsterInstance],
resource_instances_to_override: Optional[Dict[str, "InitializedResource"]],
emit_persistent_events: Optional[bool],
):
if emit_persistent_events:
check.invariant(
execution_plan,
"If emit_persistent_events is enabled, then execution_plan must be provided",
)
resource_instances_to_override = check.opt_dict_param(
resource_instances_to_override, "resource_instances_to_override")
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(
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:
if resource_name in resource_instances_to_override:
# use the given resource instances instead of re-initiating it from resource def
resource_def = ResourceDefinition.hardcoded_resource(
resource_instances_to_override[resource_name])
else:
resource_def = resource_defs[resource_name]
if not resource_name in resource_keys_to_init:
continue
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_def.resource_fn.__name__),
),
resource_instance_dict=resource_instances,
required_resource_keys=resource_def.required_resource_keys,
instance=instance,
pipeline_def_for_backwards_compat=execution_plan.
pipeline_def if execution_plan else None,
)
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
resource_managers.append(manager)
if emit_persistent_events and resource_keys_to_init:
yield DagsterEvent.resource_init_success(execution_plan,
resource_log_manager,
resource_instances,
resource_init_times)
yield ScopedResourcesBuilder(resource_instances)
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(
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
def resource_initialization_event_generator(execution_plan, environment_config,
pipeline_run, log_manager,
resource_keys_to_init):
check.inst_param(execution_plan, 'execution_plan', ExecutionPlan)
check.inst_param(environment_config, 'environment_config',
EnvironmentConfig)
check.inst_param(pipeline_run, 'pipeline_run', PipelineRun)
check.inst_param(log_manager, 'log_manager', DagsterLogManager)
check.set_param(resource_keys_to_init,
'resource_keys_to_init',
of_type=str)
if execution_plan.step_key_for_single_step_plans():
step = execution_plan.get_step_by_key(
execution_plan.step_key_for_single_step_plans())
resource_log_manager = log_manager.with_tags(**step.logging_tags)
else:
resource_log_manager = log_manager
resource_instances = {}
pipeline_def = execution_plan.pipeline_def
mode_definition = pipeline_def.get_mode_definition(pipeline_run.mode)
resource_managers = deque()
generator_closed = False
resource_init_times = {}
try:
if resource_keys_to_init:
yield DagsterEvent.resource_init_start(
execution_plan,
resource_log_manager,
resource_keys_to_init,
)
for resource_name, resource_def in sorted(
mode_definition.resource_defs.items()):
if not resource_name in resource_keys_to_init:
continue
resource_context = InitResourceContext(
pipeline_def=pipeline_def,
resource_def=resource_def,
resource_config=environment_config.resources.get(
resource_name, {}).get('config'),
run_id=pipeline_run.run_id,
# Add tags with information about the resource
log_manager=resource_log_manager.with_tags(
resource_name=resource_name,
resource_fn_name=str(resource_def.resource_fn.__name__),
),
)
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
resource_managers.append(manager)
if resource_keys_to_init:
yield DagsterEvent.resource_init_success(execution_plan,
resource_log_manager,
resource_instances,
resource_init_times)
yield ScopedResourcesBuilder(resource_instances)
except GeneratorExit:
# Shouldn't happen, but avoid runtime-exception in case this generator gets GC-ed
# (see https://amir.rachum.com/blog/2017/03/03/generator-cleanup/).
generator_closed = True
raise
except DagsterUserCodeExecutionError as dagster_user_error:
yield DagsterEvent.resource_init_failure(
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
finally:
if not generator_closed:
error = None
while len(resource_managers) > 0:
manager = resource_managers.pop()
try:
for event in manager.generate_teardown_events():
yield event
except DagsterUserCodeExecutionError as dagster_user_error:
error = dagster_user_error
if error:
yield DagsterEvent.resource_teardown_failure(
execution_plan,
resource_log_manager,
resource_keys_to_init,
serializable_error_info_from_exc_info(
error.original_exc_info),
)
def _core_resource_initialization_event_generator(
execution_plan,
environment_config,
pipeline_run,
resource_keys_to_init,
resource_log_manager,
resource_managers,
instance,
resource_instances_to_override,
):
pipeline_def = execution_plan.pipeline_def
resource_instances = {}
mode_definition = pipeline_def.get_mode_definition(pipeline_run.mode)
resource_init_times = {}
try:
if resource_keys_to_init:
yield DagsterEvent.resource_init_start(
execution_plan,
resource_log_manager,
resource_keys_to_init,
)
resource_dependencies = _resolve_resource_dependencies(
mode_definition.resource_defs)
for level in toposort(resource_dependencies):
for resource_name in level:
if resource_name in resource_instances_to_override:
# use the given resource instances instead of re-initiating it from resource def
resource_def = ResourceDefinition.hardcoded_resource(
resource_instances_to_override[resource_name])
else:
resource_def = mode_definition.resource_defs[resource_name]
if not resource_name in resource_keys_to_init:
continue
resource_context = InitResourceContext(
resource_def=resource_def,
resource_config=environment_config.resources.get(
resource_name, {}).get("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_def.resource_fn.__name__),
),
resource_instance_dict=resource_instances,
required_resource_keys=resource_def.required_resource_keys,
instance_for_backwards_compat=instance,
pipeline_def_for_backwards_compat=pipeline_def,
)
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
resource_managers.append(manager)
if resource_keys_to_init:
yield DagsterEvent.resource_init_success(execution_plan,
resource_log_manager,
resource_instances,
resource_init_times)
yield ScopedResourcesBuilder(resource_instances)
except DagsterUserCodeExecutionError as dagster_user_error:
yield DagsterEvent.resource_init_failure(
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
