feat: Adding storage profiles to create job and config (#9)

feat: Adding storage profiles to create job and config

src/deadline/client/api/__init__.py

@@ -16,6 +16,7 @@
     "list_queues",
     "list_jobs",
     "list_fleets",
+    "list_storage_profiles_for_queue",
     "get_queue_boto3_session",
 ]
 
@@ -152,3 +153,21 @@ def list_fleets(config=None, **kwargs):
 
     deadline = get_boto3_client("deadline", config=config)
     return _call_paginated_deadline_list_api(deadline.list_fleets, "fleets", **kwargs)
+
+
+def list_storage_profiles_for_queue(config=None, **kwargs):
+    """
+    Calls the deadline:ListStorageProfilesForQueue API call, applying the filter for user membership
+    depending on the configuration. If the response is paginated, it repeated
+    calls the API to get all the storage profiles.
+    """
+    if "principalId" not in kwargs:
+        user_id, _ = get_user_and_identity_store_id(config=config)
+        if user_id:
+            kwargs["principalId"] = user_id
+
+    deadline = get_boto3_client("deadline", config=config)
+
+    return _call_paginated_deadline_list_api(
+        deadline.list_storage_profiles_for_queue, "storageProfiles", **kwargs
+    )

src/deadline/client/api/_submit_job_bundle.py

@@ -99,6 +99,10 @@ def create_job_from_job_bundle(
         "templateType": file_type,
     }
 
+    storage_profile_id = get_setting("defaults.storage_profile_id", config=config)
+    if storage_profile_id:
+        create_job_args["storageProfileId"] = storage_profile_id
+
     # The job parameters
     job_bundle_parameters = read_job_bundle_parameters(job_bundle_dir)
 

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

@@ -123,6 +123,10 @@ def bundle_submit(job_bundle_dir, asset_loading_method, parameter, **args):
             "templateType": file_type,
         }
 
+        storage_profile_id = get_setting("defaults.storage_profile_id", config=config)
+        if storage_profile_id:
+            create_job_args["storageProfileId"] = storage_profile_id
+
         # The job parameters
         job_bundle_parameters = read_job_bundle_parameters(job_bundle_dir)
 

src/deadline/client/cli/groups/config_group.py

@@ -27,6 +27,9 @@ def cli_config():
     defaults.queue_id:
         The default queue ID to use for job submissions or CLI operations.
 
+    defaults.storage_profile_id:
+        The default storage profile ID to use for job submission or CLI operations.
+
     settings.job_history_dir:
         The directory in which to create new job bundles for
         submitting to Amazon Deadline Cloud, to produce a history of job submissions.

src/deadline/client/config/config_file.py

@@ -73,6 +73,11 @@
         "depend": "defaults.farm_id",
         "section_format": "{}",
     },
+    "defaults.storage_profile_id": {
+        "default": "",
+        "depend": "defaults.queue_id",
+        "section_format": "{}",
+    },
     "settings.auto_accept": {
         "default": "false",
     },

src/deadline/client/ui/dialogs/deadline_config_dialog.py

@@ -10,6 +10,7 @@
 
 __all__ = ["DeadlineConfigDialog"]
 
+import sys
 import threading
 from configparser import ConfigParser
 from logging import getLogger, root
@@ -152,6 +153,9 @@ class DeadlineWorkstationConfigWidget(QWidget):
     # Emitted when an async refresh_queues_list thread completes,
     # provides (aws_profile_name, farm_id, [(queue_id, queue_name), ...])
     _queue_list_update = Signal(str, str, list)
+    # Emitted when an async refresh_storage_profiles_name_list thread completes,
+    # provides (aws_profile_name, farm_id, queue_id, [storage_profile_id, ...])
+    _storage_profile_list_update = Signal(str, str, list)
     # This signal is sent when any background refresh thread catches an exception,
     # provides (operation_name, BaseException)
     _background_exception = Signal(str, BaseException)
@@ -238,6 +242,18 @@ def _build_farm_settings_ui(self, group, layout):
         self.default_queue_box.background_exception.connect(self.handle_background_exception)
         layout.addRow(default_queue_box_label, self.default_queue_box)
 
