Skip to content

Commit fc2b657

Browse files
committed
fix: getting cpu work left from a single source of truth
1 parent 7b80cd8 commit fc2b657

6 files changed

Lines changed: 298 additions & 104 deletions

File tree

src/DIRAC/WorkloadManagementSystem/Agent/JobAgent.py

Lines changed: 13 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
from DIRAC.RequestManagementSystem.Client.ReqClient import ReqClient
2828
from DIRAC.RequestManagementSystem.Client.Request import Request
2929
from DIRAC.RequestManagementSystem.private.RequestValidator import RequestValidator
30-
from DIRAC.Resources.Computing.BatchSystems.TimeLeft.TimeLeft import TimeLeft
3130
from DIRAC.Resources.Computing.ComputingElementFactory import ComputingElementFactory
3231
from DIRAC.WorkloadManagementSystem.Client import JobStatus, PilotStatus
3332
from DIRAC.WorkloadManagementSystem.Client.JobManagerClient import JobManagerClient
@@ -81,10 +80,9 @@ def __init__(self, agentName, loadName, baseAgentName=False, properties=None):
8180
self.defaultWrapperLocation = "DIRAC/WorkloadManagementSystem/JobWrapper/JobWrapperTemplate.py"
8281

8382
# Timeleft
84-
self.initTimes = os.times()
8583
self.initTimeLeft = 0.0
8684
self.timeLeft = self.initTimeLeft
87-
self.timeLeftUtil = None
85+
self.initTime = time.time()
8886
self.pilotInfoReportedFlag = False
8987

9088
# Attributes related to the processed jobs, it should take the following form:
@@ -109,16 +107,11 @@ def initialize(self):
109107
if not result["OK"]:
110108
return result
111109

112-
result = self._getCEDict(self.computingElement)
113-
if not result["OK"]:
114-
return result
115-
ceDict = result["Value"][0]
116-
117-
self.initTimeLeft = ceDict.get("CPUTime", self.initTimeLeft)
118-
self.initTimeLeft = gConfig.getValue("/Resources/Computing/CEDefaults/MaxCPUTime", self.initTimeLeft)
110+
# Read initial CPU work left from config (seeded by pilot via dirac-wms-get-queue-cpu-time)
111+
self.initTimeLeft = gConfig.getValue("/LocalSite/CPUTimeLeft", self.initTimeLeft)
119112
self.timeLeft = self.initTimeLeft
120113

121-
self.initTimes = os.times()
114+
self.initTime = time.time()
122115
# Localsite options
123116
self.siteName = siteName()
124117
self.pilotReference = gConfig.getValue("/LocalSite/PilotReference", self.pilotReference)
@@ -136,9 +129,6 @@ def initialize(self):
136129
self.logLevel = self.am_getOption("DefaultLogLevel", self.logLevel)
137130
self.defaultWrapperLocation = self.am_getOption("JobWrapperTemplate", self.defaultWrapperLocation)
138131

139-
# Utilities
140-
self.timeLeftUtil = TimeLeft()
141-
142132
# Some innerCEs may want to make use of CGroup2 support, so we prepare it globally here
143133
res = CG2Manager().setUp()
144134
if res["OK"]:
@@ -403,24 +393,19 @@ def _checkCEAvailability(self, computingElement):
403393
return S_OK()
404394

