python/ray/workflow/step_executor.py
def _workflow_step_executor(
    func: Callable,
    context: "WorkflowStepContext",
    step_id: "StepID",
    baked_inputs: "_BakedWorkflowInputs",
    runtime_options: "WorkflowStepRuntimeOptions",
    inplace: bool = False,
) -> Tuple[Any, Any]:
    """Executor function for workflow step.

    Args:
        step_id: ID of the step.
        func: The workflow step function.
        baked_inputs: The processed inputs for the step.
        context: Workflow step context. Used to access correct storage etc.
        runtime_options: Parameters for workflow step execution.
        inplace: Execute the workflow inplace.

    Returns:
        Workflow step output.
    """
    # Part 1: update the context for the step
    workflow_context.update_workflow_step_context(context, step_id)
    context = workflow_context.get_workflow_step_context()
    step_type = runtime_options.step_type
    context.checkpoint_context.checkpoint = runtime_options.checkpoint

    # Part 2: resolve inputs
    args, kwargs = baked_inputs.resolve()

    # Part 3: execute the step
    store = workflow_storage.get_workflow_storage()
    try:
        step_prerun_metadata = {"start_time": time.time()}
        store.save_step_prerun_metadata(step_id, step_prerun_metadata)
        with workflow_context.workflow_execution():
            persisted_output, volatile_output = _wrap_run(
                func, runtime_options, *args, **kwargs
            )
        step_postrun_metadata = {"end_time": time.time()}
        store.save_step_postrun_metadata(step_id, step_postrun_metadata)
    except Exception as e:
        # Always checkpoint the exception.
        commit_step(store, step_id, None, exception=e)
        raise e

    # Part 4: save outputs
    if step_type == StepType.READONLY_ACTOR_METHOD:
        if isinstance(volatile_output, Workflow):
            raise TypeError(
                "Returning a Workflow from a readonly virtual actor is not allowed."
            )
        assert not isinstance(persisted_output, Workflow)
    else:
        # TODO(suquark): Validate checkpoint options before
        # commit the step.
        store = workflow_storage.get_workflow_storage()
        if CheckpointMode(runtime_options.checkpoint) == CheckpointMode.SYNC:
            commit_step(
                store,
                step_id,
                persisted_output,
                exception=None,
            )
        if isinstance(persisted_output, Workflow):
            sub_workflow = persisted_output
            outer_most_step_id = context.outer_most_step_id
            assert volatile_output is None
            if step_type == StepType.FUNCTION:
                # Passing down outer most step so inner nested steps would
                # access the same outer most step.
                if not context.outer_most_step_id:
                    # The current workflow step returns a nested workflow, and
                    # there is no outer step for the current step. So the
                    # current step is the outer most step for the inner nested
                    # workflow steps.
                    outer_most_step_id = workflow_context.get_current_step_id()
            if inplace:
                _step_options = sub_workflow.data.step_options
                if (
                    _step_options.step_type != StepType.WAIT
                    and runtime_options.ray_options != _step_options.ray_options
                ):
                    logger.warning(
                        f"Workflow step '{sub_workflow.step_id}' uses "
                        f"a Ray option different to its caller step '{step_id}' "
                        f"and will be executed inplace. Ray assumes it still "
                        f"consumes the same resource as the caller. This may result "
                        f"in oversubscribing resources."
                    )
                return (
                    InplaceReturnedWorkflow(
                        sub_workflow, {"outer_most_step_id": outer_most_step_id}
                    ),
                    None,
                )
            # Execute sub-workflow. Pass down "outer_most_step_id".
            with workflow_context.fork_workflow_step_context(
                outer_most_step_id=outer_most_step_id
            ):
                result = execute_workflow(sub_workflow)
            # When virtual actor returns a workflow in the method,
            # the volatile_output and persisted_output will be put together
            persisted_output = result.persisted_output
            volatile_output = result.volatile_output
        elif context.last_step_of_workflow:
            # advance the progress of the workflow
            store.advance_progress(step_id)
        _record_step_status(step_id, WorkflowStatus.SUCCESSFUL)
    logger.info(get_step_status_info(WorkflowStatus.SUCCESSFUL))
    if isinstance(volatile_output, Workflow):
        # This is the case where a step method is called in the virtual actor.
        # We need to run the method to get the final result.
        assert step_type == StepType.ACTOR_METHOD
        volatile_output = volatile_output.run_async(
            workflow_context.get_current_workflow_id()
        )
        volatile_output = WorkflowStaticRef.from_output(step_id, volatile_output)
    return persisted_output, volatile_output
