pycue: add cloneJob() to duplicate an existing job for resubmission

pycue/opencue/api.py

@@ -18,6 +18,9 @@
 from __future__ import print_function
 from __future__ import division
 
+import logging
+from xml.etree import ElementTree as Et
+
 from opencue_proto import comment_pb2
 from opencue_proto import criterion_pb2
 from opencue_proto import cue_pb2
@@ -59,6 +62,9 @@
 from . import util
 
 
+logger = logging.getLogger("opencue")
+
+
 __protobufs = [comment_pb2, criterion_pb2, cue_pb2, department_pb2, depend_pb2, facility_pb2,
                filter_pb2, host_pb2, job_pb2, renderPartition_pb2, report_pb2, service_pb2,
                show_pb2, subscription_pb2, task_pb2]
@@ -403,6 +409,107 @@ def launchSpecAndWait(spec):
     return [Job(j) for j in jobSeq.jobs]
 
 
+def _addTextElement(parent, tag, text):
+    """Convenience method to create a sub element with text."""
+    element = Et.SubElement(parent, tag)
+    element.text = text
+    return element
+
+
+def cloneJob(job, name=None, user=None, frame_range=None, layer_frame_ranges=None):
+    """Duplicates an existing job's submission-time configuration into a new job spec
+    and launches it.
+
+    This reconstructs a launchable spec from the job and layer data available through the
+    API (commands, services, frame ranges, core/memory/gpu requirements, tags, and limits),
+    so a job can be re-run without needing its original outline/pyoutline submission script.
+    Only ``Render``, ``Util``, and ``Post`` layers can be cloned this way; any ``PreProcess``
+    layers are skipped, as are layers with no command to re-run.
+
+    :type job: opencue.wrappers.job.Job
+    :param job: the job to clone
+    :type name: str
+    :param name: name for the cloned job, defaults to "<original-name>_clone"
+    :type user: str
+    :param user: submitting user for the new job, defaults to the original job's user
+    :type frame_range: str
+    :param frame_range: frame range override applied to every cloned layer
+    :type layer_frame_ranges: dict[str, str]
+    :param layer_frame_ranges: per-layer frame range overrides, keyed by layer name; takes
+        precedence over `frame_range` for the layers it names
+    :rtype: list[opencue.wrappers.job.Job]
+    :return: the newly launched job(s)
+    """
+    layer_frame_ranges = layer_frame_ranges or {}
+    layer_type_names = {
+        job_pb2.RENDER: 'Render',
+        job_pb2.UTIL: 'Util',
+        job_pb2.POST: 'Post',
+    }
+
+    root = Et.Element('spec')
+    _addTextElement(root, 'facility', job.data.facility)
+    _addTextElement(root, 'show', job.data.show)
+    _addTextElement(root, 'shot', job.data.shot)
+    _addTextElement(root, 'user', user or job.data.user)
+
+    jobEl = Et.SubElement(root, 'job', {'name': name or '%s_clone' % job.data.name})
+    _addTextElement(jobEl, 'priority', str(job.data.priority))
+    if job.data.max_cores:
+        _addTextElement(jobEl, 'maxcores', str(job.data.max_cores))
+    if job.data.max_gpus:
+        _addTextElement(jobEl, 'maxgpus', str(job.data.max_gpus))
+    if job.data.os:
+        _addTextElement(jobEl, 'os', job.data.os)
+
+    layersEl = Et.SubElement(jobEl, 'layers')
+    for layer in job.getLayers():
+        typeName = layer_type_names.get(layer.data.type)
+        if typeName is None:
+            logger.warning(
+                "cloneJob: skipping layer %s, %s layers can't be cloned directly.",
+                layer.data.name, Layer.LayerType(layer.data.type).name)
+            continue
+        if not layer.data.command:
+            logger.warning(
+                "cloneJob: skipping layer %s, it has no command to re-run.", layer.data.name)
+            continue
+
+        layerEl = Et.SubElement(layersEl, 'layer', {'name': layer.data.name, 'type': typeName})
+        _addTextElement(layerEl, 'cmd', layer.data.command)
+        _addTextElement(
+            layerEl, 'range',
+            layer_frame_ranges.get(layer.data.name) or frame_range or layer.data.range)
+        _addTextElement(layerEl, 'chunk', str(layer.data.chunk_size or 1))
+        if layer.data.min_cores:
+            _addTextElement(layerEl, 'cores', '%0.1f' % layer.data.min_cores)
+        _addTextElement(layerEl, 'threadable', 'True' if layer.data.is_threadable else 'False')
+        if layer.data.min_memory:
+            _addTextElement(layerEl, 'memory', '%sm' % (layer.data.min_memory / 1024.0))
+        if layer.data.min_gpus or layer.data.min_gpu_memory:
+            _addTextElement(layerEl, 'gpus', str(int(round(layer.data.min_gpus)) or 1))
+            _addTextElement(
+                layerEl, 'gpu_memory',
+                '%sm' % (layer.data.min_gpu_memory / 1024.0) if layer.data.min_gpu_memory
+                else '1g')
+        if layer.data.timeout:
+            _addTextElement(layerEl, 'timeout', str(layer.data.timeout))
+        if layer.data.timeout_llu:
+            _addTextElement(layerEl, 'timeout_llu', str(layer.data.timeout_llu))
+        if layer.data.tags:
+            _addTextElement(layerEl, 'tags', '|'.join(layer.data.tags))
+        if layer.data.limits:
+            limitsEl = Et.SubElement(layerEl, 'limits')
+            for limitName in layer.data.limits:
+                _addTextElement(limitsEl, 'limit', limitName)
+
+        servicesEl = Et.SubElement(layerEl, 'services')
+        for serviceName in layer.data.services:
+            _addTextElement(servicesEl, 'service', serviceName)
+
+    return launchSpecAndWait(Et.tostring(root, encoding='unicode'))
+
+
 #
 # Job Names
 #

