Source code for matador.compute.queueing

# coding: utf-8
# Distributed under the terms of the MIT License.

""" This file implements a simple agnostic interface to various
queueing systems.


import os
import abc

[docs]class QueueManager(abc.ABC): """Abstract base class for queue managers.""" token = None def __repr__(self): return "{} with env: {}".format(type(self), self.env) def __init__(self): self.env = get_queue_env(self.token) @property def walltime(self): """Return the allotted walltime in seconds, returning None if not available. """ try: return self.get_walltime() except Exception: None @property def ntasks(self): return self.get_ntasks() @property def max_memory(self): """Return the allotted memory in MB, returning None if not available. """ try: return self.get_max_memory() except Exception: return None @property def array_id(self): return self.get_array_id()
[docs] @abc.abstractmethod def get_walltime(self): pass
[docs] @abc.abstractmethod def get_max_memory(self): pass
[docs] @abc.abstractmethod def get_array_id(self): pass
[docs]def get_queue_env(token): """Read os.environment variables for either PBS or SLURM prefixes, and return a dictionary of those vars only. Parameter: token (str): choose one of either SLURM or PBS explicitly. Returns: (dict, str): dictionary of keys from the detected/specified queue, and a string containing either "slurm" or "pbs". """ return {key: os.environ[key] for key in os.environ if key.lower().startswith(token)}
[docs]def get_queue_manager(): """Detects whether PBS, SLURM or neither is being used by probing the environment variables SLURM_NTASKS and PBS_TASKNUM. Returns: str or None: either "slurm", "pbs" or None. Raises: SystemExit: if both SLURM and PBS were found. """ queue_mgr = [] if os.environ.get("SLURM_NTASKS") is not None: queue_mgr.append("slurm") if os.environ.get("PBS_TASKNUM") is not None: queue_mgr.append("pbs") if len(queue_mgr) > 1: raise RuntimeError("Both SLURM and PBS were found... aah!") if not queue_mgr: return None if queue_mgr[0] == "slurm": from matador.compute.slurm import SlurmQueueManager return SlurmQueueManager() if queue_mgr[0] == "pbs": from matador.compute.pbs import PBSQueueManager return PBSQueueManager()