perf: avoid S3 HEAD requests when downloading job outputs

src/deadline/job_attachments/download.py

@@ -21,6 +21,7 @@
 from boto3.s3.transfer import ProgressCallbackInvoker
 from botocore.client import BaseClient
 from botocore.exceptions import BotoCoreError, ClientError
+from s3transfer.subscribers import BaseSubscriber as _BaseSubscriber
 
 from .asset_manifests.base_manifest import BaseAssetManifest, BaseManifestPath as RelativeFilePath
 from .asset_manifests.hash_algorithms import HashAlgorithm
@@ -77,6 +78,22 @@
 
 download_logger = getLogger("deadline.job_attachments.download")
 
+
+class _FileSizeSubscriber(_BaseSubscriber):
+    """Subscriber that provides file size to skip HEAD requests."""
+
+    def __init__(self, size):
+        self._size = size
+
+    def on_queued(self, future, **kwargs):
+        future.meta.provide_transfer_size(self._size)
+        # Provide a dummy etag to skip HEAD request if the method exists (added in s3transfer 0.6.0).
+        # For downloads from CAS, we don't need etag validation since files are content-addressed.
+        # Older s3transfer versions don't have this method, so we check before calling.
+        if hasattr(future.meta, "provide_object_etag"):
+            future.meta.provide_object_etag("dummy-etag")
+
+
 S3_DOWNLOAD_MAX_CONCURRENCY = 10
 WINDOWS_MAX_PATH_LENGTH = 260
 TEMP_DOWNLOAD_ADDED_CHARS_LENGTH = 9
@@ -549,7 +566,7 @@ def handler(bytes_downloaded):
             if not should_continue:
                 future.cancel()
 
-    subscribers = [ProgressCallbackInvoker(handler)]
+    subscribers = [_FileSizeSubscriber(file_bytes), ProgressCallbackInvoker(handler)]
 
     future = transfer_manager.download(
         bucket=s3_bucket,

test/unit/deadline_job_attachments/test_download.py

@@ -1784,7 +1784,7 @@ def test_download_file_error_message_on_access_denied(self):
         s3_client = boto3.client("s3")
         stubber = Stubber(s3_client)
         stubber.add_client_error(
-            "head_object",
+            "get_object",
             service_error_code="AccessDenied",
             service_message="Access Denied",
             http_status_code=403,
@@ -1871,6 +1871,58 @@ def test_download_file_error_message_on_timeout(self):
             mock_lock.assert_not_called()
             mock_collision_dict.assert_not_called()
 
+    def test_download_file_does_not_make_head_request_when_size_known(self):
+        """
+        Test that download_file does not make a HEAD request when file size is known from manifest.
+        """
+        mock_s3_client = MagicMock()
+        mock_transfer_manager = MagicMock()
+        mock_lock = MagicMock()
+        mock_collision_dict = MagicMock()
+
+        file_size = 12345
+        file_path = ManifestPathv2023_03_03(
+            path="inputs/input1.txt", hash="input1", size=file_size, mtime=1234000000
+        )
+
+        with patch(
+            f"{deadline.__package__}.job_attachments.download.get_s3_client",
+            return_value=mock_s3_client,
+        ), patch(
+            f"{deadline.__package__}.job_attachments.download.get_s3_transfer_manager",
+            return_value=mock_transfer_manager,
+        ), patch(f"{deadline.__package__}.job_attachments.download.Path.mkdir"), patch(
+            f"{deadline.__package__}.job_attachments.download.os.utime"
+        ):
+            download_file(
+                file_path,
+                HashAlgorithm.XXH128,
+                "/home/username/assets",
+                mock_lock,
+                mock_collision_dict,
+                "test-bucket",
+                "rootPrefix/Data",
+                mock_s3_client,
+            )
+            mock_s3_client.head_object.assert_not_called()
+
+    def test_file_size_subscriber_without_provide_object_etag(self):
+        """
+        Test that _FileSizeSubscriber works with older s3transfer versions
+        that don't have provide_object_etag method.
+        """
+        from deadline.job_attachments.download import _FileSizeSubscriber
+
+        mock_future = MagicMock()
+        mock_future.meta.provide_transfer_size = MagicMock()
+        # Simulate older s3transfer where provide_object_etag doesn't exist
+        del mock_future.meta.provide_object_etag
+
+        subscriber = _FileSizeSubscriber(12345)
+        subscriber.on_queued(mock_future)
+
+        mock_future.meta.provide_transfer_size.assert_called_once_with(12345)
+
     @pytest.mark.skipif(
         sys.platform == "win32",
         reason="This test is for Linux path only.",