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()