diff --git a/python/loom/client/tasks.py b/python/loom/client/tasks.py index de55f69780606d34e979ee3521c89d2638c9fbcb..8a3a4a96aaca4063e89edb5307f7937b91a1076e 100644 --- a/python/loom/client/tasks.py +++ b/python/loom/client/tasks.py @@ -168,8 +168,15 @@ def py_call(obj, inputs=(), request=cpu1): return task -def py_task(): +def py_task(label=None): def make_py_call(fn): + def py_task_builder(*args): + task = py_call(fn, args) + if label is not None: + task.label = label + else: + task.label = fn.__name__ + return task assert callable(fn) - return lambda *args: py_call(fn, args) + return py_task_builder return make_py_call diff --git a/tests/client/py_test.py b/tests/client/py_test.py index 9a4319d5e2de0971bf59b7b62777d14af4fb6585..967f0b205262d16a238789fd0828335aa6037a03 100644 --- a/tests/client/py_test.py +++ b/tests/client/py_test.py @@ -21,7 +21,7 @@ def test_py_call(loom_env): d = tasks.const("12345") p = tasks.py_call(f, (c, d)) q = tasks.py_call(g) - result1, result2 = loom_env.submit((p, q), "report") + result1, result2 = loom_env.submit((p, q)) assert result1 == b"ABC, 3, 12345, 5" assert result2 == b"Test" @@ -33,7 +33,7 @@ def test_py_task(loom_env): def t1(): return "ABC" - @tasks.py_task() + @tasks.py_task(label="Merging task") def t2(a, b): return a.read() + b.read() @@ -41,7 +41,7 @@ def test_py_task(loom_env): a = tasks.const("1234") b = t1() c = t2(a, b) - result = loom_env.submit(c) + result = loom_env.submit(c, "report") assert result == b"1234ABC"