[pycue] Improve GPU API methods. (#1050)

AcademySoftwareFoundation · Oct 20, 2021 · 553f7ae · 553f7ae
1 parent db938da
commit 553f7ae
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 11 deletions.
diff --git a/cuegui/cuegui/ServiceDialog.py b/cuegui/cuegui/ServiceDialog.py
@@ -168,7 +168,7 @@ def save(self):
  service.setMinCores(self.min_cores.value())
  service.setMaxCores(self.max_cores.value())
  service.setMinMemory(self.min_memory.value() * 1024)
- service.setMinGpu(self.min_gpu_memory.value() * 1024)
+ service.setMinGpuMemory(self.min_gpu_memory.value() * 1024)
  service.setTimeout(self.timeout.value())
  service.setTimeoutLLU(self.timeout_llu.value())
  service.setTags(self._tags_w.get_tags())

diff --git a/pycue/opencue/wrappers/service.py b/pycue/opencue/wrappers/service.py
@@ -161,21 +161,53 @@ def setMinMemory(self, minMemory):
  """
  self.data.min_memory = minMemory
 
- def minGpu(self):
- """Returns the minimum gpu of the service.
+ def minGpus(self):
+ """Returns the minimum gpus of the service.
 
  :rtype: int
- :return: min gpu
+ :return: min gpus
  """
- return self.data.min_gpu
+ return self.data.min_gpus
 
- def setMinGpu(self, minGpu):
- """Sets the minimum gpu of the service.
+ def setMinGpus(self, minGpus):
+ """Sets the minimum gpus of the service.
 
  :type: int
- :param: new min gpu
+ :param: new min gpus
  """
- self.data.min_gpu = minGpu
+ self.data.min_gpus = minGpus
+
+ def maxGpus(self):
+ """Returns the maximum gpus of the service.
+
+ :rtype: int
+ :return: max gpus
+ """
+ return self.data.max_gpus
+
+ def setMaxGpus(self, maxGpus):
+ """Sets the maximum gpus of the service.
+
+ :type: int
+ :param: new max gpus
+ """
+ self.data.max_gpus = maxGpus
+
+ def minGpuMemory(self):
+ """Returns the minimum gpu memory of the service.
+
+ :rtype: int
+ :return: min gpu memory
+ """
+ return self.data.min_gpu_memory
+
+ def setMinGpuMemory(self, minGpuMemory):
+ """Sets the minimum gpu memory of the service.
+
+ :type: int
+ :param: new min gpu memory
+ """
+ self.data.min_gpu_memory = minGpuMemory
 
  def tags(self):
  """Returns the list of tags for the service.

diff --git a/pycue/opencue/wrappers/show.py b/pycue/opencue/wrappers/show.py
@@ -176,7 +176,7 @@ def setDefaultMaxGpus(self, maxgpus):
  :return: response is empty
  """
  response = self.stub.SetDefaultMaxGpus(show_pb2.ShowSetDefaultMaxGpusRequest(
- show=self.data, max_gpu=maxgpus),
+ show=self.data, max_gpus=maxgpus),
  timeout=Cuebot.Timeout)
  return response
 
@@ -189,7 +189,7 @@ def setDefaultMinGpus(self, mingpus):
  :return: response is empty
  """
  response = self.stub.SetDefaultMinGpus(show_pb2.ShowSetDefaultMinGpusRequest(
- show=self.data, min_gpu=mingpus),
+ show=self.data, min_gpus=mingpus),
  timeout=Cuebot.Timeout)
  return response
 

diff --git a/pycue/tests/wrappers/service_test.py b/pycue/tests/wrappers/service_test.py
@@ -28,6 +28,9 @@
 
 
 TEST_SERVICE_NAME = 'testService'
+TEST_MIN_GPUS = 2
+TEST_MAX_GPUS = 7
+TEST_MIN_GPU_MEMORY = 8 * 1024 * 1024 * 1024
 
 
 @mock.patch('opencue.cuebot.Cuebot.getStub')
