def get_output_context(
execution_plan: "ExecutionPlan",
pipeline_def: "PipelineDefinition",
resolved_run_config: "ResolvedRunConfig",
step_output_handle: "StepOutputHandle",
run_id: Optional[str],
log_manager: Optional["DagsterLogManager"],
step_context: Optional["StepExecutionContext"],
resources: Optional["Resources"],
version: Optional[str],
) -> "OutputContext":
"""
Args:
run_id (str): The run ID of the run that produced the output, not necessarily the run that
the context will be used in.
"""
step = execution_plan.get_step_by_key(step_output_handle.step_key)
# get config
solid_config = resolved_run_config.solids[step.solid_handle.to_string()]
outputs_config = solid_config.outputs
if outputs_config:
output_config = outputs_config.get_output_manager_config(
step_output_handle.output_name)
else:
output_config = None
step_output = execution_plan.get_step_output(step_output_handle)
output_def = pipeline_def.get_solid(
step_output.solid_handle).output_def_named(step_output.name)
io_manager_key = output_def.io_manager_key
resource_config = resolved_run_config.resources[io_manager_key].config
if step_context:
check.invariant(
not resources,
"Expected either resources or step context to be set, but "
"received both. If step context is provided, resources for IO manager will be "
"retrieved off of that.",
)
resources = build_resources_for_manager(io_manager_key, step_context)
return OutputContext(
step_key=step_output_handle.step_key,
name=step_output_handle.output_name,
pipeline_name=pipeline_def.name,
run_id=run_id,
metadata=output_def.metadata,
mapping_key=step_output_handle.mapping_key,
config=output_config,
solid_def=pipeline_def.get_solid(step.solid_handle).definition,
dagster_type=output_def.dagster_type,
log_manager=log_manager,
version=version,
step_context=step_context,
resource_config=resource_config,
resources=resources,
)
def get_output_context(
execution_plan: "ExecutionPlan",
pipeline_def: "PipelineDefinition",
environment_config: "EnvironmentConfig",
step_output_handle: "StepOutputHandle",
run_id: Optional[str] = None,
log_manager: Optional["DagsterLogManager"] = None,
step_context: Optional["StepExecutionContext"] = None,
) -> "OutputContext":
"""
Args:
run_id (str): The run ID of the run that produced the output, not necessarily the run that
the context will be used in.
"""
step = execution_plan.get_step_by_key(step_output_handle.step_key)
# get config
solid_config = environment_config.solids[step.solid_handle.to_string()]
outputs_config = solid_config.outputs
if outputs_config:
output_config = outputs_config.get_output_manager_config(step_output_handle.output_name)
else:
output_config = None
step_output = execution_plan.get_step_output(step_output_handle)
output_def = pipeline_def.get_solid(step_output.solid_handle).output_def_named(step_output.name)
io_manager_key = output_def.io_manager_key
resource_config = environment_config.resources[io_manager_key].config
resources = build_resources_for_manager(io_manager_key, step_context) if step_context else None
return OutputContext(
step_key=step_output_handle.step_key,
name=step_output_handle.output_name,
pipeline_name=pipeline_def.name,
run_id=run_id,
metadata=output_def.metadata,
mapping_key=step_output_handle.mapping_key,
config=output_config,
solid_def=pipeline_def.get_solid(step.solid_handle).definition,
dagster_type=output_def.dagster_type,
log_manager=log_manager,
version=(
_step_output_version(
pipeline_def, execution_plan, environment_config, step_output_handle
)
if MEMOIZED_RUN_TAG in pipeline_def.tags
else None
),
step_context=step_context,
resource_config=resource_config,
resources=resources,
)
def _get_value(self, context, step_output_data):
step_output_handle = step_output_data.step_output_handle
manager = context.get_io_manager(step_output_handle)
manager_key = context.execution_plan.get_manager_key(step_output_handle)
res = manager.load_input(
context.for_input_manager(
name=None,
config=None,
metadata=None,
dagster_type=self.solid.output_def_named(step_output_data.output_name).dagster_type,
source_handle=step_output_handle,
resource_config=context.environment_config.resources[manager_key].config,
resources=build_resources_for_manager(manager_key, context),
)
)
return res
def _get_value(self, context, step_output_data):
step_output_handle = step_output_data.step_output_handle
# output capture dictionary will only have values in the in process case, but will not have
# values from steps launched via step launcher.
if self.output_capture and step_output_handle in self.output_capture:
return self.output_capture[step_output_handle]
manager = context.get_io_manager(step_output_handle)
manager_key = context.execution_plan.get_manager_key(step_output_handle, self.pipeline_def)
res = manager.load_input(
context.for_input_manager(
name=None,
config=None,
metadata=None,
dagster_type=self.solid.output_def_named(step_output_data.output_name).dagster_type,
source_handle=step_output_handle,
resource_config=context.environment_config.resources[manager_key].config,
resources=build_resources_for_manager(manager_key, context),
)
)
return res
def get_output_context(
execution_plan: "ExecutionPlan",
environment_config: EnvironmentConfig,
step_output_handle: StepOutputHandle,
run_id: str,
log_manager: Optional[DagsterLogManager] = None,
step_context: Optional[SystemStepExecutionContext] = None,
) -> OutputContext:
"""
Args:
run_id (str): The run ID of the run that produced the output, not necessarily the run that
the context will be used in.
"""
from dagster.core.execution.plan.plan import ExecutionPlan
check.inst_param(execution_plan, "execution_plan", ExecutionPlan)
check.inst_param(environment_config, "environment_config",
EnvironmentConfig)
check.inst_param(step_output_handle, "step_output_handle",
StepOutputHandle)
check.opt_str_param(run_id, "run_id")
step = execution_plan.get_step_by_key(step_output_handle.step_key)
# get config
solid_config = environment_config.solids.get(step.solid_handle.to_string())
outputs_config = solid_config.outputs
if outputs_config:
output_config = outputs_config.get_output_manager_config(
step_output_handle.output_name)
else:
output_config = None
pipeline_def = execution_plan.pipeline.get_definition()
step_output = execution_plan.get_step_output(step_output_handle)
output_def = pipeline_def.get_solid(
step_output.solid_handle).output_def_named(step_output.name)
io_manager_key = output_def.io_manager_key
resource_config = environment_config.resources[io_manager_key].get(
"config", {})
resources = build_resources_for_manager(
io_manager_key, step_context) if step_context else None
return OutputContext(
step_key=step_output_handle.step_key,
name=step_output_handle.output_name,
pipeline_name=pipeline_def.name,
run_id=run_id,
metadata=output_def.metadata,
mapping_key=step_output_handle.mapping_key,
config=output_config,
solid_def=pipeline_def.get_solid(step.solid_handle).definition,
dagster_type=output_def.dagster_type,
log_manager=log_manager,
version=_step_output_version(execution_plan, step_output_handle)
if MEMOIZED_RUN_TAG in execution_plan.pipeline.get_definition().tags
else None,
step_context=step_context,
resource_config=resource_config,
resources=resources,
)
