From 5d3b57fb0e19ed76625263110a12aee4b78bb428 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20N=C3=B6the?= <maximilian.noethe@tu-dortmund.de>
Date: Sat, 14 Apr 2018 06:36:53 +0100
Subject: [PATCH] Add job info for Torque and PBS (#78)

---
 dask_drmaa/core.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/dask_drmaa/core.py b/dask_drmaa/core.py
index dd5b7e8..0879b82 100644
--- a/dask_drmaa/core.py
+++ b/dask_drmaa/core.py
@@ -34,7 +34,10 @@ def get_session():
 worker_bin_path = os.path.join(sys.exec_prefix, 'bin', 'dask-worker')
 
 # All JOB_ID and TASK_ID environment variables
-_drm_info = drmaa.Session().drmsInfo
+_session = drmaa.Session()
+_drm_info = _session.drmsInfo
+_drmaa_implementation = _session.drmaaImplementation
+
 if "SLURM" in _drm_info:
     JOB_PARAM = "%j"
     JOB_ID = "$SLURM_JOB_ID"
@@ -47,6 +50,10 @@ def get_session():
     JOB_PARAM = "$JOB_ID"
     JOB_ID = "$JOB_ID"
     TASK_ID = "$SGE_TASK_ID"
+elif "Torque" == _drm_info or "PBS" in _drmaa_implementation:
+    JOB_PARAM = "$PBS_JOBID"
+    JOB_ID = "$PBS_JOBID"
+    TASK_ID = "$PBS_TASKNUM"
 else:
     JOB_PARAM = ""
     JOB_ID = ""