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)