def test_wait_basics_2(workflow_start_regular_shared): # Test "workflow.wait" running in the top level DAG, # or running "workflow.wait" directly. @workflow.step def sleep_identity(x: int): time.sleep(x) return x @workflow.step def identity(x): return x ws = [ sleep_identity.step(1), sleep_identity.step(5), sleep_identity.step(2), ] w = workflow.wait(ws, num_returns=2, timeout=3) ready, remaining = identity.step(w).run() assert ready == [1, 2] ws = [ sleep_identity.step(2), sleep_identity.step(5), sleep_identity.step(1), ] w = workflow.wait(ws, num_returns=2, timeout=3) ready, remaining = w.run() assert ready == [2, 1]
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 test_wait_failure_recovery_2(workflow_start_regular_shared): # Test failing "workflow.wait" and its input steps. @workflow.step def sleep_identity(x: int): # block the step by a global mark while not utils.check_global_mark(): time.sleep(0.1) time.sleep(x) return x @workflow.step def identity(x): return x ws = [ sleep_identity.step(2), sleep_identity.step(5), sleep_identity.step(1), ] w = workflow.wait(ws, num_returns=2, timeout=None) utils.unset_global_mark() _ = identity.step(w).run_async(workflow_id="wait_failure_recovery_2") # wait util "workflow.wait" has been running time.sleep(10) workflow.cancel("wait_failure_recovery_2") time.sleep(2) utils.set_global_mark() ready, unready = ray.get(workflow.resume("wait_failure_recovery_2")) assert ready == [2, 1]
def test_wait_recovery_step_id(workflow_start_regular_shared): # This test ensures workflow reuse the original directory and # step id for "workflow.wait" during recovery. @workflow.step def identity(x: int): # block the step by a global mark assert utils.check_global_mark() return x w = workflow.wait([identity.step(42)], num_returns=1, timeout=None) utils.unset_global_mark() with pytest.raises(RaySystemError): _ = w.run(workflow_id="test_wait_recovery_step_id") utils.set_global_mark() ready, unready = ray.get(workflow.resume("test_wait_recovery_step_id")) assert ready == [42] from ray.workflow import storage, workflow_storage global_storage = storage.get_global_storage() wf_storage = workflow_storage.WorkflowStorage("test_wait_recovery_step_id", global_storage) index = wf_storage.gen_step_id("workflow.wait") # no new step id assert index <= 1
def recursive_wait(s): ready, unready = s if len(unready) == 2 and not isinstance(unready[0], Workflow): ready_2, unready = unready print(ready, (ready_2, unready)) ready += ready_2 if not unready: return ready w = workflow.wait(unready) return recursive_wait.step([ready, w])
def test_wait_basics(workflow_start_regular_shared): # This tests basic usage of 'workflow.wait': # 1. It returns ready tasks precisely and preserves the original order. # 2. All steps would see the same waiting result. # 3. We can pass remaining pending workflows to another workflow, # and they can be resolved like normal workflows. @workflow.step def return_ready(wait_results): ready, unready = wait_results return ready @workflow.step def join(a, b): return a, b wait_result = wait_multiple_steps.step() a = return_ready.step(wait_result) b = return_ready.step(wait_result) ready1, ready2 = join.step(a, b).run() assert ready1 == ready2 == [1, 3, 2] @workflow.step def get_all(ready, unready): return ready, unready @workflow.step def filter_all(wait_results): ready, unready = wait_results return get_all.step(ready, unready) @workflow.step def composite(): w = wait_multiple_steps.step() return filter_all.step(w) ready, unready = composite.step().run() assert ready == [1, 3, 2] assert unready == [10, 12] with pytest.raises(TypeError): workflow.wait([1, 2])
def wait_multiple_steps(): @workflow.step def sleep_identity(x: int): time.sleep(x) return x ws = [ sleep_identity.step(1), sleep_identity.step(3), sleep_identity.step(10), sleep_identity.step(2), sleep_identity.step(12), ] return workflow.wait(ws, num_returns=4, timeout=5)