[PyCue] Supoprt gpus and gpu_memory

Co-authored-by: Lars van der Bijl <[email protected]>

pycue/opencue/wrappers/group.py

@@ -67,6 +67,24 @@ def setMinCores(self, value):
         self.stub.SetMinCores(job_pb2.GroupSetMinCoresRequest(group=self.data, min_cores=value),
                               timeout=Cuebot.Timeout)
 
+    def setMaxGpus(self, value):
+        """Sets the maximum gpus of everything in the group.
+
+        :type  value: int
+        :param value: new maximum number of gpus
+        """
+        self.stub.SetMaxGpus(job_pb2.GroupSetMaxGpusRequest(group=self.data, max_gpus=value),
+                             timeout=Cuebot.Timeout)
+
+    def setMinGpus(self, value):
+        """Sets the minimum gpus of everything the group.
+
+        :type  value: int
+        :param value: new minimum number of gpus
+        """
+        self.stub.SetMinGpus(job_pb2.GroupSetMinGpusRequest(group=self.data, min_gpus=value),
+                             timeout=Cuebot.Timeout)
+
     def setDefaultJobPriority(self, value):
         """Sets the default job priority for everything in the group.
 
@@ -97,6 +115,26 @@ def setDefaultJobMaxCores(self, value):
             job_pb2.GroupSetDefJobMaxCoresRequest(group=self.data, max_cores=value),
             timeout=Cuebot.Timeout)
 
+    def setDefaultJobMinGpus(self, value):
+        """Sets the default job minimum gpus for everything in the group.
+
+        :type  value: int
+        :param value: new default job minimum gpus
+        """
+        self.stub.SetDefaultJobMinGpus(
+            job_pb2.GroupSetDefJobMinGpusRequest(group=self.data, min_gpus=value),
+            timeout=Cuebot.Timeout)
+
+    def setDefaultJobMaxGpus(self, value):
+        """Sets the default job maximum gpus for everything in the group.
+
+        :type  value: int
+        :param value: new default job maximum gpus
+        """
+        self.stub.SetDefaultJobMaxGpus(
+            job_pb2.GroupSetDefJobMaxGpusRequest(group=self.data, max_gpus=value),
+            timeout=Cuebot.Timeout)
+
     def getGroups(self):
         """Returns child groups of this group.
 

pycue/opencue/wrappers/job.py

@@ -126,6 +126,20 @@ def setMaxCores(self, maxCores):
         self.stub.SetMaxCores(job_pb2.JobSetMaxCoresRequest(job=self.data, val=maxCores),
                               timeout=Cuebot.Timeout)
 
+    def setMinGpus(self, minGpus):
+        """Sets the minimum procs value
+        :type  minGpus: int
+        :param minGpus: New minimum cores value"""
+        self.stub.SetMinGpus(job_pb2.JobSetMinGpusRequest(job=self.data, val=minGpus),
+                             timeout=Cuebot.Timeout)
+
+    def setMaxGpus(self, maxGpus):
+        """Sets the maximum procs value
+        :type  maxGpus: int
+        :param maxGpus: New maximum cores value"""
+        self.stub.SetMaxGpus(job_pb2.JobSetMaxGpusRequest(job=self.data, val=maxGpus),
+                             timeout=Cuebot.Timeout)
+
     def setPriority(self, priority):
         """Sets the job priority.
 
@@ -211,7 +225,7 @@ def setAutoEating(self, value):
         self.stub.SetAutoEat(job_pb2.JobSetAutoEatRequest(job=self.data, value=value),
                              timeout=Cuebot.Timeout)
 
-    def addRenderPartition(self, hostname, threads, max_cores, num_mem, max_gpu):
+    def addRenderPartition(self, hostname, threads, max_cores, num_mem, max_gpus, max_gpu_memory):
         """Adds a render partition to the job.
 
         :type  hostname: str
@@ -222,16 +236,19 @@ def addRenderPartition(self, hostname, threads, max_cores, num_mem, max_gpu):
         :param max_cores: max cores enabled for the partition
         :type  num_mem: int
         :param num_mem: amount of memory reserved for the partition
-        :type  max_gpu: int
-        :param max_gpu: max gpu cores enabled for the partition
+        :type  max_gpus: int
+        :param max_gpus: max gpu cores enabled for the partition
+        :type  max_gpu_memory: int
+        :param max_gpu_memory: amount of gpu memory reserved for the partition
         """
         self.stub.AddRenderPartition(
             job_pb2.JobAddRenderPartRequest(job=self.data,
                                             host=hostname,
                                             threads=threads,
                                             max_cores=max_cores,
                                             max_memory=num_mem,
-                                            max_gpu=max_gpu,
+                                            max_gpus=max_gpus,
+                                            max_gpu_memory=max_gpu_memory,
                                             username=os.getenv("USER", "unknown")))
 
     def getWhatDependsOnThis(self):
@@ -492,6 +509,20 @@ def maxCores(self):
         """
         return self.data.max_cores
 
+    def minGpus(self):
+        """Returns the minimum number of gpus the job needs.
+        :rtype:  int
+        :return: job's min gpus
+        """
+        return self.data.min_gpus
+
+    def maxGpus(self):
+        """Returns the maximum number of gpus the job will use.
+        :rtype:  int
+        :return: job's max gpus
+        """
+        return self.data.max_gpus
+
     def os(self):
         """Returns the job's operating system.
 
@@ -823,6 +854,18 @@ def setMaxCores(self, maxCores):
         """
         self.asJob().setMaxCores(maxCores)
 
