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 = pipeline_initialization_manager(
plan,
pipeline_run.run_config,
pipeline_run,
DagsterInstance.ephemeral(),
)
for _ in initialization_manager.generate_setup_events():
pass
pipeline_context = initialization_manager.get_object()
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_lost_steps():
plan = create_execution_plan(define_diamond_pipeline())
# run to completion - but step was in unknown state so exception thrown
with pytest.raises(DagsterUnknownStepStateError):
with plan.start(
retries=Retries(RetryMode.DISABLED)) as active_execution:
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
# called by verify_complete when success / fail event not observed
active_execution.mark_unknown_state(step_1.key)
# failure assumed for start step - so rest should skip
steps_to_skip = active_execution.get_steps_to_skip()
while steps_to_skip:
_ = [
active_execution.mark_skipped(step.key)
for step in steps_to_skip
]
steps_to_skip = active_execution.get_steps_to_skip()
assert active_execution.is_complete
def test_retries_deferred_active_execution():
pipeline_def = define_diamond_pipeline()
plan = create_execution_plan(pipeline_def)
with plan.start(retries=Retries(RetryMode.DEFERRED)) as active_execution:
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
assert step_1.key == "return_two.compute"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_up_for_retry(step_1.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress, retries are deferred
assert not active_execution.is_complete
steps = active_execution.get_steps_to_skip()
# skip split of diamond
assert len(steps) == 2
_ = [active_execution.mark_skipped(step.key) for step in steps]
assert not active_execution.is_complete
steps = active_execution.get_steps_to_skip()
# skip end of diamond
assert len(steps) == 1
active_execution.mark_skipped(steps[0].key)
assert active_execution.is_complete
def initialize_step_context(scratch_dir):
pipeline_def = ExecutionTargetHandle.for_pipeline_fn(
define_basic_pipeline).build_pipeline_definition()
pipeline_run = PipelineRun(
pipeline_name='foo_pipeline',
run_id=str(uuid.uuid4()),
environment_dict=make_environment_dict(scratch_dir, 'external'),
mode='external',
)
plan = create_execution_plan(pipeline_def,
pipeline_run.environment_dict,
mode='external')
initialization_manager = pipeline_initialization_manager(
plan,
pipeline_run.environment_dict,
pipeline_run,
DagsterInstance.ephemeral(),
)
for _ in initialization_manager.generate_setup_events():
pass
pipeline_context = initialization_manager.get_object()
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_retries_active_execution():
pipeline_def = define_diamond_pipeline()
plan = create_execution_plan(pipeline_def)
with plan.start(retries=Retries(RetryMode.ENABLED)) as active_execution:
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
assert step_1.key == "return_two.compute"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_up_for_retry(step_1.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
assert steps[0].key == "return_two.compute"
active_execution.mark_up_for_retry(step_1.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
assert steps[0].key == "return_two.compute"
active_execution.mark_success(step_1.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 2
step_2 = steps[0]
step_3 = steps[1]
assert step_2.key == "add_three.compute"
assert step_3.key == "mult_three.compute"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_2.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
# uh oh failure
active_execution.mark_failed(step_3.key)
# cant progres to 4th step
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0
assert not active_execution.is_complete
steps = active_execution.get_steps_to_skip()
assert len(steps) == 1
step_4 = steps[0]
assert step_4.key == "adder.compute"
active_execution.mark_skipped(step_4.key)
assert active_execution.is_complete
def for_cli(broker=None, backend=None, include=None, config_source=None):
return CeleryConfig(
retries=Retries(RetryMode.DISABLED),
broker=broker,
backend=backend,
include=include,
config_source=config_source,
)
def test_incomplete_execution_plan():
plan = create_execution_plan(define_diamond_pipeline())
with pytest.raises(DagsterIncompleteExecutionPlanError):
with plan.start(retries=Retries(RetryMode.DISABLED)) as active_execution:
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
active_execution.mark_success(step_1.key)
def test_incomplete_execution_plan():
plan = create_execution_plan(define_diamond_pipeline())
with pytest.raises(
DagsterInvariantViolationError,
match="Execution of pipeline finished without completing the execution plan.",
):
with plan.start(retries=Retries(RetryMode.DISABLED)) as active_execution:
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
active_execution.mark_success(step_1.key)
def test_failing_execution_plan():
pipeline_def = define_diamond_pipeline()
plan = create_execution_plan(pipeline_def)
with plan.start(retries=Retries(RetryMode.DISABLED)) as active_execution:
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
assert step_1.key == "return_two"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_1.key)
active_execution.mark_step_produced_output(StepOutputHandle(step_1.key, "result"))
steps = active_execution.get_steps_to_execute()
assert len(steps) == 2
step_2 = steps[0]
step_3 = steps[1]
assert step_2.key == "add_three"
assert step_3.key == "mult_three"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_2.key)
active_execution.mark_step_produced_output(StepOutputHandle(step_2.key, "result"))
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
# uh oh failure
active_execution.mark_failed(step_3.key)
active_execution.mark_step_produced_output(StepOutputHandle(step_3.key, "result"))
# cant progres to 4th step
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0
assert not active_execution.is_complete
steps = active_execution.get_steps_to_abandon()
assert len(steps) == 1
step_4 = steps[0]
assert step_4.key == "adder"
active_execution.mark_abandoned(step_4.key)
assert active_execution.is_complete
def test_retry_deferral():
events = execute_plan(
create_execution_plan(define_retry_limit_pipeline()),
pipeline_run=PipelineRun(pipeline_name='retry_limits', run_id='42'),
retries=Retries(RetryMode.DEFERRED),
instance=DagsterInstance.local_temp(),
)
events_by_type = defaultdict(list)
for ev in events:
events_by_type[ev.event_type].append(ev)
assert len(events_by_type[DagsterEventType.STEP_START]) == 2
assert len(events_by_type[DagsterEventType.STEP_UP_FOR_RETRY]) == 2
assert DagsterEventType.STEP_RESTARTED not in events
assert DagsterEventType.STEP_SUCCESS not in events
def test_failing_execution_plan():
pipeline_def = define_diamond_pipeline()
plan = create_execution_plan(pipeline_def)
active_execution = plan.start(retries=Retries(RetryMode.DISABLED))
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
assert step_1.key == 'return_two.compute'
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_1.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 2
step_2 = steps[0]
step_3 = steps[1]
assert step_2.key == 'add_three.compute'
assert step_3.key == 'mult_three.compute'
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_2.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
# uh oh failure
active_execution.mark_failed(step_3.key)
# cant progres to 4th step
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0
assert not active_execution.is_complete
steps = active_execution.get_steps_to_skip()
assert len(steps) == 1
step_4 = steps[0]
assert step_4.key == 'adder.compute'
active_execution.mark_skipped(step_4.key)
assert active_execution.is_complete
def test_retries_disabled_active_execution():
pipeline_def = define_diamond_pipeline()
plan = create_execution_plan(pipeline_def)
active_execution = plan.start(retries=Retries(RetryMode.DISABLED))
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
assert step_1.key == "return_two.compute"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
with pytest.raises(check.CheckError):
active_execution.mark_up_for_retry(step_1.key)
def test_priorities():
@solid(tags={"priority": 5})
def pri_5(_):
pass
@solid(tags={"priority": 4})
def pri_4(_):
pass
@solid(tags={"priority": 3})
def pri_3(_):
pass
@solid(tags={"priority": 2})
def pri_2(_):
pass
@solid(tags={"priority": -1})
def pri_neg_1(_):
pass
@solid
def pri_none(_):
pass
@pipeline
def priorities():
pri_neg_1()
pri_3()
pri_2()
pri_none()
pri_5()
pri_4()
sort_key_fn = lambda step: int(step.tags.get("priority", 0)) * -1
plan = create_execution_plan(priorities)
with plan.start(Retries(RetryMode.DISABLED),
sort_key_fn) as active_execution:
steps = active_execution.get_steps_to_execute()
assert steps[0].key == "pri_5.compute"
assert steps[1].key == "pri_4.compute"
assert steps[2].key == "pri_3.compute"
assert steps[3].key == "pri_2.compute"
assert steps[4].key == "pri_none.compute"
assert steps[5].key == "pri_neg_1.compute"
_ = [active_execution.mark_skipped(step.key) for step in steps]
def test_active_execution_plan():
plan = create_execution_plan(define_diamond_pipeline())
with plan.start(retries=Retries(RetryMode.DISABLED)) as active_execution:
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_1 = steps[0]
assert step_1.key == "return_two.compute"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_1.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 2
step_2 = steps[0]
step_3 = steps[1]
assert step_2.key == "add_three.compute"
assert step_3.key == "mult_three.compute"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_2.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
active_execution.mark_success(step_3.key)
steps = active_execution.get_steps_to_execute()
assert len(steps) == 1
step_4 = steps[0]
assert step_4.key == "adder.compute"
steps = active_execution.get_steps_to_execute()
assert len(steps) == 0 # cant progress
assert not active_execution.is_complete
active_execution.mark_success(step_4.key)
assert active_execution.is_complete
def test_priorities():
@solid(tags={'priority': 5})
def pri_5(_):
pass
@solid(tags={'priority': 4})
def pri_4(_):
pass
@solid(tags={'priority': 3})
def pri_3(_):
pass
@solid(tags={'priority': 2})
def pri_2(_):
pass
@solid(tags={'priority': -1})
def pri_neg_1(_):
pass
@solid
def pri_none(_):
pass
@pipeline
def priorities():
pri_neg_1()
pri_3()
pri_2()
pri_none()
pri_5()
pri_4()
sort_key_fn = lambda step: int(step.tags.get('priority', 0)) * -1
plan = create_execution_plan(priorities)
active_execution = plan.start(Retries(RetryMode.DISABLED), sort_key_fn)
steps = active_execution.get_steps_to_execute()
assert steps[0].key == 'pri_5.compute'
assert steps[1].key == 'pri_4.compute'
assert steps[2].key == 'pri_3.compute'
assert steps[3].key == 'pri_2.compute'
assert steps[4].key == 'pri_none.compute'
assert steps[5].key == 'pri_neg_1.compute'
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()
active_execution = plan.start(retries=Retries(RetryMode.DISABLED))
step = active_execution.get_next_step()
step_context = pipeline_context.for_step(step)
return step_context