[ENH] Include tests for the queue command #110
2 changed files with 565 additions and 9 deletions
|
|
@ -84,7 +84,7 @@ def run(
|
|||
# Convert str to Path
|
||||
if isinstance(workdir, str):
|
||||
workdir = Path(workdir)
|
||||
if not isinstance(elements, List) and elements is not None:
|
||||
if not isinstance(elements, list) and elements is not None:
|
||||
elements = [elements]
|
||||
# Get datagrabber to use
|
||||
datagrabber_object = _get_datagrabber(datagrabber)
|
||||
|
|
@ -186,7 +186,7 @@ def queue(
|
|||
jobdir = cwd / "junifer_jobs" / jobname
|
||||
logger.info(f"Creating job in {str(jobdir.absolute())}")
|
||||
if jobdir.exists():
|
||||
if overwrite is not True:
|
||||
if not overwrite:
|
||||
raise_error(
|
||||
f"Job folder for {jobname} already exists. "
|
||||
"This error is raised to prevent overwriting job files "
|
||||
|
|
@ -234,7 +234,7 @@ def queue(
|
|||
elements = dg.get_elements()
|
||||
|
||||
# TODO: Fix typing of elements
|
||||
if not isinstance(elements, List):
|
||||
if not isinstance(elements, list):
|
||||
elements = [elements] # type: ignore
|
||||
|
||||
typing.cast(List[Union[str, Tuple]], elements)
|
||||
|
|
@ -258,7 +258,7 @@ def queue(
|
|||
**kwargs,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown queue kind: {kind}")
|
||||
raise_error(f"Unknown queue kind: {kind}")
|
||||
|
||||
logger.info("Queue done")
|
||||
|
||||
|
|
|
|||
|
|
@ -5,15 +5,19 @@
|
|||
# Synchon Mandal <s.mandal@fz-juelich.de>
|
||||
# License: AGPL
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
import junifer.testing.registry # noqa: F401
|
||||
from junifer.api.functions import collect, run
|
||||
from junifer.api.functions import collect, queue, run
|
||||
from junifer.datagrabber.base import BaseDataGrabber
|
||||
from junifer.pipeline.registry import build
|
||||
|
||||
|
||||
# Define datagrabber
|
||||
datagrabber = {
|
||||
"kind": "OasisVBMTestingDatagrabber",
|
||||
|
|
@ -144,10 +148,562 @@ def test_run_and_collect(tmp_path: Path) -> None:
|
|||
assert uri.exists()
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="HTCondor not installed on system.")
|
||||
def test_queue_condor() -> None:
|
||||
"""Test job queueing in HTCondor."""
|
||||
pass
|
||||
def test_queue_correct_yaml_config(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Test proper YAML config generation for queueing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
queue(
|
||||
config={
|
||||
"with": "junifer.testing.registry",
|
||||
"workdir": str(Path(tmp_path).resolve()),
|
||||
"datagrabber": datagrabber,
|
||||
"markers": markers,
|
||||
"storage": storage,
|
||||
"env": {
|
||||
"kind": "conda",
|
||||
"name": "junifer",
|
||||
},
|
||||
"mem": "8G",
|
||||
},
|
||||
kind="HTCondor",
|
||||
jobname="yaml_config_gen_check",
|
||||
)
|
||||
|
||||
generated_config_yaml_path = Path(
|
||||
tmp_path / "junifer_jobs" / "yaml_config_gen_check" / "config.yaml"
|
||||
)
|
||||
with open(generated_config_yaml_path, "r") as f:
|
||||
yaml_config = yaml.unsafe_load(f)
|
||||
# Check for correct YAML config generation
|
||||
assert all(
|
||||
key in yaml_config.keys()
|
||||
for key in [
|
||||
"with",
|
||||
"workdir",
|
||||
"datagrabber",
|
||||
"markers",
|
||||
"storage",
|
||||
"env",
|
||||
"mem",
|
||||
]
|
||||
)
|
||||
assert "queue" not in yaml_config.keys()
|
||||
|
||||
|
||||
def test_queue_invalid_job_queue(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Test queue function for invalid job queue.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
|
||||
"""
|
||||
with pytest.raises(ValueError, match="Unknown queue kind"):
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
queue(
|
||||
config={"elements": ["sub-001"]},
|
||||
kind="ABC",
|
||||
)
|
||||
|
||||
|
||||
def test_queue_assets_disallow_overwrite(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Test overwrite prevention of queue files.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
|
||||
"""
|
||||
with pytest.raises(ValueError, match="Either delete the directory"):
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
# First generate assets
|
||||
queue(
|
||||
config={"elements": ["sub-001"]},
|
||||
kind="HTCondor",
|
||||
jobname="prevent_overwrite",
|
||||
)
|
||||
# Re-run to trigger error
|
||||
queue(
|
||||
config={"elements": ["sub-001"]},
|
||||
kind="HTCondor",
|
||||
jobname="prevent_overwrite",
|
||||
)
|
||||
|
||||
|
||||
def test_queue_assets_allow_overwrite(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Test overwriting of queue files.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
# First generate assets
|
||||
queue(
|
||||
config={"elements": ["sub-001"]},
|
||||
kind="HTCondor",
|
||||
jobname="allow_overwrite",
|
||||
)
|
||||
with caplog.at_level(logging.INFO):
|
||||
# Re-run to overwrite
|
||||
queue(
|
||||
config={"elements": ["sub-001"]},
|
||||
kind="HTCondor",
|
||||
jobname="allow_overwrite",
|
||||
overwrite=True,
|
||||
)
|
||||
assert "Deleting existing job directory" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"with_",
|
||||
[
|
||||
"a.py",
|
||||
["a.py"],
|
||||
["a.py", "b"],
|
||||
],
|
||||
)
|
||||
def test_queue_with_imports(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
with_: Union[str, List[str]],
|
||||
) -> None:
|
||||
"""Test queue with `with` imports.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
with_ : str or list of str
|
||||
The parametrized imports.
|
||||
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
# Create test file, keeping it simple without conditionals
|
||||
(tmp_path / "a.py").touch()
|
||||
with caplog.at_level(logging.DEBUG):
|
||||
queue(
|
||||
config={"with": with_},
|
||||
kind="HTCondor",
|
||||
jobname="with_import_check",
|
||||
elements="sub-001",
|
||||
)
|
||||
assert "Copying" in caplog.text
|
||||
assert "Queue done" in caplog.text
|
||||
|
||||
# Check that file is copied
|
||||
assert Path(
|
||||
tmp_path / "junifer_jobs" / "with_import_check" / "a.py"
|
||||
).is_file()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"elements",
|
||||
[
|
||||
"sub-001",
|
||||
["sub-001"],
|
||||
["sub-001", "sub-002"],
|
||||
("sub-001", "ses-001"),
|
||||
[("sub-001", "ses-001")],
|
||||
[("sub-001", "ses-001"), ("sub-001", "ses-002")],
|
||||
],
|
||||
)
|
||||
def test_queue_with_elements(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
elements: Union[str, List[Union[str, Tuple[str]]], Tuple[str]],
|
||||
) -> None:
|
||||
|
Done. Done.
|
||||
"""Test queue with elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
elements : str of list of str
|
||||
The parametrized elements for the queue.
|
||||
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
with caplog.at_level(logging.INFO):
|
||||
queue(
|
||||
|
Can you add elements that are tuples? Like Can you add elements that are tuples? Like `[("sub-001", "ses-001"),("sub-001", "ses-002")]`
Done. Done.
|
||||
config={},
|
||||
kind="HTCondor",
|
||||
elements=elements,
|
||||
)
|
||||
assert "Queue done" in caplog.text
|
||||
|
||||
|
||||
def test_queue_without_elements(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Test queue without elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
with caplog.at_level(logging.INFO):
|
||||
queue(
|
||||
config={"datagrabber": datagrabber},
|
||||
kind="SLURM",
|
||||
)
|
||||
assert "Queue done" in caplog.text
|
||||
|
||||
|
||||
def test_queue_condor_invalid_python_env(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Test invalid Python environment check for HTCondor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
|
||||
"""
|
||||
with pytest.raises(ValueError, match="Unknown env kind"):
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
queue(
|
||||
config={"elements": "sub-001"},
|
||||
kind="HTCondor",
|
||||
env={"kind": "galaxy"},
|
||||
)
|
||||
|
||||
|
||||
def test_queue_condor_conda_python(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Test conda Python environment check for HTCondor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
with caplog.at_level(logging.INFO):
|
||||
queue(
|
||||
config={"elements": "sub-001"},
|
||||
kind="HTCondor",
|
||||
jobname="conda_env_check",
|
||||
env={"kind": "conda", "name": "conda-env"},
|
||||
)
|
||||
assert "Copying" in caplog.text
|
||||
assert (
|
||||
Path(
|
||||
tmp_path
|
||||
/ "junifer_jobs"
|
||||
/ "conda_env_check"
|
||||
/ "run_conda.sh"
|
||||
)
|
||||
.stat()
|
||||
.st_mode
|
||||
== 33261
|
||||
)
|
||||
|
||||
|
||||
def test_queue_condor_venv_python(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Test venv Python environment check for HTCondor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
with caplog.at_level(logging.INFO):
|
||||
queue(
|
||||
config={"elements": "sub-001"},
|
||||
kind="HTCondor",
|
||||
env={"kind": "venv", "name": "venv-env"},
|
||||
)
|
||||
# TODO: needs implementation for testing
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"elements, env, mem, cpus, disk, collect",
|
||||
[
|
||||
(
|
||||
["sub-001"],
|
||||
{"kind": "conda", "name": "conda-env"},
|
||||
"4G",
|
||||
4,
|
||||
"4G",
|
||||
True,
|
||||
),
|
||||
(
|
||||
["sub-001"],
|
||||
{"kind": "venv", "name": "venv-env"},
|
||||
"8G",
|
||||
8,
|
||||
"8G",
|
||||
False,
|
||||
),
|
||||
(["sub-001"], {"kind": "local"}, "12G", 12, "12G", True),
|
||||
],
|
||||
)
|
||||
def test_queue_condor_assets_generation(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
elements: str,
|
||||
env: str,
|
||||
mem: str,
|
||||
cpus: int,
|
||||
disk: str,
|
||||
collect: bool,
|
||||
) -> None:
|
||||
"""Test HTCondor generated assets.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
elements : str
|
||||
The parametrized element names.
|
||||
env : dict
|
||||
The parametrized env names.
|
||||
mem : str
|
||||
The parametrized memory size.
|
||||
cpus : int
|
||||
The parametrized CPUs.
|
||||
disk : str
|
||||
The parametrized disk size.
|
||||
collect : bool
|
||||
The parametrized collect option.
|
||||
|
||||
"""
|
||||
jobname = "condor_assets_gen_check"
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
with caplog.at_level(logging.INFO):
|
||||
queue(
|
||||
config={"elements": elements},
|
||||
kind="HTCondor",
|
||||
jobname=jobname,
|
||||
env=env,
|
||||
mem=mem,
|
||||
cpus=cpus,
|
||||
disk=disk,
|
||||
collect=collect,
|
||||
)
|
||||
|
||||
# Check log directory creation
|
||||
assert Path(tmp_path / "junifer_jobs" / jobname / "logs").is_dir()
|
||||
|
||||
run_submit_file_path = Path(
|
||||
tmp_path / "junifer_jobs" / jobname / f"run_{jobname}.submit"
|
||||
)
|
||||
# Check junifer run submit file
|
||||
assert run_submit_file_path.is_file()
|
||||
# Read run submit file to check if resources are correct
|
||||
with open(run_submit_file_path, "r") as f:
|
||||
for line in f.read().splitlines():
|
||||
if "request_cpus" in line:
|
||||
assert int(line.split("=")[1].strip()) == cpus
|
||||
if "request_memory" in line:
|
||||
assert line.split("=")[1].strip() == mem
|
||||
if "request_disk" in line:
|
||||
assert line.split("=")[1].strip() == disk
|
||||
|
||||
# Check junifer collect submit file
|
||||
assert Path(
|
||||
tmp_path
|
||||
/ "junifer_jobs"
|
||||
/ jobname
|
||||
/ f"collect_{jobname}.submit"
|
||||
).is_file()
|
||||
|
||||
dag_file_path = Path(
|
||||
tmp_path / "junifer_jobs" / jobname / f"{jobname}.dag"
|
||||
)
|
||||
# Check junifer dag file
|
||||
assert dag_file_path.is_file()
|
||||
# Read dag file to check if collect job is found
|
||||
element_count = 0
|
||||
has_collect_job = False
|
||||
with open(dag_file_path, "r") as f:
|
||||
for line in f.read().splitlines():
|
||||
if "JOB" in line:
|
||||
element_count += 1
|
||||
if "collect" in line:
|
||||
has_collect_job = True
|
||||
|
||||
if collect:
|
||||
assert len(elements) == element_count - 1
|
||||
assert has_collect_job is True
|
||||
else:
|
||||
assert len(elements) == element_count
|
||||
assert has_collect_job is False
|
||||
|
||||
# Check submit log
|
||||
assert (
|
||||
"HTCondor job files created, to submit the job, run"
|
||||
in caplog.text
|
||||
)
|
||||
|
||||
|
||||
def test_queue_condor_extra_preamble(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Test HTCondor extra preamble addition.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
|
||||
"""
|
||||
jobname = "condor_extra_preamble_check"
|
||||
extra_preamble = "FOO = BAR"
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
queue(
|
||||
config={"elements": ["sub-001"]},
|
||||
kind="HTCondor",
|
||||
jobname=jobname,
|
||||
extra_preamble=extra_preamble,
|
||||
)
|
||||
|
||||
# Check extra preamble in run submit file
|
||||
run_submit_file_path = Path(
|
||||
tmp_path / "junifer_jobs" / jobname / f"run_{jobname}.submit"
|
||||
)
|
||||
with open(run_submit_file_path, "r") as f:
|
||||
for line in f.read().splitlines():
|
||||
if "FOO" in line:
|
||||
assert line.strip() == extra_preamble
|
||||
|
||||
# Check extra preamble in collect submit file
|
||||
collect_submit_file_path = Path(
|
||||
tmp_path / "junifer_jobs" / jobname / f"collect_{jobname}.submit"
|
||||
)
|
||||
with open(collect_submit_file_path, "r") as f:
|
||||
for line in f.read().splitlines():
|
||||
if "FOO" in line:
|
||||
assert line.strip() == extra_preamble
|
||||
|
||||
|
||||
def test_queue_condor_submission_fail(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Test HTCondor job submission failure.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tmp_path : pathlib.Path
|
||||
The path to the test directory.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
The monkeypatch object.
|
||||
caplog : pytest.LogCaptureFixture
|
||||
The logcapturefixture object.
|
||||
|
||||
"""
|
||||
with pytest.raises(
|
||||
FileNotFoundError,
|
||||
match="No such file or directory: 'condor_submit_dag'",
|
||||
):
|
||||
with monkeypatch.context() as m:
|
||||
m.chdir(tmp_path)
|
||||
with caplog.at_level(logging.INFO):
|
||||
queue(
|
||||
config={"elements": ["sub-001"]},
|
||||
kind="HTCondor",
|
||||
jobname="condor_job_submission_fail",
|
||||
submit=True,
|
||||
)
|
||||
# Check submit log
|
||||
assert "Submitting HTCondor job" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="SLURM not installed on system.")
|
||||
|
|
|
|||
Loading…
Reference in a new issue
Can you assert that the file is present in the jobdir?