Executor function for workflow step.

    Args:
        step_id: ID of the step.
        func: The workflow step function.
        baked_inputs: The processed inputs for the step.
        context: Workflow step context. Used to access correct storage etc.
        runtime_options: Parameters for workflow step execution.
        inplace: Execute the workflow inplace.

    Returns:
        Workflow step output.
    
Executor function for workflow step.
Executor function for workflow step .
def _workflow_step_executor(
    func: Callable,
    context: "WorkflowStepContext",
    step_id: "StepID",
    baked_inputs: "_BakedWorkflowInputs",
    runtime_options: "WorkflowStepRuntimeOptions",
    inplace: bool = False,
) -> Tuple[Any, Any]:
    workflow_context.update_workflow_step_context(context, step_id)
    context = workflow_context.get_workflow_step_context()
    step_type = runtime_options.step_type
    context.checkpoint_context.checkpoint = runtime_options.checkpoint
    args, kwargs = baked_inputs.resolve()
    store = workflow_storage.get_workflow_storage()
    try:
        step_prerun_metadata = {"start_time": time.time()}
        store.save_step_prerun_metadata(step_id, step_prerun_metadata)
        with workflow_context.workflow_execution():
            persisted_output, volatile_output = _wrap_run(
                func, runtime_options, *args, **kwargs
            )
        step_postrun_metadata = {"end_time": time.time()}
        store.save_step_postrun_metadata(step_id, step_postrun_metadata)
    except Exception as e:
        commit_step(store, step_id, None, exception=e)
        raise e
    if step_type == StepType.READONLY_ACTOR_METHOD:
        if isinstance(volatile_output, Workflow):
            raise TypeError(
                "Returning a Workflow from a readonly virtual actor is not allowed."
            )
        assert not isinstance(persisted_output, Workflow)
    else:
        store = workflow_storage.get_workflow_storage()
        if CheckpointMode(runtime_options.checkpoint) == CheckpointMode.SYNC:
            commit_step(
                store,
                step_id,
                persisted_output,
                exception=None,
            )
        if isinstance(persisted_output, Workflow):
            sub_workflow = persisted_output
            outer_most_step_id = context.outer_most_step_id
            assert volatile_output is None
            if step_type == StepType.FUNCTION:
                if not context.outer_most_step_id:
                    outer_most_step_id = workflow_context.get_current_step_id()
            if inplace:
                _step_options = sub_workflow.data.step_options
                if (
                    _step_options.step_type != StepType.WAIT
                    and runtime_options.ray_options != _step_options.ray_options
                ):
                    logger.warning(
                        f"Workflow step '{sub_workflow.step_id}' uses "
                        f"a Ray option different to its caller step '{step_id}' "
                        f"and will be executed inplace. Ray assumes it still "
                        f"consumes the same resource as the caller. This may result "
                        f"in oversubscribing resources."
                    )
                return (
                    InplaceReturnedWorkflow(
                        sub_workflow, {"outer_most_step_id": outer_most_step_id}
                    ),
                    None,
                )
            with workflow_context.fork_workflow_step_context(
                outer_most_step_id=outer_most_step_id
            ):
                result = execute_workflow(sub_workflow)
            persisted_output = result.persisted_output
            volatile_output = result.volatile_output
        elif context.last_step_of_workflow:
            store.advance_progress(step_id)
        _record_step_status(step_id, WorkflowStatus.SUCCESSFUL)
    logger.info(get_step_status_info(WorkflowStatus.SUCCESSFUL))
    if isinstance(volatile_output, Workflow):
        assert step_type == StepType.ACTOR_METHOD
        volatile_output = volatile_output.run_async(
            workflow_context.get_current_workflow_id()
        )
        volatile_output = WorkflowStaticRef.from_output(step_id, volatile_output)
    return persisted_output, volatile_output
