junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/api/queue_context/htcondor_adapter.py

@@ -5,15 +5,38 @@
 
 import shutil
 import textwrap
+from enum import StrEnum
 from pathlib import Path
-from typing import Optional
+from typing import Any
 
 from ...typing import Elements
-from ...utils import logger, make_executable, raise_error, run_ext_cmd
-from .queue_context_adapter import QueueContextAdapter
+from ...utils import make_executable, raise_error, run_ext_cmd
+from .queue_context_adapter import (
+    EnvKind,
+    EnvShell,
+    QueueContextAdapter,
+    QueueContextEnv,
+    logger,
+)
 
 
-__all__ = ["HTCondorAdapter"]
+__all__ = ["HTCondorAdapter", "HTCondorCollect"]
+
+
+class HTCondorCollect(StrEnum):
+    """Accepted HTCondor collect commands.
+
+    * ``"yes"``: Submit "collect" task and run even if some of the jobs
+        fail.
+    * ``"on_success_only"``: Submit "collect" task and run only if all jobs
+        succeed.
+    * ``"no"``: Do not submit "collect" task.
+
+    """
+
+    Yes = "yes"
+    No = "no"
+    OnSuccessOnly = "on_success_only"
 
 
 class HTCondorAdapter(QueueContextAdapter):
@@ -27,14 +50,14 @@ class HTCondorAdapter(QueueContextAdapter):
         The path to the job directory.
     yaml_config_path : pathlib.Path
         The path to the YAML config file.
-    elements : list of str or tuple
+    elements : ``Elements``
         Element(s) to process. Will be used to index the DataGrabber.
     pre_run : str or None, optional
-        Extra bash commands to source before the run (default None).
+        Extra shell commands to source before the run (default None).
     pre_collect : str or None, optional
-        Extra bash commands to source before the collect (default None).
-    env : dict, optional
-        The Python environment configuration. If None, will run without a
+        Extra shell commands to source before the collect (default None).
+    env : :class:`.QueueContextEnv` or None, optional
+        The environment configuration. If None, will run without a
         virtual environment of any kind (default None).
     verbose : str, optional
         The level of verbosity (default "info").
@@ -49,25 +72,12 @@ class HTCondorAdapter(QueueContextAdapter):
         The size of disk (HDD or SSD) to use (default "1G").
     extra_preamble : str or None, optional
         Extra commands to pass to HTCondor (default None).
-    collect : {"yes", "on_success_only", "no"}, optional
+    collect_task : :class:`.HTCondorCollect`, optional
         Whether to submit "collect" task for junifer (default "yes").
-        Valid options are:
-
-            * "yes": Submit "collect" task and run even if some of the jobs
-                fail.
-            * "on_success_only": Submit "collect" task and run only if all jobs
-                succeed.
-            * "no": Do not submit "collect" task.
-
     submit : bool, optional
         Whether to submit the jobs. In any case, .dag files will be created
         for submission (default False).
 
-    Raises
-    ------
-    ValueError
-        If ``collect`` is invalid or if ``env`` is invalid.
-
     See Also
     --------
     QueueContextAdapter :
