Differential D8501 Diff 40096 python_modules/dagster-test/dagster_test/toys/branches.py

Changeset View

Standalone View

python_modules/dagster-test/dagster_test/toys/branches.py

	from time import sleep			from time import sleep

	from dagster import Field, Int, Output, OutputDefinition, PresetDefinition, pipeline, solid			from dagster import Field, MultiOut, Out, Output, graph, op


	@solid(			@op(
	config_schema={"sleep_secs": Field([int], is_required=False, default_value=[0, 0])},			config_schema={"sleep_secs": Field([int], is_required=False, default_value=[0, 0])},
	output_defs=[OutputDefinition(Int, "out_1"), OutputDefinition(Int, "out_2")],			out=MultiOut(outs=[Out(dagster_type=int, name="out_1"), Out(dagster_type=int, name="out_2")]),
	)			)
	def root(context):			def root(context):
	sleep_secs = context.solid_config["sleep_secs"]			sleep_secs = context.solid_config["sleep_secs"]
	yield Output(sleep_secs[0], "out_1")			yield Output(sleep_secs[0], "out_1")
	yield Output(sleep_secs[1], "out_2")			yield Output(sleep_secs[1], "out_2")
				alangenfeldUnsubmitted Not Done Inline Actions nit: mypy switch to tuple output typing alangenfeld: nit: mypy switch to tuple output typing


	@solid			@op
	def branch_solid(context, sec):			def branch_solid(context, sec):
	if sec < 0:			if sec < 0:
	sleep(-sec)			sleep(-sec)
	raise Exception("fail")			raise Exception("fail")
	context.log.info("Sleeping for {} seconds".format(sec))			context.log.info("Sleeping for {} seconds".format(sec))
	sleep(sec)			sleep(sec)
	return sec			return sec


	def branch(name, arg, solid_num):			def branch(name, arg, solid_num):
	out = arg			out = arg
	for i in range(solid_num):			for i in range(solid_num):
	out = branch_solid.alias(f"{name}_{i}")(out)			out = branch_solid.alias(f"{name}_{i}")(out)

	return out			return out


	@pipeline(			@graph(description="Demo fork-shaped pipeline that has two-path parallel structure of solids.")
	description=("Demo fork-shaped pipeline that has two-path parallel structure of solids."),			def branch_graph():
	preset_defs=[			out_1, out_2 = root()
	PresetDefinition(			branch("branch_1", out_1, 3)
	"sleep_failed",			branch("branch_2", out_2, 5)
	{

				sleep_failed_branch_job = branch_graph.to_job(
				name="sleep_failed_branch_job",
				default_config={
	"intermediate_storage": {"filesystem": {}},			"intermediate_storage": {"filesystem": {}},
	"execution": {"multiprocess": {}},			"execution": {"multiprocess": {}},
	"solids": {"root": {"config": {"sleep_secs": [-10, 30]}}},			"solids": {"root": {"config": {"sleep_secs": [-10, 30]}}},
	},			},
	),			)
	PresetDefinition(
	"sleep",			sleep_branch_job = branch_graph.to_job(
	{			name="sleep_branch_job",
				default_config={
	"intermediate_storage": {"filesystem": {}},			"intermediate_storage": {"filesystem": {}},
	"execution": {"multiprocess": {}},			"execution": {"multiprocess": {}},
	"solids": {"root": {"config": {"sleep_secs": [0, 10]}}},			"solids": {"root": {"config": {"sleep_secs": [0, 10]}}},
	},			},
	),
	],
	)			)
	def branch_pipeline():
	out_1, out_2 = root()
	branch("branch_1", out_1, 3)
	branch("branch_2", out_2, 5)