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 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
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=[])
environment_config = EnvironmentConfig.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,
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 = execution_plan.pipeline_def.get_mode_definition(
environment_config.mode)
generator = resource_initialization_event_generator(
resource_defs=resource_defs,
resource_configs=environment_config.resources,
log_manager=log_manager,
execution_plan=execution_plan,
pipeline_run=pipeline_run,
resource_keys_to_init={"a"},
instance=instance,
resource_instances_to_override=None,
emit_persistent_events=True,
)
next(generator)
generator.close()
generator = PipelineExecutionContextManager( # pylint: disable=protected-access
execution_plan,
{},
pipeline_run,
instance,
resource_initialization_manager,
).get_generator()
next(generator)
generator.close()
def resolve_memoized_execution_plan(execution_plan):
"""
Returns:
ExecutionPlan: Execution plan configured to only run unmemoized steps.
"""
pipeline_def = execution_plan.pipeline.get_definition()
environment_config = execution_plan.environment_config
pipeline_def = execution_plan.pipeline.get_definition()
mode_def = pipeline_def.get_mode_definition(environment_config.mode)
step_keys_to_execute = set()
for step in execution_plan.steps:
for output_name in step.step_output_dict.keys():
step_output_handle = StepOutputHandle(step.key, output_name)
io_manager_key = execution_plan.get_manager_key(step_output_handle)
# TODO: https://github.com/dagster-io/dagster/issues/3302
# The following code block is HIGHLY experimental. It initializes an IO manager
# outside of the resource initialization context, and will ignore any exit hooks defined
# for the IO manager, and will not work if the IO manager requires resource keys
# for initialization.
resource_config = (
environment_config.resources[io_manager_key]["config"]
if "config" in environment_config.resources[io_manager_key]
else {}
)
resource_def = mode_def.resource_defs[io_manager_key]
resource_context = InitResourceContext(
resource_config,
resource_def,
pipeline_run=PipelineRun(
pipeline_name=pipeline_def.name, run_id="", mode=environment_config.mode
),
)
io_manager = resource_def.resource_fn(resource_context)
context = get_output_context(
execution_plan, environment_config, step_output_handle, None
)
if not io_manager.has_output(context):
step_keys_to_execute.add(step_output_handle.step_key)
return execution_plan.build_subset_plan(list(step_keys_to_execute))
def _check_invocation_requirements(
resource_def: "ResourceDefinition",
init_context: Optional["InitResourceContext"]
) -> "InitResourceContext":
from dagster.core.execution.context.init import InitResourceContext, build_init_resource_context
if resource_def.required_resource_keys and init_context is None:
raise DagsterInvalidInvocationError(
"Resource has required resources, but no context was provided. Use the "
"`build_init_resource_context` function to construct a context with the required "
"resources.")
if init_context is not None and resource_def.required_resource_keys:
resources_dict = cast(
"InitResourceContext",
init_context,
).resources._asdict() # type: ignore[attr-defined]
for resource_key in resource_def.required_resource_keys:
if resource_key not in resources_dict:
raise DagsterInvalidInvocationError(
f'Resource requires resource "{resource_key}", but no resource '
"with that key was found on the context.")
# Check config requirements
if not init_context and resource_def.config_schema.as_field().is_required:
raise DagsterInvalidInvocationError(
"Resource has required config schema, but no context was provided. "
"Use the `build_init_resource_context` function to create a context with config."
)
resource_config = _resolve_bound_config(
init_context.resource_config if init_context else None, resource_def)
# Construct a context if None was provided. This will initialize an ephemeral instance, and
# console log manager.
init_context = init_context or build_init_resource_context()
return InitResourceContext(
resource_config=resource_config,
resources=init_context.resources,
resource_def=resource_def,
instance=init_context.instance,
log_manager=init_context.log,
)
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
pipeline = gen_basic_resource_pipeline()
instance = DagsterInstance.ephemeral()
execution_plan = create_execution_plan(pipeline)
pipeline_run = instance.create_run_for_pipeline(
pipeline=pipeline, execution_plan=execution_plan)
log_manager = DagsterLogManager(run_id=pipeline_run.run_id,
logging_tags={},
loggers=[])
environment_config = EnvironmentConfig.build(pipeline)
execution_plan = create_execution_plan(pipeline)
resource_name, resource_def = next(
iter(pipeline.get_default_mode().resource_defs.items()))
resource_context = InitResourceContext(
pipeline_def=pipeline,
resource_def=resource_def,
resource_config=None,
run_id=make_new_run_id(),
)
generator = single_resource_event_generator(resource_context,
resource_name, resource_def)
next(generator)
generator.close()
generator = resource_initialization_event_generator(
execution_plan, environment_config, pipeline_run, log_manager, {'a'})
next(generator)
generator.close()
generator = pipeline_initialization_event_generator(
pipeline,
{},
pipeline_run,
instance,
execution_plan,
resource_initialization_manager,
)
next(generator)
generator.close()
def resolve_memoized_execution_plan(execution_plan):
"""
Returns:
ExecutionPlan: Execution plan configured to only run unmemoized steps.
"""
# pylint: disable=comparison-with-callable
pipeline_def = execution_plan.pipeline.get_definition()
step_output_versions = execution_plan.resolve_step_output_versions()
if all(version is None for version in step_output_versions.values()):
raise DagsterInvariantViolationError(
"While creating a memoized pipeline run, no steps have versions. At least one step "
"must have a version.")
environment_config = execution_plan.environment_config
pipeline_def = execution_plan.pipeline.get_definition()
mode_def = pipeline_def.get_mode_definition(environment_config.mode)
step_keys_to_execute = []
for step_output_handle in step_output_versions.keys():
manager_key = execution_plan.get_manager_key(step_output_handle)
# TODO: https://github.com/dagster-io/dagster/issues/3302
# The following code block is HIGHLY experimental. It initializes an asset store outside of
# the resource initialization context, and will ignore any exit hooks defined for the asset
# store.
resource_config = (
environment_config.resources[manager_key]["config"]
if "config" in environment_config.resources[manager_key] else {})
resource_def = mode_def.resource_defs[manager_key]
resource_context = InitResourceContext(resource_config, pipeline_def,
resource_def, "")
object_manager = resource_def.resource_fn(resource_context)
context = get_output_context(execution_plan, environment_config,
step_output_handle, None)
if not object_manager.has_asset(
AssetStoreContext.from_output_context(context)):
step_keys_to_execute.append(step_output_handle.step_key)
return execution_plan.build_subset_plan(step_keys_to_execute)