@@ -99,6 +102,27 @@ def testUpdate(self, getStubMock):
  stubMock.Update.assert_called_with(
  service_pb2.ServiceUpdateRequest(service=wrapper.data), timeout=mock.ANY)
 
+ def testGpus(self, getStubMock):
+ stubMock = mock.Mock()
+ stubMock.GetService.return_value = service_pb2.ServiceGetServiceResponse(
+ service=service_pb2.Service(name=TEST_SERVICE_NAME))
+ getStubMock.return_value = stubMock
+
+ wrapper = opencue.wrappers.service.Service()
+ service = wrapper.getService(name=TEST_SERVICE_NAME)
+ self.assertEqual(service.minGpus(), 0)
+ self.assertEqual(service.maxGpus(), 0)
+ self.assertEqual(service.minGpuMemory(), 0)
+ service.setMinGpus(TEST_MIN_GPUS)
+ service.setMaxGpus(TEST_MAX_GPUS)
+ service.setMinGpuMemory(TEST_MIN_GPU_MEMORY)
+ self.assertEqual(service.minGpus(), TEST_MIN_GPUS)
+ self.assertEqual(service.maxGpus(), TEST_MAX_GPUS)
+ self.assertEqual(service.minGpuMemory(), TEST_MIN_GPU_MEMORY)
+
+ stubMock.GetService.assert_called_with(
+ service_pb2.ServiceGetServiceRequest(name=TEST_SERVICE_NAME), timeout=mock.ANY)
+ self.assertEqual(service.name(), TEST_SERVICE_NAME)
 
 if __name__ == '__main__':
  unittest.main()
diff --git a/pycue/tests/wrappers/show_test.py b/pycue/tests/wrappers/show_test.py
@@ -44,6 +44,8 @@
 TEST_SUBSCRIPTION_BURST = 1200
 TEST_MIN_CORES = 42
 TEST_MAX_CORES = 47
+TEST_MIN_GPUS = 2
+TEST_MAX_GPUS = 7
 TEST_ENABLE_VALUE = False
 TEST_GROUP_NAME = 'group'
 TEST_GROUP_DEPT = 'lighting'
@@ -197,6 +199,30 @@ def testSetDefaultMinCores(self, getStubMock):
  show_pb2.ShowSetDefaultMinCoresRequest(show=show.data, min_cores=TEST_MIN_CORES),
  timeout=mock.ANY)
 
+ def testSetDefaultMaxGpus(self, getStubMock):
+ stubMock = mock.Mock()
+ stubMock.SetDefaultMaxGpus.return_value = show_pb2.ShowSetDefaultMaxGpusResponse()
+ getStubMock.return_value = stubMock
+
+ show = opencue.wrappers.show.Show(show_pb2.Show(name=TEST_SHOW_NAME))
+ show.setDefaultMaxGpus(TEST_MAX_GPUS)
+
+ stubMock.SetDefaultMaxGpus.assert_called_with(
+ show_pb2.ShowSetDefaultMaxGpusRequest(show=show.data, max_gpus=TEST_MAX_GPUS),
+ timeout=mock.ANY)
+
+ def testSetDefaultMinGpus(self, getStubMock):
+ stubMock = mock.Mock()
+ stubMock.SetDefaultMinGpus.return_value = show_pb2.ShowSetDefaultMinGpusResponse()
+ getStubMock.return_value = stubMock
+
+ show = opencue.wrappers.show.Show(show_pb2.Show(name=TEST_SHOW_NAME))
+ show.setDefaultMinGpus(TEST_MIN_GPUS)
+
+ stubMock.SetDefaultMinGpus.assert_called_with(
+ show_pb2.ShowSetDefaultMinGpusRequest(show=show.data, min_gpus=TEST_MIN_GPUS),
+ timeout=mock.ANY)
+
  def testFindFilter(self, getStubMock):
  stubMock = mock.Mock()
  stubMock.FindFilter.return_value = show_pb2.ShowFindFilterResponse(