mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
karajan1001
b37524603a
Followed from #7592 1. Seperate remove method to a new file. 2. Add queued, failed, processed flags to remove. 3. Add new unit tests for queue status --queue/fail/success 4. Implement methods to remove done tasks. 5. add some new unit and functional test for celery_queue.remove 6. bump into dvc-task version 0.0.13
150 lines
4.4 KiB
Python
150 lines
4.4 KiB
Python
import pytest
|
|
from funcy import first
|
|
|
|
from dvc.exceptions import InvalidArgumentError
|
|
|
|
|
|
def to_dict(tasks):
|
|
status_dict = {}
|
|
for task in tasks:
|
|
status_dict[task["name"]] = task["status"]
|
|
return status_dict
|
|
|
|
|
|
@pytest.fixture
|
|
def queued_tasks(tmp_dir, dvc, scm, exp_stage):
|
|
queue_length = 3
|
|
name_list = []
|
|
for i in range(queue_length):
|
|
name = f"queued{i}"
|
|
name_list.append(name)
|
|
dvc.experiments.run(
|
|
exp_stage.addressing,
|
|
params=[f"foo={i+2*queue_length}"],
|
|
queue=True,
|
|
name=name,
|
|
)
|
|
return ["queued0", "queued1", "queued2"]
|
|
|
|
|
|
@pytest.fixture
|
|
def success_tasks(tmp_dir, dvc, scm, test_queue, exp_stage):
|
|
queue_length = 3
|
|
name_list = []
|
|
for i in range(queue_length):
|
|
name = f"success{i}"
|
|
name_list.append(name)
|
|
dvc.experiments.run(
|
|
exp_stage.addressing, params=[f"foo={i}"], queue=True, name=name
|
|
)
|
|
dvc.experiments.run(run_all=True)
|
|
return name_list
|
|
|
|
|
|
@pytest.fixture
|
|
def failed_tasks(tmp_dir, dvc, scm, test_queue, failed_exp_stage, capsys):
|
|
queue_length = 3
|
|
name_list = []
|
|
for i in range(queue_length):
|
|
name = f"failed{i}"
|
|
name_list.append(name)
|
|
dvc.experiments.run(
|
|
failed_exp_stage.addressing,
|
|
params=[f"foo={i+queue_length}"],
|
|
queue=True,
|
|
name=name,
|
|
)
|
|
dvc.experiments.run(run_all=True)
|
|
output = capsys.readouterr()
|
|
assert "Failed to reproduce experiment" in output.err
|
|
return name_list
|
|
|
|
|
|
@pytest.mark.parametrize("follow", [True, False])
|
|
def test_celery_logs(
|
|
tmp_dir,
|
|
scm,
|
|
dvc,
|
|
failed_exp_stage,
|
|
test_queue,
|
|
follow,
|
|
capsys,
|
|
):
|
|
dvc.experiments.run(failed_exp_stage.addressing, queue=True)
|
|
dvc.experiments.run(run_all=True)
|
|
|
|
queue = dvc.experiments.celery_queue
|
|
done_result = first(queue.iter_done())
|
|
name = done_result.entry.stash_rev
|
|
captured = capsys.readouterr()
|
|
queue.logs(name, follow=follow)
|
|
captured = capsys.readouterr()
|
|
assert "failed to reproduce 'failed-copy-file'" in captured.out
|
|
|
|
|
|
def test_queue_status(dvc, failed_tasks, success_tasks, queued_tasks):
|
|
assert len(dvc.experiments.stash_revs) == 3
|
|
assert len(dvc.experiments.celery_queue.failed_stash) == 3
|
|
status = to_dict(dvc.experiments.celery_queue.status())
|
|
assert len(status) == 9
|
|
for task in failed_tasks:
|
|
assert status[task] == "Failed"
|
|
for task in success_tasks:
|
|
assert status[task] == "Success"
|
|
for task in queued_tasks:
|
|
assert status[task] == "Queued"
|
|
|
|
|
|
def test_queue_remove(dvc, failed_tasks, success_tasks, queued_tasks):
|
|
assert len(dvc.experiments.stash_revs) == 3
|
|
assert len(dvc.experiments.celery_queue.failed_stash) == 3
|
|
assert len(dvc.experiments.celery_queue.status()) == 9
|
|
|
|
with pytest.raises(InvalidArgumentError):
|
|
dvc.experiments.celery_queue.remove(failed_tasks[:2] + ["non-exist"])
|
|
assert len(dvc.experiments.celery_queue.status()) == 9
|
|
|
|
to_remove = failed_tasks[:2] + success_tasks[1:] + queued_tasks[1:2]
|
|
assert set(dvc.experiments.celery_queue.remove(to_remove)) == set(
|
|
to_remove
|
|
)
|
|
|
|
assert len(dvc.experiments.stash_revs) == 2
|
|
assert len(dvc.experiments.celery_queue.failed_stash) == 1
|
|
status = to_dict(dvc.experiments.celery_queue.status())
|
|
assert set(status) == set(
|
|
queued_tasks[:1]
|
|
+ queued_tasks[2:]
|
|
+ success_tasks[:1]
|
|
+ failed_tasks[2:]
|
|
)
|
|
assert status[queued_tasks[0]] == "Queued"
|
|
assert status[queued_tasks[2]] == "Queued"
|
|
|
|
assert (
|
|
dvc.experiments.celery_queue.remove([], queued=True)
|
|
== queued_tasks[:1] + queued_tasks[2:]
|
|
)
|
|
|
|
assert len(dvc.experiments.stash_revs) == 0
|
|
assert len(dvc.experiments.celery_queue.failed_stash) == 1
|
|
assert len(dvc.experiments.celery_queue.status()) == 2
|
|
|
|
assert (
|
|
dvc.experiments.celery_queue.remove([], failed=True)
|
|
== failed_tasks[2:]
|
|
)
|
|
|
|
assert len(dvc.experiments.stash_revs) == 0
|
|
assert len(dvc.experiments.celery_queue.failed_stash) == 0
|
|
assert len(dvc.experiments.celery_queue.status()) == 1
|
|
|
|
assert (
|
|
dvc.experiments.celery_queue.remove([], success=True)
|
|
== success_tasks[:1]
|
|
)
|
|
|
|
assert len(dvc.experiments.stash_revs) == 0
|
|
assert len(dvc.experiments.celery_queue.failed_stash) == 0
|
|
assert len(dvc.experiments.celery_queue.status()) == 0
|