def wait(
self,
num_returns: int = 1,
timeout: Optional[float] = None
) -> Tuple[List[Workflow], List[Workflow]]:
"""Return a list of workflows that are ready and a list of workflows that
are not. See `api.wait()` for details.
Args:
num_returns (int): The number of workflows that should be returned.
timeout (float): The maximum amount of time in seconds to wait
before returning.
Returns:
A list of workflows that are ready and a list of the remaining
workflows.
"""
if self.workflow_refs:
raise ValueError("Currently, we do not support wait operations "
"on dynamic workflow refs. They are typically "
"generated by virtual actors.")
refs_map = {w.ref: w for w in self.workflow_outputs}
ready_ids, remaining_ids = ray.wait(list(refs_map.keys()),
num_returns=num_returns,
timeout=timeout)
ready_workflows = [Workflow.from_ref(refs_map[i]) for i in ready_ids]
remaining_workflows = [
Workflow.from_ref(refs_map[i]) for i in remaining_ids
]
return ready_workflows, remaining_workflows
def step(self, *args, **kwargs):
flattened_args = signature.flatten_args(self._signature, args, kwargs)
actor_id = workflow_context.get_current_workflow_id()
if not self.readonly:
if self._method_name == "__init__":
state_ref = None
else:
ws = WorkflowStorage(actor_id, get_global_storage())
state_ref = WorkflowRef(ws.get_entrypoint_step_id())
# This is a hack to insert a positional argument.
flattened_args = [signature.DUMMY_TYPE, state_ref] + flattened_args
workflow_inputs = serialization_context.make_workflow_inputs(
flattened_args)
if self.readonly:
_actor_method = _wrap_readonly_actor_method(
actor_id, self._original_class, self._method_name)
else:
_actor_method = _wrap_actor_method(self._original_class,
self._method_name)
workflow_data = WorkflowData(
func_body=_actor_method,
inputs=workflow_inputs,
name=self._name,
step_options=self._options,
user_metadata=self._user_metadata,
)
wf = Workflow(workflow_data)
return wf
def _build_workflow(*args, **kwargs) -> Workflow:
flattened_args = signature.flatten_args(self._func_signature, args, kwargs)
def prepare_inputs():
ensure_ray_initialized()
return serialization_context.make_workflow_inputs(flattened_args)
nonlocal step_options
if step_options is None:
step_options = WorkflowStepRuntimeOptions.make(
step_type=StepType.FUNCTION
)
# We could have "checkpoint=None" when we use @workflow.step
# with arguments. Avoid this by updating it here.
step_options.checkpoint = _inherit_checkpoint_option(
step_options.checkpoint
)
workflow_data = WorkflowData(
func_body=self._func,
inputs=None,
step_options=step_options,
name=self._name,
user_metadata=self._user_metadata,
)
return Workflow(workflow_data, prepare_inputs)
def _reconstruct_wait_step(
reader: workflow_storage.WorkflowStorage,
result: workflow_storage.StepInspectResult,
input_map: Dict[StepID, Any],
):
input_workflows = []
step_options = result.step_options
wait_options = step_options.ray_options.get("wait_options", {})
for i, _step_id in enumerate(result.workflows):
# Check whether the step has been loaded or not to avoid
# duplication
if _step_id in input_map:
r = input_map[_step_id]
else:
r = _construct_resume_workflow_from_step(reader, _step_id, input_map)
input_map[_step_id] = r
if isinstance(r, Workflow):
input_workflows.append(r)
else:
assert isinstance(r, StepID)
# TODO (Alex): We should consider caching these outputs too.
output = reader.load_step_output(r)
# Simulate a workflow with a workflow reference so it could be
# used directly by 'workflow.wait'.
static_ref = WorkflowStaticRef(step_id=r, ref=ray.put(output))
wf = Workflow.from_ref(static_ref)
input_workflows.append(wf)
from ray import workflow
return workflow.wait(input_workflows, **wait_options)
def step(method_name, method, *args, **kwargs):
readonly = getattr(method, "__virtual_actor_readonly__", False)
flattened_args = self.flatten_args(method_name, args, kwargs)
actor_id = workflow_context.get_current_workflow_id()
if not readonly:
if method_name == "__init__":
state_ref = None
else:
ws = WorkflowStorage(actor_id, get_global_storage())
state_ref = WorkflowRef(ws.get_entrypoint_step_id())
# This is a hack to insert a positional argument.
flattened_args = [signature.DUMMY_TYPE, state_ref
] + flattened_args
workflow_inputs = serialization_context.make_workflow_inputs(
flattened_args)
if readonly:
_actor_method = _wrap_readonly_actor_method(
actor_id, self.cls, method_name)
step_type = StepType.READONLY_ACTOR_METHOD
else:
_actor_method = _wrap_actor_method(self.cls, method_name)
step_type = StepType.ACTOR_METHOD
# TODO(suquark): Support actor options.
workflow_data = WorkflowData(
func_body=_actor_method,
step_type=step_type,
inputs=workflow_inputs,
max_retries=1,
catch_exceptions=False,
ray_options={},
name=None,
)
wf = Workflow(workflow_data)
return wf
def _build_workflow(*args, **kwargs) -> Workflow:
flattened_args = signature.flatten_args(self._func_signature, args, kwargs)
def prepare_inputs():
ensure_ray_initialized()
return serialization_context.make_workflow_inputs(flattened_args)
workflow_data = WorkflowData(
func_body=self._func,
inputs=None,
step_options=step_options,
name=self._name,
user_metadata=self._user_metadata,
)
return Workflow(workflow_data, prepare_inputs)
def _build_workflow(*args, **kwargs) -> Workflow:
flattened_args = signature.flatten_args(self._func_signature, args,
kwargs)
def prepare_inputs():
ensure_ray_initialized()
return serialization_context.make_workflow_inputs(
flattened_args)
workflow_data = WorkflowData(
func_body=self._func,
step_type=StepType.FUNCTION,
inputs=None,
max_retries=self._max_retries,
catch_exceptions=self._catch_exceptions,
ray_options=self._ray_options,
name=self._name,
)
return Workflow(workflow_data, prepare_inputs)
def _construct_resume_workflow_from_step(
workflow_id: str, step_id: StepID) -> Union[Workflow, Any]:
"""Try to construct a workflow (step) that recovers the workflow step.
If the workflow step already has an output checkpointing file, we return
the workflow step id instead.
Args:
workflow_id: The ID of the workflow.
step_id: The ID of the step we want to recover.
Returns:
A workflow that recovers the step, or the output of the step
if it has been checkpointed.
"""
reader = workflow_storage.WorkflowStorage(workflow_id)
# Step 1: construct dependency of the DAG (BFS)
inpsect_results = {}
dependency_map = defaultdict(list)
num_in_edges = {}
dag_visit_queue = deque([step_id])
while dag_visit_queue:
s: StepID = dag_visit_queue.popleft()
if s in inpsect_results:
continue
r = reader.inspect_step(s)
inpsect_results[s] = r
if not r.is_recoverable():
raise WorkflowStepNotRecoverableError(s)
if r.output_object_valid:
deps = []
elif isinstance(r.output_step_id, str):
deps = [r.output_step_id]
else:
deps = r.workflows
for w in deps:
dependency_map[w].append(s)
num_in_edges[s] = len(deps)
dag_visit_queue.extend(deps)
# Step 2: topological sort to determine the execution order (Kahn's algorithm)
execution_queue: List[StepID] = []
start_nodes = deque(k for k, v in num_in_edges.items() if v == 0)
while start_nodes:
n = start_nodes.popleft()
execution_queue.append(n)
for m in dependency_map[n]:
num_in_edges[m] -= 1
assert num_in_edges[m] >= 0, (m, n)
if num_in_edges[m] == 0:
start_nodes.append(m)
# Step 3: recover the workflow by the order of the execution queue
with serialization.objectref_cache():
# "input_map" is a context storing the input which has been loaded.
# This context is important for deduplicate step inputs.
input_map: Dict[StepID, Any] = {}
for _step_id in execution_queue:
result = inpsect_results[_step_id]
if result.output_object_valid:
input_map[_step_id] = reader.load_step_output(_step_id)
continue
if isinstance(result.output_step_id, str):
input_map[_step_id] = input_map[result.output_step_id]
continue
# Process the wait step as a special case.
if result.step_options.step_type == StepType.WAIT:
wait_input_workflows = []
for w in result.workflows:
output = input_map[w]
if isinstance(output, Workflow):
wait_input_workflows.append(output)
else:
# Simulate a workflow with a workflow reference so it could be
# used directly by 'workflow.wait'.
static_ref = WorkflowStaticRef(step_id=w,
ref=ray.put(output))
wait_input_workflows.append(
Workflow.from_ref(static_ref))
recovery_workflow = ray.workflow.wait(
wait_input_workflows,
**result.step_options.ray_options.get("wait_options", {}),
)
else:
args, kwargs = reader.load_step_args(
_step_id,
workflows=[input_map[w] for w in result.workflows],
workflow_refs=list(map(WorkflowRef, result.workflow_refs)),
)
func: Callable = reader.load_step_func_body(_step_id)
# TODO(suquark): Use an alternative function when "workflow.step"
# is fully deprecated.
recovery_workflow = ray.workflow.step(func).step(
*args, **kwargs)
# override step_options
recovery_workflow._step_id = _step_id
recovery_workflow.data.step_options = result.step_options
input_map[_step_id] = recovery_workflow
# Step 4: return the output of the requested step
return input_map[step_id]
def wait(workflows: List[Workflow],
num_returns: int = 1,
timeout: Optional[float] = None) -> Workflow[WaitResult]:
"""Return a list of result of workflows that are ready and a list of
workflows that are pending.
Examples:
>>> tasks = [task.step() for _ in range(3)]
>>> wait_step = workflow.wait(tasks, num_returns=1)
>>> print(wait_step.run())
([result_1], [<Workflow object>, <Workflow object>])
>>> tasks = [task.step() for _ in range(2)] + [forever.step()]
>>> wait_step = workflow.wait(tasks, num_returns=3, timeout=10)
>>> print(wait_step.run())
([result_1, result_2], [<Workflow object>])
If timeout is set, the function returns either when the requested number of
workflows are ready or when the timeout is reached, whichever occurs first.
If it is not set, the function simply waits until that number of workflows
is ready and returns that exact number of workflows.
This method returns two lists. The first list consists of workflows
references that correspond to workflows that are ready. The second
list corresponds to the rest of the workflows (which may or may not be
ready).
Ordering of the input list of workflows is preserved. That is, if A
precedes B in the input list, and both are in the ready list, then A will
precede B in the ready list. This also holds true if A and B are both in
the remaining list.
This method will issue a warning if it's running inside an async context.
Args:
workflows (List[Workflow]): List of workflows that may
or may not be ready. Note that these workflows must be unique.
num_returns (int): The number of workflows that should be returned.
timeout (float): The maximum amount of time in seconds to wait before
returning.
Returns:
A list of ready workflow results that are ready and a list of the
remaining workflows.
"""
from ray.workflow import serialization_context
from ray.workflow.common import WorkflowData
for w in workflows:
if not isinstance(w, Workflow):
raise TypeError("The input of workflow.wait should be a list "
"of workflows.")
wait_inputs = serialization_context.make_workflow_inputs(workflows)
step_options = WorkflowStepRuntimeOptions.make(
step_type=StepType.WAIT,
# Pass the options through Ray options. "num_returns" conflicts with
# the "num_returns" for Ray remote functions, so we need to wrap it
# under "wait_options".
ray_options={
"wait_options": {
"num_returns": num_returns,
"timeout": timeout,
}
},
)
workflow_data = WorkflowData(func_body=None,
inputs=wait_inputs,
step_options=step_options,
name="workflow.wait",
user_metadata={})
return Workflow(workflow_data)
