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)