+    def setMinGpus(self, minGpus):
+        """Sets the minimum gpus value
+        :type  minGpus: int
+        :param minGpus: New minimum gpus value"""
+        self.asJob().setMinGpus(minGpus)
+
+    def setMaxGpus(self, maxGpus):
+        """Sets the maximum gpus value
+        :type  maxGpus: int
+        :param maxGpus: New maximum gpus value"""
+        self.asJob().setMaxGpus(maxGpus)
+
     def setPriority(self, priority):
         """Sets the job priority.
 

pycue/opencue/wrappers/layer.py

@@ -140,14 +140,30 @@ def setMinCores(self, cores):
             job_pb2.LayerSetMinCoresRequest(layer=self.data, cores=cores/100.0),
             timeout=Cuebot.Timeout)
 
-    def setMinGpu(self, gpu):
+    def setMaxGpus(self, max_gpus):
+        """Sets the maximum number of gpus that this layer requires.
+        :type  max_gpus: int
+        :param max_gpus: gpu cores"""
+        return self.stub.SetMaxGpus(
+            job_pb2.LayerSetMaxGpusRequest(layer=self.data, max_gpus=max_gpus),
+            timeout=Cuebot.Timeout)
+
+    def setMinGpus(self, min_gpus):
+        """Sets the minimum number of gpus that this layer requires.
+        :type  min_gpus: int
+        :param min_gpus: gou cores"""
+        return self.stub.SetMinGpus(
+            job_pb2.LayerSetMinGpusRequest(layer=self.data, min_gpus=min_gpus),
+            timeout=Cuebot.Timeout)
+
+    def setMinGpuMemory(self, gpu_memory):
         """Sets the minimum number of gpu memory that this layer requires.
 
-        :type  gpu: int
-        :param gpu: gpu value
+        :type  gpu_memory: int
+        :param gpu_memory: gpu_memory value
         """
-        return self.stub.SetMinGpu(
-            job_pb2.LayerSetMinGpuRequest(layer=self.data, gpu=gpu),
+        return self.stub.SetMinGpuMemory(
+            job_pb2.LayerSetMinGpuMemoryRequest(layer=self.data, gpu_memory=gpu_memory),
             timeout=Cuebot.Timeout)
 
     def setMinMemory(self, memory):
@@ -401,6 +417,12 @@ def coresReserved(self):
         """
         return self.data.layer_stats.reserved_cores
 
+    def gpusReserved(self):
+        """Returns the number of gpus reserved on this layer
+        :rtype: float
+        :return: gpus reserved"""
+        return self.data.layer_stats.reserved_gpus
+
     def minCores(self):
         """Returns the minimum number of cores that frames in this layer require.
 
@@ -409,6 +431,12 @@ def minCores(self):
         """
         return self.data.min_cores
 
+    def minGpus(self):
+        """Returns the minimum number of gpus that frames in this layer require
+        :rtype:  int
+        :return: Minimum number of gpus required"""
+        return self.data.min_gpus
+
     def minMemory(self):
         """Returns the minimum amount of memory that frames in this layer require.
 

pycue/opencue/wrappers/show.py

@@ -167,6 +167,32 @@ def setDefaultMinCores(self, mincores):
             timeout=Cuebot.Timeout)
         return response
 
+    def setDefaultMaxGpus(self, maxgpus):
+        """Sets the default maximum number of gpus
+        that new jobs are launched with.
+        :type: float
+        :param: value to set maxGpu to
+        :rtype: show_pb2.ShowSetDefaultMaxGpuResponse
+        :return: response is empty
+        """
+        response = self.stub.SetDefaultMaxGpus(show_pb2.ShowSetDefaultMaxGpusRequest(
+            show=self.data, max_gpu=maxgpus),
+            timeout=Cuebot.Timeout)
+        return response
+
+    def setDefaultMinGpus(self, mingpus):
+        """Sets the default minimum number of gpus
+        all new jobs are launched with.
+        :type: float
+        :param: value to set minGpus to
+        :rtype: show_pb2.ShowSetDefaultMinGpusResponse
+        :return: response is empty
+        """
+        response = self.stub.SetDefaultMinGpus(show_pb2.ShowSetDefaultMinGpusRequest(
+            show=self.data, min_gpu=mingpus),
+            timeout=Cuebot.Timeout)
+        return response
+
     def findFilter(self, name):
         """Finds a filter by name.
 

pycue/tests/wrappers/layer_test.py

@@ -201,18 +201,18 @@ def testSetMaxCores(self, getStubMock):
             job_pb2.LayerSetMaxCoresRequest(layer=layer.data, cores=testCoresActual),
             timeout=mock.ANY)
 
-    def testSetMinGpu(self, getStubMock):
+    def testSetMinGpuMemory(self, getStubMock):
         stubMock = mock.Mock()
-        stubMock.SetMinGpu.return_value = job_pb2.LayerSetMinGpuResponse()
+        stubMock.SetMinGpuMemory.return_value = job_pb2.LayerSetMinGpuResponse()
         getStubMock.return_value = stubMock
 
         testCores = 100
         layer = opencue.wrappers.layer.Layer(
             job_pb2.Layer(name=TEST_LAYER_NAME))
-        layer.setMinGpu(testCores)
+        layer.setMinGpuMemory(testCores)
 
-        stubMock.SetMinGpu.assert_called_with(
-            job_pb2.LayerSetMinGpuRequest(layer=layer.data, gpu=testCores),
+        stubMock.SetMinGpuMemory.assert_called_with(
+            job_pb2.LayerSetMinGpuMemoryRequest(layer=layer.data, gpu_memory=testCores),
             timeout=mock.ANY)
 
     def testSetMinMemory(self, getStubMock):