Also test drmaaImplementation

dask · Apr 10, 2018 · 07b20dc · 07b20dc
1 parent 5ba9f15
commit 07b20dc
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/dask_drmaa/core.py b/dask_drmaa/core.py
@@ -34,7 +34,9 @@ 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"
@@ -48,7 +50,7 @@ def get_session():
     JOB_PARAM = "$JOB_ID"
     JOB_ID = "$JOB_ID"
     TASK_ID = "$SGE_TASK_ID"
-elif _drm_info == "Torque":
+elif "Torque" == _drm_info or "PBS" in _drmaa_implementation:
     JOB_PARAM = "$PBS_JOBID"
     JOB_ID = "$PBS_JOBID"
     TASK_ID = "$PBS_TASKNUM"