fix: Add OCIO configuration to Job Environments (#166)

* fix: Add OCIO configuration to Job Environments

Signed-off-by: Haejung Choi <chaejung@amazon.com>
---------

Signed-off-by: Haejung Choi <chaejung@amazon.com>
Co-authored-by: Haejung Choi <chaejung@amazon.com>

job_bundle_output_tests/ocio/expected_job_bundle/parameter_values.yaml

@@ -3,6 +3,8 @@ parameterValues:
   value: 1-100
 - name: NukeScriptFile
   value: /normalized/job/bundle/dir/ocio.nk
+- name: OCIOConfigPath
+  value: /normalized/job/bundle/dir/config.ocio
 - name: ProxyMode
   value: 'false'
 - name: deadline:targetTaskRunStatus

job_bundle_output_tests/ocio/expected_job_bundle/template.yaml

@@ -16,6 +16,11 @@ parameterDefinitions:
       patterns:
       - '*'
   description: The Nuke script file to render.
+- name: OCIOConfigPath
+  type: PATH
+  objectType: FILE
+  dataFlow: IN
+  description: The OCIO config file used by this job.
 - name: Frames
   type: STRING
   description: The frames to render. E.g. 1-3,8,11-15
@@ -126,3 +131,7 @@ steps:
         cancelation:
           mode: NOTIFY_THEN_TERMINATE
         timeout: 518400
+jobEnvironments:
+- name: Add OCIO Path to Environment Variable
+  variables:
+    OCIO: '{{Param.OCIOConfigPath}}'

src/deadline/nuke_adaptor/NukeClient/nuke_client.py

@@ -58,8 +58,8 @@ def ensure_output_dir():
                 os.makedirs(output_dir)
 
         def verify_ocio_config():
-            """If using a custom OCIO config, update the internal search paths if necessary"""
-            if nuke_ocio.is_custom_config_enabled():
+            """If using an OCIO config, update the internal search paths if necessary"""
+            if nuke_ocio.is_OCIO_enabled():
                 self._map_ocio_config()
 
         nuke.addBeforeRender(verify_ocio_config)

src/deadline/nuke_submitter/assets.py

@@ -7,7 +7,7 @@
 from collections.abc import Generator
 from os.path import commonpath, dirname, join, normpath, samefile
 from sys import platform
-
+from typing import Optional
 import nuke
 
 from deadline.client.job_bundle.submission import AssetReferences
@@ -81,17 +81,38 @@ def get_scene_asset_references() -> AssetReferences:
             for filename in get_node_filenames(node):
                 asset_references.output_directories.add(dirname(filename))
 
-    # if using a custom OCIO config, add the config file and associated search directories
-    if nuke_ocio.is_custom_config_enabled():
-        ocio_config_path = nuke_ocio.get_custom_config_path()
-        ocio_config_search_paths = nuke_ocio.get_config_absolute_search_paths(ocio_config_path)
+    if nuke_ocio.is_OCIO_enabled():
+        # Determine and add the config file and associated search directories
+        ocio_config_path = get_ocio_config_path()
+        # Add the references
+        if ocio_config_path is not None:
+            if os.path.isfile(ocio_config_path):
+                asset_references.input_filenames.add(ocio_config_path)
+
+                ocio_config_search_paths = nuke_ocio.get_config_absolute_search_paths(
+                    ocio_config_path
+                )
+                for search_path in ocio_config_search_paths:
+                    asset_references.input_directories.add(search_path)
+            else:
+                raise DeadlineOperationError(
+                    "OCIO config file specified(%s) is not an existing file. Please check and update the config file before proceeding."
+                    % ocio_config_path
+                )
 
-        asset_references.input_filenames.add(ocio_config_path)
+    return asset_references
 
-        for search_path in ocio_config_search_paths:
-            asset_references.input_directories.add(search_path)
 
-    return asset_references
+def get_ocio_config_path() -> Optional[str]:
+    # if using a custom OCIO environment variable
+    if nuke_ocio.is_env_config_enabled():
+        return nuke_ocio.get_env_config_path()
+    elif nuke_ocio.is_custom_config_enabled():
+        return nuke_ocio.get_custom_config_path()
+    elif nuke_ocio.is_stock_config_enabled():
+        return nuke_ocio.get_stock_config_path()
+    else:
+        return None
 
 
 def find_all_write_nodes() -> set:

src/deadline/nuke_submitter/deadline_submitter_for_nuke.py

@@ -16,6 +16,7 @@
     SubmitJobToDeadlineDialog,
     JobBundlePurpose,
 )
