class ExecutionContextManager(ABC):
def __init__(
self,
event_generator,
raise_on_error=False,
):
self._manager = EventGenerationManager(generator=event_generator,
object_cls=self.context_type,
require_object=raise_on_error)
@abstractproperty
def context_type(self):
pass
def prepare_context(self): # ode to Preparable
return self._manager.generate_setup_events()
def get_context(self):
return self._manager.get_object()
def shutdown_context(self):
return self._manager.generate_teardown_events()
def get_generator(self):
return self._manager.generator
def __init__(
self,
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm=None,
system_storage_data=None,
intermediate_storage=None,
raise_on_error=False,
):
scoped_resources_builder_cm = check.opt_callable_param(
scoped_resources_builder_cm,
"scoped_resources_builder_cm",
default=resource_initialization_manager,
)
generator = self.event_generator(
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm,
system_storage_data,
intermediate_storage,
raise_on_error,
)
self._manager = EventGenerationManager(generator, self.context_type,
raise_on_error)
def __init__(
self,
event_generator,
raise_on_error=False,
):
self._manager = EventGenerationManager(generator=event_generator,
object_cls=self.context_type,
require_object=raise_on_error)
def pipeline_initialization_manager(
pipeline_def,
environment_dict,
pipeline_run,
instance,
execution_plan,
scoped_resources_builder_cm=None,
system_storage_data=None,
raise_on_error=False,
):
scoped_resources_builder_cm = check.opt_callable_param(
scoped_resources_builder_cm,
'scoped_resources_builder_cm',
default=resource_initialization_manager,
)
generator = pipeline_initialization_event_generator(
pipeline_def,
environment_dict,
pipeline_run,
instance,
execution_plan,
scoped_resources_builder_cm,
system_storage_data,
raise_on_error,
)
return EventGenerationManager(generator, SystemPipelineExecutionContext,
raise_on_error)
def resource_initialization_manager(
execution_plan, environment_config, pipeline_run, log_manager, resource_keys_to_init,
):
generator = resource_initialization_event_generator(
execution_plan, environment_config, pipeline_run, log_manager, resource_keys_to_init,
)
return EventGenerationManager(generator, ScopedResourcesBuilder)
def resource_initialization_manager(
resource_defs: Dict[str, ResourceDefinition],
resource_configs: Dict[str, ResourceConfig],
log_manager: DagsterLogManager,
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],
pipeline_def_for_backwards_compat: Optional[PipelineDefinition],
):
generator = resource_initialization_event_generator(
resource_defs=resource_defs,
resource_configs=resource_configs,
log_manager=log_manager,
execution_plan=execution_plan,
pipeline_run=pipeline_run,
resource_keys_to_init=resource_keys_to_init,
instance=instance,
resource_instances_to_override=resource_instances_to_override,
emit_persistent_events=emit_persistent_events,
pipeline_def_for_backwards_compat=pipeline_def_for_backwards_compat,
)
return EventGenerationManager(generator, ScopedResourcesBuilder)
class ExecutionContextManager(six.with_metaclass(ABCMeta)):
def __init__(
self,
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm=None,
system_storage_data=None,
intermediate_storage=None,
raise_on_error=False,
):
scoped_resources_builder_cm = check.opt_callable_param(
scoped_resources_builder_cm,
"scoped_resources_builder_cm",
default=resource_initialization_manager,
)
generator = self.event_generator(
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm,
system_storage_data,
intermediate_storage,
raise_on_error,
)
self._manager = EventGenerationManager(generator, self.context_type,
raise_on_error)
@abstractproperty
def context_type(self):
pass
@abstractmethod
def construct_context(
self,
context_creation_data,
scoped_resources_builder,
system_storage_data,
intermediate_storage,
log_manager,
raise_on_error,
):
pass
def prepare_context(self): # ode to Preperable
return self._manager.generate_setup_events()
def get_context(self):
return self._manager.get_object()
def shutdown_context(self):
return self._manager.generate_teardown_events()
def event_generator(
self,
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm,
system_storage_data=None,
intermediate_storage=None,
raise_on_error=False,
):
execution_plan = check.inst_param(execution_plan, "execution_plan",
ExecutionPlan)
pipeline_def = execution_plan.pipeline.get_definition()
run_config = check.dict_param(run_config, "run_config", key_type=str)
pipeline_run = check.inst_param(pipeline_run, "pipeline_run",
PipelineRun)
instance = check.inst_param(instance, "instance", DagsterInstance)
scoped_resources_builder_cm = check.callable_param(
scoped_resources_builder_cm, "scoped_resources_builder_cm")
system_storage_data = check.opt_inst_param(system_storage_data,
"system_storage_data",
SystemStorageData)
intermediate_storage = check.opt_inst_param(
intermediate_storage, "intermediate_storage_data",
IntermediateStorage)
raise_on_error = check.bool_param(raise_on_error, "raise_on_error")
execution_context = None
resources_manager = None
try:
context_creation_data = create_context_creation_data(
execution_plan,
run_config,
pipeline_run,
instance,
)
log_manager = create_log_manager(context_creation_data)
resources_manager = scoped_resources_builder_cm(
execution_plan,
context_creation_data.environment_config,
context_creation_data.pipeline_run,
log_manager,
context_creation_data.resource_keys_to_init,
)
for event in resources_manager.generate_setup_events():
yield event
scoped_resources_builder = check.inst(
resources_manager.get_object(), ScopedResourcesBuilder)
system_storage_data = create_system_storage_data(
context_creation_data, system_storage_data,
scoped_resources_builder)
if intermediate_storage or context_creation_data.intermediate_storage_def:
intermediate_storage = create_intermediate_storage(
context_creation_data,
intermediate_storage,
scoped_resources_builder,
)
else:
# remove this as part of https://github.com/dagster-io/dagster/issues/2705
intermediate_storage = system_storage_data.intermediate_storage
execution_context = self.construct_context(
context_creation_data=context_creation_data,
scoped_resources_builder=scoped_resources_builder,
system_storage_data=system_storage_data,
log_manager=log_manager,
intermediate_storage=intermediate_storage,
raise_on_error=raise_on_error,
)
_validate_plan_with_context(execution_context, execution_plan)
yield execution_context
for event in resources_manager.generate_teardown_events():
yield event
except DagsterError as dagster_error:
if execution_context is None:
user_facing_exc_info = (
# pylint does not know original_exc_info exists is is_user_code_error is true
# pylint: disable=no-member
dagster_error.original_exc_info
if dagster_error.is_user_code_error else sys.exc_info())
error_info = serializable_error_info_from_exc_info(
user_facing_exc_info)
yield DagsterEvent.pipeline_init_failure(
pipeline_name=pipeline_def.name,
failure_data=PipelineInitFailureData(error=error_info),
log_manager=_create_context_free_log_manager(
instance, pipeline_run, pipeline_def),
)
if resources_manager:
for event in resources_manager.generate_teardown_events():
yield event
else:
# pipeline teardown failure
raise dagster_error
if raise_on_error:
raise dagster_error
def get_generator(self):
return self._manager.generator
def single_resource_generation_manager(context, resource_name, resource_def):
generator = single_resource_event_generator(context, resource_name,
resource_def)
return EventGenerationManager(generator, InitializedResource)
class ExecutionContextManager(ABC):
def __init__(
self,
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm=None,
intermediate_storage=None,
raise_on_error=False,
resource_instances_to_override=None,
output_capture=None,
):
scoped_resources_builder_cm = check.opt_callable_param(
scoped_resources_builder_cm,
"scoped_resources_builder_cm",
default=resource_initialization_manager,
)
generator = self.event_generator(
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm,
intermediate_storage,
raise_on_error,
resource_instances_to_override,
output_capture,
)
self._manager = EventGenerationManager(generator, self.context_type,
raise_on_error)
@abstractproperty
def context_type(self):
pass
@abstractmethod
def construct_context(
self,
context_creation_data,
scoped_resources_builder,
intermediate_storage,
log_manager,
raise_on_error,
output_capture,
):
pass
def prepare_context(self): # ode to Preparable
return self._manager.generate_setup_events()
def get_context(self):
return self._manager.get_object()
def shutdown_context(self):
return self._manager.generate_teardown_events()
def event_generator(
self,
execution_plan,
run_config,
pipeline_run,
instance,
scoped_resources_builder_cm,
intermediate_storage=None,
raise_on_error=False,
resource_instances_to_override=None,
output_capture=None,
):
execution_plan = check.inst_param(execution_plan, "execution_plan",
ExecutionPlan)
pipeline_def = execution_plan.pipeline.get_definition()
run_config = check.dict_param(run_config, "run_config", key_type=str)
pipeline_run = check.inst_param(pipeline_run, "pipeline_run",
PipelineRun)
instance = check.inst_param(instance, "instance", DagsterInstance)
scoped_resources_builder_cm = check.callable_param(
scoped_resources_builder_cm, "scoped_resources_builder_cm")
intermediate_storage = check.opt_inst_param(
intermediate_storage, "intermediate_storage_data",
IntermediateStorage)
raise_on_error = check.bool_param(raise_on_error, "raise_on_error")
resource_instances_to_override = check.opt_dict_param(
resource_instances_to_override, "resource_instances_to_override")
execution_context = None
resources_manager = None
try:
context_creation_data = create_context_creation_data(
execution_plan,
run_config,
pipeline_run,
instance,
)
log_manager = create_log_manager(context_creation_data)
resource_defs = execution_plan.pipeline_def.get_mode_definition(
context_creation_data.environment_config.mode).resource_defs
resources_manager = scoped_resources_builder_cm(
resource_defs=resource_defs,
resource_configs=context_creation_data.environment_config.
resources,
log_manager=log_manager,
execution_plan=execution_plan,
pipeline_run=context_creation_data.pipeline_run,
resource_keys_to_init=context_creation_data.
resource_keys_to_init,
instance=instance,
resource_instances_to_override=resource_instances_to_override,
emit_persistent_events=True,
)
yield from resources_manager.generate_setup_events()
scoped_resources_builder = check.inst(
resources_manager.get_object(), ScopedResourcesBuilder)
intermediate_storage = create_intermediate_storage(
context_creation_data,
intermediate_storage,
scoped_resources_builder,
)
execution_context = self.construct_context(
context_creation_data=context_creation_data,
scoped_resources_builder=scoped_resources_builder,
log_manager=log_manager,
intermediate_storage=intermediate_storage,
raise_on_error=raise_on_error,
output_capture=output_capture,
)
_validate_plan_with_context(execution_context, execution_plan)
yield execution_context
yield from resources_manager.generate_teardown_events()
except DagsterError as dagster_error:
if execution_context is None:
user_facing_exc_info = (
# pylint does not know original_exc_info exists is is_user_code_error is true
# pylint: disable=no-member
dagster_error.original_exc_info
if dagster_error.is_user_code_error else sys.exc_info())
error_info = serializable_error_info_from_exc_info(
user_facing_exc_info)
yield DagsterEvent.pipeline_init_failure(
pipeline_name=pipeline_def.name,
failure_data=PipelineInitFailureData(error=error_info),
log_manager=_create_context_free_log_manager(
instance, pipeline_run, pipeline_def),
)
if resources_manager:
yield from resources_manager.generate_teardown_events()
else:
# pipeline teardown failure
raise dagster_error
if raise_on_error:
raise dagster_error
def get_generator(self):
return self._manager.generator
def test_event_generation_manager():
def basic_generator():
yield "A"
yield "B"
yield 2
yield "C"
with pytest.raises(CheckError, match="Not a generator"):
EventGenerationManager(None, int)
with pytest.raises(CheckError, match="was supposed to be a type"):
EventGenerationManager(basic_generator(), None)
with pytest.raises(CheckError, match="Called `get_object` before `generate_setup_events`"):
basic_manager = EventGenerationManager(basic_generator(), int)
basic_manager.get_object()
with pytest.raises(CheckError, match="generator never yielded object of type bool"):
basic_manager = EventGenerationManager(basic_generator(), bool)
list(basic_manager.generate_setup_events())
basic_manager.get_object()
basic_manager = EventGenerationManager(basic_generator(), int)
setup_events = list(basic_manager.generate_setup_events())
assert setup_events == ["A", "B"]
result = basic_manager.get_object()
assert result == 2
teardown_events = list(basic_manager.generate_teardown_events())
assert teardown_events == ["C"]
def test_event_generation_manager():
def basic_generator():
yield 'A'
yield 'B'
yield 2
yield 'C'
with pytest.raises(CheckError, match='Not a generator'):
EventGenerationManager(None, int)
with pytest.raises(CheckError, match='was supposed to be a type'):
EventGenerationManager(basic_generator(), None)
with pytest.raises(
CheckError,
match='Called `get_object` before `generate_setup_events`'):
basic_manager = EventGenerationManager(basic_generator(), int)
basic_manager.get_object()
with pytest.raises(CheckError,
match='generator never yielded object of type bool'):
basic_manager = EventGenerationManager(basic_generator(), bool)
list(basic_manager.generate_setup_events())
basic_manager.get_object()
basic_manager = EventGenerationManager(basic_generator(), int)
setup_events = list(basic_manager.generate_setup_events())
assert setup_events == ['A', 'B']
result = basic_manager.get_object()
assert result == 2
teardown_events = list(basic_manager.generate_teardown_events())
assert teardown_events == ['C']
