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
async def _write_step_inputs(wf_storage: workflow_storage.WorkflowStorage,
step_id: StepID, inputs: WorkflowData) -> None:
"""Save workflow inputs."""
metadata = inputs.to_metadata()
with serialization_context.workflow_args_keeping_context():
# TODO(suquark): in the future we should write to storage directly
# with plasma store object in memory.
args_obj = ray.get(inputs.inputs.args)
workflow_id = wf_storage._workflow_id
storage = wf_storage._storage
save_tasks = [
# TODO (Alex): Handle the json case better?
wf_storage._put(wf_storage._key_step_input_metadata(step_id), metadata,
True),
wf_storage._put(wf_storage._key_step_user_metadata(step_id),
inputs.user_metadata, True),
serialization.dump_to_storage(
wf_storage._key_step_function_body(step_id),
inputs.func_body,
workflow_id,
storage,
),
serialization.dump_to_storage(wf_storage._key_step_args(step_id),
args_obj, workflow_id, storage),
]
await asyncio.gather(*save_tasks)
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 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 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)