+from deadline.nuke_util import ocio as nuke_ocio
 from nuke import Node
 from PySide2.QtCore import Qt  # pylint: disable=import-error
 from PySide2.QtWidgets import (  # pylint: disable=import-error; type: ignore
@@ -25,7 +26,12 @@
 )
 
 from ._version import version, version_tuple as adaptor_version_tuple
-from .assets import get_nuke_script_file, get_scene_asset_references, find_all_write_nodes
+from .assets import (
+    get_nuke_script_file,
+    get_scene_asset_references,
+    find_all_write_nodes,
+    get_ocio_config_path,
+)
 from .data_classes import RenderSubmitterUISettings
 from .ui.components.scene_settings_tab import SceneSettingsWidget
 from deadline.client.job_bundle.submission import AssetReferences
@@ -112,6 +118,28 @@ def _add_gizmo_dir_to_job_template(job_template: dict[str, Any]) -> None:
     )
 
 
+def _add_ocio_path_to_job_template(job_template: dict[str, Any]) -> None:
+    if "jobEnvironments" not in job_template:
+        job_template["jobEnvironments"] = []
+
+    # This needs to be prepended rather than appended
+    # as it must run before the "Nuke" environment.
+    job_template["jobEnvironments"].insert(
+        0,
+        {
+            "name": "Add OCIO Path to Environment Variable",
+            "variables": {"OCIO": "{{Param.OCIOConfigPath}}"},
+        },
+    )
+
+
+def _remove_ocio_path_from_job_template(job_template: dict[str, Any]) -> None:
+    for index, param in enumerate(job_template["parameterDefinitions"]):
+        if param["name"] == "OCIOConfigPath":
+            job_template["parameterDefinitions"].pop(index)
+            break
+
+
 def _get_job_template(settings: RenderSubmitterUISettings) -> dict[str, Any]:
     # Load the default Nuke job template, and then fill in scene-specific
     # values it needs.
@@ -144,6 +172,12 @@ def _get_job_template(settings: RenderSubmitterUISettings) -> dict[str, Any]:
     # Set the View parameter allowed values
     parameter_def_map["View"]["allowedValues"] = ["All Views"] + sorted(nuke.views())
 
+    # if OCIO is disabled, remove OCIO path from the template
+    if nuke_ocio.is_OCIO_enabled():
+        _add_ocio_path_to_job_template(job_template)
+    else:
+        _remove_ocio_path_from_job_template(job_template)
+
     # If this developer option is enabled, merge the adaptor_override_environment
     if settings.include_adaptor_wheels:
         with open(Path(__file__).parent / "adaptor_override_environment.yaml") as f:
