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

hpcflow/sdk/submission/schedulers/sge.py

@@ -2,22 +2,34 @@
 An interface to SGE.
 """
 
-from pathlib import Path
+from __future__ import annotations
+from collections.abc import Sequence
 import re
-from typing import Dict, List, Tuple
+from typing import TYPE_CHECKING
+from typing_extensions import override
+from hpcflow.sdk.typing import hydrate
 from hpcflow.sdk.core.errors import (
     IncompatibleSGEPEError,
     NoCompatibleSGEPEError,
     UnknownSGEPEError,
 )
 from hpcflow.sdk.log import TimeIt
-from hpcflow.sdk.submission.jobscript_info import JobscriptElementState
-from hpcflow.sdk.submission.schedulers import Scheduler
+from hpcflow.sdk.submission.enums import JobscriptElementState
+from hpcflow.sdk.submission.schedulers import QueuedScheduler
 from hpcflow.sdk.submission.schedulers.utils import run_cmd
-from hpcflow.sdk.submission.shells.base import Shell
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Mapping
+    from typing import Any, ClassVar
+    from ...config.types import SchedulerConfigDescriptor
+    from ...core.element import ElementResources
+    from ..jobscript import Jobscript
+    from ..types import VersionInfo
+    from ..shells.base import Shell
 
-class SGEPosix(Scheduler):
+
+@hydrate
+class SGEPosix(QueuedScheduler):
     """
     A scheduler that uses SGE.
 
