def test_run_dynamic_task_job(n_jobs, storage):
dynamic = DynamicTask1(
parent=Task1(name="test-dynamic-task", resource_spec=None).output())
# Case - case independently executed
run_job(dynamic.parent.src_task, storage, n_jobs=n_jobs)
task = generate_task(dynamic, storage)
assert isinstance(task, WriteValue)
assert task.value_to_write == "test-dynamic-task"
run_job(task, storage)
assert is_completed(task, storage)
assert is_completed(dynamic, storage)
_flush(storage)
# Case - executed at the same time
assert not is_completed(dynamic,
storage), "confirm initialization of storage"
run_job(dynamic, storage, n_jobs=n_jobs)
assert is_completed(dynamic, storage)
def test_run_dynamic_task_job():
dynamic = DynamicTask1(parent=Task1(name="test-dynamic-task").output())
# Case - case independently executed
with tempfile.TemporaryDirectory() as dirname:
storage = LocalStorage(dirname)
run_job(dynamic.parent.src_task, storage)
task = generate_task(dynamic, storage)
assert isinstance(task, WriteValue)
assert task.value_to_write == "test-dynamic-task"
run_job(task, storage)
assert is_completed(task, storage)
assert is_completed(dynamic, storage)
assert not is_completed(dynamic, storage)
# Case - executed at the same time
with tempfile.TemporaryDirectory() as dirname:
run_job(dynamic, storage)
assert is_completed(dynamic, storage)
def test_run_job(n_jobs: int, storage):
task = Task2(resource_spec=None, parent=Task1(resource_spec=None).output())
assert not is_completed(task.input().src_task, storage)
assert not is_completed(task, storage)
run_job(task, storage, n_jobs=n_jobs)
assert is_completed(task.input().src_task, storage)
assert is_completed(task, storage)
def test_run_with_context(method: str, storage):
task = Task2(resource_spec=None, parent=Task1(resource_spec=None).output())
assert not is_completed(task.input().src_task, storage)
assert not is_completed(task, storage)
run_job(task, storage, n_jobs=2, context=get_context(method))
assert is_completed(task.input().src_task, storage)
assert is_completed(task, storage)
def test_with_ephemeral_output_task(task, n_jobs, storage):
run_job(task, storage, n_jobs=2)
assert is_completed(task, storage)
assert task.input().ephemeral
assert not is_completed(
task.input().src_task,
storage), "ephemeral output is already removed automatically"
def test_run_job():
task = Task2(parent=Task1().output())
with tempfile.TemporaryDirectory() as dirname:
storage = LocalStorage(dirname)
assert not is_completed(task.input().src_task, storage)
assert not is_completed(task, storage)
run_job(task, storage)
assert is_completed(task.input().src_task, storage)
assert is_completed(task, storage)
def _to_ref_map(
tasks: Dict[str, AbstractTask],
only_incomplete: bool = False,
storage: Optional[Storage] = None,
) -> Tuple[Dict[str, Set[str]], Dict[str, bool]]:
"""Get reference count dictionary
"""
if only_incomplete:
assert storage is not None, "storage must be given when only_incomplete = True"
# key = Output.key, value set of task_ids who uses the output
ref: Dict[str, Set[str]] = defaultdict(set)
# key = Output.key, value where an output is all ephemeral or not.
ephemeral_map: Dict[str, bool] = defaultdict(lambda: True)
while len(tasks) > 0:
next_tasks: Dict[str, AbstractTask] = {}
for task_id, task in tasks.items():
if only_incomplete:
assert isinstance(storage, Storage)
if is_completed(task, storage):
continue
inputs = flatten(task.input())
for inp in inputs:
for key in inp.physical_key_list():
ref[key].add(task_id)
ephemeral_map[key] = ephemeral_map[key] and inp.ephemeral
dependent_tasks = {inp.src_task.task_id: inp.src_task for inp in inputs}
for dependent_task_id, task in dependent_tasks.items():
next_tasks[dependent_task_id] = task
tasks = next_tasks
return ref, ephemeral_map
def test_complex_input_task(storage):
task = ComplexInputTask()
run_job(task, storage, n_jobs=1)
assert is_completed(task, storage)
