feat: support ksp bundling

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

README.md

@@ -31,6 +31,24 @@ This package provides a KeyShot plugin script that creates jobs for AWS Deadline
     - e.g. System install `%PROGRAMFILES%/KeyShot/Scripts`
 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. Submit `the scene BIP file and all external file references`
+    - The open scene file and all external files referenced within will be
+    submitted 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. Submit `only the scene BIP file`
+    - Only the open scene file will be submitted. 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:

pyproject.toml

@@ -140,7 +140,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

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,
             )
             output = result.stdout
     except subprocess.CalledProcessError as e:
@@ -357,55 +347,139 @@ 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():
-    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",
                 "value": "PNG",
             },
         ],
         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)

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"))

test/keyshot_submitter/test_keyshot_submitter.py

@@ -4,7 +4,8 @@
 import json
 import os
 import tempfile
-
+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")
@@ -30,7 +31,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"],
     )
 
@@ -43,10 +43,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
 
 
@@ -62,7 +61,6 @@ def test_construct_parameter_values():
         auto_detected_input_filenames=[],
         input_directories=[],
         output_directories=[],
-        auto_detected_output_directories=[],
         referenced_paths=[],
     )
 
@@ -90,7 +88,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"],
     )
 
@@ -119,7 +116,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()
@@ -189,7 +185,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"],
     )
 
@@ -222,7 +217,6 @@ def test_settings_apply_submitter_settings():
                             "directories": [
                                 "test_directory_2",
                                 "test_directory_3",
-                                "test_directory_4",
                             ]
                         },
                         "referencedPaths": ["test_ref_path_2"],
@@ -240,5 +234,47 @@ 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_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)