@@ -43,36 +55,34 @@ class SGEPosix(Scheduler):
 
     """
 
-    _app_attr = "app"
-
     #: Default args for shebang line.
-    DEFAULT_SHEBANG_ARGS = ""
+    DEFAULT_SHEBANG_ARGS: ClassVar[str] = ""
     #: Default submission command.
-    DEFAULT_SUBMIT_CMD = "qsub"
+    DEFAULT_SUBMIT_CMD: ClassVar[str] = "qsub"
     #: Default command to show the queue state.
-    DEFAULT_SHOW_CMD = ["qstat"]
+    DEFAULT_SHOW_CMD: ClassVar[Sequence[str]] = ("qstat",)
     #: Default cancel command.
-    DEFAULT_DEL_CMD = "qdel"
+    DEFAULT_DEL_CMD: ClassVar[str] = "qdel"
     #: Default job control directive prefix.
-    DEFAULT_JS_CMD = "#$"
+    DEFAULT_JS_CMD: ClassVar[str] = "#$"
     #: Default prefix to enable array processing.
-    DEFAULT_ARRAY_SWITCH = "-t"
+    DEFAULT_ARRAY_SWITCH: ClassVar[str] = "-t"
     #: Default shell variable with array ID.
-    DEFAULT_ARRAY_ITEM_VAR = "SGE_TASK_ID"
+    DEFAULT_ARRAY_ITEM_VAR: ClassVar[str] = "SGE_TASK_ID"
     #: Default switch to control CWD.
-    DEFAULT_CWD_SWITCH = "-cwd"
+    DEFAULT_CWD_SWITCH: ClassVar[str] = "-cwd"
     #: Default command to get the login nodes.
-    DEFAULT_LOGIN_NODES_CMD = ["qconf", "-sh"]
+    DEFAULT_LOGIN_NODES_CMD: ClassVar[Sequence[str]] = ("qconf", "-sh")
 
     #: Maps scheduler state codes to :py:class:`JobscriptElementState` values.
-    state_lookup = {
+    state_lookup: ClassVar[Mapping[str, JobscriptElementState]] = {
         "qw": JobscriptElementState.pending,
         "hq": JobscriptElementState.waiting,
         "hR": JobscriptElementState.waiting,
         "r": JobscriptElementState.running,
         "t": JobscriptElementState.running,
         "Rr": JobscriptElementState.running,
-        "Rt": JobscriptElementState.running,
+        # "Rt": JobscriptElementState.running,
         "s": JobscriptElementState.errored,
         "ts": JobscriptElementState.errored,
         "S": JobscriptElementState.errored,
@@ -93,17 +103,22 @@ class SGEPosix(Scheduler):
         "dT": JobscriptElementState.cancelled,
     }
 
-    def __init__(self, cwd_switch=None, *args, **kwargs):
+    def __init__(self, cwd_switch: str | None = None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.cwd_switch = cwd_switch or self.DEFAULT_CWD_SWITCH
 
     @classmethod
+    @override
     @TimeIt.decorator
-    def process_resources(cls, resources, scheduler_config: Dict) -> None:
-        """Perform scheduler-specific processing to the element resources.
-
-        Note: this mutates `resources`.
+    def process_resources(
+        cls, resources: ElementResources, scheduler_config: SchedulerConfigDescriptor
+    ) -> None:
+        """
+        Perform scheduler-specific processing to the element resources.
 
+        Note
+        ----
+        This mutates `resources`.
         """
         if resources.num_nodes is not None:
             raise ValueError(
@@ -116,7 +131,7 @@ class SGEPosix(Scheduler):
         if resources.SGE_parallel_env is not None:
             # check user-specified `parallel_env` is valid and compatible with
             # `num_cores`:
-            if resources.num_cores > 1:
+            if resources.num_cores and resources.num_cores > 1:
                 raise ValueError(
                     f"An SGE parallel environment should not be specified if `num_cores` "
                     f"is 1 (`SGE_parallel_env` was specified as "
@@ -126,105 +141,91 @@ class SGEPosix(Scheduler):
             try:
                 env = para_envs[resources.SGE_parallel_env]
             except KeyError:
-                raise UnknownSGEPEError(
-                    f"The SGE parallel environment {resources.SGE_parallel_env!r} is not "
-                    f"specified in the configuration. Specified parallel environments "
-                    f"are {list(para_envs.keys())!r}."
-                )
+                raise UnknownSGEPEError(resources.SGE_parallel_env, para_envs)
             if not cls.is_num_cores_supported(resources.num_cores, env["num_cores"]):
                 raise IncompatibleSGEPEError(
-                    f"The SGE parallel environment {resources.SGE_parallel_env!r} is not "
-                    f"compatible with the number of cores requested: "
-                    f"{resources.num_cores!r}."
+                    resources.SGE_parallel_env, resources.num_cores
                 )
         else:
             # find the first compatible PE:
-            pe_match = -1  # pe_name might be `None`
             for pe_name, pe_info in para_envs.items():
                 if cls.is_num_cores_supported(resources.num_cores, pe_info["num_cores"]):
-                    pe_match = pe_name
+                    resources.SGE_parallel_env = pe_name
                     break
-            if pe_match != -1:
-                resources.SGE_parallel_env = pe_name
             else:
-                raise NoCompatibleSGEPEError(
-                    f"No compatible SGE parallel environment could be found for the "
-                    f"specified `num_cores` ({resources.num_cores!r})."
-                )
+                raise NoCompatibleSGEPEError(resources.num_cores)
 
-    def get_login_nodes(self):
+    def get_login_nodes(self) -> list[str]:
         """Return a list of hostnames of login/administrative nodes as reported by the
         scheduler."""
-        stdout, stderr = run_cmd(self.login_nodes_cmd)
+        get_login = self.login_nodes_cmd
+        assert get_login is not None and len(get_login) >= 1
+        stdout, stderr = run_cmd(get_login)
         if stderr:
             print(stderr)
-        nodes = stdout.strip().split("\n")
-        return nodes
-
-    def _format_core_request_lines(self, resources):
-        lns = []
-        if resources.num_cores > 1:
-            lns.append(
-                f"{self.js_cmd} -pe {resources.SGE_parallel_env} {resources.num_cores}"
-            )
+        return stdout.strip().split("\n")
+
+    def __format_core_request_lines(self, resources: ElementResources) -> Iterator[str]:
+        if resources.num_cores and resources.num_cores > 1:
+            yield f"{self.js_cmd} -pe {resources.SGE_parallel_env} {resources.num_cores}"
         if resources.max_array_items:
-            lns.append(f"{self.js_cmd} -tc {resources.max_array_items}")
-        return lns
+            yield f"{self.js_cmd} -tc {resources.max_array_items}"
 
-    def _format_array_request(self, num_elements):
+    def __format_array_request(self, num_elements: int) -> str:
         return f"{self.js_cmd} {self.array_switch} 1-{num_elements}"
 
-    def _format_std_stream_file_option_lines(self, is_array, sub_idx):
+    def __format_std_stream_file_option_lines(
+        self, is_array: bool, sub_idx: int
+    ) -> Iterator[str]:
         # note: we can't modify the file names
-        base = f"./artifacts/submissions/{sub_idx}"
-        return [
-            f"{self.js_cmd} -o {base}",
-            f"{self.js_cmd} -e {base}",
-        ]
+        yield f"{self.js_cmd} -o ./artifacts/submissions/{sub_idx}"
+        yield f"{self.js_cmd} -e ./artifacts/submissions/{sub_idx}"
 
-    def format_options(self, resources, num_elements, is_array, sub_idx):
+    @override
+    def format_options(
+        self, resources: ElementResources, num_elements: int, is_array: bool, sub_idx: int
+    ) -> str:
         """
         Format the options to the jobscript command.
         """
-        opts = []
+        opts: list[str] = []
         opts.append(self.format_switch(self.cwd_switch))
-        opts.extend(self._format_core_request_lines(resources))
+        opts.extend(self.__format_core_request_lines(resources))
         if is_array:
-            opts.append(self._format_array_request(num_elements))
+            opts.append(self.__format_array_request(num_elements))
 
-        opts.extend(self._format_std_stream_file_option_lines(is_array, sub_idx))
+        opts.extend(self.__format_std_stream_file_option_lines(is_array, sub_idx))
 
         for opt_k, opt_v in self.options.items():
-            if isinstance(opt_v, list):
-                for i in opt_v:
-                    opts.append(f"{self.js_cmd} {opt_k} {i}")
+            if opt_v is None:
+                opts.append(f"{self.js_cmd} {opt_k}")
+            elif isinstance(opt_v, list):
+                opts.extend(f"{self.js_cmd} {opt_k} {i}" for i in opt_v)
             elif opt_v:
                 opts.append(f"{self.js_cmd} {opt_k} {opt_v}")
-            elif opt_v is None:
-                opts.append(f"{self.js_cmd} {opt_k}")
 
         return "\n".join(opts) + "\n"
 
+    @override
     @TimeIt.decorator
-    def get_version_info(self):
-        vers_cmd = self.show_cmd + ["-help"]
-        stdout, stderr = run_cmd(vers_cmd)
+    def get_version_info(self) -> VersionInfo:
+        stdout, stderr = run_cmd([*self.show_cmd, "-help"])
         if stderr:
             print(stderr)
-        version_str = stdout.split("\n")[0].strip()
-        name, version = version_str.split()
-        out = {
+        first_line, *_ = stdout.split("\n")
+        name, version, *_ = first_line.strip().split()
+        return {
             "scheduler_name": name,
             "scheduler_version": version,
         }
-        return out
 
+    @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 command to use to submit a job to the scheduler.
 
@@ -234,8 +235,8 @@ class SGEPosix(Scheduler):
         """
         cmd = [self.submit_cmd, "-terse"]
 
-        dep_job_IDs = []
-        dep_job_IDs_arr = []
+        dep_job_IDs: list[str] = []
+        dep_job_IDs_arr: list[str] = []
         for job_ID, is_array_dep in deps.values():
             if is_array_dep:  # array dependency
                 dep_job_IDs_arr.append(str(job_ID))
@@ -253,60 +254,59 @@ class SGEPosix(Scheduler):
         cmd.append(js_path)
         return cmd
 
+    __SGE_JOB_ID_RE: ClassVar[re.Pattern] = re.compile(r"^\d+")
+
     def parse_submission_output(self, stdout: str) -> str:
         """Extract scheduler reference for a newly submitted jobscript"""
-        match = re.search(r"^\d+", stdout)
-        if match:
-            job_ID = match.group()
-        else:
+        if not (match := self.__SGE_JOB_ID_RE.search(stdout)):
             raise RuntimeError(f"Could not parse Job ID from scheduler output {stdout!r}")
-        return job_ID
+        return match.group()
 
-    def get_job_statuses(self):
+    def get_job_statuses(
+        self,
+    ) -> Mapping[str, Mapping[int | None, JobscriptElementState]]:
         """Get information about all of this user's jobscripts that currently listed by
         the scheduler."""
-        cmd = self.show_cmd + ["-u", "$USER", "-g", "d"]  # "-g d": separate arrays items
-        stdout, stderr = run_cmd(cmd, logger=self.app.submission_logger)
+        cmd = [*self.show_cmd, "-u", "$USER", "-g", "d"]  # "-g d": separate arrays items
+        stdout, stderr = run_cmd(cmd, logger=self._app.submission_logger)
         if stderr:
             raise ValueError(
                 f"Could not get query SGE jobs. Command was: {cmd!r}; stderr was: "
                 f"{stderr}"
             )
         elif not stdout:
-            info = {}
-        else:
-            info = {}
-            lines = stdout.split("\n")
-            # assuming a job name with spaces means we can't split on spaces to get
-            # anywhere beyond the job name, so get the column index of the state heading
-            # and assume the state is always left-aligned with the heading:
-            state_idx = lines[0].index("state")
-            task_id_idx = lines[0].index("ja-task-ID")
-            for ln in lines[2:]:
-                if not ln:
-                    continue
-                ln_s = ln.split()
-                base_job_ID = ln_s[0]
-
-                # states can be one or two chars (for our limited purposes):
-                state_str = ln[state_idx : state_idx + 2].strip()
-                state = self.state_lookup[state_str]
-
-                arr_idx = ln[task_id_idx:].strip()
-                if arr_idx:
-                    arr_idx = int(arr_idx) - 1  # zero-index
-                else:
-                    arr_idx = None
-
-                if base_job_ID not in info:
-                    info[base_job_ID] = {}
-
-                info[base_job_ID][arr_idx] = state
+            return {}
+
+        info: dict[str, dict[int | None, JobscriptElementState]] = {}
+        lines = stdout.split("\n")
+        # assuming a job name with spaces means we can't split on spaces to get
+        # anywhere beyond the job name, so get the column index of the state heading
+        # and assume the state is always left-aligned with the heading:
+        state_idx = lines[0].index("state")
+        task_id_idx = lines[0].index("ja-task-ID")
+        for ln in lines[2:]:
+            if not ln:
+                continue
+            base_job_ID, *_ = ln.split()
+
+            # states can be one or two chars (for our limited purposes):
+            state_str = ln[state_idx : state_idx + 2].strip()
+            state = self.state_lookup[state_str]
+
+            arr_idx_s = ln[task_id_idx:].strip()
+            arr_idx = (
+                int(arr_idx_s) - 1  # We are using zero-indexed info
+                if arr_idx_s
+                else None
+            )
+
+            info.setdefault(base_job_ID, {})[arr_idx] = state
         return info
 
+    @override
     def get_job_state_info(
-        self, js_refs: List[str] = None
-    ) -> Dict[str, Dict[int, JobscriptElementState]]:
+        self, *, js_refs: Sequence[str] | 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.
 
@@ -316,23 +316,29 @@ class SGEPosix(Scheduler):
         """
         info = self.get_job_statuses()
         if js_refs:
-            info = {k: v for k, v in info.items() if k in js_refs}
+            return {k: v for k, v in info.items() if k in js_refs}
         return info
 
-    def cancel_jobs(self, js_refs: List[str], jobscripts: List = None):
+    @override
+    def cancel_jobs(
+        self,
+        js_refs: list[str],
+        jobscripts: list[Jobscript] | None = None,
+        num_js_elements: int = 0,  # Ignored!
+    ):
         """
         Cancel submitted jobs.
         """
         cmd = [self.del_cmd] + js_refs
-        self.app.submission_logger.info(
+        self._app.submission_logger.info(
             f"cancelling {self.__class__.__name__} jobscripts with command: {cmd}."
         )
-        stdout, stderr = run_cmd(cmd, logger=self.app.submission_logger)
+        stdout, stderr = run_cmd(cmd, logger=self._app.submission_logger)
         if stderr:
             raise ValueError(
                 f"Could not get query SGE {self.__class__.__name__}. Command was: "
                 f"{cmd!r}; stderr was: {stderr}"
             )
-        self.app.submission_logger.info(
+        self._app.submission_logger.info(
             f"jobscripts cancel command executed; stdout was: {stdout}."
         )