diff --git a/src/deadline/client/cli/_incremental_download.py b/src/deadline/client/cli/_incremental_download.py index 46f3b024f..2a617d4dc 100644 --- a/src/deadline/client/cli/_incremental_download.py +++ b/src/deadline/client/cli/_incremental_download.py @@ -126,13 +126,13 @@ def _get_download_candidate_jobs( region=region, ) } - print(f"DEBUG: Got {len(download_candidate_jobs)} active jobs") + print_function_callback(f"DEBUG: Got {len(download_candidate_jobs)} active jobs") download_candidate_jobs = { job_id: _datetimes_to_str(job) for job_id, job in download_candidate_jobs.items() if job["taskRunStatusCounts"]["SUCCEEDED"] > 0 } - print( + print_function_callback( f"DEBUG: Filtered down to {len(download_candidate_jobs)} active jobs based on SUCCEEDED task filter" ) @@ -157,14 +157,16 @@ def _get_download_candidate_jobs( }, region=region, ) - print( + print_function_callback( f"DEBUG: Got {len(recently_ended_jobs)} jobs with job[endedAt] >= {starting_timestamp.astimezone().isoformat()}" ) # Filter to jobs where the count of SUCCEEDED tasks is positive. recently_ended_jobs = [ job for job in recently_ended_jobs if job["taskRunStatusCounts"]["SUCCEEDED"] > 0 ] - print(f"DEBUG: Filtered down to {len(recently_ended_jobs)} jobs based on SUCCEEDED task filter") + print_function_callback( + f"DEBUG: Filtered down to {len(recently_ended_jobs)} jobs based on SUCCEEDED task filter" + ) download_candidate_jobs.update( {job["jobId"]: _datetimes_to_str(job) for job in recently_ended_jobs} ) @@ -599,13 +601,21 @@ def _get_job_sessions( for job_id in job_ids: # Use the greater of the bootstrap command timestamp and the session ended timestamps # recorded in the checkpoint. - session_ended_threshold = job_session_ended_timestamp.get(job_id) + saved_session_ended_timestamp = job_session_ended_timestamp.get(job_id) + session_ended_threshold = saved_session_ended_timestamp if session_ended_threshold is None: session_ended_threshold = checkpoint.downloads_started_timestamp - # For all jobs that are not NEW (including re-queued jobs) - i.e. completed and updated jobs - # Use an eventual consistency window to accept a little extra - if job_id not in categorized_job_ids.added: + # For all jobs that are not brand new - i.e. completed, updated, and re-queued jobs - + # use an eventual consistency window to accept a little extra. A brand new job is one + # in the 'added' category with no previously saved session_ended_timestamp. A re-queued + # job is also categorized as 'added' but carries a saved session_ended_timestamp from + # when it was previously tracked, so it must still get the window to avoid missing + # sessions near the re-queue boundary. + is_brand_new_job = ( + job_id in categorized_job_ids.added and saved_session_ended_timestamp is None + ) + if not is_brand_new_job: session_ended_threshold = session_ended_threshold - timedelta( seconds=checkpoint.eventual_consistency_max_seconds ) @@ -849,7 +859,7 @@ def _filter_session_actions_without_manifests_from_job_sessions( filtered_session_action_list = [ session_action for session_action in session.get("sessionActions", []) - if any(item != {} for item in session_action["manifests"]) + if any(item != {} for item in session_action.get("manifests", [])) ] filtered_count += len(filtered_session_action_list) if total_count != filtered_count: @@ -891,7 +901,10 @@ def _update_checkpoint_jobs_list( if job.session_ended_timestamp is not None } for job_id, session_list in job_sessions.items(): - max_session_ended_timestamp = None + # Seed with the value carried over from the previous checkpoint so that a job whose + # current sessions are all still running (no endedAt yet) keeps its saved timestamp + # instead of having it overwritten with None. + max_session_ended_timestamp = job_session_ended_timestamps.get(job_id) for session in session_list: if "endedAt" in session: if max_session_ended_timestamp is None: @@ -915,7 +928,6 @@ def _update_checkpoint_jobs_list( job_session_completed_indexes.setdefault(job_id, {})[session["sessionId"]] = max( session_action["sessionActionIndex"] for session_action in session_actions ) - job_session_ended_timestamps[job_id] = max_session_ended_timestamp # These categories keep the download_candidate_jobs job as is. for job_id in ( diff --git a/test/unit/deadline_client/cli/test_cli_queue_incremental_download.py b/test/unit/deadline_client/cli/test_cli_queue_incremental_download.py index db5e626e3..8e5155515 100644 --- a/test/unit/deadline_client/cli/test_cli_queue_incremental_download.py +++ b/test/unit/deadline_client/cli/test_cli_queue_incremental_download.py @@ -4,12 +4,13 @@ Tests for the CLI queue incremental output download command. """ +import json import os -import sys import pytest -from unittest.mock import patch -from datetime import datetime, timedelta +from unittest.mock import patch, MagicMock +from datetime import datetime, timedelta, timezone +import boto3 from freezegun import freeze_time from click.testing import CliRunner from deadline.client.cli import main @@ -22,6 +23,7 @@ MOCK_STORAGE_PROFILE_ID, MOCK_FLEET_ID, MOCK_WORKER_ID, + MOCK_BUCKET_NAME, ) from ..mock_deadline_job_apis import ( mock_search_jobs_for_set, @@ -30,9 +32,23 @@ ) from deadline.job_attachments._incremental_downloads.incremental_download_state import ( EVENTUAL_CONSISTENCY_MAX_SECONDS, + IncrementalDownloadState, + IncrementalDownloadJob, +) +from deadline.job_attachments.asset_manifests.hash_algorithms import HashAlgorithm +from deadline.job_attachments.asset_manifests.v2023_03_03.asset_manifest import ( + AssetManifest, + ManifestPath, ) from deadline.job_attachments.models import StorageProfileOperatingSystemFamily import deadline.client.api +import deadline.client.cli._incremental_download as mod +from deadline.client.cli._incremental_download import ( + CategorizedJobIds, + _update_checkpoint_jobs_list, + _filter_session_actions_without_manifests_from_job_sessions, + _get_job_sessions, +) ISO_FREEZE_TIME_MINUS_5MIN = "2025-05-26 11:55:00+00:00" ISO_FREEZE_TIME_MINUS_1MIN = "2025-05-26 11:59:00+00:00" @@ -63,9 +79,6 @@ def deadline_telemetry_client_mock(): yield m -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_requires_queue_with_job_attachments( fresh_deadline_config, deadline_mock, checkpoint_dir ): @@ -101,9 +114,6 @@ def test_incremental_output_download_requires_queue_with_job_attachments( ) -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_pid_lock_already_held_error( fresh_deadline_config, deadline_mock, @@ -148,9 +158,6 @@ def test_incremental_output_download_pid_lock_already_held_error( assert os.path.exists(pid_lock_file) -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_storage_profile_options_mutually_exclusive( fresh_deadline_config, deadline_mock, @@ -187,9 +194,6 @@ def test_incremental_output_download_storage_profile_options_mutually_exclusive( ), result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) @pytest.mark.parametrize("storage_profile_id", [None, MOCK_STORAGE_PROFILE_ID]) def test_incremental_output_download_bootstrap_and_completion( fresh_deadline_config, @@ -407,9 +411,6 @@ def test_incremental_output_download_bootstrap_and_completion( assert "inactive: 0" in result.output, result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_storage_profile_path_mapping( fresh_deadline_config, tmp_path, @@ -569,9 +570,6 @@ def mock_get_storage_profile_for_queue(farmId: str, queueId: str, storageProfile ), result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_bootstrap_retire_job_without_attachments( fresh_deadline_config, deadline_mock, checkpoint_dir ): @@ -705,9 +703,6 @@ def test_incremental_output_download_bootstrap_retire_job_without_attachments( assert "inactive: 1" in result.output, result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_job_unchanged( fresh_deadline_config, deadline_mock, checkpoint_dir ): @@ -803,9 +798,6 @@ def test_incremental_output_download_job_unchanged( assert "unchanged: 1" in result.output, result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_job_canceled( fresh_deadline_config, deadline_mock, checkpoint_dir ): @@ -910,9 +902,6 @@ def test_incremental_output_download_job_canceled( assert "inactive: 1" in result.output, result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_job_completed_then_requeued( fresh_deadline_config, deadline_mock, checkpoint_dir ): @@ -1053,9 +1042,6 @@ def test_incremental_output_download_job_completed_then_requeued( assert "added: 1" in result.output, result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_dry_run(fresh_deadline_config, deadline_mock, checkpoint_dir): """Test a new job through bootstrap, completion, and retirement.""" mock_jobs = create_fake_job_list(1) @@ -1119,9 +1105,6 @@ def test_incremental_output_download_dry_run(fresh_deadline_config, deadline_moc assert "This is a DRY RUN so the checkpoint was not saved" in result.output, result.output -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Incremental output download requires Python >= 3.9" -) def test_incremental_output_download_stats_telemetry( fresh_deadline_config, deadline_mock, @@ -1295,3 +1278,360 @@ def fake_download_all_manifests( assert result.exit_code == 0, result.output assert "WARNING: THE FOLLOWING FILES WILL NOT BE DOWNLOADED" in result.output, result.output assert "/etc/cron.d/evil" in result.output, result.output + + +# --------------------------------------------------------------------------- +# Unit tests for _update_checkpoint_jobs_list (checkpoint session-ended timestamps) +# --------------------------------------------------------------------------- + + +def _make_categorized_job_ids(**kwargs): + """Build a CategorizedJobIds with all categories reset to fresh empty sets. + + CategorizedJobIds defines its sets as class attributes, so instances share + them unless reassigned. Reset every category to avoid cross-test contamination. + """ + cats = CategorizedJobIds() + for name in ( + "added", + "updated", + "unchanged", + "completed", + "inactive", + "missing_storage_profile", + "attachments_free", + ): + setattr(cats, name, set(kwargs.get(name, set()))) + return cats + + +def test_update_checkpoint_jobs_list_does_not_corrupt_session_ended_timestamp(): + """Each job's session_ended_timestamp must reflect its own sessions. + + Regression for the stale-variable bug: a leftover max_session_ended_timestamp + from the last job of the first loop was written to every job in a second loop, + corrupting the checkpoint timestamps. + """ + t0 = datetime(2025, 1, 1, tzinfo=timezone.utc) + t_a = datetime(2025, 1, 2, tzinfo=timezone.utc) + t_b = datetime(2025, 1, 3, tzinfo=timezone.utc) + + checkpoint = IncrementalDownloadState( + local_storage_profile_id=None, downloads_started_timestamp=t0 + ) + download_candidate_jobs = { + "job-a": {"jobId": "job-a", "name": "A"}, + "job-b": {"jobId": "job-b", "name": "B"}, + } + job_sessions = { + "job-a": [ + {"sessionId": "s-a", "endedAt": t_a, "sessionActions": [{"sessionActionIndex": 1}]} + ], + "job-b": [ + {"sessionId": "s-b", "endedAt": t_b, "sessionActions": [{"sessionActionIndex": 1}]} + ], + } + cats = _make_categorized_job_ids(added={"job-a", "job-b"}) + + _update_checkpoint_jobs_list(checkpoint, download_candidate_jobs, cats, job_sessions) + + result = {job.job_id: job.session_ended_timestamp for job in checkpoint.jobs} + assert result["job-a"] == t_a, result + assert result["job-b"] == t_b, result + + +def test_update_checkpoint_running_only_job_keeps_saved_timestamp(): + """A job whose current sessions are all still running (no endedAt) must keep + the session_ended_timestamp saved in the previous checkpoint rather than have + it overwritten with None.""" + t0 = datetime(2025, 1, 1, tzinfo=timezone.utc) + t_saved = datetime(2025, 1, 2, tzinfo=timezone.utc) + + prior_job = IncrementalDownloadJob({"jobId": "job-a", "name": "A"}, t_saved, {}) + checkpoint = IncrementalDownloadState( + local_storage_profile_id=None, downloads_started_timestamp=t0, jobs=[prior_job] + ) + download_candidate_jobs = {"job-a": {"jobId": "job-a", "name": "A"}} + # Session is still running: no "endedAt" field. + job_sessions = {"job-a": [{"sessionId": "s-a", "sessionActions": [{"sessionActionIndex": 1}]}]} + cats = _make_categorized_job_ids(updated={"job-a"}) + + _update_checkpoint_jobs_list(checkpoint, download_candidate_jobs, cats, job_sessions) + + result = {job.job_id: job.session_ended_timestamp for job in checkpoint.jobs} + assert result["job-a"] == t_saved, result + + +def test_filter_session_actions_tolerates_missing_manifests_key(): + """A session action without a 'manifests' key must not raise KeyError.""" + job_sessions = { + "job-a": [ + { + "sessionId": "s-a", + "sessionActions": [{"sessionActionId": "sa-0"}], # no "manifests" key + } + ] + } + download_candidate_jobs = {"job-a": {"jobId": "job-a", "name": "A"}} + + # Must not raise (previously raised KeyError on session_action["manifests"]). + _filter_session_actions_without_manifests_from_job_sessions( + job_sessions, download_candidate_jobs + ) + # The action lacked any output manifests, so it is filtered out. + assert job_sessions["job-a"][0].get("sessionActions", []) == [] + + +def test_get_job_sessions_requeued_job_uses_eventual_consistency_window(): + """A requeued job reappears categorized as 'added' but carries a saved + session_ended_timestamp. It must get the eventual-consistency window applied, + like updated/completed jobs, so sessions near the requeue boundary aren't missed.""" + t0 = datetime(2025, 1, 1, tzinfo=timezone.utc) + t_saved = datetime(2025, 1, 5, tzinfo=timezone.utc) + + prior_job = IncrementalDownloadJob({"jobId": "job-a"}, t_saved, {}) + checkpoint = IncrementalDownloadState( + local_storage_profile_id=None, downloads_started_timestamp=t0, jobs=[prior_job] + ) + download_candidate_jobs = {"job-a": {"jobId": "job-a", "name": "A"}} + cats = _make_categorized_job_ids(added={"job-a"}) + + captured_thresholds = {} + + def fake_retrieve_sessions_for_job( + deadline_client, farm_id, queue_id, job_id, session_ended_threshold, output_job_sessions + ): + captured_thresholds[job_id] = session_ended_threshold + + with ( + patch.object(mod, "get_session_client", return_value=MagicMock()), + patch.object(mod, "_retrieve_sessions_for_job", side_effect=fake_retrieve_sessions_for_job), + ): + _get_job_sessions( + MagicMock(), # boto3_session + MagicMock(), # boto3_session_for_s3 + MOCK_FARM_ID, + {"queueId": MOCK_QUEUE_ID}, + {}, # checkpoint_job_session_completed_indexes + cats, + checkpoint, + download_candidate_jobs, + ) + + expected = t_saved - timedelta(seconds=checkpoint.eventual_consistency_max_seconds) + assert captured_thresholds["job-a"] == expected, captured_thresholds + + +def test_incremental_output_download_json_mode_emits_no_debug_lines( + fresh_deadline_config, deadline_mock, checkpoint_dir +): + """In --json mode the command must not emit raw DEBUG print() lines that would + corrupt the JSON output stream.""" + mock_jobs = create_fake_job_list(1) + mock_jobs[0]["name"] = "Mock Job" + mock_jobs[0]["jobId"] = MOCK_JOB_ID + mock_jobs[0]["taskRunStatus"] = "READY" + mock_jobs[0]["taskRunStatusCounts"] = {"SUCCEEDED": 1, "READY": 1} + mock_jobs[0]["attachments"] = { + "manifests": [ + {"rootPath": "/", "rootPathFormat": "posix", "outputRelativeDirectories": ["."]} + ], + "fileSystem": "VIRTUAL", + } + del mock_jobs[0]["endedAt"] + deadline_mock.search_jobs = mock_search_jobs_for_set(MOCK_FARM_ID, MOCK_QUEUE_ID, mock_jobs) + deadline_mock.get_job = mock_get_job_for_set(MOCK_FARM_ID, MOCK_QUEUE_ID, mock_jobs) + + runner = CliRunner() + with freeze_time(ISO_FREEZE_TIME): + result = runner.invoke( + main, + [ + "queue", + "sync-output", + "--ignore-storage-profiles", + "--json", + "--farm-id", + MOCK_FARM_ID, + "--queue-id", + MOCK_QUEUE_ID, + "--checkpoint-dir", + checkpoint_dir, + ], + ) + + assert result.exit_code == 0, result.output + assert "DEBUG" not in result.output, result.output + + +def _put_manifest_and_data_objects_in_s3(root_path: str) -> dict[str, bytes]: + """Upload asset manifests and CAS data objects to the moto S3 bucket so that + sync-output performs real S3 downloads. Returns {relative_path: content}. + + The keys follow the job attachments layout under the queue's rootPrefix + ("MockRootPrefix" from the deadline_mock fixture's get_queue response): + manifests at MockRootPrefix/Manifests/ and CAS data at + MockRootPrefix/Data/.xxh128. + """ + files = { + "output/file1.txt": b"content of file one", + "output/file2.txt": b"content of file two, a bit longer", + "output/deeper/file3.txt": b"third file content", + } + hashes = {path: f"fakehash{i}" for i, path in enumerate(files)} + + def make_manifest(paths: list[str]) -> str: + return AssetManifest( + hash_alg=HashAlgorithm.XXH128, + paths=[ + ManifestPath(path=path, hash=hashes[path], size=len(files[path]), mtime=1) + for path in paths + ], + total_size=sum(len(files[path]) for path in paths), + ).encode() + + s3 = boto3.client("s3", region_name="us-west-2") + s3.put_object( + Bucket=MOCK_BUCKET_NAME, + Key="MockRootPrefix/Manifests/manifest_action_0_output.xxh128", + Body=make_manifest(["output/file1.txt", "output/file2.txt"]).encode("utf-8"), + ) + s3.put_object( + Bucket=MOCK_BUCKET_NAME, + Key="MockRootPrefix/Manifests/manifest_action_1_output.xxh128", + Body=make_manifest(["output/deeper/file3.txt"]).encode("utf-8"), + ) + for path, content in files.items(): + s3.put_object( + Bucket=MOCK_BUCKET_NAME, + Key=f"MockRootPrefix/Data/{hashes[path]}.xxh128", + Body=content, + ) + return files + + +def test_incremental_output_download_json_mode_with_real_s3_download( + fresh_deadline_config, deadline_mock, checkpoint_dir, tmp_path +): + """In --json mode, stdout must stay byte-clean even while sync-output performs + real S3 file downloads (via moto) with the real ProgressTracker firing progress + callbacks. A contrast run without --json verifies that the same download emits + progress output, proving the callbacks fired and were suppressed rather than + the download being skipped.""" + root_path = str(tmp_path / "job_root") + + mock_jobs = create_fake_job_list(1) + mock_jobs[0]["name"] = "Mock Job" + mock_jobs[0]["jobId"] = MOCK_JOB_ID + mock_jobs[0]["taskRunStatus"] = "READY" + mock_jobs[0]["taskRunStatusCounts"] = {"SUCCEEDED": 2, "READY": 1} + mock_jobs[0]["attachments"] = { + "manifests": [ + { + "rootPath": root_path, + "rootPathFormat": "posix", + "outputRelativeDirectories": ["output"], + } + ], + "fileSystem": "VIRTUAL", + } + del mock_jobs[0]["endedAt"] + deadline_mock.search_jobs = mock_search_jobs_for_set(MOCK_FARM_ID, MOCK_QUEUE_ID, mock_jobs) + deadline_mock.get_job = mock_get_job_for_set(MOCK_FARM_ID, MOCK_QUEUE_ID, mock_jobs) + + # One running session (no endedAt, so it passes the session threshold filter) + # with two SUCCEEDED task-run session actions carrying output manifests. + deadline_mock.list_sessions.return_value = { + "sessions": [ + { + "sessionId": MOCK_SESSION_ID, + "fleetId": MOCK_FLEET_ID, + "workerId": MOCK_WORKER_ID, + "startedAt": "2025-05-26T11:40:00+00:00", + "lifecycleStatus": "STARTED", + } + ] + } + deadline_mock.list_session_actions.return_value = { + "sessionActions": [ + { + "sessionActionId": MOCK_SESSION_ACTION_ID_1, + "status": "SUCCEEDED", + "startedAt": "2025-05-26T11:41:00+00:00", + "endedAt": "2025-05-26T11:42:00+00:00", + "definition": {"taskRun": {"taskId": "task-abc-0", "stepId": "step-abc"}}, + "manifests": [{"outputManifestPath": "manifest_action_0_output.xxh128"}], + }, + { + "sessionActionId": MOCK_SESSION_ACTION_ID_2, + "status": "SUCCEEDED", + "startedAt": "2025-05-26T11:43:00+00:00", + "endedAt": "2025-05-26T11:44:00+00:00", + "definition": {"taskRun": {"taskId": "task-abc-1", "stepId": "step-abc"}}, + "manifests": [{"outputManifestPath": "manifest_action_1_output.xxh128"}], + }, + ] + } + + files = _put_manifest_and_data_objects_in_s3(root_path) + + # RUN 1: --json mode, real downloads through moto S3. + runner = CliRunner() + with freeze_time(ISO_FREEZE_TIME): + result = runner.invoke( + main, + [ + "queue", + "sync-output", + "--ignore-storage-profiles", + "--json", + "--farm-id", + MOCK_FARM_ID, + "--queue-id", + MOCK_QUEUE_ID, + "--checkpoint-dir", + checkpoint_dir, + ], + ) + assert result.exit_code == 0, result.output + + # The files must have actually been downloaded from S3, byte-for-byte. + for rel_path, content in files.items(): + local_file = os.path.join(root_path, *rel_path.split("/")) + assert os.path.isfile(local_file), f"missing downloaded file: {local_file}" + with open(local_file, "rb") as f: + assert f.read() == content, rel_path + + # stdout must be byte-clean for scripting: the command currently emits no JSON + # payload of its own, so any output at all is leakage. If a legitimate JSON + # payload is added later, every line must still parse as JSON. + for line in result.output.splitlines(): + if line.strip(): + json.loads(line) + assert result.output == "", repr(result.output) + + # RUN 2 (contrast): the same download without --json emits the DEBUG lines and + # the ProgressTracker's 100% progress message, proving the progress callbacks + # fired during RUN 1 and were suppressed rather than the download not happening. + contrast_checkpoint_dir = str(tmp_path / "contrast_checkpoint") + with freeze_time(ISO_FREEZE_TIME): + result = runner.invoke( + main, + [ + "queue", + "sync-output", + "--ignore-storage-profiles", + "--farm-id", + MOCK_FARM_ID, + "--queue-id", + MOCK_QUEUE_ID, + "--checkpoint-dir", + contrast_checkpoint_dir, + ], + ) + assert result.exit_code == 0, result.output + assert "DEBUG: Got" in result.output, result.output + assert "Downloading 3 files from S3..." in result.output, result.output + # The 100% progress callback message, e.g. "Downloaded 70 B / 70 B of 3 files (...)" + assert "of 3 files" in result.output, result.output + assert "Downloaded files: 3" in result.output, result.output