def initialize_step_context(scratch_dir, instance):
pipeline_run = PipelineRun(
pipeline_name="foo_pipeline",
run_id=str(uuid.uuid4()),
run_config=make_run_config(scratch_dir, "external"),
mode="external",
)
plan = create_execution_plan(reconstructable(define_basic_pipeline),
pipeline_run.run_config,
mode="external")
initialization_manager = PipelineExecutionContextManager(
plan,
pipeline_run.run_config,
pipeline_run,
instance,
)
for _ in initialization_manager.prepare_context():
pass
pipeline_context = initialization_manager.get_context()
step_context = pipeline_context.for_step(
plan.get_step_by_key("return_two"))
return step_context
def initialize_step_context(scratch_dir):
pipeline_run = PipelineRun(
pipeline_name='foo_pipeline',
run_id=str(uuid.uuid4()),
run_config=make_run_config(scratch_dir, 'external'),
mode='external',
)
plan = create_execution_plan(reconstructable(define_basic_pipeline),
pipeline_run.run_config,
mode='external')
initialization_manager = PipelineExecutionContextManager(
plan,
pipeline_run.run_config,
pipeline_run,
DagsterInstance.ephemeral(),
)
for _ in initialization_manager.prepare_context():
pass
pipeline_context = initialization_manager.get_context()
active_execution = plan.start(retries=Retries(RetryMode.DISABLED))
step = active_execution.get_next_step()
step_context = pipeline_context.for_step(step)
return step_context
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 = 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,
pipeline_def_for_backwards_compat=pipeline_def,
)
next(generator)
generator.close()
generator = PipelineExecutionContextManager( # pylint: disable=protected-access
InMemoryPipeline(pipeline_def),
execution_plan,
{},
pipeline_run,
instance,
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(
pipeline_def=pipeline_def,
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 = PipelineExecutionContextManager( # pylint: disable=protected-access
execution_plan,
{},
pipeline_run,
instance,
resource_initialization_manager,
).get_generator()
next(generator)
generator.close()