hpcflow-new2 0.2.0a188__py3-none-any.whl → 0.2.0a190__py3-none-any.whl

hpcflow/sdk/submission/schedulers/__init__.py

@@ -2,18 +2,36 @@
 Job scheduler models.
 """
 
-from pathlib import Path
+from __future__ import annotations
+from abc import ABC, abstractmethod
 import sys
 import time
-from typing import Any, List, Tuple
-
-
-class NullScheduler:
+from typing import Generic, TypeVar, TYPE_CHECKING
+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.
 
-    Keyword Args
-    ------------
+    Parameters
+    ----------
     shell_args: str
         Arguments to pass to the shell. Pre-quoted.
     shebang_args: str
@@ -22,48 +40,65 @@ class NullScheduler:
         Options to the scheduler.
     """
 
+    # This would be in the docstring except it renders really wrongly!
+    # Type Parameters
+    # ---------------
+    # T
+    #     The type of a jobscript reference.
+
     #: Default value for arguments to the shell.
-    DEFAULT_SHELL_ARGS = ""
+    DEFAULT_SHELL_ARGS: ClassVar[str] = ""
     #: Default value for arguments on the shebang line.
-    DEFAULT_SHEBANG_ARGS = ""
+    DEFAULT_SHEBANG_ARGS: ClassVar[str] = ""
 
     def __init__(
         self,
-        shell_args=None,
-        shebang_args=None,
-        options=None,
+        shell_args: str | None = None,
+        shebang_args: str | None = None,
+        options: dict | None = None,
     ):
         self.shebang_args = shebang_args or self.DEFAULT_SHEBANG_ARGS
         self.shell_args = shell_args or self.DEFAULT_SHELL_ARGS
         self.options = options or {}
 
     @property
-    def unique_properties(self):
+    def unique_properties(self) -> tuple[str, ...]:
         """
         Unique properties, for hashing.
         """
         return (self.__class__.__name__,)
 
-    def __eq__(self, other) -> bool:
-        if type(self) != type(other):
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, self.__class__):
             return False
-        else:
-            return self.__dict__ == other.__dict__
+        return self.__dict__ == other.__dict__
+
+    @abstractmethod
+    def process_resources(
+        self, resources: ElementResources, scheduler_config: SchedulerConfigDescriptor
+    ) -> None:
+        """
+        Perform scheduler-specific processing to the element resources.
 
-    def get_version_info(self):
+        Note
+        ----
+        This mutates `resources`.
+        """
+
+    def get_version_info(self) -> VersionInfo:
         """
         Get the version of the scheduler.
         """
         return {}
 
-    def parse_submission_output(self, stdout: str) -> None:
+    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, core_range: List[int]):
+    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.
         """
@@ -71,8 +106,45 @@ class NullScheduler:
         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.
+        """
 
-class Scheduler(NullScheduler):
+    @abstractmethod
+    def get_job_state_info(
+        self, *, js_refs: Sequence[JSRefType] | None = None, num_js_elements: int = 0
+    ) -> Mapping[str, Mapping[int | None, 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,
+        num_js_elements: int = 0,  # Ignored!
+    ) -> None:
+        """
+        Cancel one or more jobscripts.
+        """
+
+
+@hydrate
+class QueuedScheduler(Scheduler[str]):
     """
     Base class for schedulers that use a job submission system.
 
@@ -86,7 +158,7 @@ class Scheduler(NullScheduler):
         The delete command, if overridden from default.
     js_cmd: str
         The job script command, if overridden from default.
-    login_nodes_cmd: str
+    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.
@@ -95,25 +167,37 @@ class Scheduler(NullScheduler):
     """
 
     #: Default command for logging into nodes.
-    DEFAULT_LOGIN_NODES_CMD = None
+    DEFAULT_LOGIN_NODES_CMD: ClassVar[Sequence[str] | None] = None
     #: Default pattern for matching the names of login nodes.
-    DEFAULT_LOGIN_NODE_MATCH = "*login*"
+    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,
-        submit_cmd=None,
-        show_cmd=None,
-        del_cmd=None,
-        js_cmd=None,
-        login_nodes_cmd=None,
-        array_switch=None,
-        array_item_var=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)
 
