diff --git a/src/deadline_worker_agent/aws/deadline/__init__.py b/src/deadline_worker_agent/aws/deadline/__init__.py index 6380ba37..27bb427f 100644 --- a/src/deadline_worker_agent/aws/deadline/__init__.py +++ b/src/deadline_worker_agent/aws/deadline/__init__.py @@ -13,7 +13,7 @@ from deadline.client.api import TelemetryClient from deadline.client import version as deadline_client_lib_version -from deadline.job_attachments.progress_tracker import SummaryStatistics +from deadline.job_attachments.upload import SummaryStatistics from openjd.model import version as openjd_model_version from openjd.sessions import version as openjd_sessions_version @@ -984,3 +984,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: return cast(F, wrapper) return inner + + +def hash_summary_telemetry_callback(hash_summary: SummaryStatistics): + _get_deadline_telemetry_client().record_hashing_summary(hash_summary) diff --git a/src/deadline_worker_agent/sessions/actions/scripts/attachment_upload.py b/src/deadline_worker_agent/sessions/actions/scripts/attachment_upload.py index 6983d665..c01574d7 100644 --- a/src/deadline_worker_agent/sessions/actions/scripts/attachment_upload.py +++ b/src/deadline_worker_agent/sessions/actions/scripts/attachment_upload.py @@ -42,6 +42,7 @@ record_attachment_upload_latencies_telemetry_event, record_attachment_upload_telemetry_event, record_success_fail_telemetry_event, + hash_summary_telemetry_callback, ) from deadline_worker_agent.sessions.attachment_models import ( WorkerManifestProperties, @@ -180,6 +181,8 @@ def snapshot( # when the code reaches here, it's guaranteed output_relative_directories contains value include=[glob.escape(subdir) + "/**" for subdir in output_relative_directories], name="output", + hash_cache_dir=config_file.get_cache_directory(), + telemetry_callback=hash_summary_telemetry_callback, ) if output_manifest: diff --git a/test/unit/sessions/actions/scripts/test_attachment_upload.py b/test/unit/sessions/actions/scripts/test_attachment_upload.py index c0a59caa..3353739a 100644 --- a/test/unit/sessions/actions/scripts/test_attachment_upload.py +++ b/test/unit/sessions/actions/scripts/test_attachment_upload.py @@ -1,4 +1,5 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +from deadline_worker_agent.aws.deadline import hash_summary_telemetry_callback import json import pytest @@ -7,6 +8,7 @@ from typing import Optional, Generator import deadline_worker_agent.sessions.actions.scripts.attachment_upload as attachment_upload_mod +from deadline.client.config import config_file from deadline.job_attachments.asset_manifests.decode import decode_manifest from deadline.job_attachments.progress_tracker import ProgressStatus, ProgressTracker from deadline.job_attachments.models import JobAttachmentS3Settings @@ -212,6 +214,8 @@ def test_snapshot_mixed_output_directories(self, mock_manifest_snapshot: Mock): diff="/base/manifest.json", include=["output/**"], name="output", + hash_cache_dir=config_file.get_cache_directory(), + telemetry_callback=hash_summary_telemetry_callback, ) mock_manifest_snapshot.assert_any_call( root="/local/path4", @@ -219,6 +223,8 @@ def test_snapshot_mixed_output_directories(self, mock_manifest_snapshot: Mock): diff=None, include=["./**"], name="output", + hash_cache_dir=config_file.get_cache_directory(), + telemetry_callback=hash_summary_telemetry_callback, ) @patch("deadline_worker_agent.sessions.actions.scripts.attachment_upload._manifest_snapshot") @@ -255,6 +261,8 @@ def test_snapshot(self, mock_manifest_snapshot: Mock): diff="/base/manifest.json", include=["output/**"], name="output", + hash_cache_dir=config_file.get_cache_directory(), + telemetry_callback=hash_summary_telemetry_callback, ) mock_manifest_snapshot.assert_any_call( root="/local/path2", @@ -262,6 +270,8 @@ def test_snapshot(self, mock_manifest_snapshot: Mock): diff=None, include=["output/**"], name="output", + hash_cache_dir=config_file.get_cache_directory(), + telemetry_callback=hash_summary_telemetry_callback, ) @patch("deadline_worker_agent.sessions.actions.scripts.attachment_upload._manifest_snapshot") @@ -299,6 +309,8 @@ def test_snapshot_with_special_characters_in_directory_names( "[[]data_][*]dir2[?]/**", # Complex case with multiple special chars ], name="output", + hash_cache_dir=config_file.get_cache_directory(), + telemetry_callback=hash_summary_telemetry_callback, ) @patch.dict(