hpcflow 0.1.9__py3-none-any.whl → 0.2.0a271__py3-none-any.whl

hpcflow/sdk/submission/schedulers/__init__.py

@@ -0,0 +1,281 @@
+"""
+Job scheduler models.
+"""
+
+from __future__ import annotations
+from abc import ABC, abstractmethod
+import sys
+import time
+from typing import Generic, TypeVar, TYPE_CHECKING
+import warnings
+from typing_extensions import override
+from hpcflow.sdk.typing import hydrate
+from hpcflow.sdk.core.app_aware import AppAware
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+    from typing import Any, ClassVar
+    from ..shells import Shell
+    from ..jobscript import Jobscript
+    from ..enums import JobscriptElementState
+    from ..types import VersionInfo
+    from ...config.types import SchedulerConfigDescriptor
+    from ...core.element import ElementResources
+
+#: The type of a jobscript reference.
+JSRefType = TypeVar("JSRefType")
+
+
+@hydrate
+class Scheduler(ABC, Generic[JSRefType], AppAware):
+    """
+    Abstract base class for schedulers.
+
+    Note
+    ----
+    Do not make immediate subclasses of this class other than
+    :py:class:`DirectScheduler` and :py:class:`QueuedScheduler`;
+    subclass those two instead. Code (e.g., in :py:class:`Jobscript`)
+    assumes that this model is followed and does not check it.
+
+    Parameters
+    ----------
+    shebang_executable: list[str]
+        If specified, this will be used in the jobscript's shebang line instead of the
+        shell's `executable` and `executable_args` attributes.
+    """
+
+    # This would be in the docstring except it renders really wrongly!
+    # Type Parameters
+    # ---------------
+    # T
+    #     The type of a jobscript reference.
+
+    def __init__(self, shebang_executable: list[str] | None = None):
+        self.shebang_executable = shebang_executable
+
+    @property
+    def unique_properties(self) -> tuple[str, ...]:
+        """
+        Unique properties, for hashing.
+        """
+        return (self.__class__.__name__,)
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+        return self.__dict__ == other.__dict__
+
+    @abstractmethod
+    def process_resources(
+        self, resources: ElementResources, scheduler_config: SchedulerConfigDescriptor
+    ) -> None:
+        """
+        Perform scheduler-specific processing to the element resources.
+
+        Note
+        ----
+        This mutates `resources`.
+        """
+
+    def get_version_info(self) -> VersionInfo:
+        """
+        Get the version of the scheduler.
+        """
+        return {}
+
+    def parse_submission_output(self, stdout: str) -> str | None:
+        """
+        Parse the output from a submission to determine the submission ID.
+        """
+        return None
+
+    @staticmethod
+    def is_num_cores_supported(num_cores: int | None, core_range: Sequence[int]) -> bool:
+        """
+        Test whether particular number of cores is supported in given range of cores.
+        """
+        step = core_range[1] if core_range[1] is not None else 1
+        upper = core_range[2] + 1 if core_range[2] is not None else sys.maxsize
+        return num_cores in range(core_range[0], upper, step)
+
+    @abstractmethod
+    def get_submit_command(
+        self,
+        shell: Shell,
+        js_path: str,
+        deps: dict[Any, tuple[Any, ...]],
+    ) -> list[str]:
+        """
+        Get a command for submitting a jobscript.
+        """
+
+    @abstractmethod
+    def get_job_state_info(
+        self, *, js_refs: Sequence[JSRefType] | None = None
+    ) -> Mapping[str, JobscriptElementState | Mapping[int, JobscriptElementState]]:
+        """
+        Get the state of one or more jobscripts.
+        """
+
+    @abstractmethod
+    def wait_for_jobscripts(self, js_refs: list[JSRefType]) -> None:
+        """
+        Wait for one or more jobscripts to complete.
+        """
+
+    @abstractmethod
+    def cancel_jobs(
+        self,
+        js_refs: list[JSRefType],
+        jobscripts: list[Jobscript] | None = None,
+    ) -> None:
+        """
+        Cancel one or more jobscripts.
+        """
+
+    @abstractmethod
+    def get_std_out_err_filename(self, js_idx: int, *args, **kwargs) -> str:
+        """File name of combined standard output and error streams."""
+
+    @abstractmethod
+    def get_stdout_filename(self, js_idx: int, *args, **kwargs) -> str:
+        """File name of the standard output stream file."""
+
+    @abstractmethod
+    def get_stderr_filename(self, js_idx: int, *args, **kwargs) -> str:
+        """File name of the standard error stream file."""
+
+
+@hydrate
+class QueuedScheduler(Scheduler[str]):
+    """
+    Base class for schedulers that use a job submission system.
+
+    Parameters
+    ----------
+    directives: dict
+        Scheduler directives. Each item is written verbatim in the jobscript as a
+        scheduler directive, and is not processed in any way. If a value is `None`, the
+        key is considered a flag-like directive. If a value is a list, multiple directives
+        will be printed to the jobscript with the same key, but different values.
+    options: dict
+        Deprecated. Please use `directives` instead.
+    submit_cmd: str
+        The submission command, if overridden from default.
+    show_cmd: str
+        The show command, if overridden from default.
+    del_cmd: str
+        The delete command, if overridden from default.
+    js_cmd: str
+        The job script command, if overridden from default.
+    login_nodes_cmd: list[str]
+        The login nodes command, if overridden from default.
+    array_switch: str
+        The switch to enable array jobs, if overridden from default.
+    array_item_var: str
+        The variable for array items, if overridden from default.
+    """
+
+    #: Default command for logging into nodes.
+    DEFAULT_LOGIN_NODES_CMD: ClassVar[Sequence[str] | None] = None
+    #: Default pattern for matching the names of login nodes.
+    DEFAULT_LOGIN_NODE_MATCH: ClassVar[str] = "*login*"
+    #: Default command for submitting a job.
+    DEFAULT_SUBMIT_CMD: ClassVar[str]
+    #: Default command for listing current submitted jobs.
+    DEFAULT_SHOW_CMD: ClassVar[Sequence[str]]
+    #: Default command for deleting a job.
+    DEFAULT_DEL_CMD: ClassVar[str]
+    #: Default marker for job control metadata in a job script.
+    DEFAULT_JS_CMD: ClassVar[str]
+    #: Default switch for enabling array mode.
+    DEFAULT_ARRAY_SWITCH: ClassVar[str]
+    #: Default shell variable containin the current array index.
+    DEFAULT_ARRAY_ITEM_VAR: ClassVar[str]
+
+    def __init__(
+        self,
+        directives: dict | None = None,
+        options: dict | None = None,
+        submit_cmd: str | None = None,
+        show_cmd: Sequence[str] | None = None,
+        del_cmd: str | None = None,
+        js_cmd: str | None = None,
+        login_nodes_cmd: Sequence[str] | None = None,
+        array_switch: str | None = None,
+        array_item_var: str | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        if options:
+            warnings.warn(
+                f"{self.__class__.__name__!r}: Please use `directives` instead of "
+                f"`options`, which will be removed in a future release.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            directives = options
+
+        self.directives = directives or {}
+        self.submit_cmd: str = submit_cmd or self.DEFAULT_SUBMIT_CMD
+        self.show_cmd = show_cmd or self.DEFAULT_SHOW_CMD
+        self.del_cmd = del_cmd or self.DEFAULT_DEL_CMD
+        self.js_cmd = js_cmd or self.DEFAULT_JS_CMD
+        self.login_nodes_cmd = login_nodes_cmd or self.DEFAULT_LOGIN_NODES_CMD
+        self.array_switch = array_switch or self.DEFAULT_ARRAY_SWITCH
+        self.array_item_var = array_item_var or self.DEFAULT_ARRAY_ITEM_VAR
+
+    @property
+    def unique_properties(self) -> tuple[str, str, Any, Any]:
+        return (self.__class__.__name__, self.submit_cmd, self.show_cmd, self.del_cmd)
+
+    def format_switch(self, switch: str) -> str:
+        """
+        Format a particular switch to use the JS command.
+        """
+        return f"{self.js_cmd} {switch}"
+
+    def is_jobscript_active(self, job_ID: str) -> bool:
+        """Query if a jobscript is running/pending."""
+        return bool(self.get_job_state_info(js_refs=[job_ID]))
+
+    @override
+    def wait_for_jobscripts(self, js_refs: list[str]) -> None:
+        """
+        Wait for jobscripts to update their state.
+        """
+        while js_refs:
+            info: Mapping[str, Any] = self.get_job_state_info(js_refs=js_refs)
+            if not info:
+                break
+            js_refs = list(info)
+            time.sleep(2)
+
+    @abstractmethod
+    def format_directives(
+        self,
+        resources: ElementResources,
+        num_elements: int,
+        is_array: bool,
+        sub_idx: int,
+        js_idx: int,
+    ) -> str:
+        """
+        Render directives in a way that the scheduler can handle.
+        """
+
+    def get_std_out_err_filename(
+        self, js_idx: int, job_ID: str, array_idx: int | None = None
+    ):
+        """File name of combined standard output and error streams.
+
+        Notes
+        -----
+        We use the standard output stream filename format for the combined output and
+        error streams file.
+
+        """
+        return self.get_stdout_filename(js_idx=js_idx, job_ID=job_ID, array_idx=array_idx)

hpcflow/sdk/submission/schedulers/direct.py

@@ -0,0 +1,233 @@
+"""
+A direct job "scheduler" that just runs immediate subprocesses.
+"""
+
+from __future__ import annotations
+import shutil
+import signal
+from typing import overload, cast, TYPE_CHECKING
+from typing_extensions import override, TypeAlias
+import psutil
+
+from hpcflow.sdk.typing import hydrate
+from hpcflow.sdk.submission.enums import JobscriptElementState
+from hpcflow.sdk.submission.schedulers import Scheduler
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping, Sequence
+    from typing import Any, ClassVar
+    from ...config.types import SchedulerConfigDescriptor
+    from ..jobscript import Jobscript
+    from ..shells.base import Shell
+
+DirectRef: TypeAlias = "tuple[int, list[str]]"
+
+
+def _is_process_cmdline_equal(proc: psutil.Process, cmdline: list[str]) -> bool:
+    """Check if the `cmdline` of a psutil `Process` is equal to the specified
+    `cmdline`."""
+    try:
+        if proc.cmdline() == cmdline:
+            return True
+        else:
+            return False
+    except (psutil.NoSuchProcess, psutil.ZombieProcess):
+        # process no longer exists or, on unix, process has completed but still has a
+        # record
+        return False
+
+
+class DirectScheduler(Scheduler[DirectRef]):
+    """
+    A direct scheduler, that just runs jobs immediately as direct subprocesses.
+
+    The correct subclass (:py:class:`DirectPosix` or :py:class:`DirectWindows`) should
+    be used to create actual instances.
+
+    """
+
+    @classmethod
+    @override
+    def process_resources(
+        cls, resources, scheduler_config: SchedulerConfigDescriptor
+    ) -> None:
+        """Perform scheduler-specific processing to the element resources.
+
+        Note
+        ----
+        This mutates `resources`.
+        """
+        return
+
+    @override
+    def get_submit_command(
+        self,
+        shell: Shell,
+        js_path: str,
+        deps: dict[Any, tuple[Any, ...]],
+    ) -> list[str]:
+        """
+        Get the concrete submission command.
+        """
+        return shell.get_direct_submit_command(js_path)
+
+    @staticmethod
+    def __kill_processes(
+        procs: list[psutil.Process],
+        sig: signal.Signals = signal.SIGTERM,
+        timeout: float | None = None,
+        on_terminate: Callable[[psutil.Process], object] | None = None,
+    ):
+        all_procs: list[psutil.Process] = []
+        for process in procs:
+            all_procs.append(process)
+            all_procs.extend(process.children(recursive=True))
+
+        for process in all_procs:
+            try:
+                process.send_signal(sig)
+            except psutil.NoSuchProcess:
+                pass
+        _, alive = psutil.wait_procs(all_procs, timeout=timeout, callback=on_terminate)
+        for process in alive:
+            process.kill()
+
+    @staticmethod
+    def __get_jobscript_processes(js_refs: list[DirectRef]) -> list[psutil.Process]:
+        procs: list[psutil.Process] = []
+        for p_id, p_cmdline in js_refs:
+            try:
+                proc_i = psutil.Process(p_id)
+            except psutil.NoSuchProcess:
+                # process might have completed already
+                continue
+            if _is_process_cmdline_equal(proc_i, p_cmdline):
+                procs.append(proc_i)
+        return procs
+
+    @overload
+    @override
+    @classmethod
+    def wait_for_jobscripts(cls, js_refs: list[DirectRef]) -> None: ...
+
+    @overload
+    @classmethod
+    def wait_for_jobscripts(
+        cls,
+        js_refs: list[DirectRef],
+        *,
+        callback: Callable[[psutil.Process], None],
+    ) -> list[psutil.Process]: ...
+
+    @classmethod
+    def wait_for_jobscripts(
+        cls,
+        js_refs: list[DirectRef],
+        *,
+        callback: Callable[[psutil.Process], None] | None = None,
+    ) -> list[psutil.Process] | None:
+        """Wait until the specified jobscripts have completed."""
+        procs = cls.__get_jobscript_processes(js_refs)
+        (gone, alive) = psutil.wait_procs(procs, callback=callback)
+        assert not alive
+        return gone if callback else None
+
+    @override
+    def get_job_state_info(
+        self, *, js_refs: Sequence[DirectRef] | None = None
+    ) -> Mapping[str, JobscriptElementState]:
+        """Query the scheduler to get the states of all of this user's jobs, optionally
+        filtering by specified job IDs.
+
+        Jobs that are not in the scheduler's status output will not appear in the output
+        of this method."""
+        info: dict[str, JobscriptElementState] = {}
+        for p_id, p_cmdline in js_refs or ():
+            if self.is_jobscript_active(p_id, p_cmdline):
+                # as far as the "scheduler" is concerned, all elements are running:
+                info[str(p_id)] = JobscriptElementState.running
+
+        return info
+
+    @override
+    def cancel_jobs(
+        self,
+        js_refs: list[DirectRef],
+        jobscripts: list[Jobscript] | None = None,
+    ):
+        """
+        Cancel some jobs.
+        """
+
+        js_proc_id: dict[int, Jobscript]
+
+        def callback(proc: psutil.Process):
+            try:
+                js_proc_id[proc.pid]
+            except KeyError:
+                # child process of one of the jobscripts
+                self._app.submission_logger.debug(
+                    f"jobscript child process ({proc.pid}) killed"
+                )
+                return
+
+        procs = self.__get_jobscript_processes(js_refs)
+        self._app.submission_logger.info(
+            f"cancelling {self.__class__.__name__} jobscript processes: {procs}."
+        )
+        js_proc_id = {i.pid: jobscripts[idx] for idx, i in enumerate(procs) if jobscripts}
+        self.__kill_processes(procs, timeout=3, on_terminate=callback)
+        print(f"Cancelled {len(procs)} jobscript{'s' if len(procs) > 1 else ''}.")
+        self._app.submission_logger.info("jobscripts cancel command executed.")
+
+    def is_jobscript_active(self, process_ID: int, process_cmdline: list[str]):
+        """Query if a jobscript is running.
+
+        Note that a "running" jobscript might be waiting on upstream jobscripts to
+        complete.
+
+        """
+        try:
+            proc = psutil.Process(process_ID)
+        except psutil.NoSuchProcess:
+            return False
+        return _is_process_cmdline_equal(proc, process_cmdline)
+
+    def get_std_out_err_filename(self, js_idx: int, **kwargs) -> str:
+        """File name of combined standard output and error streams."""
+        return f"js_{js_idx}_std.log"
+
+    def get_stdout_filename(self, js_idx: int, **kwargs) -> str:
+        """File name of the standard output stream file."""
+        return f"js_{js_idx}_stdout.log"
+
+    def get_stderr_filename(self, js_idx: int, **kwargs) -> str:
+        """File name of the standard error stream file."""
+        return f"js_{js_idx}_stderr.log"
+
+
+@hydrate
+class DirectPosix(DirectScheduler):
+    """
+    A direct scheduler for POSIX systems.
+
+    """
+
+
+@hydrate
+class DirectWindows(DirectScheduler):
+    """
+    A direct scheduler for Windows.
+
+    """
+
+    @override
+    def get_submit_command(
+        self, shell: Shell, js_path: str, deps: dict[Any, tuple[Any, ...]]
+    ) -> list[str]:
+        cmd = super().get_submit_command(shell, js_path, deps)
+        # `Start-Process` (see `Jobscript._launch_direct_js_win`) seems to resolve the
+        # executable, which means the process's `cmdline` might look different to what we
+        # record; so let's resolve it ourselves:
+        cmd[0] = cast("str", shutil.which(cmd[0]))
+        return cmd