hpcpy 0.7.3__py3-none-any.whl

hpcpy/client/direct.py

@@ -0,0 +1,210 @@
+"""Direct execution client (no scheduler)."""
+
+import subprocess
+import shlex
+import os
+import re
+from pathlib import Path
+from typing import Union
+
+from hpcpy.client.base import BaseClient
+from hpcpy.constants.direct import COMMANDS, DIRECTIVES, STATUSES
+from hpcpy.utilities import shell
+from hpcpy.job import Job
+import hpcpy.constants as hc
+
+
+class DirectClient(BaseClient):
+    """Direct execution interface for systems without a scheduler.
+
+    Runs scripts directly as local processes, using the process ID as the
+    job ID. Status is tracked via ps(1).
+
+    Parameters
+    ----------
+    *args
+        Positional arguments forwarded to the base class.
+    **kwargs
+        Keyword arguments forwarded to the base class.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(
+            cmd_templates=COMMANDS,
+            directive_templates=DIRECTIVES,
+            statuses=STATUSES,
+            status_attribute="short",
+            *args,
+            **kwargs,
+        )
+
+    def status(self, job_id):
+        """Get the status of a directly-executed process.
+
+        Parameters
+        ----------
+        job_id : str
+            Process ID.
+
+        Returns
+        -------
+        tuple
+            Generic status code and native status dictionary.
+        """
+        cmd = self.cmd_templates["status"].format(job_id=job_id)
+
+        # ps returns non-zero when the process does not exist, so use check=False
+        result = shell(cmd, check=False)
+        raw = result.stdout.decode("utf8").strip()
+
+        return self._parse_status(raw, job_id)
+
+    def submit(
+        self,
+        job_script: Union[str, Path],
+        directives: list = None,
+        render: bool = False,
+        dry_run: bool = False,
+        variables: dict = dict(),
+        **context,
+    ):
+        """Submit a job by running it directly as a local process.
+
+        Parameters
+        ----------
+        job_script : Union[str, Path]
+            Path to the script to execute.
+        directives : list, optional
+            Ignored for direct execution.
+        render : bool, optional
+            Render the job script from a template, by default False
+        dry_run : bool, optional
+            Return the command rather than executing it, by default False
+        variables : dict, optional
+            Environment variables to set for the process.
+        **context :
+            Additional key/value pairs for job script rendering.
+
+        Returns
+        -------
+        Union[Job, str]
+            Submitted job object, or the command string if dry_run=True.
+        """
+        # Sanitize the variables, if there are any, then create string version and add to context
+        variables_sanitized = self._sanitize_dict(variables)
+        variables_str = []
+        for key, value in variables_sanitized.items():
+            variables_str.append(f"{key}={value}")
+
+        context["variables_str"] = (
+            " ".join(variables_str) + " " if variables_str else ""
+        )
+
+        # Delegate command assembly (directive rendering, job script rendering) to the parent
+        cmd = super().submit(
+            job_script=job_script,
+            directives=directives if isinstance(directives, list) else [],
+            render=render,
+            dry_run=True,
+            variables=self._sanitize_dict(variables),
+            **context,
+        )
+
+        if dry_run:
+            return cmd
+
+        _env = os.environ.copy()
+        if isinstance(variables_sanitized, dict) and variables_sanitized:
+            _env.update({k: str(v) for k, v in variables_sanitized.items()})
+
+        proc = subprocess.Popen(shlex.split(cmd), env=_env)
+        job_id = str(proc.pid)
+
+        job = Job(job_id, client=None, auto_update=False)
+        job.set_client(self)
+        job._auto_update = True
+
+        return job
+
+    def _parse_status(self, raw, job_id):
+        """Extract the status from the raw ps output.
+
+        Parameters
+        ----------
+        raw : str
+            Raw output from ps(1). Empty string means the process has finished.
+        job_id : str
+            Process ID.
+
+        Returns
+        -------
+        tuple
+            Generic status code and native status dictionary.
+        """
+        # Empty output means the process no longer exists
+        if not raw:
+            return hc.STATUS_FINISHED, {"pid": job_id, "state": "FINISHED"}
+
+        # ps -o stat= returns one or more characters; the first is the main state
+        state_char = raw[0]
+        generic_status = None
+
+        for s in self.statuses:
+            if state_char == s.short:
+                generic_status = s.status
+                break
+
+        return generic_status, {"pid": job_id, "state": raw}
+
+    def _sanitize_dict(self, user_dict: dict) -> dict:
+        """Sanitize user-supplied dict for use with subprocess.run.
+
+        Parameters
+        ----------
+        user_dict : dict
+            Key/value pairs.
+
+        Returns
+        -------
+        dict
+            Sanitized version of the dict.
+
+        Raises
+        ------
+        ValueError
+            When a key is invalid.
+            When a value is invalid.
+            When a null byte is provided.
+        """
+        VALID_KEY = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+        sanitized = dict()
+
+        for key, value in user_dict.items():
+
+            # Validate key
+            if not isinstance(key, str):
+                raise ValueError(f"Key {key!r} must be a string")
+
+            if not VALID_KEY.match(key):
+                raise ValueError(
+                    f"Key {key!r} is invalid: must start with a letter or underscore "
+                    f"and contain only alphanumeric characters and underscores"
+                )
+
+            # Validate and sanitize value
+            if not isinstance(value, (str, int, float, bool)):
+                raise ValueError(
+                    f"Value for key {key!r} must be a string, int, float, or bool; "
+                    f"got {type(value).__name__}"
+                )
+
+            str_value = str(value)
+
+            # Reject null bytes — they silently truncate strings in most shells
+            if "\x00" in str_value:
+                raise ValueError(f"Value for key {key!r} contains a null byte")
+
+            # Add to dict and quote for safety
+            sanitized[key] = shlex.quote(str_value)
+
+        return sanitized

hpcpy/client/pbs.py

@@ -0,0 +1,223 @@
+"""PBS implementation."""
+
+from hpcpy.client.base import BaseClient
+from hpcpy.constants.pbs import COMMANDS, DIRECTIVES, STATUSES, DELAY_DIRECTIVE_FMT
+from datetime import datetime, timedelta
+from typing import Union
+import json
+from pathlib import Path
+from shlex import quote
+
+
+class PBSClient(BaseClient):
+    """PBS interface.
+
+    Parameters
+    ----------
+    *args
+        Positional arguments forwarded to the base class.
+    **kwargs
+        Keyword arguments forwarded to the base class.
+    """
+
+    def __init__(self, *args, **kwargs):
+        # Set up the templates
+        super().__init__(
+            cmd_templates=COMMANDS,
+            directive_templates=DIRECTIVES,
+            statuses=STATUSES,
+            status_attribute="short",
+            *args,
+            **kwargs,
+        )
+
+    def status(self, job_id):
+        """Get the status of a job.
+
+        Parameters
+        ----------
+        job_id : str
+            Job ID.
+
+        Returns
+        -------
+        str
+            Generic status code.
+        """
+        # Get the raw response
+        raw = super().status(job_id=job_id)
+
+        # Parse the status as per this implementation
+        generic_status, native_full = self._parse_status(raw, job_id)
+
+        # Return the generic status
+        return generic_status, native_full
+
+    def _render_variables(self, variables):
+        """Render the variables flag for PBS.
+
+        Parameters
+        ----------
+        variables : dict
+            Dictionary of variables
+
+        Returns
+        -------
+        str
+            String formatted variables for PBS
+        """
+        formatted = list()
+
+        for k, v in variables.items():
+
+            # Fix for broken quotes #52
+            if isinstance(v, str) and " " in v:
+                v = quote(v)
+                line = f'"{k}={v}"'
+                formatted.append(line)
+            else:
+                formatted.append(f"{k}={v}")
+
+        formatted = ",".join(formatted)
+        return f"-v {formatted}"
+
+    def submit(
+        self,
+        job_script: Union[str, Path],
+        directives: list = None,
+        render: bool = False,
+        dry_run: bool = False,
+        depends_on: list = None,
+        delay: Union[datetime, timedelta] = None,
+        queue: str = None,
+        walltime: timedelta = None,
+        storage: list = None,
+        variables: dict = None,
+        **context,
+    ):
+        """Submit a job to the scheduler.
+
+        Parameters
+        ----------
+        job_script : Union[str, Path]
+            Path to the script.
+        directives : list, optional
+            List of complete directives to submit, by default list()
+        render : bool, optional
+            Render the job script from a template, by default False
+        dry_run : bool, optional
+            Return rather than executing the command, by default False
+        depends_on : list, optional
+            List of job IDs with successful exit on which this job depends, by default list()
+        delay: Union[datetime, timedelta]
+            Delay the start of this job until specific date or interval, by default None
+        queue: str, optional
+            Queue on which to submit the job, by default None
+        walltime: timedelta, optional
+            Walltime expressed as a timedelta, by default None
+        storage: list, optional
+            List of storage mounts to apply, by default None
+        variables: dict, optional
+            Key/value environment variable pairs added to the qsub command.
+        **context:
+            Additional key/value pairs to be added to command/jobscript interpolation
+
+        Returns
+        -------
+        Job : hpcpy.job.Job
+            Job object.
+        """
+        # Initialise the directives to an empty list
+        directives = directives if isinstance(directives, list) else []
+
+        # Add job depends
+        if depends_on:
+
+            # Normalise to a list of strs
+            depends_on = super()._normalise_depends_on(depends_on)
+
+            directives = self._interpolate_directive(
+                directives,
+                "depends_on",
+                depends_on_str=":".join(depends_on),
+            )
+
+        # Add delay (specified time or delta)
+        if delay:
+            delay_directive = self._assemble_delay_directive(delay, DELAY_DIRECTIVE_FMT)
+            directives.append(delay_directive)
+
+        # Add queue
+        if queue:
+            directives = self._interpolate_directive(directives, "queue", queue=queue)
+            context["queue"] = queue
+
+        # Add walltime
+        if walltime:
+            _walltime = str(walltime)
+            directives.append(f"-l walltime={_walltime}")
+            context["walltime"] = _walltime
+
+        # Add storage
+        if storage:
+            storage_str = "+".join(storage)
+            directives.append(f"-l storage={storage_str}")
+            context["storage"] = storage
+            context["storage_str"] = storage_str
+
+        # Add variables
+        if isinstance(variables, dict) and len(variables) > 0:
+            directives.append(self._render_variables(variables))
+
+        # Call the super
+        job_or_cmd = super().submit(
+            job_script=job_script,
+            directives=directives,
+            render=render,
+            dry_run=dry_run,
+            **context,
+        )
+
+        # Return the command if requested
+        if dry_run:
+            return job_or_cmd
+
+        # Point to this client from the job and switch on auto_update
+        job_or_cmd.set_client(self)
+        job_or_cmd._auto_update = True
+
+        # Return the job object
+        return job_or_cmd
+
+    def _parse_status(self, raw, job_id):
+        """Extract the status from the raw response.
+
+        Parameters
+        ----------
+        raw : str
+            Raw response from the scheduler.
+        job_id : str
+            Job ID required to extract status from the response.
+
+        Returns
+        -------
+        str
+            Status code.
+        """
+        # Convert to JSON
+        parsed = json.loads(raw)
+
+        # Get the status out of the job ID
+        native_full = parsed.get("Jobs").get(job_id)
+        native_status = native_full.get("job_state")
+
+        # Set the generic status attribute
+        generic_status = None
+
+        for s in self.statuses:
+            if native_status == s.short:
+                generic_status = s.status
+                break
+
+        # Return the generic status
+        return generic_status, native_full

hpcpy/client/slurm.py

@@ -0,0 +1,203 @@
+"""SLURM Client."""
+
+from hpcpy.client.base import BaseClient
+from hpcpy.constants.slurm import COMMANDS, STATUSES, DIRECTIVES, DELAY_DIRECTIVE_FMT
+from datetime import datetime, timedelta
+from typing import Union
+import json
+from pathlib import Path
+import os
+
+
+class SlurmClient(BaseClient):
+    """SLURM interface.
+
+    Parameters
+    ----------
+    *args
+        Positional arguments forwarded to the base class.
+    **kwargs
+        Keyword arguments forwarded to the base class.
+    """
+
+    def __init__(self, *args, **kwargs):
+        # Set up the templates
+        super().__init__(
+            cmd_templates=COMMANDS,
+            directive_templates=DIRECTIVES,
+            statuses=STATUSES,
+            status_attribute="short",
+            *args,
+            **kwargs,
+        )
+
+    def status(self, job_id):
+        """Get the status of a job.
+
+        Parameters
+        ----------
+        job_id : str
+            Job ID.
+
+        Returns
+        -------
+        str
+            Generic status code.
+        """
+        # Get the raw response
+        raw = super().status(job_id=job_id)
+
+        # Parse the status as per this implementation
+        generic_status, native_full = self._parse_status(raw, job_id)
+
+        # Return the generic status
+        return generic_status, native_full
+
+    def submit(
+        self,
+        job_script: Union[str, Path],
+        directives: list = None,
+        render: bool = False,
+        dry_run: bool = False,
+        depends_on: list = None,
+        delay: Union[datetime, timedelta] = None,
+        queue: str = None,
+        walltime: timedelta = None,
+        variables: dict = dict(),
+        **context,
+    ):
+        """Submit a job to the scheduler.
+
+        Parameters
+        ----------
+        job_script : Union[str, Path]
+            Path to the script.
+        directives : list, optional
+            List of complete directives to submit, by default list()
+        render : bool, optional
+            Render the job script from a template, by default False
+        dry_run : bool, optional
+            Return rather than executing the command, by default False
+        depends_on : list, optional
+            List of job IDs with successful exit on which this job depends, by default list()
+        delay: Union[datetime, timedelta]
+            Delay the start of this job until specific date or interval, by default None
+        queue: str, optional
+            Queue on which to submit the job, by default None
+        walltime: timedelta, optional
+            Walltime expressed as a timedelta, by default None
+        variables: dict, optional
+            Key/value environment variable pairs added to the qsub command.
+        **context:
+            Additional key/value pairs to be added to command/jobscript interpolation
+
+        Returns
+        -------
+        Job : hpcpy.job.Job
+            Job object.
+        """
+        # Get a logger from the parent
+        self._logger.debug(f"Submitting {job_script}")
+
+        # Initialise directives
+        directives = directives if isinstance(directives, list) else []
+
+        # Add job depends
+        if depends_on:
+
+            self._logger.debug("Job dependency specified.")
+
+            # Normalise depends_on to strs
+            depends_on = super()._normalise_depends_on(depends_on)
+
+            directives = self._interpolate_directive(
+                directives,
+                "depends_on",
+                depends_on_str=":".join(depends_on),
+            )
+
+        # Add delay (specified time or delta)
+        if delay:
+            self._logger.debug("Delay specified.")
+            delay_directive = self._assemble_delay_directive(delay, DELAY_DIRECTIVE_FMT)
+            directives.append(delay_directive)
+
+        # Add queue
+        if queue:
+            self._logger.debug(f"Queue specified ({queue})")
+            directives = self._interpolate_directive(directives, "queue", queue=queue)
+            context["queue"] = queue
+
+        # Add walltime
+        if walltime:
+            walltime_str = str(int(walltime.total_seconds() / 60.0))
+            self._logger.debug(f"Walltime specified ({walltime_str})")
+            directives = self._interpolate_directive(
+                directives, "walltime", walltime_str=walltime_str
+            )
+
+        # Update the environment with the variables (which is how Slurm does job variables)
+        self._logger.debug("Updating environment.")
+        _env = os.environ.copy()
+        _env.update(variables)
+
+        # Call the super submit
+        self._logger.debug("Submitting via parent class.")
+        job_or_cmd = super().submit(
+            job_script=job_script,
+            directives=directives,
+            render=render,
+            dry_run=dry_run,
+            env=_env,
+            **context,
+        )
+
+        # Return the command from the super
+        if dry_run:
+            self._logger.debug("Dry run requested, returning fully-formed command.")
+            return job_or_cmd
+
+        # Point to this client in the job object and switch on auto_update
+        job_or_cmd.set_client(self)
+        job_or_cmd._auto_update = True
+
+        # Get the job ID out of the return string
+        self._logger.debug(f"job_id={job_or_cmd.id}")
+
+        # Return the job object
+        return job_or_cmd
+
+    def _parse_status(self, raw, job_id):
+        """Extract the statue from the raw response.
+
+        Parameters
+        ----------
+        raw : str
+            Raw response from the scheduler.
+        job_id : str
+            Job ID.
+
+        Returns
+        -------
+        generic_status : str
+            Generic status code.
+        native_full : dict
+            Full native status dictionary.
+        """
+        # Parse the response
+        parsed = json.loads(raw)
+
+        # Get the status from the first job
+        native_full = parsed.get("jobs")[0]
+        native_status = native_full.get("job_state")[0]
+
+        # Set the default generic status to None
+        generic_status = None
+
+        # Look up the generic status
+        for s in self.statuses:
+            if native_status == s.long:
+                generic_status = s.status
+                break
+
+        return generic_status, native_full

hpcpy/constants/__init__.py

@@ -0,0 +1,21 @@
+"""Common constants."""
+
+from pathlib import Path
+
+# Location for rendered job scripts
+JOB_SCRIPT_DIR = Path.home() / ".hpcpy" / "job_scripts"
+JOB_SCRIPT_DIR.mkdir(parents=True, exist_ok=True)
+
+# Generic Statuses (actually just PBS statuses, but used as a generic status)
+STATUS_CYCLE_HARVESTING = "U"
+STATUS_EXITING = "E"
+STATUS_FINISHED = "F"
+STATUS_HAS_SUBJOB = "B"
+STATUS_HELD = "H"
+STATUS_MOVED = "M"
+STATUS_MOVING = "T"
+STATUS_QUEUED = "Q"
+STATUS_RUNNING = "R"
+STATUS_SUBJOB_COMPLETED = "X"
+STATUS_SUSPENDED = "S"
+STATUS_WAITING = "W"

hpcpy/constants/direct.py

@@ -0,0 +1,55 @@
+"""Constants for direct execution (no scheduler)."""
+
+from hpcpy.status import Status
+import hpcpy.constants as hc
+
+# Commands
+COMMANDS = dict(
+    submit="{variables_str}bash {job_script}",
+    status="ps -p {job_id} -o stat=",
+    delete="kill {job_id}",
+    hold="kill -s STOP {job_id}",
+    release="kill -s CONT {job_id}",
+)
+
+# Directives (none for direct execution)
+DIRECTIVES = dict()
+
+# Statuses - mapped from ps(1) stat codes
+STATUSES = [
+    Status("R", "RUNNING", "Process is running.", generic=hc.STATUS_RUNNING),
+    Status(
+        "S",
+        "SLEEPING",
+        "Process is sleeping (interruptible).",
+        generic=hc.STATUS_RUNNING,
+    ),
+    Status(
+        "D",
+        "DISK_SLEEP",
+        "Process is in uninterruptible disk sleep.",
+        generic=hc.STATUS_RUNNING,
+    ),
+    Status("I", "IDLE", "Process is idle.", generic=hc.STATUS_RUNNING),
+    Status(
+        "T", "STOPPED", "Process has been stopped by a signal.", generic=hc.STATUS_HELD
+    ),
+    Status(
+        "Z",
+        "ZOMBIE",
+        "Process is a zombie (awaiting reaping).",
+        generic=hc.STATUS_EXITING,
+    ),
+    Status(
+        "F",
+        "FINISHED",
+        "Job has finished, or never existed",
+        generic=hc.STATUS_FINISHED,
+    ),
+    Status(
+        "E",
+        "EXITING",
+        "Job has finished, or never existed",
+        generic=hc.STATUS_EXITING,
+    ),
+]