fix: make TimeLeft backward compatible

DIRACGrid · Dec 7, 2023 · 2ffe69a · 2ffe69a
1 parent 1b00fd0
commit 2ffe69a
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 4 deletions.
diff --git a/src/DIRAC/Resources/Computing/BatchSystems/TimeLeft/MJFResourceUsage.py b/src/DIRAC/Resources/Computing/BatchSystems/TimeLeft/MJFResourceUsage.py
@@ -19,8 +19,6 @@ def __init__(self, jobID, parameters):
         """Standard constructor"""
         super().__init__("MJF", jobID, parameters)
 
-        self.queue = os.environ.get("QUEUE")
-
         self.log.verbose(f"jobID={self.jobID}, queue={self.queue}")
         self.startTime = time.time()
 

diff --git a/src/DIRAC/Resources/Computing/BatchSystems/TimeLeft/TimeLeft.py b/src/DIRAC/Resources/Computing/BatchSystems/TimeLeft/TimeLeft.py
@@ -9,6 +9,7 @@
     With this information the utility can calculate in normalized units the
     CPU time remaining for a given slot.
 """
+import os
 import shlex
 
 import DIRAC
@@ -136,6 +137,47 @@ def _getBatchSystemPlugin(self):
         jobID = batchSystemInfo.get("JobID")
         parameters = batchSystemInfo.get("Parameters")
 
+        ###########################################################################################
+        # TODO: remove the following block in v9.0
+        # This is a temporary fix for the case where the batch system is not set in the configuration
+        if not type:
+            self.log.warn(
+                "Batch system info not set in local configuration: trying to guess from environment variables."
+            )
+            self.log.warn("Consider updating your pilot version before switching to v9.0.")
+            batchSystemInfo = {
+                "LSF": {
+                    "JobID": "LSB_JOBID",
+                    "Parameters": {
+                        "BinaryPath": "LSB_BINDIR",
+                        "Host": "LSB_HOSTS",
+                        "InfoPath": "LSB_ENVDIR",
+                        "Queue": "LSB_QUEUE",
+                    },
+                },
+                "PBS": {"JobID": "PBS_JOBID", "Parameters": {"BinaryPath": "PBS_O_PATH", "Queue": "PBS_O_QUEUE"}},
+                "SGE": {"JobID": "SGE_TASK_ID", "Parameters": {"BinaryPath": "SGE_BINARY_PATH", "Queue": "QUEUE"}},
+                "SLURM": {"JobID": "SLURM_JOB_ID", "Parameters": {}},
+                "HTCondor": {"JobID": "HTCONDOR_JOBID", "Parameters": {"InfoPath": "_CONDOR_JOB_AD"}},
+            }
+            type = None
+            for batchSystem, attributes in batchSystemInfo.items():
+                if attributes["JobID"] in os.environ:
+                    type = batchSystem
+                    jobID = os.environ[attributes["JobID"]]
+                    parameters = {}
+                    for parameterName, parameterVariable in attributes["Parameters"].items():
+                        parameters[parameterName] = os.environ.get(parameterVariable)
+                    break
+
+            if not type and "MACHINEFEATURES" in os.environ and "JOBFEATURES" in os.environ:
+                # Only use MJF if legacy batch system information not available for now
+                type = "MJF"
+                jobID = os.environ.get("JOB_ID")
+                parameters = {"Queue": os.environ.get("QUEUE")}
+
+        ###########################################################################################
+
         if not type or type == "Unknown":
             self.log.warn(f"Batch system type for site {DIRAC.siteName()} is not currently supported")
             return S_ERROR(DErrno.ERESUNK, "Current batch system is not supported")

diff --git a/src/DIRAC/Resources/Computing/BatchSystems/TimeLeft/test/Test_TimeLeft.py b/src/DIRAC/Resources/Computing/BatchSystems/TimeLeft/test/Test_TimeLeft.py
@@ -1,7 +1,5 @@
 """ Test TimeLeft utility
 """
-import pytest
-
 from DIRAC import S_OK
 from DIRAC.Resources.Computing.BatchSystems.TimeLeft.TimeLeft import TimeLeft
 
@@ -27,6 +25,23 @@ def test_batchSystemNotDefined(mocker):
     assert "Current batch system is not supported" in res["Message"]
 
 
+def test_batchSystemNotDefinedInConfigButInEnvironmentVariables(mocker, monkeypatch):
+    """Test batch system not defined but present in environment variables (should fail from v9.0)"""
+    mocker.patch(
+        "DIRAC.Resources.Computing.BatchSystems.TimeLeft.HTCondorResourceUsage.runCommand",
+        return_value=S_OK("9000 800"),
+    )
+    mocker.patch("DIRAC.gConfig.getSections", return_value={})
+    monkeypatch.setenv("HTCONDOR_JOBID", "12345.0")
+    monkeypatch.setenv("_CONDOR_JOB_AD", "/path/to/config")
+
+    tl = TimeLeft()
+    tl.cpuPower = 10
+    res = tl.getTimeLeft()
+    assert res["OK"]
+    assert res["Value"] == 82000
+
+
 def test_getScaledCPU(mocker):
     """Test getScaledCPU()"""
     mocker.patch(