def _workflow_step_executor ( func : Callable , context : "WorkflowStepContext" , step_id : "StepID" , baked_inputs : "_BakedWorkflowInputs" , runtime_options : "WorkflowStepRuntimeOptions" , inplace : bool = False , ) -> Tuple [ Any , Any ] : workflow_context . update_workflow_step_context ( context , step_id ) context = workflow_context . get_workflow_step_context ( ) step_type = runtime_options . step_type context . checkpoint_context . checkpoint = runtime_options . checkpoint args , kwargs = baked_inputs . resolve ( ) store = workflow_storage . get_workflow_storage ( ) try : step_prerun_metadata = { "start_time" : time . time ( ) } store . save_step_prerun_metadata ( step_id , step_prerun_metadata ) with workflow_context . workflow_execution ( ) : persisted_output , volatile_output = _wrap_run ( func , runtime_options , * args , ** kwargs ) step_postrun_metadata = { "end_time" : time . time ( ) } store . save_step_postrun_metadata ( step_id , step_postrun_metadata ) except Exception as e : commit_step ( store , step_id , None , exception = e ) raise e if step_type == StepType . READONLY_ACTOR_METHOD : if isinstance ( volatile_output , Workflow ) : raise TypeError ( "Returning a Workflow from a readonly virtual actor is not allowed." ) assert not isinstance ( persisted_output , Workflow ) else : store = workflow_storage . get_workflow_storage ( ) if CheckpointMode ( runtime_options . checkpoint ) == CheckpointMode . SYNC : commit_step ( store , step_id , persisted_output , exception = None , ) if isinstance ( persisted_output , Workflow ) : sub_workflow = persisted_output outer_most_step_id = context . outer_most_step_id assert volatile_output is None if step_type == StepType . FUNCTION : if not context . outer_most_step_id : outer_most_step_id = workflow_context . get_current_step_id ( ) if inplace : _step_options = sub_workflow . data . step_options if ( _step_options . step_type != StepType . WAIT and runtime_options . ray_options != _step_options . ray_options ) : logger . warning ( f"Workflow step '{sub_workflow.step_id}' uses " f"a Ray option different to its caller step '{step_id}' " f"and will be executed inplace. Ray assumes it still " f"consumes the same resource as the caller. This may result " f"in oversubscribing resources." ) return ( InplaceReturnedWorkflow ( sub_workflow , { "outer_most_step_id" : outer_most_step_id } ) , None , ) with workflow_context . fork_workflow_step_context ( outer_most_step_id = outer_most_step_id ) : result = execute_workflow ( sub_workflow ) persisted_output = result . persisted_output volatile_output = result . volatile_output elif context . last_step_of_workflow : store . advance_progress ( step_id ) _record_step_status ( step_id , WorkflowStatus . SUCCESSFUL ) logger . info ( get_step_status_info ( WorkflowStatus . SUCCESSFUL ) ) if isinstance ( volatile_output , Workflow ) : assert step_type == StepType . ACTOR_METHOD volatile_output = volatile_output . run_async ( workflow_context . get_current_workflow_id ( ) ) volatile_output = WorkflowStaticRef . from_output ( step_id , volatile_output ) return persisted_output , volatile_output
Executor function for workflow step.
Executor function for workflow step .
"""Executor function for workflow step.

    Args:
        step_id: ID of the step.
        func: The workflow step function.
        baked_inputs: The processed inputs for the step.
        context: Workflow step context. Used to access correct storage etc.
        runtime_options: Parameters for workflow step execution.
        inplace: Execute the workflow inplace.

    Returns:
        Workflow step output.
    """ # Part 1: update the context for the step # Part 2: resolve inputs # Part 3: execute the step # Always checkpoint the exception. # Part 4: save outputs # TODO(suquark): Validate checkpoint options before # commit the step. # Passing down outer most step so inner nested steps would # access the same outer most step. # The current workflow step returns a nested workflow, and # there is no outer step for the current step. So the # current step is the outer most step for the inner nested # workflow steps. # Execute sub-workflow. Pass down "outer_most_step_id". # When virtual actor returns a workflow in the method, # the volatile_output and persisted_output will be put together # advance the progress of the workflow # This is the case where a step method is called in the virtual actor. # We need to run the method to get the final result.
Workflow step output.
Workflow step output .
func
Callable
The workflow step function.
The workflow step function .
context
"WorkflowStepContext"
Workflow step context. Used to access correct storage etc.
Workflow step context . Used to access correct storage etc .
step_id
"StepID"
ID of the step.
ID of the step .
baked_inputs
"_BakedWorkflowInputs"
The processed inputs for the step.
The processed inputs for the step .
runtime_options
"WorkflowStepRuntimeOptions"
Parameters for workflow step execution.
Parameters for workflow step execution .
inplace
bool
Execute the workflow inplace.
Execute the workflow inplace .