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

@@ -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)
 =========

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

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"] )

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 )