feat(asset-cli): added snapshot subcommmand for deadline asset

Signed-off-by: Tang <[email protected]>

src/deadline/client/api/_submit_job_bundle.py

@@ -215,7 +215,6 @@ def create_job_from_job_bundle(
         # Extend input_filenames with all the files in the input_directories
         missing_directories: set[str] = set()
         for directory in asset_references.input_directories:
-
             if not os.path.isdir(directory):
                 if require_paths_exist:
                     missing_directories.add(directory)
@@ -356,12 +355,6 @@ def _default_create_job_result_callback() -> bool:
         raise DeadlineOperationError("CreateJob response was empty, or did not contain a Job ID.")
 
 
-def get_input_paths(root_directory: str, missing_directories: set, require_paths_exist: bool):
-    """
-    Collect all input paths from specified directory
-    """
-
-
 def wait_for_create_job_to_complete(
     farm_id: str,
     queue_id: str,

src/deadline/client/cli/_common.py

@@ -17,7 +17,6 @@
 
 import click
 from contextlib import ExitStack
-
 from deadline.job_attachments.progress_tracker import ProgressReportMetadata
 
 from ..config import config_file

src/deadline/client/cli/_groups/asset_group.py

@@ -32,22 +32,23 @@ def cli_asset():
 
 @cli_asset.command(name="snapshot")
 @click.option("--root-dir", required=True, help="The root directory to snapshot. ")
-@click.option("--manifest-out-dir", help="Destination path to directory for created manifest. ")
+@click.option("--manifest-out", help="Destination path to directory for created manifest. ")
 @click.option(
-    "--r",
+    "--recursive",
+    "-r",
     help="Flag to recursively snapshot subdirectories. ",
     is_flag=True,
     show_default=True,
     default=False,
 )
 @_handle_error
-def asset_snapshot(r, **args):
+def asset_snapshot(recursive, **args):
     """
     Creates manifest of files specified root directory.
     """
     root_dir = args.pop("root_dir")
     root_dir_basename = os.path.basename(root_dir) + "_"
-    out_dir = args.pop("manifest_out_dir")
+    out_dir = args.pop("manifest_out")
 
     if not os.path.isdir(root_dir):
         misconfigured_directories_msg = f"Specified root directory {root_dir} does not exist. "
@@ -68,7 +69,7 @@ def asset_snapshot(r, **args):
         for file in files:
             file_full_path = str(os.path.join(root, file))
             inputs.append(file_full_path)
-        if not r:
+        if not recursive:
             break
 
     # Placeholder Asset Manager
@@ -126,7 +127,7 @@ def asset_upload(**args):
 
 @cli_asset.command(name="diff")
 @click.option("--root-dir", help="The root directory to compare changes to. ")
-@click.option("--manifest", help="The path to manifest of working directory to show changes of. ")
+@click.option("--manifest", help="The path to manifest folder of directory to show changes of. ")
 @click.option(
     "--print",
     help="Pretty prints diff information. ",

src/deadline/client/cli/_groups/bundle_group.py

@@ -280,42 +280,3 @@ def bundle_gui_submit(job_bundle_dir, browse, **args):
             click.echo(f"Job ID: {response['jobId']}")
         else:
             click.echo("Job submission canceled.")
-
-
-# class _ProgressBarCallbackManager:
-#     """
-#     Manages creation, update, and deletion of a progress bar. On first call of the callback, the progress bar is created. The progress bar is closed
-#     on the final call (100% completion)
-#     """
-
-#     BAR_NOT_CREATED = 0
-#     BAR_CREATED = 1
-#     BAR_CLOSED = 2
-
-#     def __init__(self, length: int, label: str):
-#         self._length = length
-#         self._label = label
-#         self._bar_status = self.BAR_NOT_CREATED
-#         self._exit_stack = ExitStack()
-
-#     def callback(self, upload_metadata: ProgressReportMetadata) -> bool:
-#         if self._bar_status == self.BAR_CLOSED:
-#             # from multithreaded execution this can be called after completion somtimes.
-#             return sigint_handler.continue_operation
-#         elif self._bar_status == self.BAR_NOT_CREATED:
-#             # Note: click doesn't export the return type of progressbar(), so we suppress mypy warnings for
-#             # not annotating the type of hashing_progress.
-#             self._upload_progress = click.progressbar(length=self._length, label=self._label)  # type: ignore[var-annotated]
-#             self._exit_stack.enter_context(self._upload_progress)
-#             self._bar_status = self.BAR_CREATED
-
-#         total_progress = int(upload_metadata.progress)
-#         new_progress = total_progress - self._upload_progress.pos
-#         if new_progress > 0:
-#             self._upload_progress.update(new_progress)
-
-#         if total_progress == self._length or not sigint_handler.continue_operation:
-#             self._bar_status = self.BAR_CLOSED
-#             self._exit_stack.close()
-
-#         return sigint_handler.continue_operation

src/deadline/job_attachments/upload.py

@@ -179,7 +179,7 @@ def upload_assets(
 
         if manifest_write_dir:
             self._write_local_manifest(
-                manifest_write_dir, manifest_name, full_manifest_key, manifest
+                manifest_write_dir, manifest_name, full_manifest_key, manifest, None
             )
 
         self.upload_bytes_to_s3(
@@ -224,35 +224,31 @@ def _write_local_manifest(
         manifest_name: str,
         full_manifest_key: str,
         manifest: BaseAssetManifest,
+        root_dir_name: Optional[str],
     ) -> None:
         """
         Writes a manifest file locally in a 'manifests' sub-directory.
         Also creates/appends to a file mapping the local manifest name to the full S3 key in the same directory.
         """
-        # local_manifest_file = Path(manifest_write_dir, "manifests", manifest_name)
-        # logger.info(f"Creating local manifest file: {local_manifest_file}\n")
-        # local_manifest_file.parent.mkdir(parents=True, exist_ok=True)
-        # with open(local_manifest_file, "w") as file:
-        #     file.write(manifest.encode())
-        self._write_local_input_manifest(manifest_write_dir, manifest_name, manifest)
-
-        # manifest_map_file = Path(manifest_write_dir, "manifests", "manifest_s3_mapping")
-        # mapping = {"local_file": manifest_name, "s3_key": full_manifest_key}
-        # with open(manifest_map_file, "a") as mapping_file:
-        #     mapping_file.write(f"{mapping}\n")
+        self._write_local_input_manifest(manifest_write_dir, manifest_name, manifest, root_dir_name)
+
         self._write_local_manifest_s3_mapping(manifest_write_dir, manifest_name, full_manifest_key)
 
     def _write_local_input_manifest(
         self,
         manifest_write_dir: str,
         manifest_name: str,
         manifest: BaseAssetManifest,
-        root_dir_name: Optional[str]
+        root_dir_name: Optional[str],
     ):
         """
         Creates 'manifests' sub-directory and writes a local input manifest file
         """
-        local_manifest_file = Path(manifest_write_dir, root_dir_name + "manifests", manifest_name)
+        input_manifest_folder_name = "manifests"
+        if root_dir_name is not None:
+            input_manifest_folder_name = root_dir_name + input_manifest_folder_name
+
+        local_manifest_file = Path(manifest_write_dir, input_manifest_folder_name, manifest_name)
         logger.info(f"Creating local manifest file: {local_manifest_file}\n")
         local_manifest_file.parent.mkdir(parents=True, exist_ok=True)
         with open(local_manifest_file, "w") as file:

test/unit/deadline_client/api/test_job_bundle_submission.py

@@ -479,7 +479,7 @@ def test_create_job_from_job_bundle_job_attachments(
     ) as client_mock, patch.object(
         _submit_job_bundle.api, "get_queue_user_boto3_session"
     ), patch.object(
-        api._submit_job_bundle, "_hash_attachments", return_value=(None, None)
+        api._submit_job_bundle, "hash_attachments", return_value=(None, None)
     ) as mock_hash_attachments, patch.object(
         S3AssetManager,
         "prepare_paths_for_upload",
@@ -601,7 +601,7 @@ def test_create_job_from_job_bundle_empty_job_attachments(
     ) as client_mock, patch.object(
         _submit_job_bundle.api, "get_queue_user_boto3_session"
     ), patch.object(
-        api._submit_job_bundle, "_hash_attachments", return_value=(None, None)
+        api._submit_job_bundle, "hash_attachments", return_value=(None, None)
     ) as mock_hash_attachments, patch.object(
         S3AssetManager,
         "prepare_paths_for_upload",
@@ -923,7 +923,7 @@ def test_create_job_from_job_bundle_with_single_asset_file(
     ) as client_mock, patch.object(
         _submit_job_bundle.api, "get_queue_user_boto3_session"
     ), patch.object(
-        api._submit_job_bundle, "_hash_attachments", return_value=(None, None)
+        api._submit_job_bundle, "hash_attachments", return_value=(None, None)
     ) as mock_hash_attachments, patch.object(
         S3AssetManager,
         "prepare_paths_for_upload",

test/unit/deadline_client/cli/test_cli_asset.py

@@ -0,0 +1,108 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+"""
+Tests for the CLI asset commands.
+"""
+import os
+import json
+from click.testing import CliRunner
+import pytest
+import tempfile
+
+from deadline.client.cli import main
+
+
+@pytest.fixture
+def temp_dir():
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        yield tmpdirname
+
+
+def test_cli_asset_snapshot_valid_root_dir(temp_dir):
+    root_dir = os.path.join(temp_dir, "root_dir")
+    os.makedirs(root_dir)
+    file_path = os.path.join(root_dir, "file.txt")
+    with open(file_path, "w") as f:
+        f.write("test file content")
+
+    runner = CliRunner()
+    runner.invoke(main, ["asset", "snapshot", "--root-dir", root_dir])
+
+    # Check manifest file details to match correct content
+    # since manifest file name is hashed depending on source location, we have to list out manifest
+    manifest_folder_path = os.path.join(root_dir, f"{os.path.basename(root_dir)}_manifests")
+    manifest_file_name = os.listdir(manifest_folder_path)[0]
+    manifest_file_path = os.path.join(manifest_folder_path, manifest_file_name)
+
+    with open(manifest_file_path, "r") as f:
+        manifest_data = json.load(f)
+        print(manifest_data["paths"][0]["hash"])
+        assert manifest_data["paths"][0]["path"] == "file.txt"
+        assert manifest_data["paths"][0]["hash"] == "0741993e50c8bc250cefba3959c81eb8"
+
+
+def test_cli_asset_snapshot_invalid_root_dir(temp_dir):
+    invalid_root_dir = os.path.join(temp_dir, "invalid_root_dir")
+    runner = CliRunner()
+    result = runner.invoke(main, ["asset", "snapshot", "--root-dir", invalid_root_dir])
+    assert result.exit_code == 1
+
+
+def test_cli_asset_snapshot_recursive(temp_dir):
+    root_dir = os.path.join(temp_dir, "root_dir")
+    os.makedirs(os.path.join(root_dir, "subdir1", "subdir2"))
+
+    # Create a file in the root directory
+    root_file_path = os.path.join(root_dir, "root_file.txt")
+    with open(root_file_path, "w") as f:
+        f.write("root file content")
+
+    # Create a file in the subdirectory
+    subdir_file_path = os.path.join(root_dir, "subdir1", "subdir2", "subdir_file.txt")
+    os.makedirs(os.path.dirname(subdir_file_path), exist_ok=True)
+    with open(subdir_file_path, "w") as f:
+        f.write("subdir file content")
+
+    runner = CliRunner()
+    runner.invoke(main, ["asset", "snapshot", "--root-dir", root_dir, "--recursive"])
+
+    # Check manifest file details to match correct content
+    # since manifest file name is hashed depending on source location, we have to list out manifest
+    manifest_folder_path = os.path.join(root_dir, f"{os.path.basename(root_dir)}_manifests")
+    root_manifest_file_name = [file for file in os.listdir(manifest_folder_path)][0]
+    root_manifest_file_path = os.path.join(manifest_folder_path, root_manifest_file_name)
+    with open(root_manifest_file_path, "r") as f:
+        manifest_data = json.load(f)
+        assert manifest_data["paths"][0]["path"] == "root_file.txt"
+        assert manifest_data["paths"][0]["hash"] == "a5fc4a07191e2c90364319d2fd503cc1"
+        assert manifest_data["paths"][1]["path"] == os.path.normcase(
+            os.path.join("subdir1", "subdir2", "subdir_file.txt")
+        )
+        assert manifest_data["paths"][1]["hash"] == "a3ede7fa4c2694d59ff09ed553fcc806"
+
+
+def test_cli_asset_snapshot_valid_manifest_out_dir(temp_dir):
+    root_dir = os.path.join(temp_dir, "root_dir")
+    os.makedirs(root_dir)
+    manifest_out_dir = os.path.join(temp_dir, "manifest_out")
+    os.makedirs(manifest_out_dir)
+    file_path = os.path.join(root_dir, "file.txt")
+    with open(file_path, "w") as f:
+        f.write("test file content")
+
+    runner = CliRunner()
+    runner.invoke(
+        main, ["asset", "snapshot", "--root-dir", root_dir, "--manifest-out", manifest_out_dir]
+    )
+
+    # Check manifest file details to match correct content
+    # since manifest file name is hashed depending on source location, we have to list out manifest
+    manifest_folder_path = os.path.join(manifest_out_dir, f"{os.path.basename(root_dir)}_manifests")
+    manifest_file_name = os.listdir(manifest_folder_path)[0]
+    manifest_file_path = os.path.join(manifest_folder_path, manifest_file_name)
+
+    with open(manifest_file_path, "r") as f:
+        manifest_data = json.load(f)
+        print(manifest_data["paths"][0]["hash"])
+        assert manifest_data["paths"][0]["path"] == "file.txt"
+        assert manifest_data["paths"][0]["hash"] == "0741993e50c8bc250cefba3959c81eb8"

test/unit/deadline_client/cli/test_cli_bundle.py

@@ -118,7 +118,7 @@ def test_cli_bundle_submit(fresh_deadline_config, temp_job_bundle_dir):
     ) as get_boto3_client_mock, patch.object(
         _queue_parameters, "get_boto3_client"
     ) as qp_boto3_client_mock, patch.object(
-        _submit_job_bundle, "_hash_attachments", return_value=[]
+        _submit_job_bundle, "hash_attachments", return_value=[]
     ), patch.object(
         _submit_job_bundle.api, "get_queue_user_boto3_session"
     ), patch.object(
@@ -339,7 +339,7 @@ def test_cli_bundle_asset_load_method(fresh_deadline_config, temp_job_bundle_dir
     ) as bundle_boto3_client_mock, patch.object(
         _queue_parameters, "get_boto3_client"
     ) as qp_boto3_client_mock, patch.object(
-        _submit_job_bundle, "_hash_attachments", return_value=(attachment_mock, {})
+        _submit_job_bundle, "hash_attachments", return_value=(attachment_mock, {})
     ), patch.object(
         _submit_job_bundle, "_upload_attachments", return_value={}
     ), patch.object(
@@ -633,7 +633,7 @@ def test_cli_bundle_accept_upload_confirmation(fresh_deadline_config, temp_job_b
     with patch.object(
         _submit_job_bundle.api, "get_boto3_client"
     ) as get_boto3_client_mock, patch.object(
-        _submit_job_bundle, "_hash_attachments", return_value=[SummaryStatistics(), "test"]
+        _submit_job_bundle, "hash_attachments", return_value=[SummaryStatistics(), "test"]
     ), patch.object(
         _submit_job_bundle, "_upload_attachments"
     ), patch.object(
@@ -711,7 +711,7 @@ def test_cli_bundle_reject_upload_confirmation(fresh_deadline_config, temp_job_b
     ) as get_boto3_client_mock, patch.object(
         _queue_parameters, "get_boto3_client"
     ) as qp_boto3_client_mock, patch.object(
-        _submit_job_bundle, "_hash_attachments", return_value=[SummaryStatistics(), "test"]
+        _submit_job_bundle, "hash_attachments", return_value=[SummaryStatistics(), "test"]
     ), patch.object(
         _submit_job_bundle, "_upload_attachments"
     ) as upload_attachments_mock, patch.object(