-        self.submit_cmd = submit_cmd or self.DEFAULT_SUBMIT_CMD
+        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
@@ -122,27 +206,36 @@ class Scheduler(NullScheduler):
         self.array_item_var = array_item_var or self.DEFAULT_ARRAY_ITEM_VAR
 
     @property
-    def unique_properties(self):
+    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):
+    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):
+    def is_jobscript_active(self, job_ID: str) -> bool:
         """Query if a jobscript is running/pending."""
-        return bool(self.get_job_state_info([job_ID]))
+        return bool(self.get_job_state_info(js_refs=[job_ID]))
 
-    def wait_for_jobscripts(self, js_refs: List[Any]) -> None:
+    @override
+    def wait_for_jobscripts(self, js_refs: list[str]) -> None:
         """
         Wait for jobscripts to update their state.
         """
         while js_refs:
-            info = self.get_job_state_info(js_refs)
+            info: Mapping[str, Any] = self.get_job_state_info(js_refs=js_refs)
             print(info)
             if not info:
                 break
-            js_refs = list(info.keys())
+            js_refs = list(info)
             time.sleep(2)
+
+    @abstractmethod
+    def format_options(
+        self, resources: ElementResources, num_elements: int, is_array: bool, sub_idx: int
+    ) -> str:
+        """
+        Render options in a way that the scheduler can handle.
+        """

hpcflow/sdk/submission/schedulers/direct.py

@@ -2,19 +2,28 @@
 A direct job "scheduler" that just runs immediate subprocesses.
 """
 
-from pathlib import Path
+from __future__ import annotations
 import shutil
 import signal
-from typing import Callable, Dict, List, Optional, Tuple
-
+from typing import overload, cast, TYPE_CHECKING
+from typing_extensions import override, TypeAlias
 import psutil
-from hpcflow.sdk.submission.jobscript_info import JobscriptElementState
 
-from hpcflow.sdk.submission.schedulers import NullScheduler
-from hpcflow.sdk.submission.shells.base import Shell
+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]]"
 
 
-class DirectScheduler(NullScheduler):
+class DirectScheduler(Scheduler[DirectRef]):
     """
     A direct scheduler, that just runs jobs immediately as direct subprocesses.
 
@@ -31,53 +40,55 @@ class DirectScheduler(NullScheduler):
         Options to the jobscript command.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
     @classmethod
-    def process_resources(cls, resources, scheduler_config: Dict) -> None:
+    @override
+    def process_resources(
+        cls, resources, scheduler_config: SchedulerConfigDescriptor
+    ) -> None:
         """Perform scheduler-specific processing to the element resources.
 
-        Note: this mutates `resources`.
-
+        Note
+        ----
+        This mutates `resources`.
         """
         return
 
+    @override
     def get_submit_command(
         self,
         shell: Shell,
         js_path: str,
-        deps: List[Tuple],
-    ) -> List[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.SIGTERM,
-        timeout=None,
-        on_terminate=None,
+    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 = []
-        for i in procs:
-            all_procs.append(i)
-            all_procs.extend(i.children(recursive=True))
+        all_procs: list[psutil.Process] = []
+        for process in procs:
+            all_procs.append(process)
+            all_procs.extend(process.children(recursive=True))
 
-        for i in all_procs:
+        for process in all_procs:
             try:
-                i.send_signal(sig)
+                process.send_signal(sig)
             except psutil.NoSuchProcess:
                 pass
-        gone, alive = psutil.wait_procs(all_procs, timeout=timeout, callback=on_terminate)
-        for p in alive:
-            p.kill()
+        _, 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[Tuple[int, List[str]]]):
-        procs = []
+    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)
@@ -89,71 +100,94 @@ class DirectScheduler(NullScheduler):
                 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[Tuple[int, List[str]]],