@@ -77,144 +87,67 @@ class HTCondorAdapter(QueueContextAdapter):
 
     """
 
-    def __init__(
-        self,
-        job_name: str,
-        job_dir: Path,
-        yaml_config_path: Path,
-        elements: Elements,
-        pre_run: Optional[str] = None,
-        pre_collect: Optional[str] = None,
-        env: Optional[dict[str, str]] = None,
-        verbose: str = "info",
-        verbose_datalad: Optional[str] = None,
-        cpus: int = 1,
-        mem: str = "8G",
-        disk: str = "1G",
-        extra_preamble: Optional[str] = None,
-        collect: str = "yes",
-        submit: bool = False,
-    ) -> None:
-        """Initialize the class."""
-        self._job_name = job_name
-        self._job_dir = job_dir
-        self._yaml_config_path = yaml_config_path
-        self._elements = elements
-        self._pre_run = pre_run
-        self._pre_collect = pre_collect
-        self._check_env(env)
-        self._verbose = verbose
-        self._verbose_datalad = verbose_datalad
-        self._cpus = cpus
-        self._mem = mem
-        self._disk = disk
-        self._extra_preamble = extra_preamble
-        self._collect = self._check_collect(collect)
-        self._submit = submit
-
-        self._log_dir = self._job_dir / "logs"
-        self._pre_run_path = self._job_dir / "pre_run.sh"
-        self._pre_collect_path = self._job_dir / "pre_collect.sh"
-        self._submit_run_path = self._job_dir / f"run_{self._job_name}.submit"
-        self._submit_collect_path = (
-            self._job_dir / f"collect_{self._job_name}.submit"
-        )
-        self._dag_path = self._job_dir / f"{self._job_name}.dag"
-
-    def _check_env(self, env: Optional[dict[str, str]]) -> None:
-        """Check value of env parameter on init.
-
-        Parameters
-        ----------
-        env : dict or None
-            The value of env parameter.
-
-        Raises
-        ------
-        ValueError
-            If ``env.kind`` is invalid or
-            if ``env.shell`` is invalid.
-
-        """
-        # Set env related variables
-        if env is None:
-            env = {"kind": "local"}
-        # Check env kind
-        valid_env_kinds = ["conda", "venv", "local"]
-        if env["kind"] not in valid_env_kinds:
-            raise_error(
-                f"Invalid value for `env.kind`: {env['kind']}, "
-                f"must be one of {valid_env_kinds}"
+    job_name: str
+    job_dir: Path
+    yaml_config_path: Path
+    elements: Elements
+    pre_run: str | None = None
+    pre_collect: str | None = None
+    env: QueueContextEnv | None = None
+    verbose: str = "info"
+    verbose_datalad: str | None = None
+    cpus: int = 1
+    mem: str = "8G"
+    disk: str = "1G"
+    extra_preamble: str | None = None
+    collect_task: HTCondorCollect = HTCondorCollect.Yes
+    submit: bool = False
+
+    def model_post_init(self, context: Any):  # noqa: D102
+        if self.env is None:
+            self.env = QueueContextEnv(
+                kind=EnvKind.Local, shell=EnvShell.Bash, name=""
             )
+        if self.env["kind"] == EnvKind.Local:
+            # No virtual environment
+            self._executable = "junifer"
+            self._arguments = ""
         else:
-            # Check shell
-            shell = env.get("shell", "bash")
-            valid_shells = ["bash", "zsh"]
-            if shell not in valid_shells:
-                raise_error(
-                    f"Invalid value for `env.shell`: {shell}, "
-                    f"must be one of {valid_shells}"
-                )
-            self._shell = shell
-            # Set variables
-            if env["kind"] == "local":
-                # No virtual environment
-                self._executable = "junifer"
-                self._arguments = ""
-            else:
-                self._executable = f"run_{env['kind']}.{self._shell}"
-                self._arguments = f"{env['name']} junifer"
-                self._exec_path = self._job_dir / self._executable
-
-    def _check_collect(self, collect: str) -> str:
-        """Check value of collect parameter on init.
-
-        Parameters
-        ----------
-        collect : str
-            The value of collect parameter.
-
-        Returns
-        -------
-        str
-            The checked value of collect parameter.
-
-        Raises
-        ------
-        ValueError
-            If ``collect`` is invalid.
-
-        """
-        valid_options = ["yes", "no", "on_success_only"]
-        if collect not in valid_options:
-            raise_error(
-                f"Invalid value for `collect`: {collect}, "
-                f"must be one of {valid_options}"
-            )
-        else:
-            return collect
+            if self.env["name"] is None:
+                raise_error("`env.name` is required")
+            self._executable = f"run_{self.env['kind']}.{self.env['shell']}"
+            self._arguments = f"{self.env['name']} junifer"
+            self._exec_path = self.job_dir / self._executable
+        self._log_dir = self.job_dir / "logs"
+        self._pre_run_path = self.job_dir / "pre_run.sh"
+        self._pre_collect_path = self.job_dir / "pre_collect.sh"
+        self._submit_run_path = self.job_dir / f"run_{self.job_name}.submit"
+        self._submit_collect_path = (
+            self.job_dir / f"collect_{self.job_name}.submit"
+        )
+        self._dag_path = self.job_dir / f"{self.job_name}.dag"
 
-    def pre_run(self) -> str:
+    def pre_run_cmds(self) -> str:
         """Return pre-run commands."""
         fixed = (
-            f"#!/usr/bin/env {self._shell}\n\n"
+            f"#!/usr/bin/env {self.env['shell']}\n\n"
             "# This script is auto-generated by junifer.\n\n"
             "# Force datalad to run in non-interactive mode\n"
             "DATALAD_UI_INTERACTIVE=false\n"
         )
-        var = self._pre_run or ""
+        var = self.pre_run or ""
         return fixed + "\n" + var
 
-    def run(self) -> str:
+    def run_cmds(self) -> str:
         """Return run commands."""
-        verbose_args = f"--verbose {self._verbose} "
-        if self._verbose_datalad is not None:
+        verbose_args = f"--verbose {self.verbose} "
+        if self.verbose_datalad is not None:
             verbose_args = (
-                f"{verbose_args} --verbose-datalad {self._verbose_datalad} "
+                f"{verbose_args} --verbose-datalad {self.verbose_datalad} "
             )
         junifer_run_args = (
             "run "
-            f"{self._yaml_config_path.resolve()!s} "
+            f"{self.yaml_config_path.resolve()!s} "
             f"{verbose_args}"
             "--element $(element)"
         )
@@ -227,11 +160,11 @@ class HTCondorAdapter(QueueContextAdapter):
             "universe = vanilla\n"
             "getenv = True\n\n"
             "# Resources\n"
-            f"request_cpus = {self._cpus}\n"
-            f"request_memory = {self._mem}\n"
-            f"request_disk = {self._disk}\n\n"
+            f"request_cpus = {self.cpus}\n"
+            f"request_memory = {self.mem}\n"
+            f"request_disk = {self.disk}\n\n"
             "# Executable\n"
-            f"initial_dir = {self._job_dir.resolve()!s}\n"
+            f"initial_dir = {self.job_dir.resolve()!s}\n"
             f"executable = $(initial_dir)/{self._executable}\n"
             f"transfer_executable = False\n\n"
             f"arguments = {self._arguments} {junifer_run_args}\n\n"
@@ -240,31 +173,31 @@ class HTCondorAdapter(QueueContextAdapter):
             f"output = {log_dir_prefix}.out\n"
             f"error = {log_dir_prefix}.err\n"
         )
-        var = self._extra_preamble or ""
+        var = self.extra_preamble or ""
         return fixed + "\n" + var + "\n" + "queue"
 
-    def pre_collect(self) -> str:
+    def pre_collect_cmds(self) -> str:
         """Return pre-collect commands."""
         fixed = (
-            f"#!/usr/bin/env {self._shell}\n\n"
+            f"#!/usr/bin/env {self.env['shell']}\n\n"
             "# This script is auto-generated by junifer.\n"
         )
-        var = self._pre_collect or ""
+        var = self.pre_collect or ""
         # Add commands if collect="yes"
-        if self._collect == "yes":
+        if self.collect_task == "yes":
             var += 'if [ "${1}" == "4" ]; then\n    exit 1\nfi\n'
         return fixed + "\n" + var
 
-    def collect(self) -> str:
+    def collect_cmds(self) -> str:
         """Return collect commands."""
-        verbose_args = f"--verbose {self._verbose} "
-        if self._verbose_datalad is not None:
+        verbose_args = f"--verbose {self.verbose} "
+        if self.verbose_datalad is not None:
             verbose_args = (
-                f"{verbose_args} --verbose-datalad {self._verbose_datalad} "
+                f"{verbose_args} --verbose-datalad {self.verbose_datalad} "
             )
 
         junifer_collect_args = (
-            f"collect {self._yaml_config_path.resolve()!s} {verbose_args}"
+            f"collect {self.yaml_config_path.resolve()!s} {verbose_args}"
         )
         log_dir_prefix = f"{self._log_dir.resolve()!s}/junifer_collect"
         fixed = (
@@ -273,11 +206,11 @@ class HTCondorAdapter(QueueContextAdapter):
             "universe = vanilla\n"
             "getenv = True\n\n"
             "# Resources\n"
-            f"request_cpus = {self._cpus}\n"
-            f"request_memory = {self._mem}\n"
-            f"request_disk = {self._disk}\n\n"
+            f"request_cpus = {self.cpus}\n"
+            f"request_memory = {self.mem}\n"
+            f"request_disk = {self.disk}\n\n"
             "# Executable\n"
-            f"initial_dir = {self._job_dir.resolve()!s}\n"
+            f"initial_dir = {self.job_dir.resolve()!s}\n"
             f"executable = $(initial_dir)/{self._executable}\n"
             "transfer_executable = False\n\n"
             f"arguments = {self._arguments} {junifer_collect_args}\n\n"
@@ -286,13 +219,13 @@ class HTCondorAdapter(QueueContextAdapter):
             f"output = {log_dir_prefix}.out\n"
             f"error = {log_dir_prefix}.err\n"
         )
-        var = self._extra_preamble or ""
+        var = self.extra_preamble or ""
         return fixed + "\n" + var + "\n" + "queue"
 
-    def dag(self) -> str:
+    def dag_cmds(self) -> str:
         """Return HTCondor DAG commands."""
         fixed = ""
-        for idx, element in enumerate(self._elements):
+        for idx, element in enumerate(self.elements):
             # Stringify elements if tuple for operation
             str_element = (
                 ",".join(element) if isinstance(element, tuple) else element
@@ -307,15 +240,15 @@ class HTCondorAdapter(QueueContextAdapter):
                 f'log_element="{log_element}"\n\n'  # double quoted
             )
         var = ""
-        if self._collect == "yes":
+        if self.collect_task == "yes":
             var += (
                 f"FINAL collect {self._submit_collect_path}\n"
                 f"SCRIPT PRE collect {self._pre_collect_path.as_posix()} "
                 "$DAG_STATUS\n"
             )
-        elif self._collect == "on_success_only":
+        elif self.collect_task == "on_success_only":
             var += f"JOB collect {self._submit_collect_path}\nPARENT "
-            for idx, _ in enumerate(self._elements):
+            for idx, _ in enumerate(self.elements):
                 var += f"run{idx} "
             var += "CHILD collect\n"
 
@@ -326,7 +259,7 @@ class HTCondorAdapter(QueueContextAdapter):
         logger.info("Creating HTCondor job")
         # Create logs
         logger.info(
-            f"Creating logs directory under {self._job_dir.resolve()!s}"
+            f"Creating logs directory under {self.job_dir.resolve()!s}"
         )
         self._log_dir.mkdir(exist_ok=True, parents=True)
         # Copy executable if not local
@@ -341,46 +274,50 @@ class HTCondorAdapter(QueueContextAdapter):
             make_executable(self._exec_path)
         # Create pre run
         logger.info(
-            f"Writing {self._pre_run_path.name} to {self._job_dir.resolve()!s}"
+            f"Writing {self._pre_run_path.name} to {self.job_dir.resolve()!s}"
         )
         self._pre_run_path.touch()
-        self._pre_run_path.write_text(textwrap.dedent(self.pre_run()))
+        self._pre_run_path.write_text(textwrap.dedent(self.pre_run_cmds()))
         make_executable(self._pre_run_path)
         # Create run
         logger.debug(
             f"Writing {self._submit_run_path.name} to "
-            f"{self._job_dir.resolve()!s}"
+            f"{self.job_dir.resolve()!s}"
         )
         self._submit_run_path.touch()
-        self._submit_run_path.write_text(textwrap.dedent(self.run()))
+        self._submit_run_path.write_text(textwrap.dedent(self.run_cmds()))
         # Create pre collect
         logger.info(
             f"Writing {self._pre_collect_path.name} to "
-            f"{self._job_dir.resolve()!s}"
+            f"{self.job_dir.resolve()!s}"
         )
         self._pre_collect_path.touch()
-        self._pre_collect_path.write_text(textwrap.dedent(self.pre_collect()))
+        self._pre_collect_path.write_text(
+            textwrap.dedent(self.pre_collect_cmds())
+        )
         make_executable(self._pre_collect_path)
         # Create collect
         logger.debug(
             f"Writing {self._submit_collect_path.name} to "
-            f"{self._job_dir.resolve()!s}"
+            f"{self.job_dir.resolve()!s}"
         )
         self._submit_collect_path.touch()
-        self._submit_collect_path.write_text(textwrap.dedent(self.collect()))
+        self._submit_collect_path.write_text(
+            textwrap.dedent(self.collect_cmds())
+        )
         # Create DAG
         logger.debug(
-            f"Writing {self._dag_path.name} to {self._job_dir.resolve()!s}"
+            f"Writing {self._dag_path.name} to {self.job_dir.resolve()!s}"
         )
         self._dag_path.touch()
-        self._dag_path.write_text(textwrap.dedent(self.dag()))
+        self._dag_path.write_text(textwrap.dedent(self.dag_cmds()))
         # Submit if required
         condor_submit_dag_cmd = [
             "condor_submit_dag",
             "-include_env HOME",
             f"{self._dag_path.resolve()!s}",
         ]
-        if self._submit:
+        if self.submit:
             run_ext_cmd(name="condor_submit_dag", cmd=condor_submit_dag_cmd)
         else:
             logger.info(

junifer/api/queue_context/queue_context_adapter.py

@@ -3,24 +3,68 @@
 # Authors: Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
+import sys
+
+
+if sys.version_info < (3, 12):  # pragma: no cover
+    from typing_extensions import TypedDict
+else:
+    from typing import TypedDict
+
 from abc import ABC, abstractmethod
+from enum import StrEnum
+from typing import Required
+
+import structlog
+from pydantic import BaseModel, ConfigDict
 
 from ...utils import raise_error
 
 
-__all__ = ["QueueContextAdapter"]
+__all__ = ["EnvKind", "EnvShell", "QueueContextAdapter", "QueueContextEnv"]
+
+_log = structlog.get_logger("junifer")
+logger = _log.bind(pkg="api")
+
+
+class EnvKind(StrEnum):
+    """Accepted Python environment kind."""
+
+    Venv = "venv"
+    Conda = "conda"
+    Local = "local"
 
 
-class QueueContextAdapter(ABC):
+class EnvShell(StrEnum):
+    """Accepted environment shell."""
+
+    Bash = "bash"
+    Zsh = "zsh"
+
+
+class QueueContextEnv(TypedDict, total=False):
+    """Accepted environment configuration for queue context."""
+
+    kind: Required[EnvKind]
+    name: str
+    shell: Required[EnvShell]
+
+
+class QueueContextAdapter(BaseModel, ABC):
     """Abstract base class for queue context adapter.
 
