From 40a85e479b65457e2eb0801d06f5ffccc9ef6dc9 Mon Sep 17 00:00:00 2001 From: lucaseck <117225985+lucaseck@users.noreply.github.com> Date: Thu, 4 Jul 2024 15:23:16 -0500 Subject: [PATCH] feat: support ksp bundling Signed-off-by: lucaseck <117225985+lucaseck@users.noreply.github.com> --- README.md | 18 +++ pyproject.toml | 2 +- .../Submit to AWS Deadline Cloud.py | 150 +++++++++++++----- test/keyshot_submitter/mock_lux.py | 5 +- .../test_keyshot_submitter.py | 63 ++++++-- 5 files changed, 183 insertions(+), 55 deletions(-) diff --git a/README.md b/README.md index 87c25ef..b03aea2 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,24 @@ This package provides a KeyShot plugin script that creates jobs for AWS Deadline - You can navigate to the folder by going to Finder, clicking the menu for Go -> Go to Folder, and typing in the folder path. 3. Launch KeyShot. The submitter can be launched within KeyShot from `Window > Scripting Console > Scripts > Submit to AWS Deadline Cloud > Run` +#### Submission Modes + +There are two submission modes for the KeyShot submitter which a dialog will ask you to select from before opening the submitter UI. + +1. Attach `The scene BIP file and all external files references` + - The open scene file and all external files referenced within will be included + as job attachments. The submitter will export the open scene to a + KSP(KeyShot Package) which turns all file paths in the scene into relative + paths and creates a flattened directory with all of the external files directly + beside the scene file. The KSP is then unzipped, and the new scene file with + the relative paths and all external files will be submitted with the job. + The temporary directory used to save the KSP will be deleted after each + submission. +1. Attach `Only the scene BIP file` + - Only the open scene file will be attached to the submission. The expectation is that any + external files referenced within the scene will be available to the workers + through network storage or some other method. + ## Adaptor The KeyShot Adaptor implements the [OpenJD][openjd-adaptor-runtime] interface that allows render workloads to launch KeyShot and feed it commands. This gives the following benefits: diff --git a/pyproject.toml b/pyproject.toml index 3dba1d3..ccdb221 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,7 +139,7 @@ source = [ [tool.coverage.report] show_missing = true -fail_under = 10 +fail_under = 18 [tool.semantic_release] # Can be removed or set to true once we are v1 diff --git a/src/deadline/keyshot_submitter/Submit to AWS Deadline Cloud.py b/src/deadline/keyshot_submitter/Submit to AWS Deadline Cloud.py index 1743e4b..82cb1eb 100644 --- a/src/deadline/keyshot_submitter/Submit to AWS Deadline Cloud.py +++ b/src/deadline/keyshot_submitter/Submit to AWS Deadline Cloud.py @@ -7,15 +7,20 @@ import os import platform import shlex +import shutil import subprocess import tempfile +import glob from dataclasses import dataclass from pathlib import Path -from typing import Any, Optional +from typing import Any, Optional, Tuple import lux RENDER_SUBMITTER_SETTINGS_FILE_EXT = ".deadline_render_settings.json" +SUBMISSION_MODE_KEY = "submission_mode" +# Unique ID required to allow KeyShot to save selections for a dialog +DEADLINE_CLOUD_DIALOG_ID = "e309ce79-3ee8-446a-8308-10d16dfcbb42" @dataclass @@ -26,7 +31,6 @@ class Settings: output_directories: list[str] referenced_paths: list[str] auto_detected_input_filenames: list[str] - auto_detected_output_directories: list[str] def output_sticky_settings(self): return { @@ -88,11 +92,7 @@ def apply_submitter_settings(self, output: dict): if asset_references.get("inputs", {}).get("directories"): self.input_directories = asset_references["inputs"]["directories"] if asset_references.get("outputs", {}).get("directories"): - # Persist output directories that were not autodetected (i.e. were manually added) - self.output_directories = list( - set(asset_references["outputs"]["directories"]) - - set(self.auto_detected_output_directories) - ) + self.output_directories = asset_references["outputs"]["directories"] if asset_references.get("referencedPaths"): self.referenced_paths = asset_references["referencedPaths"] @@ -259,18 +259,7 @@ def construct_asset_references(settings: Settings) -> dict: list(set([*settings.input_filenames, *settings.auto_detected_input_filenames])) ), }, - "outputs": { - "directories": sorted( - list( - set( - [ - *settings.output_directories, - *settings.auto_detected_output_directories, - ] - ) - ) - ), - }, + "outputs": {"directories": sorted(settings.output_directories)}, "referencedPaths": sorted(settings.referenced_paths), } } @@ -346,6 +335,7 @@ def gui_submit(bundle_directory: str) -> Optional[dict[str, Any]]: check=True, capture_output=True, text=True, + creationflags=subprocess.CREATE_NO_WINDOW, # type: ignore[attr-defined] ) output = result.stdout except subprocess.CalledProcessError as e: @@ -357,7 +347,83 @@ def gui_submit(bundle_directory: str) -> Optional[dict[str, Any]]: return None +def options_dialog() -> dict[str, Any]: + """ + Builds and displays a dialog within KeyShot to get the submission options + reuired before the main gui submission window is opened outside of KeyShot. + Options: + Option 1: Dropdown to select whether to submit just the scene file itself + or all external file references as well by packing/unpacking a + KSP bundle before submission. + Returns a dictionary of the selected option values in the format: + {'SUBMISSION_MODE_KEY': [1, 'only the scene BIP file']} + """ + dialog_items = [ + ( + SUBMISSION_MODE_KEY, + lux.DIALOG_ITEM, + "What files would you like to attach to the job?", + 0, + ["The scene BIP file and all external files references", "Only the scene BIP file"], + ) + ] + selections = lux.getInputDialog( + title="AWS Deadline Cloud Submission Options", + values=dialog_items, + id=DEADLINE_CLOUD_DIALOG_ID, + ) + + return selections + + +def save_ksp_bundle(directory: str, bundle_name: str) -> str: + """ + Saves out the current scene and any file references to a ksp bundle in a + directory. + Returns the file path where the bundle was saved to. + """ + full_ksp_path = os.path.join(directory, bundle_name) + + if not lux.savePackage(path=full_ksp_path): + raise RuntimeError("KSP was not able to be saved!") + + return full_ksp_path + + +def get_ksp_bundle_files(directory: str) -> Tuple[str, list[str]]: + """ + Creates a ksp bundle from the current scene containing the scene file and + any external file references. The bundle is unpacked into a directory passed + in. + Returns the scene file and a list of the external files from the directory + where the ksp was extracted to. + """ + + ksp_dir = os.path.join(directory, "ksp") + unpack_dir = os.path.join(directory, "unpack") + + ksp_archive = save_ksp_bundle(ksp_dir, "temp_deadline_cloud.ksp") + shutil.unpack_archive(ksp_archive, unpack_dir, "zip") + input_filenames = [ + os.path.join(unpack_dir, file) + for file in os.listdir(unpack_dir) + if not file.endswith(".bip") + ] + + bip_files = glob.glob(os.path.join(unpack_dir, "*.bip")) + + if not bip_files: + raise RuntimeError("No .bip files found in the KSP bundle.") + elif len(bip_files) > 1: + raise RuntimeError("Multiple .bip files found in the KSP bundle.") + + bip_file = bip_files[0] + + return bip_file, input_filenames + + def main(lux): + if lux.isSceneChanged(): result = lux.getInputDialog( title="Unsaved changes", @@ -370,24 +436,27 @@ def main(lux): else: lux.saveFile() - scene_file = lux.getSceneInfo()["file"] - external_files = lux.getExternalFiles() + dialog_selections = options_dialog() + + if not dialog_selections: + print("Job submission canceled.") + return + + scene_info = lux.getSceneInfo() + scene_file = scene_info["file"] + scene_name, _ = os.path.splitext(scene_info["name"]) current_frame = lux.getAnimationFrame() frame_count = lux.getAnimationInfo().get("frames") settings = Settings( parameter_values=[ - { - "name": "KeyShotFile", - "value": scene_file, - }, { "name": "Frames", "value": f"1-{frame_count}" if frame_count else f"{current_frame}", }, { "name": "OutputFilePath", - "value": f"{scene_file}.%d.png", + "value": os.path.join(os.path.dirname(scene_file), f"{scene_name}.%d.png"), }, { "name": "OutputFormat", @@ -395,29 +464,34 @@ def main(lux): }, ], input_filenames=[], - auto_detected_input_filenames=[*external_files, scene_file], + auto_detected_input_filenames=[], input_directories=[], output_directories=[], - auto_detected_output_directories=[str(os.path.dirname(scene_file))], referenced_paths=[], ) - _, filename = os.path.split(scene_file) - sticky_settings = load_sticky_settings(scene_file) if sticky_settings: settings.apply_sticky_settings(sticky_settings) - job_template = construct_job_template(filename) - asset_references = construct_asset_references(settings) - parameter_values = construct_parameter_values(settings) + with tempfile.TemporaryDirectory() as bundle_temp_dir: + # {'submission_mode': [0, 'the scene BIP file and all external files references']} + if not dialog_selections[SUBMISSION_MODE_KEY][0]: + temp_scene_file, input_filenames = get_ksp_bundle_files(bundle_temp_dir) + settings.auto_detected_input_filenames = input_filenames + settings.parameter_values.append({"name": "KeyShotFile", "value": temp_scene_file}) + else: + settings.parameter_values.append({"name": "KeyShotFile", "value": scene_file}) + + job_template = construct_job_template(scene_name) + asset_references = construct_asset_references(settings) + parameter_values = construct_parameter_values(settings) - with tempfile.TemporaryDirectory() as temp_dir: - dump_json_to_dir(job_template, temp_dir, "template.json") - dump_json_to_dir(asset_references, temp_dir, "asset_references.json") - dump_json_to_dir(parameter_values, temp_dir, "parameter_values.json") + dump_json_to_dir(job_template, bundle_temp_dir, "template.json") + dump_json_to_dir(asset_references, bundle_temp_dir, "asset_references.json") + dump_json_to_dir(parameter_values, bundle_temp_dir, "parameter_values.json") - output = gui_submit(temp_dir) + output = gui_submit(bundle_temp_dir) if output: settings.apply_submitter_settings(output) diff --git a/test/keyshot_submitter/mock_lux.py b/test/keyshot_submitter/mock_lux.py index 0c89122..25cc330 100644 --- a/test/keyshot_submitter/mock_lux.py +++ b/test/keyshot_submitter/mock_lux.py @@ -4,11 +4,12 @@ from unittest.mock import Mock module_name = "lux" -hou_module = types.ModuleType(module_name) -sys.modules[module_name] = hou_module +lux_module = types.ModuleType(module_name) +sys.modules[module_name] = lux_module this_module = sys.modules[module_name] # Usage: set mocked names here, set mocked return values/properties in unit tests. # Mocked names setattr(this_module, "getSceneInfo", Mock(name=module_name + ".getSceneInfo")) setattr(this_module, "getExternalFiles", Mock(name=module_name + ".getExternalFiles")) setattr(this_module, "getAnimationFrame", Mock(name=module_name + ".getAnimationFrame")) +setattr(this_module, "savePackage", Mock(name=module_name + ".savePackage")) diff --git a/test/keyshot_submitter/test_keyshot_submitter.py b/test/keyshot_submitter/test_keyshot_submitter.py index 1c1cc41..33ff867 100644 --- a/test/keyshot_submitter/test_keyshot_submitter.py +++ b/test/keyshot_submitter/test_keyshot_submitter.py @@ -1,12 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -from unittest.mock import Mock import mock_lux # type: ignore[import-not-found] # noqa: F401 import json import os import tempfile import pytest - +import shutil +from unittest import mock deadline = __import__("deadline.keyshot_submitter.Submit to AWS Deadline Cloud") submitter = getattr(deadline.keyshot_submitter, "Submit to AWS Deadline Cloud") @@ -32,7 +32,6 @@ def test_construct_asset_references(): auto_detected_input_filenames=["test_filename_2"], input_directories=["test_directory_1", "test_directory_2"], output_directories=["test_directory_3"], - auto_detected_output_directories=["test_directory_4"], referenced_paths=["reference_path_1", "reference_path_2"], ) @@ -45,10 +44,9 @@ def test_construct_asset_references(): assert ( asset_references["assetReferences"]["inputs"]["directories"] == settings.input_directories ) - assert asset_references["assetReferences"]["outputs"]["directories"] == [ - "test_directory_3", - "test_directory_4", - ] + assert ( + asset_references["assetReferences"]["outputs"]["directories"] == settings.output_directories + ) assert asset_references["assetReferences"]["referencedPaths"] == settings.referenced_paths @@ -64,7 +62,6 @@ def test_construct_parameter_values(): auto_detected_input_filenames=[], input_directories=[], output_directories=[], - auto_detected_output_directories=[], referenced_paths=[], ) @@ -92,7 +89,6 @@ def test_settings_serialize_correctly(): auto_detected_input_filenames=["test_directory_3"], input_directories=["test_directory_1"], output_directories=["test_directory_2"], - auto_detected_output_directories=[], referenced_paths=["test_ref_path"], ) @@ -121,7 +117,6 @@ def test_settings_apply_sticky_settings(): auto_detected_input_filenames=["test_filename_2"], input_directories=["test_directory_1"], output_directories=["test_directory_2"], - auto_detected_output_directories=[], referenced_paths=["test_ref_path"], ) initial_settings = settings.output_sticky_settings() @@ -191,7 +186,6 @@ def test_settings_apply_submitter_settings(): auto_detected_input_filenames=["test_filename_2"], input_directories=["test_directory_1"], output_directories=["test_directory_2"], - auto_detected_output_directories=["test_directory_3"], referenced_paths=["test_ref_path"], ) @@ -224,7 +218,6 @@ def test_settings_apply_submitter_settings(): "directories": [ "test_directory_2", "test_directory_3", - "test_directory_4", ] }, "referencedPaths": ["test_ref_path_2"], @@ -242,12 +235,12 @@ def test_settings_apply_submitter_settings(): ] assert sorted(settings.input_filenames) == ["test_filename_1", "test_filename_3"] assert sorted(settings.input_directories) == ["test_directory_2"] - assert sorted(settings.output_directories) == ["test_directory_2", "test_directory_4"] + assert sorted(settings.output_directories) == ["test_directory_2", "test_directory_3"] assert sorted(settings.referenced_paths) == ["test_ref_path_2"] def test_unsaved_changes_prompt(): - local_mock_lux = Mock() + local_mock_lux = mock.Mock() local_mock_lux.isSceneChanged.return_value = True local_mock_lux.getInputDialog.return_value = None # emulate clicking Cancel @@ -263,3 +256,45 @@ def test_unsaved_changes_prompt(): with pytest.raises(Exception): submitter.main(local_mock_lux) local_mock_lux.saveFile.assert_called() + + +def test_save_ksp_bundle(): + dir = os.path.normpath("/testdir/test") + bundle_name = "test_bundle.ksp" + expected_bundle_path = os.path.normpath(f"{dir}/{bundle_name}") + + output = submitter.save_ksp_bundle(dir, bundle_name) + + assert output == expected_bundle_path + mock_lux.lux_module.savePackage.assert_called_once_with(path=expected_bundle_path) + + +def test_get_ksp_bundle_files(): + + TEST_SCENE_FILE = "test_scene_file.bip" + TEST_ASSET_FILE = "test_asset.png" + TEST_KSP_BUNDLE_NAME = "test_ksp_bundle" + + with tempfile.TemporaryDirectory() as temp_dir: + to_zip_dir = os.path.join(temp_dir, "to_zip") + os.mkdir(to_zip_dir) + with open(os.path.join(to_zip_dir, TEST_SCENE_FILE), "w") as file: + file.write("test scene") + with open(os.path.join(to_zip_dir, TEST_ASSET_FILE), "w") as file: + file.write("test asset") + + # Creating a zip archive as that's what is used to unpack a ksp and a ksp + # can't be created easily outside of KeyShot + shutil.make_archive(os.path.join(temp_dir, TEST_KSP_BUNDLE_NAME), "zip", to_zip_dir) + + with mock.patch.object( + submitter, + "save_ksp_bundle", + return_value=os.path.join(temp_dir, f"{TEST_KSP_BUNDLE_NAME}.zip"), + ) as mock_save_ksp_bundle: + scene_file, input_filenames = submitter.get_ksp_bundle_files(temp_dir) + + assert scene_file == os.path.join(temp_dir, "unpack", TEST_SCENE_FILE) + assert len(input_filenames) == 1 + assert input_filenames[0] == os.path.join(temp_dir, "unpack", TEST_ASSET_FILE) + mock_save_ksp_bundle.assert_called_once_with(os.path.join(temp_dir, "ksp"), mock.ANY)