From 574f29e8f5e48dc97fa3895cde2ebdef1f57dfba Mon Sep 17 00:00:00 2001 From: Kazuki Sakamoto Date: Sun, 21 Feb 2021 22:23:40 -0800 Subject: [PATCH] [PyCue] Supoprt gpus and gpu_memory Co-authored-by: Lars van der Bijl <285658+larsbijl@users.noreply.github.com> --- pycue/opencue/wrappers/group.py | 38 ++++++++++++++++++++++ pycue/opencue/wrappers/job.py | 51 +++++++++++++++++++++++++++--- pycue/opencue/wrappers/layer.py | 38 +++++++++++++++++++--- pycue/opencue/wrappers/show.py | 26 +++++++++++++++ pycue/tests/wrappers/layer_test.py | 10 +++--- 5 files changed, 149 insertions(+), 14 deletions(-) diff --git a/pycue/opencue/wrappers/group.py b/pycue/opencue/wrappers/group.py index 59685cc8b..1c8175e57 100644 --- a/pycue/opencue/wrappers/group.py +++ b/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. diff --git a/pycue/opencue/wrappers/job.py b/pycue/opencue/wrappers/job.py index eb977a280..733c91e61 100644 --- a/pycue/opencue/wrappers/job.py +++ b/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,8 +236,10 @@ 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, @@ -231,7 +247,8 @@ def addRenderPartition(self, hostname, threads, max_cores, num_mem, max_gpu): 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. diff --git a/pycue/opencue/wrappers/layer.py b/pycue/opencue/wrappers/layer.py index 7c33b68e7..0e34985bf 100644 --- a/pycue/opencue/wrappers/layer.py +++ b/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. diff --git a/pycue/opencue/wrappers/show.py b/pycue/opencue/wrappers/show.py index 750d05645..9fff4b47d 100644 --- a/pycue/opencue/wrappers/show.py +++ b/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. diff --git a/pycue/tests/wrappers/layer_test.py b/pycue/tests/wrappers/layer_test.py index cf8fb0c33..4f5578681 100644 --- a/pycue/tests/wrappers/layer_test.py +++ b/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):