def test_raises_non_sequence_args(self):
with Flow("test"):
with pytest.raises(TypeError, match="non-sequence object"):
apply_map(lambda x: inc(x), 1)
with pytest.raises(TypeError, match="non-sequence object"):
apply_map(lambda x: inc(x), x=1)
def test_raises_use_map_in_mapped_function(self, method):
if method == "map":
def func(x):
inc.map(x)
pass_flow = False
elif method == "set_dependencies":
def func(x, flow):
flow.add_task(inc)
flow.set_dependencies(inc, keyword_tasks={"x": x}, mapped=True)
pass_flow = True
elif method == "add_edge":
def func(x, flow):
flow.add_task(inc)
flow.add_edge(x, inc, key="x", mapped=True)
pass_flow = True
flow = Flow("test")
with pytest.raises(ValueError,
match="running from inside a mapped context"):
if pass_flow:
apply_map(func, range(3), flow=flow)
else:
with flow:
apply_map(func, range(3))
def test_apply_map_can_be_used_on_apply_map_result(self, ):
# Prior to commit 4b0df740de99a1fad3182c01f4b182ba83445bcc this would introduce
# a cycle
@task
def combine(item_combination):
return item_combination[0] + item_combination[1]
# mapping function one
def apply_map_one(item):
result_one = inc(item)
result_two = inc(inc(item))
return (result_one, result_two)
# mapping function two
def apply_map_two(item):
return combine(item)
with Flow("test") as flow:
one_result = apply_map(apply_map_one, [1, 2, 3])
two_result = apply_map(apply_map_two, one_result)
state = flow.run()
assert state.is_successful()
assert state.result[one_result[0]].result == [2, 3, 4]
assert state.result[one_result[1]].result == [3, 4, 5]
assert state.result[two_result].result == [5, 7, 9]
def test_apply_map_simple(self, api):
if api == "functional":
def func(x, y, z):
a = add(x, y)
a.name = "add-a"
b = add(a, z)
b.name = "add-b"
c = add(b, 1)
c.name = "add-c"
return inc(c)
with Flow("test") as flow:
y = ranged(3)
z = edges.unmapped(1)
res = apply_map(func, range(3, 6), y=y, z=z)
else:
def func(x, y, z, flow):
a = add.copy(name="add-a").bind(x, y, flow=flow)
b = add.copy(name="add-b").bind(a, z, flow=flow)
c = add.copy(name="add-c").bind(b, 1, flow=flow)
return inc.copy().bind(c, flow=flow)
flow = Flow("test")
y = ranged.copy().bind(3, flow=flow)
z = edges.unmapped(tasks.as_task(1, flow=flow))
res = apply_map(func, range(3, 6), y=y, z=z, flow=flow)
consts = {t.name: c for t, c in flow.constants.items()}
assert consts == {
"ranged": {
"n": 3
},
"add-b": {
"y": 1
},
"add-c": {
"y": 1
}
}
for task in flow.tasks:
if task.name != "ranged":
for e in flow.edges_to(task):
assert e.mapped
state = flow.run()
assert state.result[res].result == [6, 8, 10]
def test_apply_map_inputs_added_to_subflow_before_calling_func(self):
"""We need to ensure all args to `appy_map` are added to the temporary
subflow *before* calling the mapped func. Otherwise things like
`case`/`resource_manager` statements that check the subgraph can get
confused and create new edges that shouldn't exist, leading to cycles."""
def func(cond, a, b):
with case(cond, True):
res1 = a + 1
with case(cond, False):
res2 = b + 1
return merge(res1, res2)
@tasks.task
def identity(x):
return x
with Flow("test") as flow:
cond = identity([True, False, True])
a = identity([1, 2, 3])
b = identity([4, 5, 6])
c = apply_map(func, cond, a, b)
state = flow.run()
assert state.result[c].result == [2, 6, 4]
def test_apply_map_disparate_length_args(self):
def func(x, y):
return inc(x), inc(y)
with Flow("test") as flow:
a, b = apply_map(func, range(3), range(3, 300))
state = flow.run()
assert state.result[a].result == [1, 2, 3]
assert state.result[b].result == [4, 5, 6]
def test_apply_map_return_multiple(self):
def func(x, y):
return inc(x), inc(y)
with Flow("test") as flow:
a, b = apply_map(func, range(3), range(3, 6))
state = flow.run()
assert state.result[a].result == [1, 2, 3]
assert state.result[b].result == [4, 5, 6]
def test_apply_map_flatten_works(self, input_type):
def func(a):
return a * 2
@tasks.task
def nested_list():
return [[1], [1, 2], [1, 2, 3]]
constant_nested = nested_list.run()
with Flow("test") as flow:
nested = nested_list()
if input_type == "constant":
a = apply_map(func, edges.flatten(nested))
else:
a = apply_map(func, edges.flatten(constant_nested))
state = flow.run()
assert state.result[a].result == [2, 2, 4, 2, 4, 6]
def test_imperative_args_are_added_to_flow_before_mapping(self):
# Check an edge case when mapping over tasks that haven't been added to flow yet.
@tasks.task
def data():
return range(3)
def func(a, flow):
return inc.copy().bind(a, flow=flow)
flow = Flow("test")
res = apply_map(func, data, flow=flow)
state = flow.run()
assert state.result[res].result == [1, 2, 3]
def test_apply_map_control_flow(self, api):
if api == "functional":
def func(x):
with case(is_even(x), True):
x2 = add(x, 1)
return merge(x2, x)
with Flow("test") as flow:
res = apply_map(func, range(4))
else:
def func(x, flow):
cond = is_even.copy().bind(x, flow=flow)
with case(cond, True):
x2 = add.copy().bind(x, 1, flow=flow)
return merge(x2, x, flow=flow)
flow = Flow("test")
res = apply_map(func, range(4), flow=flow)
state = flow.run()
assert state.result[res].result == [1, 1, 3, 3]
def test_tasks_have_all_non_unmapped_constant_args_as_transitive_upstream_deps(
self,
):
def func(a, b, c, d):
m = inc.copy(name="m").bind(1)
n = inc.copy(name="n").bind(a)
o = inc.copy(name="o").bind(b)
p = add.copy(name="p").bind(n, o)
q = add.copy(name="q").bind(c, d)
r = add.copy(name="r").bind(q, m)
return m, n, o, p, q, r
with Flow("test") as flow:
a = ranged.copy(name="a").bind(3)
b = inc.copy(name="b").bind(1)
c = Constant(1, name="c")
d = Constant(range(3), name="d")
m, n, o, p, q, r = apply_map(
func, a, edges.unmapped(b), c=edges.unmapped(c), d=d
)
def edge_info(task):
"""Returns a map of {upstream: (is_data_dep, is_mapped)}"""
return {
e.upstream_task: (e.key is not None, e.mapped)
for e in flow.edges_to(task)
}
assert edge_info(m) == {a: (False, True), b: (False, False), d: (False, True)}
assert edge_info(n) == {a: (True, True), b: (False, False), d: (False, True)}
assert edge_info(o) == {a: (False, True), b: (True, False), d: (False, True)}
assert edge_info(p) == {n: (True, True), o: (True, True)}
assert edge_info(q) == {a: (False, True), b: (False, False), d: (True, True)}
assert edge_info(r) == {q: (True, True), m: (True, True)}
state = flow.run()
res = {t: state.result[t].result for t in [m, n, o, p, q, r]}
sol = {
m: [2, 2, 2],
n: [1, 2, 3],
o: [3, 3, 3],
p: [4, 5, 6],
q: [1, 2, 3],
r: [3, 4, 5],
}
assert res == sol
def test_apply_map_inside_case_statement_works(self):
def func(x):
return add(x, 1), add(x, 2)
with Flow("test") as flow:
branch = Parameter("branch")
with case(branch, True):
a, b = apply_map(func, range(4))
c = add.map(a, b)
state = flow.run(branch=True)
assert state.result[a].result == [1, 2, 3, 4]
assert state.result[b].result == [2, 3, 4, 5]
assert state.result[c].result == [3, 5, 7, 9]
state = flow.run(branch=False)
assert state.result[a].is_skipped()
assert state.result[b].is_skipped()
assert state.result[c].is_skipped()
def test_apply_map_mixed_edge_types(self):
@tasks.task
def get_mixed_types():
return 3, [1, 2, 3], [[1, 2], [3]]
@tasks.task
def identity(a, b, c):
return a, b, c
def func(u, m, f):
return identity(u, m, f)
with Flow("test") as flow:
m = get_mixed_types()
a = apply_map(func, edges.unmapped(m[0]), m[1],
edges.flatten(m[2]))
state = flow.run()
assert state.result[a].result == [(3, 1, 1), (3, 2, 2), (3, 3, 3)]
def test_apply_map_inside_resource_manager_works(self):
@resource_manager
class MyResource:
def setup(self):
return 1
def cleanup(self, _):
pass
def func(x, a):
return add(x, a), add(x, 2)
with Flow("test") as flow:
with MyResource() as a:
b, c = apply_map(func, range(4), edges.unmapped(a))
d = add.map(b, c)
state = flow.run()
assert state.result[a].result == 1
assert state.result[b].result == [1, 2, 3, 4]
assert state.result[c].result == [2, 3, 4, 5]
assert state.result[d].result == [3, 5, 7, 9]
def test_raises_no_flow_found(self):
with pytest.raises(ValueError, match="flow in the current context"):
apply_map(lambda x: inc(x), range(10))