diff --git a/pysqa/queueadapter.py b/pysqa/queueadapter.py index 0792d4a..72501bd 100644 --- a/pysqa/queueadapter.py +++ b/pysqa/queueadapter.py @@ -2,9 +2,7 @@ # Copyright (c) Jan Janssen import os -from pysqa.utils.basic import BasisQueueAdapter, read_config -from pysqa.utils.modular import ModularQueueAdapter -from pysqa.utils.remote import RemoteQueueAdapter +from pysqa.utils.functs import read_config, set_queue_adapter __author__ = "Jan Janssen" __copyright__ = "Copyright 2019, Jan Janssen" @@ -43,15 +41,50 @@ class QueueAdapter(object): Queues available for auto completion QueueAdapter().queues. returns the queue name. """ def __init__(self, directory="~/.queues"): - config = read_config(file_name=os.path.join(directory, "queue.yaml")) - if config["queue_type"] in ["SGE", "TORQUE", "SLURM", "LSF", "MOAB"]: - self._adapter = BasisQueueAdapter(config=config, directory=directory) - elif config["queue_type"] in ["GENT"]: - self._adapter = ModularQueueAdapter(config=config, directory=directory) - elif config["queue_type"] in ["REMOTE"]: - self._adapter = RemoteQueueAdapter(config=config, directory=directory) + queue_yaml = os.path.join(directory, "queue.yaml") + clusters_yaml = os.path.join(directory, "clusters.yaml") + self._adapter = None + if os.path.exists(queue_yaml): + self._queue_dict = { + "default": set_queue_adapter( + config=read_config(file_name=queue_yaml), + directory=directory + ) + } + primary_queue = "default" + elif os.path.exists(clusters_yaml): + config = read_config(file_name=clusters_yaml) + self._queue_dict = { + k: set_queue_adapter( + config=read_config( + file_name=os.path.join(directory, v) + ), + directory=directory + ) + for k, v in config["cluster"].items() + } + primary_queue = config["cluster_primary"] else: raise ValueError + self._adapter = self._queue_dict[primary_queue] + + def list_clusters(self): + """ + List available computing clusters for remote submission + + Returns: + list: List of computing clusters + """ + return self._queue_dict.keys() + + def switch_cluster(self, cluster_name): + """ + Switch to a different computing cluster + + Args: + cluster_name (str): name of the computing cluster + """ + self._adapter = self._queue_dict[cluster_name] @property def config(self): @@ -251,3 +284,5 @@ def check_queue_parameters( memory_max=memory_max, active_queue=active_queue, ) + + diff --git a/pysqa/utils/basic.py b/pysqa/utils/basic.py index b352a7f..cc3208f 100644 --- a/pysqa/utils/basic.py +++ b/pysqa/utils/basic.py @@ -8,7 +8,6 @@ import subprocess import re import pandas -import yaml from pysqa.utils.queues import Queues @@ -39,7 +38,6 @@ class BasisQueueAdapter(object): Queues available for auto completion QueueAdapter().queues. returns the queue name. """ - def __init__(self, config, directory="~/.queues"): self._config = config self._fill_queue_dict(queue_lst_dict=self._config["queues"]) @@ -548,16 +546,3 @@ def _memory_spec_string_to_value( ) else: return value - - -def read_config(file_name="queue.yaml"): - """ - - Args: - file_name (str): - - Returns: - dict: - """ - with open(file_name, "r") as f: - return yaml.load(f, Loader=yaml.FullLoader) diff --git a/pysqa/utils/functs.py b/pysqa/utils/functs.py new file mode 100644 index 0000000..75fb2e7 --- /dev/null +++ b/pysqa/utils/functs.py @@ -0,0 +1,35 @@ +import yaml +from pysqa.utils.basic import BasisQueueAdapter +from pysqa.utils.remote import RemoteQueueAdapter +from pysqa.utils.modular import ModularQueueAdapter + + +def read_config(file_name="queue.yaml"): + """ + + Args: + file_name (str): + + Returns: + dict: + """ + with open(file_name, "r") as f: + return yaml.load(f, Loader=yaml.FullLoader) + + +def set_queue_adapter(config, directory): + """ + Initialize the queue adapter + + Args: + config (dict): configuration for one cluster + directory (str): directory which contains the queue configurations + """ + if config["queue_type"] in ["SGE", "TORQUE", "SLURM", "LSF", "MOAB"]: + return BasisQueueAdapter(config=config, directory=directory) + elif config["queue_type"] in ["GENT"]: + return ModularQueueAdapter(config=config, directory=directory) + elif config["queue_type"] in ["REMOTE"]: + return RemoteQueueAdapter(config=config, directory=directory) + else: + raise ValueError