From 553f7aeabaa2c7ff415ef0e8210a1e760b933291 Mon Sep 17 00:00:00 2001 From: Kazuki Sakamoto Date: Wed, 20 Oct 2021 12:01:16 -0700 Subject: [PATCH] [pycue] Improve GPU API methods. (#1050) --- cuegui/cuegui/ServiceDialog.py | 2 +- pycue/opencue/wrappers/service.py | 48 +++++++++++++++++++++++----- pycue/opencue/wrappers/show.py | 4 +-- pycue/tests/wrappers/service_test.py | 24 ++++++++++++++ pycue/tests/wrappers/show_test.py | 26 +++++++++++++++ 5 files changed, 93 insertions(+), 11 deletions(-) diff --git a/cuegui/cuegui/ServiceDialog.py b/cuegui/cuegui/ServiceDialog.py index 5f78cf670..bdeec4dd6 100644 --- 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 index 1a941772b..3878d9b3d 100644 --- 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 index 9fff4b47d..90094e757 100644 --- 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 index 92b2e3b2c..e46b29bdf 100644 --- 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 index b7cc3dd9a..b8fb26285 100644 --- 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(