def step_run_ref_to_step_context(
step_run_ref: StepRunRef, instance: DagsterInstance
) -> StepExecutionContext:
check.inst_param(instance, "instance", DagsterInstance)
pipeline = step_run_ref.recon_pipeline.subset_for_execution_from_existing_pipeline(
step_run_ref.pipeline_run.solids_to_execute
)
execution_plan = create_execution_plan(
pipeline,
step_run_ref.run_config,
mode=step_run_ref.pipeline_run.mode,
step_keys_to_execute=[step_run_ref.step_key],
)
initialization_manager = PlanExecutionContextManager(
retry_mode=step_run_ref.retry_mode.for_inner_plan(),
pipeline=pipeline,
execution_plan=execution_plan,
run_config=step_run_ref.run_config,
pipeline_run=step_run_ref.pipeline_run,
instance=instance,
)
for _ in initialization_manager.prepare_context():
pass
execution_context = initialization_manager.get_context()
execution_step = cast("ExecutionStep", execution_plan.get_step_by_key(step_run_ref.step_key))
step_execution_context = execution_context.for_step(execution_step)
# Since for_step is abstract for IPlanContext, its return type is IStepContext.
# Since we are launching from a PlanExecutionContext, the type will always be
# StepExecutionContext.
step_execution_context = cast(StepExecutionContext, step_execution_context)
return step_execution_context
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",
)
recon_pipeline = reconstructable(define_basic_pipeline)
plan = create_execution_plan(recon_pipeline,
pipeline_run.run_config,
mode="external")
initialization_manager = PlanExecutionContextManager(
pipeline=recon_pipeline,
execution_plan=plan,
run_config=pipeline_run.run_config,
pipeline_run=pipeline_run,
instance=instance,
retry_mode=RetryMode.DISABLED,
)
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 step_run_ref_to_step_context(step_run_ref):
pipeline = step_run_ref.recon_pipeline.subset_for_execution_from_existing_pipeline(
step_run_ref.pipeline_run.solids_to_execute
)
execution_plan = create_execution_plan(
pipeline, step_run_ref.run_config, mode=step_run_ref.pipeline_run.mode
).build_subset_plan([step_run_ref.step_key])
retries = step_run_ref.retries.for_inner_plan()
initialization_manager = PlanExecutionContextManager(
retries,
execution_plan,
step_run_ref.run_config,
step_run_ref.pipeline_run,
DagsterInstance.ephemeral(),
)
for _ in initialization_manager.prepare_context():
pass
execution_context = initialization_manager.get_context()
active_execution = execution_plan.start(retries=retries)
step = active_execution.get_next_step()
return execution_context.for_step(step)
def step_run_ref_to_step_context(
step_run_ref: StepRunRef, instance: DagsterInstance
) -> SystemStepExecutionContext:
check.inst_param(instance, "instance", DagsterInstance)
pipeline = step_run_ref.recon_pipeline.subset_for_execution_from_existing_pipeline(
step_run_ref.pipeline_run.solids_to_execute
)
execution_plan = create_execution_plan(
pipeline,
step_run_ref.run_config,
mode=step_run_ref.pipeline_run.mode,
step_keys_to_execute=[step_run_ref.step_key],
)
initialization_manager = PlanExecutionContextManager(
retry_mode=step_run_ref.retry_mode.for_inner_plan(),
pipeline=pipeline,
execution_plan=execution_plan,
run_config=step_run_ref.run_config,
pipeline_run=step_run_ref.pipeline_run,
instance=instance,
)
for _ in initialization_manager.prepare_context():
pass
execution_context = initialization_manager.get_context()
return execution_context.for_step(execution_plan.get_step_by_key(step_run_ref.step_key))
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 step_run_ref_to_step_context(step_run_ref, instance):
check.inst_param(instance, "instance", DagsterInstance)
pipeline = step_run_ref.recon_pipeline.subset_for_execution_from_existing_pipeline(
step_run_ref.pipeline_run.solids_to_execute)
execution_plan = create_execution_plan(
pipeline, step_run_ref.run_config,
mode=step_run_ref.pipeline_run.mode).build_subset_plan(
[step_run_ref.step_key])
retries = step_run_ref.retries.for_inner_plan()
initialization_manager = PlanExecutionContextManager(
retries,
execution_plan,
step_run_ref.run_config,
step_run_ref.pipeline_run,
instance,
)
for _ in initialization_manager.prepare_context():
pass
execution_context = initialization_manager.get_context()
return execution_context.for_step(
execution_plan.get_step_by_key(step_run_ref.step_key))
def execute(self, pipeline_context, execution_plan):
check.inst_param(pipeline_context, "pipeline_context",
PlanOrchestrationContext)
check.inst_param(execution_plan, "execution_plan", ExecutionPlan)
step_keys_to_execute = execution_plan.step_keys_to_execute
yield DagsterEvent.engine_event(
pipeline_context,
"Executing steps in process (pid: {pid})".format(pid=os.getpid()),
event_specific_data=EngineEventData.in_process(
os.getpid(), step_keys_to_execute),
)
with time_execution_scope() as timer_result:
yield from iter(
ExecuteRunWithPlanIterable(
execution_plan=pipeline_context.execution_plan,
iterator=inner_plan_execution_iterator,
execution_context_manager=PlanExecutionContextManager(
pipeline=pipeline_context.pipeline,
retry_mode=pipeline_context.retry_mode,
execution_plan=pipeline_context.execution_plan,
run_config=pipeline_context.run_config,
pipeline_run=pipeline_context.pipeline_run,
instance=pipeline_context.instance,
raise_on_error=pipeline_context.raise_on_error,
output_capture=pipeline_context.output_capture,
),
))
yield DagsterEvent.engine_event(
pipeline_context,
"Finished steps in process (pid: {pid}) in {duration_ms}".format(
pid=os.getpid(),
duration_ms=format_duration(timer_result.millis)),
event_specific_data=EngineEventData.in_process(
os.getpid(), step_keys_to_execute),
)
