Accept configuration of a cluster

Author: w8jcik

src/moldrug/cli.py

@@ -8,6 +8,7 @@
 import argparse
 import datetime
 import inspect
+import importlib
 import os
 import sys
 from typing import Union
@@ -114,7 +115,8 @@ def _translate_config(self):
         InitArgs = MainConfig.copy()
 
         # Modifying InitArgs
-        _ = [InitArgs.pop(key, None) for key in ['type', 'njobs', 'pick']]
+        for key in ['type', 'njobs', 'cluster', 'pick']:
+            InitArgs.pop(key, None)
         InitArgs['costfunc'] = self.costfunc
 
         # Getting call arguments
@@ -125,6 +127,23 @@ def _translate_config(self):
             except KeyError:
                 pass
 
+        if 'cluster' in MainConfig:
+            if 'type' not in MainConfig['cluster'] or 'kwargs' not in MainConfig['cluster']:
+                raise ValueError("The cluster configuration must contain 'type' and 'kwargs'.")
+
+            from moldrug.runner import Runner, RunnerMode  # terminates with a message if dask is not installed
+
+            try:
+                cluster_class = getattr(importlib.import_module("dask_jobqueue"), MainConfig['cluster']['type'])
+            except ImportError:
+                raise ImportError(f"Unable to import {MainConfig['cluster']['type']} from dask_jobqueue module.")
+
+            cluster = cluster_class(**MainConfig['cluster']['kwargs'])
+            cluster.scale(MainConfig.get('njobs', 1))
+
+            CallArgs['runner'] = Runner(RunnerMode.DASK_JOB_QUEUE_CLUSTER, dask_cluster=cluster)
+            del CallArgs['njobs']
+
         # Checking for follow jobs and sanity check on the arguments
         if FollowConfig:
             # Defining the possible mutable arguments with its default values depending on the type of run
@@ -145,7 +164,7 @@ def _translate_config(self):
                         InitArgs[param.name] = param.default
 
                 MutableArgs = {
-                    'njobs': CallArgs['njobs'],
+                    'njobs': MainConfig['njobs'],
                     'crem_db_path': InitArgs['crem_db_path'],
                     'maxiter': InitArgs['maxiter'],
                     'popsize': InitArgs['popsize'],

tests/test_moldrug.py

@@ -121,7 +121,8 @@ def test_single_receptor_command_line():
         # This is currently unused
         if execution_mode == "smaug":
             Config['01_grow']['cluster'] = {
-                "SLURMCluster_kwargs": {
+                "type": "SLURMCluster",
+                "kwargs": {
                     "queue": "short",
                     "cores": 12,
                     "processes": 1,
@@ -132,7 +133,8 @@ def test_single_receptor_command_line():
             }
         elif execution_mode == "elwe":
             Config['01_grow']['cluster'] = {
-                "SLURMCluster_kwargs": {
+                "type": "SLURMCluster",
+                "kwargs": {
                     "queue": "uds-hub",
                     "cores": 10,
                     "processes": 1,