pycue/tests/test_api.py

@@ -20,6 +20,7 @@
 from __future__ import division
 from __future__ import absolute_import
 import unittest
+from xml.etree import ElementTree as Et
 
 import mock
 
@@ -35,6 +36,7 @@
 from opencue_proto import show_pb2
 from opencue_proto import subscription_pb2
 import opencue.api
+import opencue.wrappers.job
 
 
 TEST_SHOW_NAME = 'pipe'
@@ -283,6 +285,96 @@ def testLaunchSpecAndWait(self, getStubMock):
             job_pb2.JobLaunchSpecAndWaitRequest(spec=spec), timeout=mock.ANY)
         self.assertEqual([TEST_JOB_NAME], [job.name() for job in jobs])
 
+    @mock.patch('opencue.cuebot.Cuebot.getStub')
+    def testCloneJob(self, getStubMock):
+        origJob = job_pb2.Job(
+            name=TEST_JOB_NAME, facility=TEST_FACILITY_NAME, show=TEST_SHOW_NAME, shot='dev',
+            user='origuser', priority=3, max_cores=10.0, max_gpus=2.0, os='linux')
+        renderLayer = job_pb2.Layer(
+            name='render', type=job_pb2.RENDER, command='render_it -f #IFRAME#',
+            range='1-10', chunk_size=2, min_cores=2.0, is_threadable=True,
+            min_memory=4 * 1048576, min_gpus=1.0, min_gpu_memory=524288,
+            timeout=60, timeout_llu=30, tags=['general', 'gpu'], limits=['mylimit'],
+            services=['shell'])
+        preLayer = job_pb2.Layer(name='pre', type=job_pb2.PRE, command='setup')
+        noCmdLayer = job_pb2.Layer(name='empty', type=job_pb2.UTIL, command='')
+
+        stubMock = mock.Mock()
+        stubMock.GetLayers.return_value = job_pb2.JobGetLayersResponse(
+            layers=job_pb2.LayerSeq(layers=[renderLayer, preLayer, noCmdLayer]))
+        stubMock.LaunchSpecAndWait.return_value = job_pb2.JobLaunchSpecAndWaitResponse(
+            jobs=job_pb2.JobSeq(jobs=[job_pb2.Job(name=TEST_JOB_NAME + '_clone')]))
+        getStubMock.return_value = stubMock
+
+        clonedJobs = opencue.api.cloneJob(opencue.wrappers.job.Job(origJob))
+
+        self.assertEqual([TEST_JOB_NAME + '_clone'], [job.name() for job in clonedJobs])
+
+        sentSpec = stubMock.LaunchSpecAndWait.call_args[0][0].spec
+        specXml = Et.fromstring(sentSpec)
+
+        self.assertEqual(TEST_FACILITY_NAME, specXml.find('facility').text)
+        self.assertEqual(TEST_SHOW_NAME, specXml.find('show').text)
+        self.assertEqual('dev', specXml.find('shot').text)
+        self.assertEqual('origuser', specXml.find('user').text)
+
+        jobEl = specXml.find('job')
+        self.assertEqual(TEST_JOB_NAME + '_clone', jobEl.get('name'))
+        self.assertEqual('3', jobEl.find('priority').text)
+        self.assertEqual('10.0', jobEl.find('maxcores').text)
+        self.assertEqual('2.0', jobEl.find('maxgpus').text)
+        self.assertEqual('linux', jobEl.find('os').text)
+
+        # Only the Render layer with a command should survive; PRE and empty-command
+        # layers are skipped.
+        layerEls = jobEl.find('layers').findall('layer')
+        self.assertEqual(1, len(layerEls))
+        layerEl = layerEls[0]
+        self.assertEqual('render', layerEl.get('name'))
+        self.assertEqual('Render', layerEl.get('type'))
+        self.assertEqual('render_it -f #IFRAME#', layerEl.find('cmd').text)
+        self.assertEqual('1-10', layerEl.find('range').text)
+        self.assertEqual('2', layerEl.find('chunk').text)
+        self.assertEqual('2.0', layerEl.find('cores').text)
+        self.assertEqual('True', layerEl.find('threadable').text)
+        self.assertEqual('4096.0m', layerEl.find('memory').text)
+        self.assertEqual('1', layerEl.find('gpus').text)
+        self.assertEqual('512.0m', layerEl.find('gpu_memory').text)
+        self.assertEqual('60', layerEl.find('timeout').text)
+        self.assertEqual('30', layerEl.find('timeout_llu').text)
+        self.assertEqual('general|gpu', layerEl.find('tags').text)
+        self.assertEqual(['mylimit'], [e.text for e in layerEl.find('limits').findall('limit')])
+        self.assertEqual(['shell'], [e.text for e in layerEl.find('services').findall('service')])
+
+    @mock.patch('opencue.cuebot.Cuebot.getStub')
+    def testCloneJobOverrides(self, getStubMock):
+        origJob = job_pb2.Job(name=TEST_JOB_NAME, show=TEST_SHOW_NAME, shot='dev', user='origuser')
+        layerA = job_pb2.Layer(
+            name='a', type=job_pb2.RENDER, command='cmd_a', range='1-10', services=['shell'])
+        layerB = job_pb2.Layer(
+            name='b', type=job_pb2.RENDER, command='cmd_b', range='1-10', services=['shell'])
+
+        stubMock = mock.Mock()
+        stubMock.GetLayers.return_value = job_pb2.JobGetLayersResponse(
+            layers=job_pb2.LayerSeq(layers=[layerA, layerB]))
+        stubMock.LaunchSpecAndWait.return_value = job_pb2.JobLaunchSpecAndWaitResponse(
+            jobs=job_pb2.JobSeq(jobs=[job_pb2.Job(name='custom-name')]))
+        getStubMock.return_value = stubMock
+
+        opencue.api.cloneJob(
+            opencue.wrappers.job.Job(origJob), name='custom-name', user='newuser',
+            frame_range='1-5', layer_frame_ranges={'b': '20-30'})
+
+        sentSpec = stubMock.LaunchSpecAndWait.call_args[0][0].spec
+        specXml = Et.fromstring(sentSpec)
+
+        self.assertEqual('newuser', specXml.find('user').text)
+        self.assertEqual('custom-name', specXml.find('job').get('name'))
+        layersByName = {
+            el.get('name'): el for el in specXml.find('job').find('layers').findall('layer')}
+        self.assertEqual('1-5', layersByName['a'].find('range').text)
+        self.assertEqual('20-30', layersByName['b'].find('range').text)
+
 
 class LayerTests(unittest.TestCase):