-    For every interface that is required, one needs to provide a concrete
+    For every queue context, one needs to provide a concrete
     implementation of this abstract class.
 
     """
 
+    model_config = ConfigDict(
+        use_enum_values=True,
+        extra="forbid",
+    )
+
     @abstractmethod
-    def pre_run(self) -> str:
+    def pre_run_cmds(self) -> str:
         """Return pre-run commands."""
         raise_error(
             msg="Concrete classes need to implement pre_run()",
@@ -28,7 +72,7 @@ class QueueContextAdapter(ABC):
         )
 
     @abstractmethod
-    def run(self) -> str:
+    def run_cmds(self) -> str:
         """Return run commands."""
         raise_error(
             msg="Concrete classes need to implement run()",
@@ -36,7 +80,7 @@ class QueueContextAdapter(ABC):
         )
 
     @abstractmethod
-    def pre_collect(self) -> str:
+    def pre_collect_cmds(self) -> str:
         """Return pre-collect commands."""
         raise_error(
             msg="Concrete classes need to implement pre_collect()",
@@ -44,7 +88,7 @@ class QueueContextAdapter(ABC):
         )
 
     @abstractmethod
-    def collect(self) -> str:
+    def collect_cmds(self) -> str:
         """Return collect commands."""
         raise_error(
             msg="Concrete classes need to implement collect()",

junifer/api/queue_context/tests/test_gnu_parallel_local_adapter.py

@@ -5,37 +5,12 @@
 
 import logging
 from pathlib import Path
-from typing import Optional, Union
 
 import pytest
 
 from junifer.api.queue_context import GnuParallelLocalAdapter
 
 
-def test_GnuParallelLocalAdapter_env_kind_error() -> None:
-    """Test error for invalid env kind."""
-    with pytest.raises(ValueError, match=r"Invalid value for `env.kind`"):
-        GnuParallelLocalAdapter(
-            job_name="check_env_kind",
-            job_dir=Path("."),
-            yaml_config_path=Path("."),
-            elements=["sub01"],
-            env={"kind": "jambalaya"},
-        )
-
-
-def test_GnuParallelLocalAdapter_env_shell_error() -> None:
-    """Test error for invalid env shell."""
-    with pytest.raises(ValueError, match=r"Invalid value for `env.shell`"):
-        GnuParallelLocalAdapter(
-            job_name="check_env_shell",
-            job_dir=Path("."),
-            yaml_config_path=Path("."),
-            elements=["sub01"],
-            env={"kind": "conda", "shell": "fish"},
-        )
-
-
 @pytest.mark.parametrize(
     "elements, expected_text",
     [
@@ -44,7 +19,7 @@ def test_GnuParallelLocalAdapter_env_shell_error() -> None:
     ],
 )
 def test_GnuParallelLocalAdapter_elements(
-    elements: list[Union[str, tuple]],
+    elements: list[str | tuple],
     expected_text: str,
 ) -> None:
     """Test GnuParallelLocalAdapter elements().
