From b580d411dd1fe2df357e98f62ec56a7713ca3449 Mon Sep 17 00:00:00 2001
From: jaime <jcernudagarcia@hawk.iit.edu>
Date: Mon, 16 Oct 2023 17:17:28 -0500
Subject: [PATCH] Added support for PBS scheduler

---
 jarvis_util/jutil_manager.py   |  1 +
 jarvis_util/shell/exec.py      |  3 ++
 jarvis_util/shell/exec_info.py |  1 +
 jarvis_util/shell/pbs_exec.py  | 82 ++++++++++++++++++++++++++++++++++
 4 files changed, 87 insertions(+)
 create mode 100644 jarvis_util/shell/pbs_exec.py

diff --git a/jarvis_util/jutil_manager.py b/jarvis_util/jutil_manager.py
index fa79468..c69674b 100644
--- a/jarvis_util/jutil_manager.py
+++ b/jarvis_util/jutil_manager.py
@@ -27,4 +27,5 @@ def __init__(self):
         self.debug_local_exec = False
         self.debug_scp = False
         self.debug_slurm = False
+        self.debug_pbs = True
 
diff --git a/jarvis_util/shell/exec.py b/jarvis_util/shell/exec.py
index 348a019..a8c63bf 100644
--- a/jarvis_util/shell/exec.py
+++ b/jarvis_util/shell/exec.py
@@ -2,6 +2,7 @@
 This module provides mechanisms to execute binaries either locally or
 remotely.
 """
+from jarvis_util.shell.pbs_exec import PbsExec
 from jarvis_util.shell.slurm_exec import SlurmExec
 from .local_exec import LocalExec
 from .pssh_exec import PsshExec
@@ -37,6 +38,8 @@ def __init__(self, cmd, exec_info=None):
             exec_type = MpiVersion(exec_info).version
         elif exec_type == ExecType.SLURM:
             exec_type = SlurmExec(cmd, exec_info)
+        elif exec_type == ExecType.PBS:
+            exec_type = PbsExec(cmd, exec_info)
 
         if exec_type == ExecType.MPICH:
             self.exec_ = MpichExec(cmd, exec_info)
diff --git a/jarvis_util/shell/exec_info.py b/jarvis_util/shell/exec_info.py
index 108ff7f..db91d70 100644
--- a/jarvis_util/shell/exec_info.py
+++ b/jarvis_util/shell/exec_info.py
@@ -24,6 +24,7 @@ class ExecType(Enum):
     OPENMPI = 'OPENMPI'
     INTEL_MPI = 'INTEL_MPI'
     SLURM = 'SLURM'
+    PBS = 'PBS'
 
 
 class ExecInfo:
diff --git a/jarvis_util/shell/pbs_exec.py b/jarvis_util/shell/pbs_exec.py
new file mode 100644
index 0000000..98b9fa6
--- /dev/null
+++ b/jarvis_util/shell/pbs_exec.py
@@ -0,0 +1,82 @@
+"""
+This module provides methods to execute a process in parallel using the
+Message Passing Interface (MPI). This module assumes MPI is installed
+on the system. This class is intended to be called from Exec,
+not by general users.
+"""
+
+from jarvis_util.jutil_manager import JutilManager
+from jarvis_util.shell.local_exec import LocalExec
+from .exec_info import ExecInfo, ExecType
+
+
+class PbsExec(LocalExec):
+    """
+    This class contains methods for executing a command
+    through the PBS scheduler
+    """
+
+    def __init__(self, cmd, exec_info):
+        """
+        Execute a command through qsub
+        :param cmd: A command (string) to execute
+        :param exec_info: Information needed by qsub
+        """
+        self.cmd = cmd
+        self.interactive = exec_info.interactive
+        self.nnodes = exec_info.nnodes
+        self.system = exec_info.system
+        self.filesystems = exec_info.filesystems
+        self.walltime = exec_info.walltime
+        self.account = exec_info.account
+        self.queue = exec_info.queue
+        super().__init__(self.pbscmd(),
+                         exec_info.mod(env=exec_info.basic_env))
+
+        def generate_qsub_command(self):
+            cmd = 'qsub'
+
+            if self.interactive:
+                cmd += ' -I'
+
+            options_map = {
+                'filesystems': 'l filesystems',
+                'walltime': 'l walltime',
+                'account': 'A',
+                'queue': 'q'
+            }
+
+            if self.nnodes and self.system:
+                cmd += f' -l select={self.nnodes}:system={self.system}'
+            elif self.nnodes:
+                cmd += f' -l select={self.nnodes}'
+            else:
+                raise ValueError("System defined without select value.")
+
+            for attr, option in options_map.items():
+                value = getattr(self, attr)
+                if value is not None:
+                    cmd += f' -{option}={value}'
+
+            cmd += f' {self.cmd}'
+            return cmd
+
+        def pbscmd(self):
+            cmd = self.generate_qsub_command()
+            jutil = JutilManager.get_instance()
+            if jutil.debug_pbs:
+                print(cmd)
+            return cmd
+
+
+class PbsExecInfo(ExecInfo):
+    def __init__(self, **kwargs):
+        super().__init__(exec_type=ExecType.PBS, **kwargs)
+        allowed_options = ['interactive', 'nnodes', 'system', 'filesystems', 'walltime', 'account', 'queue']
+        self.keys += allowed_options
+        # We use output and error file from the base Exec Info
+        for key in allowed_options:
+            if key in kwargs:
+                setattr(self, key, kwargs[key])
+            else:
+                setattr(self, key, None)