405395
#############################################################################
406-
def _computeCPUWorkLeft(self, processors=1):
396+
def _computeCPUWorkLeft(self):
407397
"""
408-
Compute CPU Work Left in hepspec06 seconds
398+
Compute CPU Work Left in hepspec06 seconds.
399+
400+
Uses a simple wall-clock countdown from the initial value (seeded by the pilot
401+
via dirac-wms-get-queue-cpu-time). The elapsed wall-clock time is multiplied by
402+
the CPU normalization factor to get the consumed CPU work.
409403
410-
:param int processors: number of processors available
411404
:return: cpu work left (cpu time left * cpu power of the cpus)
412405
"""
413-
# Sum all times but the last one (elapsed_time) and remove times at init (is this correct?)
414-
cpuTimeConsumed = sum(os.times()[:-1]) - sum(self.initTimes[:-1])
415-
result = self.timeLeftUtil.getTimeLeft(cpuTimeConsumed, processors)
416-
if not result["OK"]:
417-
self.log.warn("There were errors calculating time left using the Timeleft utility", result["Message"])
418-
self.log.warn("The time left will be calculated using os.times() and the info in our possession")
419-
self.log.info(f"Current raw CPU time consumed is {cpuTimeConsumed}")
420-
if self.cpuFactor:
421-
return self.initTimeLeft - cpuTimeConsumed * self.cpuFactor
422-
return self.timeLeft
423-
return result["Value"]
406+
elapsed = time.time() - self.initTime
407+
cpuWorkConsumed = elapsed * self.cpuFactor
408+
return self.initTimeLeft - cpuWorkConsumed
424409

425410
def _checkCPUWorkLeft(self, cpuWorkLeft):
426411
"""Check that fillingMode is enabled and time left is sufficient to continue the execution"""

src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_JobAgent.py

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
from DIRAC.Core.Security.X509Chain import X509Chain # pylint: disable=import-error
1212

1313
from DIRAC import S_ERROR, S_OK, gLogger
14-
from DIRAC.Resources.Computing.BatchSystems.TimeLeft.TimeLeft import TimeLeft
1514
from DIRAC.Resources.Computing.ComputingElementFactory import ComputingElementFactory
1615
from DIRAC.Resources.Computing.test.Test_PoolComputingElement import badJobScript, jobScript
1716
from DIRAC.WorkloadManagementSystem.Agent.JobAgent import JobAgent
@@ -150,28 +149,27 @@ def test__checkCEAvailability(mocker, ceType, mockCEReply, expectedResult):
150149

151150

152151
@pytest.mark.parametrize(
153-
"initTimeLeft, timeLeft, cpuFactor, mockTimeLeftReply, expectedTimeLeft",
152+
"initTimeLeft, cpuFactor, elapsedSeconds, expectedTimeLeft",
154153
[
155-
(100000, 75000, None, {"OK": False, "Message": "Error"}, 75000),
156-
(100000, 75000, 10, {"OK": False, "Message": "Error"}, 100000),
157-
(100000, 75000, 10, {"OK": True, "Value": 25000}, 25000),
154+
# No CPU factor: no work consumed, time left equals initial
155+
(100000, 0, 100, 100000),
156+
# With CPU factor: elapsed * cpuFactor is subtracted from initTimeLeft
157+
(100000, 10, 100, 99000),
158+
# Longer elapsed time
159+
(100000, 10, 5000, 50000),
158160
],
159161
)
160-
def test__computeCPUWorkLeft(mocker, initTimeLeft, timeLeft, cpuFactor, mockTimeLeftReply, expectedTimeLeft):
162+
def test__computeCPUWorkLeft(mocker, initTimeLeft, cpuFactor, elapsedSeconds, expectedTimeLeft):
161163
"""Test JobAgent()._computeCPUWorkLeft()"""
162164
mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobAgent.AgentModule.__init__")
163-
mocker.patch(
164-
"DIRAC.Resources.Computing.BatchSystems.TimeLeft.TimeLeft.TimeLeft.getTimeLeft", return_value=mockTimeLeftReply
165-
)
166165

167166
jobAgent = JobAgent("Test", "Test1")
168167
jobAgent.log = gLogger
169168
jobAgent.log.setLevel("DEBUG")
170-
jobAgent.timeLeftUtil = TimeLeft()
171169

172170
jobAgent.initTimeLeft = initTimeLeft
173-
jobAgent.timeLeft = timeLeft
174171
jobAgent.cpuFactor = cpuFactor
172+
jobAgent.initTime = time.time() - elapsedSeconds
175173
result = jobAgent._computeCPUWorkLeft()
176174

177175
assert abs(result - expectedTimeLeft) < 10

