11767: Correctly parse values with unit suffixes printed by squeue.

diff --git a/services/nodemanager/arvnodeman/jobqueue.py b/services/nodemanager/arvnodeman/jobqueue.py
index 66cf73acb357f4e5cd78dfcd20793cf697749be5..cf23df0022286505053bf807ce373eaa910054ca 100644 (file)
--- a/services/nodemanager/arvnodeman/jobqueue.py
+++ b/services/nodemanager/arvnodeman/jobqueue.py
@@ -117,6 +117,19 @@ class JobQueueMonitorActor(clientactor.RemotePollLoopActor):
             client, timer_actor, *args, **kwargs)
         self._calculator = server_calc
 
+    @staticmethod
+    def coerce_to_mb(x):
+        if x.endswith("M"):
+            return int(x[:-1])
+        elif x.endswith("G"):
+            return int(x[:-1]) * 2**10
+        elif x.endswith("T"):
+            return int(x[:-1]) * 2**20
+        elif x.endswith("P"):
+            return int(x[:-1]) * 2**30
+        else:
+            return int(x)
+
     def _send_request(self):
         # cpus, memory, tempory disk space, reason, job name
         squeue_out = subprocess.check_output(["squeue", "--state=PENDING", "--noheader", "--format=%c %m %d %r %j"])
@@ -128,8 +141,8 @@ class JobQueueMonitorActor(clientactor.RemotePollLoopActor):
                     "uuid": jobname,
                     "runtime_constraints": {
                         "min_cores_per_node": cpu,
-                        "min_ram_mb_per_node": ram,
-                        "min_scratch_mb_per_node": disk
+                        "min_ram_mb_per_node": self.coerce_to_mb(ram),
+                        "min_scratch_mb_per_node": self.coerce_to_mb(disk)
                     }
                 })
 

diff --git a/services/nodemanager/tests/test_jobqueue.py b/services/nodemanager/tests/test_jobqueue.py
index 08a813185ed6f52fbf90c22430435f4b25075c8d..8d9c54c33f5b07f5feb5d7d13b3e630dff8b9eac 100644 (file)
--- a/services/nodemanager/tests/test_jobqueue.py
+++ b/services/nodemanager/tests/test_jobqueue.py
@@ -150,17 +150,44 @@ class JobQueueMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
 
     @mock.patch("subprocess.check_output")
     def test_squeue_server_list(self, mock_squeue):
-        mock_squeue.return_value = """1 0 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzy
-2 0 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzz
+        mock_squeue.return_value = """1 1024 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzy
+2 1024 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzz
 """
 
         super(JobQueueMonitorActorTestCase, self).build_monitor(jobqueue.ServerCalculator(
-            [(testutil.MockSize(n), {'cores': n, 'ram': n, 'scratch': n}) for n in range(1, 3)]))
+            [(testutil.MockSize(n), {'cores': n, 'ram': n*1024, 'scratch': n}) for n in range(1, 3)]))
         self.monitor.subscribe(self.subscriber).get(self.TIMEOUT)
         self.stop_proxy(self.monitor)
         self.subscriber.assert_called_with([testutil.MockSize(1),
                                             testutil.MockSize(2)])
 
+    @mock.patch("subprocess.check_output")
+    def test_squeue_server_list_suffix(self, mock_squeue):
+        mock_squeue.return_value = """1 1024M 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzy
+1 2G 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzz
+"""
+
+        super(JobQueueMonitorActorTestCase, self).build_monitor(jobqueue.ServerCalculator(
+            [(testutil.MockSize(n), {'cores': n, 'ram': n*1024, 'scratch': n}) for n in range(1, 3)]))
+        self.monitor.subscribe(self.subscriber).get(self.TIMEOUT)
+        self.stop_proxy(self.monitor)
+        self.subscriber.assert_called_with([testutil.MockSize(1),
+                                            testutil.MockSize(2)])
+
+    def test_coerce_to_mb(self):
+        self.assertEqual(1, jobqueue.JobQueueMonitorActor.coerce_to_mb("1"))
+        self.assertEqual(512, jobqueue.JobQueueMonitorActor.coerce_to_mb("512"))
+        self.assertEqual(512, jobqueue.JobQueueMonitorActor.coerce_to_mb("512M"))
+        self.assertEqual(1024, jobqueue.JobQueueMonitorActor.coerce_to_mb("1024M"))
+        self.assertEqual(1024, jobqueue.JobQueueMonitorActor.coerce_to_mb("1G"))
+        self.assertEqual(2048, jobqueue.JobQueueMonitorActor.coerce_to_mb("2G"))
+        self.assertEqual(1025, jobqueue.JobQueueMonitorActor.coerce_to_mb("1025M"))
+        self.assertEqual(1048576, jobqueue.JobQueueMonitorActor.coerce_to_mb("1T"))
+        self.assertEqual(1073741824, jobqueue.JobQueueMonitorActor.coerce_to_mb("1P"))
+        self.assertEqual(0, jobqueue.JobQueueMonitorActor.coerce_to_mb("0"))
+        self.assertEqual(0, jobqueue.JobQueueMonitorActor.coerce_to_mb("0M"))
+        self.assertEqual(0, jobqueue.JobQueueMonitorActor.coerce_to_mb("0G"))
+
 
 if __name__ == '__main__':
     unittest.main()