@@ -239,6 +273,16 @@ def _get_parameter_values(
     parameter_values.append(
         {"name": "ProxyMode", "value": "true" if settings.is_proxy_mode else "false"}
     )
+
+    # Set the OCIO config path value
+    if nuke_ocio.is_OCIO_enabled():
+        ocio_config_path = get_ocio_config_path()
+        if ocio_config_path:
+            parameter_values.append({"name": "OCIOConfigPath", "value": ocio_config_path})
+        else:
+            raise DeadlineOperationError(
+                "OCIO is enabled but OCIO config file is not specified. Please check and update the config file before proceeding."
+            )
     if settings.include_adaptor_wheels:
         wheels_path = str(Path(__file__).parent.parent.parent.parent / "wheels")
         parameter_values.append({"name": "AdaptorWheels", "value": wheels_path})

src/deadline/nuke_submitter/default_nuke_job_template.yaml

@@ -25,6 +25,11 @@ parameterDefinitions:
     label: Gizmo Directory
   description: The directory containing Nuke Gizmo files used by this job.
   default: './gizmos'
+- name: OCIOConfigPath
+  type: PATH
+  objectType: FILE
+  dataFlow: IN
+  description: The OCIO config file used by this job.
 - name: Frames
   type: STRING
   description: The frames to render. E.g. 1-3,8,11-15

src/deadline/nuke_util/ocio.py

@@ -4,11 +4,40 @@
 
 import os
 from pathlib import PurePath
+from typing import Optional
 
 import nuke
 import PyOpenColorIO as OCIO
 
 
+def is_env_config_enabled() -> bool:
+    """True if the OCIO environment variable is specified"""
+    return "OCIO" in os.environ
+
+
+def get_env_config_path() -> Optional[str]:
+    """This is the path to the custom OCIO config used by the OCIO env var."""
+    return os.environ.get("OCIO")
+
+
+def is_stock_config_enabled() -> bool:
+    """True if the script is using a default OCIO config"""
+    return (
+        nuke.root().knob("colorManagement").value() == "OCIO"
+        and nuke.root().knob("OCIO_config").value() != "custom"
+    )
+
+
+def get_stock_config_path() -> str:
+    """This is the path to the UI defined OCIO config file."""
+    return os.path.abspath(nuke.root().knob("OCIOConfigPath").getEvaluatedValue())
+
+
+def is_OCIO_enabled() -> bool:
+    """Nuke is set to use OCIO."""
+    return nuke.root().knob("colorManagement").value() == "OCIO"
+
+
 def is_custom_config_enabled() -> bool:
     """True if the script is using a custom OCIO config"""
     return (
@@ -19,7 +48,7 @@ def is_custom_config_enabled() -> bool:
 
 def get_custom_config_path() -> str:
     """This is the path to the custom OCIO config used by the script"""
-    return nuke.root().knob("customOCIOConfigPath").getEvaluatedValue()
+    return os.path.abspath(nuke.root().knob("customOCIOConfigPath").getEvaluatedValue())
 
 
 def create_config_from_file(ocio_config_path: str) -> OCIO.Config:

test/unit/deadline_submitter_for_nuke/test_assets.py

@@ -14,6 +14,7 @@
     get_node_file_knob_paths,
     get_node_filenames,
     get_scene_asset_references,
+    get_ocio_config_path,
 )
 
 
@@ -37,6 +38,8 @@ def _activated_reading_write_node_knobs(knob_name: str):
     return_value=["/one/asset.png", "/two/asset.png"],
 )
 @patch("deadline.nuke_util.ocio.is_custom_config_enabled", return_value=False)