src/DIRAC/WorkloadManagementSystem/Client/CPUNormalization.py

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,13 @@
1616
def getCPUTime(cpuNormalizationFactor):
1717
"""Trying to get CPUTime left for execution (in seconds).
1818
19-
It will first look to get the work left looking for batch system information useing the TimeLeft utility.
20-
If it succeeds, it will convert it in real second, and return it.
19+
It will first look to get the CPU work left from the local configuration
20+
(written by the JobAgent every cycle). If not available (e.g. at pilot bootstrap),
21+
it will try to compute it from the batch system using the TimeLeft utility.
2122
22-
If it fails, it tries to get it from the static info found in CS.
23+
If both fail, it tries to get it from the static info found in CS.
2324
If it fails, it returns the default, which is a large 9999999, that we may consider as "Infinite".
2425
25-
This is a generic method, independent from the middleware of the resource if TimeLeft doesn't return a value
26-
2726
args:
2827
cpuNormalizationFactor (float): the CPU power of the current Worker Node.
2928
If not passed in, it's get from the local configuration
@@ -35,22 +34,23 @@ def getCPUTime(cpuNormalizationFactor):
3534
cpuWorkLeft = gConfig.getValue("/LocalSite/CPUTimeLeft", 0)
3635

3736
if not cpuWorkLeft:
38-
# Try and get the information from the CPU left utility
37+
# At pilot bootstrap, CPUTimeLeft is not yet in config.
38+
# Try to compute it from the batch system (sacct, qstat, etc.)
3939
result = TimeLeft().getTimeLeft()
4040
if result["OK"]:
4141
cpuWorkLeft = result["Value"]
4242

4343
if cpuWorkLeft > 0:
44-
# This is in HS06sseconds
44+
# This is in HS06*seconds
4545
# We need to convert in real seconds
4646
if not cpuNormalizationFactor: # if cpuNormalizationFactor passed in is 0, try get it from the local cfg
4747
cpuNormalizationFactor = gConfig.getValue("/LocalSite/CPUNormalizationFactor", 0.0)
4848
if cpuNormalizationFactor:
4949
cpuTimeLeft = cpuWorkLeft / cpuNormalizationFactor
5050

5151
if not cpuTimeLeft:
52-
# now we know that we have to find the CPUTimeLeft by looking in the CS
53-
# this is not granted to be correct as the CS units may not be real seconds
52+
# Try to get CPUTimeLeft from the queue configuration in the CS.
53+
# These values are wall-clock minutes from BDII, so we convert to seconds.
5454
gridCE = gConfig.getValue("/LocalSite/GridCE")
5555
ceQueue = gConfig.getValue("/LocalSite/CEQueue")
5656
if not ceQueue:
@@ -63,23 +63,27 @@ def getCPUTime(cpuNormalizationFactor):
6363
if not res["OK"]:
6464
raise RuntimeError(res["Message"])
6565
queues = res["Value"]
66-
cpuTimes = [gConfig.getValue(queueSection + "/" + queue + "/maxCPUTime", 9999999.0) for queue in queues]
67-
# These are (real, wall clock) minutes - damn BDII!
68-
cpuTimeLeft = min(cpuTimes) * 60
66+
cpuTimes = [gConfig.getValue(queueSection + "/" + queue + "/maxCPUTime", 0.0) for queue in queues]
67+
cpuTimes = [t for t in cpuTimes if t > 0]
68+
if cpuTimes:
69+
cpuTimeLeft = min(cpuTimes) * 60
6970
else:
7071
queueInfo = getQueueInfo(f"{gridCE}/{ceQueue}")
71-
cpuTimeLeft = 9999999.0
7272
if not queueInfo["OK"] or not queueInfo["Value"]:
73-
gLogger.warn("Can't find a CE/queue, defaulting CPUTime to %d" % cpuTimeLeft)
73+
gLogger.warn("Can't find a CE/queue in CS")
7474
else:
7575
queueCSSection = queueInfo["Value"]["QueueCSSection"]
76-
# These are (real, wall clock) minutes - damn BDII!
7776
cpuTimeInMinutes = gConfig.getValue(f"{queueCSSection}/maxCPUTime", 0.0)
7877
if cpuTimeInMinutes:
7978
cpuTimeLeft = cpuTimeInMinutes * 60.0
8079
gLogger.info(f"CPUTime for {queueCSSection}: {cpuTimeLeft:f}")
8180
else:
82-
gLogger.warn(f"Can't find maxCPUTime for {queueCSSection}, defaulting CPUTime to {cpuTimeLeft:f}")
81+
gLogger.warn(f"Can't find maxCPUTime for {queueCSSection}")
82+
83+
if not cpuTimeLeft:
84+
# Last resort: global default from CS, or hardcoded large value
85+
cpuTimeLeft = gConfig.getValue("/Resources/Computing/CEDefaults/MaxCPUTime", 9999999)
86+
gLogger.warn(f"Using fallback MaxCPUTime: {cpuTimeLeft}")
8387

8488
return int(cpuTimeLeft)
8589

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
"""Unit tests for CPUNormalization.getCPUTime()"""
2+
from unittest.mock import patch
3+
4+
from DIRAC import S_OK, S_ERROR
5+
6+
7+
@patch("DIRAC.WorkloadManagementSystem.Client.CPUNormalization.TimeLeft")
8+
@patch("DIRAC.WorkloadManagementSystem.Client.CPUNormalization.gConfig")
9+
class TestGetCPUTime:
10+
"""Tests for getCPUTime() fallback chain."""
11+
12+
def _import_getCPUTime(self):
13+
from DIRAC.WorkloadManagementSystem.Client.CPUNormalization import getCPUTime
14+
15+
return getCPUTime
16+
17+
def test_from_config(self, mock_gConfig, mock_TimeLeft):
18+
"""Primary path: CPUTimeLeft is in config (written by JobAgent)."""
19+
mock_gConfig.getValue.side_effect = lambda key, default=0: {
20+
"/LocalSite/CPUTimeLeft": 50000, # HS06*s
21+
"/LocalSite/CPUNormalizationFactor": 10.0,
22+
}.get(key, default)
23+
24+
result = self._import_getCPUTime()(cpuNormalizationFactor=10.0)
25+
26+
# 50000 / 10.0 = 5000 seconds
27+
assert result == 5000
28+
# TimeLeft should NOT be instantiated since config had a value
29+
mock_TimeLeft.assert_not_called()
30+
31+
def test_from_batch_system(self, mock_gConfig, mock_TimeLeft):
32+
"""Bootstrap path: no CPUTimeLeft in config, falls back to batch system."""
33+
mock_gConfig.getValue.side_effect = lambda key, default=0: {
34+
"/LocalSite/CPUTimeLeft": 0, # not set yet
35+
"/LocalSite/CPUNormalizationFactor": 10.0,
36+
}.get(key, default)
37+
mock_TimeLeft.return_value.getTimeLeft.return_value = S_OK(30000) # HS06*s
38+
39+
result = self._import_getCPUTime()(cpuNormalizationFactor=10.0)
40+
41+
# 30000 / 10.0 = 3000 seconds
42+
assert result == 3000
43+
mock_TimeLeft.return_value.getTimeLeft.assert_called_once()
44+
45+
def test_from_queue_cs(self, mock_gConfig, mock_TimeLeft):
46+
"""Fallback: batch system fails, uses queue maxCPUTime from CS."""
47+
mock_TimeLeft.return_value.getTimeLeft.return_value = S_ERROR("No batch info")
48+
49+
config_values = {
50+
"/LocalSite/CPUTimeLeft": 0,
51+
"/LocalSite/GridCE": "ce.example.com",
52+
"/LocalSite/CEQueue": "default",
53+
"/LocalSite/Site": "LCG.Example.com",
54+
}
55+
56+
def mock_getValue(key, default=0):
57+
if key in config_values:
58+
return config_values[key]
59+
# maxCPUTime in minutes
60+
if "maxCPUTime" in key:
61+
return 120.0 # 120 minutes
62+
return default
63+
64+
mock_gConfig.getValue.side_effect = mock_getValue
65+
66+
with patch(
67+
"DIRAC.WorkloadManagementSystem.Client.CPUNormalization.getQueueInfo",
68+
return_value=S_OK(
69+
{"QueueCSSection": "/Resources/Sites/LCG/LCG.Example.com/CEs/ce.example.com/Queues/default"}
70+
),
71+
):
72+
result = self._import_getCPUTime()(cpuNormalizationFactor=10.0)
73+
74+
# 120 minutes * 60 = 7200 seconds
75+
assert result == 7200
76+
77+
def test_fallback_max_cpu_time(self, mock_gConfig, mock_TimeLeft):
78+
"""Last resort: everything fails, uses /Resources/Computing/CEDefaults/MaxCPUTime."""
79+
mock_TimeLeft.return_value.getTimeLeft.return_value = S_ERROR("No batch info")
80+
81+
config_values = {
82+
"/LocalSite/CPUTimeLeft": 0,
83+
"/LocalSite/GridCE": "ce.example.com",
84+
"/LocalSite/CEQueue": "default",
85+
"/LocalSite/Site": "LCG.Example.com",
86+
"/Resources/Computing/CEDefaults/MaxCPUTime": 86400,
87+
}
88+
89+
def mock_getValue(key, default=0):
90+
if key in config_values:
91+
return config_values[key]
92+
return default
93+
94+
mock_gConfig.getValue.side_effect = mock_getValue
95+
96+
with patch(
97+
"DIRAC.WorkloadManagementSystem.Client.CPUNormalization.getQueueInfo",
98+
return_value=S_OK(
99+
{"QueueCSSection": "/Resources/Sites/LCG/LCG.Example.com/CEs/ce.example.com/Queues/default"}
100+
),
101+
):
102+
result = self._import_getCPUTime()(cpuNormalizationFactor=10.0)
103+
104+
# maxCPUTime from queue returned 0, so falls through to CEDefaults/MaxCPUTime
105+
assert result == 86400
106+
107+
def test_hardcoded_fallback(self, mock_gConfig, mock_TimeLeft):
108+
"""Absolute last resort: no CS default either, returns 9999999."""
109+
mock_TimeLeft.return_value.getTimeLeft.return_value = S_ERROR("No batch info")
110+
111+
config_values = {
112+
"/LocalSite/CPUTimeLeft": 0,
113+
"/LocalSite/GridCE": "ce.example.com",
114+
"/LocalSite/CEQueue": "default",
115+
"/LocalSite/Site": "LCG.Example.com",
116+
}
117+
118+
def mock_getValue(key, default=0):
119+
if key in config_values:
120+
return config_values[key]
121+
return default
122+
123+
mock_gConfig.getValue.side_effect = mock_getValue
124+
125+
with patch(
126+
"DIRAC.WorkloadManagementSystem.Client.CPUNormalization.getQueueInfo",
127+
return_value=S_OK(
128+
{"QueueCSSection": "/Resources/Sites/LCG/LCG.Example.com/CEs/ce.example.com/Queues/default"}
129+
),
130+
):
131+
result = self._import_getCPUTime()(cpuNormalizationFactor=10.0)
132+
133+
assert result == 9999999
134+
135+
def test_cpu_normalization_factor_from_config(self, mock_gConfig, mock_TimeLeft):
136+
"""When cpuNormalizationFactor=0, it should be read from config."""
137+
mock_gConfig.getValue.side_effect = lambda key, default=0: {
138+
"/LocalSite/CPUTimeLeft": 50000,
139+
"/LocalSite/CPUNormalizationFactor": 5.0,
140+
}.get(key, default)
141+
142+
result = self._import_getCPUTime()(cpuNormalizationFactor=0)
143+
144+
# 50000 / 5.0 = 10000 seconds
145+
assert result == 10000

0 commit comments

Comments
 (0)