-        callback: Optional[Callable] = None,
-    ) -> None:
+        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)
+        procs = cls.__get_jobscript_processes(js_refs)
         (gone, alive) = psutil.wait_procs(procs, callback=callback)
         assert not alive
-        return gone
+        return gone if callback else None
 
+    @override
     def get_job_state_info(
         self,
-        js_refs: List[Tuple[int, List[str]]],
-        num_js_elements: int,
-    ) -> Dict[int, Dict[int, JobscriptElementState]]:
+        *,
+        js_refs: Sequence[DirectRef] | None = None,
+        num_js_elements: int = 0,
+    ) -> Mapping[str, Mapping[int | None, 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 = {}
-        for p_id, p_cmdline in js_refs:
-            is_active = self.is_jobscript_active(p_id, p_cmdline)
-            if is_active:
+        info: dict[str, Mapping[int | None, 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[p_id] = {
+                info[str(p_id)] = {
                     i: JobscriptElementState.running for i in range(num_js_elements)
                 }
 
         return info
 
+    @override
     def cancel_jobs(
         self,
-        js_refs: List[Tuple[int, List[str]]],
-        jobscripts: List = None,
+        js_refs: list[DirectRef],
+        jobscripts: list[Jobscript] | None = None,
+        num_js_elements: int = 0,  # Ignored!
     ):
         """
         Cancel some jobs.
         """
 
-        def callback(proc):
+        js_proc_id: dict[int, Jobscript]
+
+        def callback(proc: psutil.Process):
             try:
                 js = js_proc_id[proc.pid]
             except KeyError:
                 # child process of one of the jobscripts
-                self.app.submission_logger.debug(
+                self._app.submission_logger.debug(
                     f"jobscript child process ({proc.pid}) killed"
                 )
                 return
+            assert hasattr(proc, "returncode")
             print(
                 f"Jobscript {js.index} from submission {js.submission.index} "
                 f"terminated (user-initiated cancel) with exit code {proc.returncode}."
             )
 
-        procs = self._get_jobscript_processes(js_refs)
-        self.app.submission_logger.info(
+        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)}
-        self._kill_processes(procs, timeout=3, on_terminate=callback)
-        self.app.submission_logger.info(f"jobscripts cancel command executed.")
+        js_proc_id = {i.pid: jobscripts[idx] for idx, i in enumerate(procs) if jobscripts}
+        self.__kill_processes(procs, timeout=3, on_terminate=callback)
+        self._app.submission_logger.info("jobscripts cancel command executed.")
 
-    def is_jobscript_active(self, process_ID: int, process_cmdline: List[str]):
+    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
@@ -165,12 +199,10 @@ class DirectScheduler(NullScheduler):
         except psutil.NoSuchProcess:
             return False
 
-        if proc.cmdline() == process_cmdline:
-            return True
-        else:
-            return False
+        return proc.cmdline() == process_cmdline
 
 
+@hydrate
 class DirectPosix(DirectScheduler):
     """
     A direct scheduler for POSIX systems.
@@ -185,14 +217,11 @@ class DirectPosix(DirectScheduler):
         Options to the jobscript command.
     """
 
-    _app_attr = "app"
     #: Default shell.
-    DEFAULT_SHELL_EXECUTABLE = "/bin/bash"
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    DEFAULT_SHELL_EXECUTABLE: ClassVar[str] = "/bin/bash"
 
 
+@hydrate
 class DirectWindows(DirectScheduler):
     """
     A direct scheduler for Windows.
@@ -205,19 +234,16 @@ class DirectWindows(DirectScheduler):
         Options to the jobscript command.
     """
 
-    _app_attr = "app"
     #: Default shell.
-    DEFAULT_SHELL_EXECUTABLE = "powershell.exe"
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    DEFAULT_SHELL_EXECUTABLE: ClassVar[str] = "powershell.exe"
 
+    @override
     def get_submit_command(
-        self, shell: Shell, js_path: str, deps: List[Tuple]
-    ) -> List[str]:
+        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] = shutil.which(cmd[0])
+        cmd[0] = cast("str", shutil.which(cmd[0]))
         return cmd