+        self.default_storage_profile_box = DeadlineStorageProfileNameListComboBox(parent=group)
+        default_storage_profile_box_label = self.labels["defaults.storage_profile_id"] = QLabel(
+            "Default Storage Profile"
+        )
+        self.default_storage_profile_box.box.currentIndexChanged.connect(
+            self.default_storage_profile_name_changed
+        )
+        self.default_storage_profile_box.background_exception.connect(
+            self.handle_background_exception
+        )
+        layout.addRow(default_storage_profile_box_label, self.default_storage_profile_box)
+
     def _build_general_settings_ui(self, group, layout):
         self.auto_accept = self._init_checkbox_setting(
             group, layout, "settings.auto_accept", "Auto Accept Confirmation Prompts"
@@ -360,6 +376,7 @@ def _fill_aws_profiles_box(self):
     def refresh_lists(self):
         self.default_farm_box.refresh_list()
         self.default_queue_box.refresh_list()
+        self.default_storage_profile_box.refresh_list()
 
     def refresh(self):
         """
@@ -372,6 +389,7 @@ def refresh(self):
             config_file.set_setting(setting_name, value, self.config)
         self.default_farm_box.set_config(self.config)
         self.default_queue_box.set_config(self.config)
+        self.default_storage_profile_box.set_config(self.config)
 
         with block_signals(self.aws_profiles_box):
             aws_profile_name = config_file.get_setting(
@@ -404,6 +422,7 @@ def refresh(self):
             refresh_callback()
 
         self.default_queue_box.refresh_selected_id()
+        self.default_storage_profile_box.refresh_selected_id()
 
         # Put an orange box around the labels for any settings that are changed
         for setting_name, label in self.labels.items():
@@ -447,6 +466,7 @@ def aws_profile_changed(self, value):
         self.changes["defaults.aws_profile_name"] = value
         self.default_farm_box.clear_list()
         self.default_queue_box.clear_list()
+        self.default_storage_profile_box.clear_list()
         self.refresh()
 
     def job_history_dir_changed(self):
@@ -462,6 +482,7 @@ def default_farm_changed(self, index):
         self.changes["defaults.farm_id"] = self.default_farm_box.box.itemData(index)
         self.refresh()
         self.default_queue_box.refresh_list()
+        self.default_storage_profile_box.refresh_list()
 
     def deadline_endpoint_url_edited(self):
         deadline_endpoint_url = self.deadline_endpoint_url_edit.text()
@@ -475,6 +496,13 @@ def deadline_endpoint_url_edited(self):
     def default_queue_changed(self, index):
         self.changes["defaults.queue_id"] = self.default_queue_box.box.itemData(index)
         self.refresh()
+        self.default_storage_profile_box.refresh_list()
+
+    def default_storage_profile_name_changed(self, index):
+        self.changes["defaults.storage_profile_id"] = self.default_storage_profile_box.box.itemData(
+            index
+        )
+        self.refresh()
 
 
 class _DeadlineResourceListComboBox(QWidget):
@@ -621,3 +649,47 @@ def list_resources(self, config: Optional[ConfigParser]):
             return sorted([(item["displayName"], item["queueId"]) for item in response["queues"]])
         else:
             return []
+
+
+class DeadlineStorageProfileNameListComboBox(_DeadlineResourceListComboBox):
+    WINDOWS_OS = "windows"
+    MAC_OS = "macos"
+    LINUX_OS = "linux"
+
+    def __init__(self, parent=None):
+        super().__init__(
+            resource_name="Storage Profile",
+            setting_name="defaults.storage_profile_id",
+            parent=parent,
+        )
+
+    def list_resources(self, config: Optional[ConfigParser]):
+        default_farm_id = config_file.get_setting("defaults.farm_id", config=config)
+        default_queue_id = config_file.get_setting("defaults.queue_id", config=config)
+        if default_farm_id and default_queue_id:
+            response = api.list_storage_profiles_for_queue(
+                config=config, farmId=default_farm_id, queueId=default_queue_id
+            )
+            storage_profiles = response.get("storageProfiles", [])
+            return sorted(
+                (item["displayName"], item["storageProfileId"])
+                for item in storage_profiles
+                if self._get_current_os() == item["osFamily"]
+            )
+        else:
+            return []
+
+    def _get_current_os(self) -> str:
+        """
+        Get a string specifying what the OS is, following the format the Deadline storage profile API expects.
+        """
+        if sys.platform.startswith("linux"):
+            return self.LINUX_OS
+
+        if sys.platform.startswith("darwin"):
+            return self.MAC_OS
+
+        if sys.platform.startswith("win"):
+            return self.WINDOWS_OS
+
+        return "Unknown"

src/deadline/client/ui/dialogs/submit_job_progress_dialog.py

@@ -71,6 +71,7 @@ class SubmitJobProgressDialog(QDialog):
     def start_submission(
         farm_id: str,
         queue_id: str,
+        storage_profile_id: str,
         job_bundle_dir: str,
         asset_manager: S3AssetManager,
         deadline_client: BaseClient,
@@ -86,6 +87,8 @@ def start_submission(
         Args:
             farm_id (str): Id of the farm to submit to
             queue_id (str): Id of the queue to submit to
+            storage_profile_id (str): Id of the storage profile to associate
+                with the job.
             job_bundle_dir (str): Path to the folder containing the job bundle to
                 submit.
             asset_manager (S3AssetManager): A job attachments S3AssetManager
@@ -100,6 +103,7 @@ def start_submission(
         job_progress_dialog = SubmitJobProgressDialog(
             farm_id,
             queue_id,
+            storage_profile_id,
             job_bundle_dir,
             asset_manager,
             deadline_client,
@@ -113,6 +117,7 @@ def __init__(
         self,
         farm_id: str,
         queue_id: str,
+        storage_profile_id: str,
         job_bundle_dir: str,
         asset_manager: S3AssetManager,
         deadline_client: BaseClient,
@@ -124,6 +129,7 @@ def __init__(
 
         self._farm_id = farm_id
         self._queue_id = queue_id
+        self._storage_profile_id = storage_profile_id
         self._job_bundle_dir = job_bundle_dir
         self._asset_manager = asset_manager
         self._deadline_client = deadline_client
@@ -201,6 +207,9 @@ def _start_submission(self):
         self._create_job_args["template"] = file_contents
         self._create_job_args["templateType"] = file_type
 
+        if self._storage_profile_id:
+            self._create_job_args["storageProfileId"] = self._storage_profile_id
+
         # The job parameters
         job_bundle_parameters = read_job_bundle_parameters(self._job_bundle_dir)
 

src/deadline/client/ui/dialogs/submit_job_to_deadline_dialog.py

@@ -261,6 +261,7 @@ def on_submit(self):
 
             farm_id = get_setting("defaults.farm_id")
             queue_id = get_setting("defaults.queue_id")
+            storage_profile_id = get_setting("defaults.storage_profile_id")
 
             queue = deadline.get_queue(farmId=farm_id, queueId=queue_id)
 
@@ -281,6 +282,7 @@ def on_submit(self):
             self.create_job_response = SubmitJobProgressDialog.start_submission(
                 farm_id,
                 queue_id,
+                storage_profile_id,
                 job_bundle_dir,
                 asset_manager,
                 deadline,

src/deadline/client/ui/widgets/shared_job_settings_tab.py

@@ -3,6 +3,7 @@
 """
 A UI Widget containing the render setup tab
 """
+import sys
 import threading
 from typing import Any, Dict, Optional
 
@@ -358,3 +359,51 @@ def get_item(self):
             return (response["queueId"], response["displayName"], response["description"])
         else:
             return ("", "", "")
+
+
+class DeadlineStorageProfileNameDisplay(_DeadlineNamedResourceDisplay):
+    WINDOWS_OS = "Windows"
+    MAC_OS = "Macos"
+    LINUX_OS = "Linux"
+
+    def __init__(self, parent=None):
+        super().__init__(
+            resource_name="Storage Profile Name",
+            setting_name="defaults.storage_profile_id",
+            parent=parent,
+        )
+
+    def get_item(self):
+        farm_id = get_setting("defaults.farm_id")
+        queue_id = get_setting("defaults.queue_id")
+        storage_profile_id = get_setting(self.setting_name)
+
+        if farm_id and queue_id and storage_profile_id:
+            deadline = api.get_boto3_client("deadline")
+            response = deadline.list_storage_profiles_for_queue(farmId=farm_id, queueId=queue_id)
+            farm_storage_profiles = response.get("storageProfiles", {})
+
+            if farm_storage_profiles:
+                storage_profile = [
+                    (item["storageProfileId"], item["displayName"], item["osFamily"])
+                    for item in farm_storage_profiles
+                    if storage_profile_id == item["storageProfileId"]
+                ]
+                return storage_profile[0]
+
+        return ("", "", "")
+
+    def _get_default_storage_profile_name(self) -> str:
+        """
+        Get a string specifying what the OS is, following the format the Deadline storage profile API expects.
+        """
+        if sys.platform.startswith("linux"):
+            return self.LINUX_OS
+
+        if sys.platform.startswith("darwin"):
+            return self.MAC_OS
+
+        if sys.platform.startswith("win"):
+            return self.WINDOWS_OS
+
+        return ""