@@ -63,7 +38,7 @@ def test_GnuParallelLocalAdapter_elements(
         yaml_config_path=Path("."),
         elements=elements,
     )
-    assert expected_text in adapter.elements()
+    assert expected_text in adapter.elements_to_run()
 
 
 @pytest.mark.parametrize(
@@ -76,7 +51,7 @@ def test_GnuParallelLocalAdapter_elements(
     ],
 )
 def test_GnuParallelLocalAdapter_pre_run(
-    pre_run: Optional[str],
+    pre_run: str | None,
     expected_text: str,
     shell: str,
 ) -> None:
@@ -100,8 +75,8 @@ def test_GnuParallelLocalAdapter_pre_run(
         env={"kind": "conda", "name": "junifer", "shell": shell},
         pre_run=pre_run,
     )
-    assert shell in adapter.pre_run()
-    assert expected_text in adapter.pre_run()
+    assert shell in adapter.pre_run_cmds()
+    assert expected_text in adapter.pre_run_cmds()
 
 
 @pytest.mark.parametrize(
@@ -114,7 +89,7 @@ def test_GnuParallelLocalAdapter_pre_run(
     ],
 )
 def test_GnuParallelLocalAdapter_pre_collect(
-    pre_collect: Optional[str],
+    pre_collect: str | None,
     expected_text: str,
     shell: str,
 ) -> None:
@@ -138,8 +113,8 @@ def test_GnuParallelLocalAdapter_pre_collect(
         env={"kind": "venv", "name": "junifer", "shell": shell},
         pre_collect=pre_collect,
     )
-    assert shell in adapter.pre_collect()
-    assert expected_text in adapter.pre_collect()
+    assert shell in adapter.pre_collect_cmds()
+    assert expected_text in adapter.pre_collect_cmds()
 
 
 def test_GnuParallelLocalAdapter_run() -> None:
@@ -150,7 +125,7 @@ def test_GnuParallelLocalAdapter_run() -> None:
         yaml_config_path=Path("."),
         elements=["sub01"],
     )
-    assert "run" in adapter.run()
+    assert "run" in adapter.run_cmds()
 
 
 def test_GnuParallelLocalAdapter_collect() -> None:
@@ -161,7 +136,7 @@ def test_GnuParallelLocalAdapter_collect() -> None:
         yaml_config_path=Path("."),
         elements=["sub01"],
     )
-    assert "collect" in adapter.collect()
+    assert "collect" in adapter.collect_cmds()
 
 
 @pytest.mark.parametrize(