Function test_restore_partially_finished_task

tests/test_journal.py:51–98

(hq_env: HqEnv, tmp_path)

def test_restore_partially_finished_task(hq_env: HqEnv, tmp_path):
    journal_path = os.path.join(tmp_path, "my.journal")
    hq_env.start_server(args=["--journal", journal_path])
    hq_env.command(
        [
            "submit",
            "--array=0-4",
            "--",
            *python(
                "import os, time;"
                f"open('{tmp_path}/task-marker' + os.environ.get('HQ_TASK_ID'), 'w').write(os.environ.get('HQ_INSTANCE_ID'));"
                "time.sleep(1.5 if os.environ.get('HQ_TASK_ID') in ('1','3','4') else 0);",
            ),
        ]
    )
    hq_env.start_worker(cpus=5)
    wait_for_job_state(hq_env, 1, "RUNNING")
    time.sleep(0.2)

    markers = sorted([name for name in os.listdir(tmp_path) if name.startswith("task-marker")])
    assert markers == ["task-marker0", "task-marker1", "task-marker2", "task-marker3", "task-marker4"]
    for marker in markers:
        path = os.path.join(tmp_path, marker)
        with open(path) as f:
            assert f.read() == "0"
        os.unlink(path)

    out = hq_env.command(["--output-mode=json", "job", "info", "1"], as_json=True)
    stats = out[0]["info"]["task_stats"]
    assert stats["running"] == 3
    assert stats["finished"] == 2
    hq_env.stop_server()

    hq_env.start_server(args=["--journal", journal_path])

    out = hq_env.command(["--output-mode=json", "job", "info", "1"], as_json=True)
    stats = out[0]["info"]["task_stats"]
    assert stats["waiting"] == 3
    assert stats["finished"] == 2

    hq_env.start_worker(cpus=4)
    wait_for_job_state(hq_env, [1], "FINISHED")
    markers = sorted([name for name in os.listdir(tmp_path) if name.startswith("task-marker")])
    assert markers == ["task-marker1", "task-marker3", "task-marker4"]
    for marker in markers:
        path = os.path.join(tmp_path, marker)
        with open(path) as f:
            assert f.read() == "1"


def test_restore_partially_failed_task(hq_env: HqEnv, tmp_path):