From 8ec280621d001f5fc799b9406f613c1e435a6164 Mon Sep 17 00:00:00 2001 From: John Haddon Date: Tue, 2 Apr 2024 14:24:03 +0100 Subject: [PATCH] TractorDispatcher : Add `taskData` to `preSpoolSignal()` This allows slots to map from Tractor tasks to the Gaffer tasks that they will execute. I would have preferred to inject this information directly as additional attributes on the `tractor.api.author.Task` objects, but that isn't possible because they use Python's `__slots__` mechanism and therefore reject additional attributes. This means we need to pass a separate dictionary of `TaskData` objects to slots, but backwards compatibility is provided for legacy slots which don't accept that argument. I deliberated a bit about whether we should just expose the protected `Dispatcher.TaskBatch` objects directly since they contain the same information as `TaskData`. But things like `TaskBatch.blindData()` are definitely meant to be private to the Dispatcher, and it'll be nice to be free to change TaskBatch in future with fewer worries about client code depending on it. --- Changes.md | 1 + python/GafferTractor/TractorDispatcher.py | 38 ++++++--- .../TractorDispatcherTest.py | 55 ++++++++++++- startup/GafferTractor/signalCompatibility.py | 78 +++++++++++++++++++ 4 files changed, 159 insertions(+), 13 deletions(-) create mode 100644 startup/GafferTractor/signalCompatibility.py diff --git a/Changes.md b/Changes.md index 2534c825c6b..e64732b8a64 100644 --- a/Changes.md +++ b/Changes.md @@ -21,6 +21,7 @@ API --- - SelectionTool : Added static `registerSelectMode()` method for registering a Python or C++ function that will modify a selected scene path location. Users can choose which mode is active when selecting. +- TractorDispatcher : The `preSpoolSignal()` now provides an additional `taskData` argument to slots, which maps from Tractor tasks to information about the Gaffer tasks they will execute. 1.4.0.0b5 (relative to 1.4.0.0b4) ========= diff --git a/python/GafferTractor/TractorDispatcher.py b/python/GafferTractor/TractorDispatcher.py index 286ebfa9452..67430fe6f4c 100644 --- a/python/GafferTractor/TractorDispatcher.py +++ b/python/GafferTractor/TractorDispatcher.py @@ -34,7 +34,7 @@ # ########################################################################## -import os +import dataclasses import IECore @@ -51,19 +51,35 @@ def __init__( self, name = "TractorDispatcher" ) : self["service"] = Gaffer.StringPlug( defaultValue = '"*"' ) self["envKey"] = Gaffer.StringPlug() - ## Emitted prior to spooling the Tractor job, to allow - # custom modifications to be applied. + ## Decribes the Gaffer task associated with a particular Tractor task. + @dataclasses.dataclass + class TaskData : + + # The `GafferDispatch.TaskPlug` being executed. The associated node + # can be accessed via `plug.node()`. + plug : Gaffer.Plug + # The Gaffer context in which the task is executed. Does not + # contain the frame number, which is provided via `frames`. + context : Gaffer.Context + # The list of frames being executed. + frames : list + + ## Emitted prior to spooling the Tractor job, to allow custom modifications to + # be applied. Slots should have the signature `slot( dispatcher, job, taskData )` : # - # Slots should have the signature `slot( dispatcher, job )`, - # where dispatcher is the TractorDispatcher and job will - # be the instance of tractor.api.author.Job that is about - # to be spooled. + # - `dispatcher` : The TractorDispatcher that is about to spool the job. + # - `job` : The `tractor.api.author.Job` that is about to be spooled. + # This may be modified in place. + # - `taskData` : A dictionary mapping from `tractor.api.author.Task` to + # TaskData, specifying the Gaffer tasks that will be executed by each + # Tractor task. For example, the Gaffer node for the first Tractor task + # can be accessed as `taskData[job.subtasks[0]].plug.node()`. @classmethod def preSpoolSignal( cls ) : return cls.__preSpoolSignal - __preSpoolSignal = Gaffer.Signals.Signal2() + __preSpoolSignal = Gaffer.Signals.Signal3() def _doDispatch( self, rootBatch ) : @@ -84,6 +100,7 @@ def _doDispatch( self, rootBatch ) : dispatchData["scriptNode"] = rootBatch.preTasks()[0].node().scriptNode() dispatchData["scriptFile"] = Gaffer.Context.current()["dispatcher:scriptFileName"] dispatchData["batchesToTasks"] = {} + dispatchData["taskData"] = {} # Create a Tractor job and set its basic properties. @@ -105,7 +122,7 @@ def _doDispatch( self, rootBatch ) : # Signal anyone who might want to make just-in-time # modifications to the job. - self.preSpoolSignal()( self, job ) + self.preSpoolSignal()( self, job, dispatchData["taskData"] ) # Save a copy of our job script to the job directory. # This isn't strictly necessary because we'll spool via @@ -197,6 +214,9 @@ def __acquireTask( self, batch, dispatchData ) : # Remember the task for next time, and return it. dispatchData["batchesToTasks"][batch] = task + dispatchData["taskData"][task] = self.TaskData( + batch.plug(), batch.context(), batch.frames() + ) return task @staticmethod diff --git a/python/GafferTractorTest/TractorDispatcherTest.py b/python/GafferTractorTest/TractorDispatcherTest.py index 05c8ce482c7..3e151fa8573 100644 --- a/python/GafferTractorTest/TractorDispatcherTest.py +++ b/python/GafferTractorTest/TractorDispatcherTest.py @@ -37,7 +37,7 @@ import unittest import unittest.mock import inspect -import sys +import warnings import imath @@ -63,7 +63,7 @@ def __dispatcher( self ) : def __job( self, nodes, dispatcher = None ) : jobs = [] - def f( dispatcher, job ) : + def f( dispatcher, job, taskData ) : jobs.append( job ) @@ -79,7 +79,51 @@ def f( dispatcher, job ) : return jobs[0] - def testPreSpoolSignal( self ) : + def testPreSpoolSignalTaskData( self ) : + + script = Gaffer.ScriptNode() + script["variables"].addChild( Gaffer.NameValuePlug( "myVariable", "test" ) ) + + script["node"] = GafferDispatchTest.LoggingTaskNode() + script["node"]["frame"] = Gaffer.StringPlug( defaultValue = "${frame}", flags = Gaffer.Plug.Flags.Default | Gaffer.Plug.Flags.Dynamic ) + script["node"]["dispatcher"]["batchSize"].setValue( 4 ) + + script["dispatcher"] = self.__dispatcher() + script["dispatcher"]["tasks"][0].setInput( script["node"]["task"] ) + script["dispatcher"]["framesMode"].setValue( script["dispatcher"].FramesMode.CustomRange ) + script["dispatcher"]["frameRange"].setValue( "1-10" ) + + taskDataChecked = False + def checkTaskData( dispatcher, job, taskData ) : + + self.assertEqual( len( job.subtasks ), 3 ) + self.assertEqual( len( taskData ), 3 ) + for subtask in job.subtasks : + + self.assertIn( subtask, taskData ) + self.assertTrue( taskData[subtask].plug.isSame( script["node"]["task"] ) ) + + self.assertIsInstance( taskData[subtask].context, Gaffer.Context ) + self.assertNotIn( "frame", taskData[subtask].context ) + self.assertEqual( taskData[subtask].context["myVariable"], "test" ) + self.assertIn( "dispatcher:jobDirectory", taskData[subtask].context ) + self.assertIn( "dispatcher:scriptFileName", taskData[subtask].context ) + + self.assertEqual( + [ taskData[subtask].frames for subtask in job.subtasks ], + [ [ 1, 2, 3, 4 ], [ 5, 6, 7, 8 ], [ 9, 10 ] ] + ) + + nonlocal taskDataChecked + taskDataChecked = True + + connection = GafferTractor.TractorDispatcher.preSpoolSignal().connect( checkTaskData, scoped = True ) + with script.context() : + script["dispatcher"]["task"].execute() + + self.assertTrue( taskDataChecked ) + + def testPreSpoolSignalCompatibility( self ) : s = Gaffer.ScriptNode() s["n"] = GafferDispatchTest.LoggingTaskNode() @@ -89,7 +133,10 @@ def f( dispatcher, job ) : spooled.append( ( dispatcher, job ) ) - c = GafferTractor.TractorDispatcher.preSpoolSignal().connect( f, scoped = True ) + # Connecting a function which doesn't have the additional `taskData` argument. + with warnings.catch_warnings() : + warnings.simplefilter( "ignore", DeprecationWarning ) + c = GafferTractor.TractorDispatcher.preSpoolSignal().connect( f, scoped = True ) dispatcher = self.__dispatcher() dispatcher["tasks"][0].setInput( s["n"]["task"] ) diff --git a/startup/GafferTractor/signalCompatibility.py b/startup/GafferTractor/signalCompatibility.py new file mode 100644 index 00000000000..def64cc22e6 --- /dev/null +++ b/startup/GafferTractor/signalCompatibility.py @@ -0,0 +1,78 @@ +########################################################################## +# +# Copyright (c) 2024, Cinesite VFX Ltd. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are +# met: +# +# * Redistributions of source code must retain the above +# copyright notice, this list of conditions and the following +# disclaimer. +# +# * Redistributions in binary form must reproduce the above +# copyright notice, this list of conditions and the following +# disclaimer in the documentation and/or other materials provided with +# the distribution. +# +# * Neither the name of John Haddon nor the names of +# any other contributors to this software may be used to endorse or +# promote products derived from this software without specific prior +# written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +########################################################################## + +import inspect +import warnings + +import GafferTractor + +# Backwards compatibility for slots written when `preSpoolSignal()` didn't +# have the `taskData` argument. + +def __slotWrapper( slot ) : + + signature = inspect.signature( slot ) + try : + # Throws if not callable with three arguments + signature.bind( None, None, None ) + # No need for a wrapper + return slot + except TypeError : + pass + + # We'll need a wrapper + + warnings.warn( + 'Slot connected to `TractorDispatcher.preSpoolSignal() should have an additional `taskData` argument', + DeprecationWarning + ) + + def call( dispatcher, *args ) : + + slot( dispatcher, *args[:-1] ) + + return call + +def __connectWrapper( originalConnect ) : + + def connect( slot, scoped = None ) : + + return originalConnect( __slotWrapper( slot ), scoped ) + + return connect + +GafferTractor.TractorDispatcher._TractorDispatcher__preSpoolSignal.connect = __connectWrapper( GafferTractor.TractorDispatcher._TractorDispatcher__preSpoolSignal.connect ) +GafferTractor.TractorDispatcher._TractorDispatcher__preSpoolSignal.connectFront = __connectWrapper( GafferTractor.TractorDispatcher._TractorDispatcher__preSpoolSignal.connectFront )