+@patch("deadline.nuke_util.ocio.is_stock_config_enabled", return_value=False)
+@patch("deadline.nuke_util.ocio.is_OCIO_enabled", return_value=False)
 @patch(
     "deadline.nuke_util.ocio.get_custom_config_path",
     return_value="/this/ocio_configs/config.ocio",
@@ -48,6 +51,8 @@ def _activated_reading_write_node_knobs(knob_name: str):
 def test_get_scene_asset_references(
     mock_get_config_absolute_search_paths: Mock,
     mock_get_custom_config_path: Mock,
+    mock_is_OCIO_enabled: Mock,
+    mock_is_stock_config_enabled: Mock,
     mock_is_custom_config_enabled: Mock,
     mock_get_node_filenames: Mock,
     mock_get_nuke_script_file: Mock,
@@ -95,10 +100,11 @@ def test_get_scene_asset_references(
 
     nuke.allNodes.return_value = []
     mock_is_custom_config_enabled.return_value = True
+    mock_is_OCIO_enabled.return_value = True
+    mock_is_stock_config_enabled.return_value = False
 
     # WHEN
     results = get_scene_asset_references()
-
     # THEN
     assert expected_script_file in results.input_filenames
     assert expected_ocio_config_path in results.input_filenames
@@ -108,6 +114,44 @@ def test_get_scene_asset_references(
     )
 
 
+@patch("deadline.nuke_util.ocio.is_env_config_enabled", return_value=True)
+@patch("deadline.nuke_util.ocio.is_custom_config_enabled", return_value=False)
+@patch("deadline.nuke_util.ocio.is_stock_config_enabled", return_value=False)
+@patch(
+    "deadline.nuke_util.ocio.get_env_config_path",
+    return_value="/this/ocio_configs/env_variable_config.ocio",
+)
+@patch(
+    "deadline.nuke_util.ocio.get_custom_config_path",
+    return_value="/this/ocio_configs/custom_config.ocio",
+)
+@patch(
+    "deadline.nuke_util.ocio.get_stock_config_path",
+    return_value="/this/ocio_configs/stock_config.ocio",
+)
+def test_get_ocio_config_path(
+    mock_get_stock_config_path,
+    mock_get_custom_config_path,
+    mock_get_env_config_path,
+    mock_is_stock_config_enabled,
+    mock_is_custom_enabled,
+    mock_is_env_config_enabled,
+):
+    env_variable_ocio_path = get_ocio_config_path()
+    assert env_variable_ocio_path == "/this/ocio_configs/env_variable_config.ocio"
+
+    mock_is_env_config_enabled.return_value = False
+    mock_is_custom_enabled.return_value = True
+    custom_ocio_path = get_ocio_config_path()
+    assert custom_ocio_path == "/this/ocio_configs/custom_config.ocio"
+
+    mock_is_env_config_enabled.return_value = False
+    mock_is_custom_enabled.return_value = False
+    mock_is_stock_config_enabled.return_value = True
+    stock_ocio_path = get_ocio_config_path()
+    assert stock_ocio_path == "/this/ocio_configs/stock_config.ocio"
+
+
 @patch("os.path.isfile", return_value=False)
 @patch("deadline.nuke_submitter.assets.get_nuke_script_file", return_value="/this/scriptfile.nk")
 def test_get_scene_asset_references_script_not_saved(

test/unit/deadline_submitter_for_nuke/test_ocio.py

@@ -0,0 +1,43 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+
+from deadline.nuke_submitter.deadline_submitter_for_nuke import (
+    _add_ocio_path_to_job_template,
+    _remove_ocio_path_from_job_template,
+)
+
+
+def test_add__dir_to_job_template():
+    job_template = {
+        "parameterDefinitions": [],
+        "steps": [
+            {
+                "name": "Render",
+                "stepEnvironments": [
+                    {
+                        "name": "OtherEnv",
+                    },
+                ],
+            }
+        ],
+    }
+
+    _add_ocio_path_to_job_template(job_template)
+
+    expected_environment = {
+        "name": "Add OCIO Path to Environment Variable",
+        "variables": {"OCIO": "{{Param.OCIOConfigPath}}"},
+    }
+
+    assert len(job_template["jobEnvironments"]) == 1
+    assert job_template["jobEnvironments"][0] == expected_environment
+
+
+def test_remove_ocio_path_from_job_template():
+    job_template = {"parameterDefinitions": [{"name": "OCIOConfigPath", "type": "PATH"}]}
+
+    _remove_ocio_path_from_job_template(job_template)
+    expected_job_template: dict = {"parameterDefinitions": []}
+
+    assert len(job_template["parameterDefinitions"]) == 0
+    assert job_template == expected_job_template