Added support for mounting file storage for oci job (#139)

lu-ohai · web-flow · commit 3a38dcf5ffa7 · 2023-04-28T17:06:54.000-07:00
diff --git a/ads/common/dsc_file_system.py b/ads/common/dsc_file_system.py
@@ -0,0 +1,219 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2023 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+import ads
+import oci
+import ipaddress
+
+from dataclasses import dataclass
+
+FILE_STORAGE_TYPE = "FILE_STORAGE"
+
+
+@dataclass
+class DSCFileSystem:
+
+    src: str = None
+    dest: str = None
+    storage_type: str = None
+    destination_directory_name: str = None
+
+    def update_to_dsc_model(self) -> dict:
+        """Updates arguments to dsc model.
+
+        Returns
+        -------
+        dict:
+            A dictionary of arguments.
+        """
+        pass
+
+    @classmethod
+    def update_from_dsc_model(cls, dsc_model) -> dict:
+        """Updates arguments from dsc model.
+
+        Parameters
+        ----------
+        dsc_model: oci.data_science.models.JobStorageMountConfigurationDetails
+            An instance of oci.data_science.models.JobStorageMountConfigurationDetails.
+
+        Returns
+        -------
+        dict
+            A dictionary of arguments.
+        """
+        pass
+
+
+@dataclass
+class OCIFileStorage(DSCFileSystem):
+
+    mount_target_id: str = None
+    mount_target: str = None
+    export_id: str = None
+    export_path: str = None
+    storage_type: str = FILE_STORAGE_TYPE        
+
+    def update_to_dsc_model(self) -> dict:
+        """Updates arguments to dsc model.
+
+        Returns
+        -------
+        dict:
+            A dictionary of arguments.
+        """
+        arguments = {
+            "destinationDirectoryName" : self.dest,
+            "storageType" : self.storage_type
+        }
+        
+        self._get_mount_target_and_export_ids(arguments)
+
+        return arguments
+
+    def _get_mount_target_and_export_ids(self, arguments: dict):
+        """Gets the mount target id and export id from src.
+
+        Parameters
+        ----------
+        arguments: dict
+            A dictionary of arguments.
+        """
+        resource_client = oci.resource_search.ResourceSearchClient(**ads.auth.default_signer())
+        src_list = self.src.split(":")
+        first_segment = src_list[0]
+        second_segment = src_list[1]
+
+        if first_segment.startswith("ocid") and "mounttarget" in first_segment:
+            arguments["mountTargetId"] = first_segment
+        else:
+            ip_resource = resource_client.search_resources(
+                search_details=oci.resource_search.models.FreeTextSearchDetails(
+                    text=first_segment,
+                    matching_context_type="NONE"
+                )
+            ).data.items
+
+            ip_resource = sorted(ip_resource, key=lambda resource_summary: resource_summary.time_created)
+
+            if not ip_resource or not hasattr(ip_resource[-1], "identifier"):
+                raise ValueError(f"Can't find the identifier from ip {first_segment}. Specify a valid `src`.")
+
+            mount_target_resource = resource_client.search_resources(
+                search_details=oci.resource_search.models.FreeTextSearchDetails(
+                    text=ip_resource[-1].identifier,
+                    matching_context_type="NONE"
+                )
+            ).data.items
+
+            mount_targets = [
+                mount_target.identifier
+                for mount_target in mount_target_resource
+                if mount_target.resource_type == "MountTarget"
+            ]
+            if len(mount_targets) == 0:
+                raise ValueError(
+                    f"No `mount_target_id` found under ip {first_segment}. Specify a valid `src`."
+                )
+            if len(mount_targets) > 1:
+                raise ValueError(
+                    f"Multiple `mount_target_id` found under ip {first_segment}. Specify the `mount_target_id` in `src` instead."
+                )
+
+            arguments["mountTargetId"] = mount_targets[0]
+        
+        if second_segment.startswith("ocid") and "export" in second_segment:
+            arguments["exportId"] = second_segment
+        else:
+            export_resource = resource_client.search_resources(
+                search_details=oci.resource_search.models.FreeTextSearchDetails(
+                    text=second_segment,
+                    matching_context_type="NONE"
+                )
+            ).data.items
+
+            exports = [
+                export.identifier
+                for export in export_resource
+                if export.resource_type == "Export"
+            ]
+            if len(exports) == 0:
+                raise ValueError(
+                    f"No `export_id` found with `export_path` {second_segment}. Specify a valid `src`."
+                )
+            if len(exports) > 1:
+                raise ValueError(
+                    f"Multiple `export_id` found with `export_path` {second_segment}. Specify the `export_id` in `src` instead."
+                )
+
+            arguments["exportId"] = exports[0]
+
+    @classmethod
+    def update_from_dsc_model(cls, dsc_model) -> dict:
+        """Updates arguments from dsc model.
+
+        Parameters
+        ----------
+        dsc_model: oci.data_science.models.JobStorageMountConfigurationDetails
+            An instance of oci.data_science.models.JobStorageMountConfigurationDetails.
+
+        Returns
+        -------
+        dict
+            A dictionary of arguments.
+        """
+        if not dsc_model.mount_target_id:
+            raise ValueError(
+                "Missing parameter `mount_target_id` from service. Check service log to see the error."
+            )
+        if not dsc_model.export_id:
+            raise ValueError(
+                "Missing parameter `export_id` from service. Check service log to see the error."
+            )
+        if not dsc_model.destination_directory_name:
+            raise ValueError(
+                "Missing parameter `destination_directory_name` from service. Check service log to see the error."
+            )
+
+        return {
+            "src" : f"{dsc_model.mount_target_id}:{dsc_model.export_id}",
+            "dest" : dsc_model.destination_directory_name
+        }
+
+
+class DSCFileSystemManager:
+
+    @classmethod
+    def initialize(cls, arguments: dict) -> dict:
+        """Initialize and update arguments to dsc model.
+
+        Parameters
+        ----------
+        arguments: dict
+            A dictionary of arguments.
+        """
+        if "src" not in arguments:
+            raise ValueError(
+                "Parameter `src` is required for mounting file storage system."
+            )
+
+        if "dest" not in arguments:
+            raise ValueError(
+                "Parameter `dest` is required for mounting file storage system."
+            )
+
+        first_segment = arguments["src"].split(":")[0]
+        # case <mount_target_id>:<export_id>
+        if first_segment.startswith("ocid"):
+            return OCIFileStorage(**arguments).update_to_dsc_model()
+
+        # case <ip_address>:<export_path>
+        try:
+            ipaddress.IPv4Network(first_segment)
+            return OCIFileStorage(**arguments).update_to_dsc_model()
+        except:
+            pass
+
+        raise ValueError("Invalid dict for mounting file systems. Specify a valid one.")
diff --git a/ads/jobs/builders/infrastructure/dsc_job.py b/ads/jobs/builders/infrastructure/dsc_job.py
@@ -7,6 +7,7 @@
 
 import datetime
 import logging
+import oci
 import os
 import time
 import traceback
@@ -34,10 +35,17 @@
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
 from ads.jobs.builders.runtimes.python_runtime import GitPythonRuntime
 
+from ads.common.dsc_file_system import (
+    OCIFileStorage,
+    DSCFileSystemManager
+)
+
 logger = logging.getLogger(__name__)
 
 SLEEP_INTERVAL = 3
 WAIT_SECONDS_AFTER_FINISHED = 90
+MAXIMUM_MOUNT_COUNT = 5
+FILE_STORAGE_TYPE = "FILE_STORAGE"
 
 
 class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
@@ -833,6 +841,13 @@ class DataScienceJob(Infrastructure):
             .with_shape_config_details(memory_in_gbs=16, ocpus=1)
             # Minimum/Default block storage size is 50 (GB).
             .with_block_storage_size(50)
+            # A list of file systems to be mounted
+            .with_storage_mount(
+                {
+                    "src" : "<mount_target_ip_address>:<export_path>",
+	            "dest" : "<destination_directory_name>"
+                }
+            )
         )
 
     """
@@ -850,6 +865,7 @@ class DataScienceJob(Infrastructure):
     CONST_OCPUS = "ocpus"
     CONST_LOG_ID = "logId"
     CONST_LOG_GROUP_ID = "logGroupId"
+    CONST_STORAGE_MOUNT = "storageMount"
 
     attribute_map = {
         CONST_PROJECT_ID: "project_id",
@@ -863,6 +879,7 @@ class DataScienceJob(Infrastructure):
         CONST_SHAPE_CONFIG_DETAILS: "shape_config_details",
         CONST_LOG_ID: "log_id",
         CONST_LOG_GROUP_ID: "log_group_id",
+        CONST_STORAGE_MOUNT: "storage_mount",
     }
 
     shape_config_details_attribute_map = {
@@ -888,6 +905,8 @@ class DataScienceJob(Infrastructure):
         v.split(".", maxsplit=1)[-1]: k for k, v in payload_attribute_map.items()
     }
 
+    storage_mount_type_dict = {FILE_STORAGE_TYPE: OCIFileStorage}
+
     @staticmethod
     def standardize_spec(spec):
         if not spec:
@@ -1212,6 +1231,46 @@ def log_group_id(self) -> str:
         """
         return self.get_spec(self.CONST_LOG_GROUP_ID)
 
+    def with_storage_mount(
+        self, *storage_mount: List[dict]
+    ) -> DataScienceJob:
+        """Sets the file systems to be mounted for the data science job.
+        A maximum number of 5 file systems are allowed to be mounted for a single data science job.
+
+        Parameters
+        ----------
+        storage_mount : List[dict]
+            A list of file systems to be mounted.
+
+        Returns
+        -------
+        DataScienceJob
+            The DataScienceJob instance (self)
+        """
+        storage_mount_list = []
+        for item in storage_mount:
+            if not isinstance(item, dict):
+                raise ValueError(
+                    "Parameter `storage_mount` should be a list of dictionaries."
+                )
+            storage_mount_list.append(item)
+        if len(storage_mount_list) > MAXIMUM_MOUNT_COUNT:
+            raise ValueError(
+                f"A maximum number of {MAXIMUM_MOUNT_COUNT} file systems are allowed to be mounted at this time for a job."
+            )
+        return self.set_spec(self.CONST_STORAGE_MOUNT, storage_mount_list)
+
+    @property
+    def storage_mount(self) -> List[dict]:
+        """Files systems that have been mounted for the data science job
+
+        Returns
+        -------
+        list
+            A list of file systems that have been mounted
+        """
+        return self.get_spec(self.CONST_STORAGE_MOUNT, [])
+
     def _prepare_log_config(self) -> dict:
         if not self.log_group_id and not self.log_id:
             return None
@@ -1285,6 +1344,42 @@ def _update_from_dsc_model(
                         sub_spec[sub_infra_attr] = sub_value
                 if sub_spec:
                     self._spec[infra_attr] = sub_spec
+
+        self._update_storage_mount_from_dsc_model(dsc_job, overwrite)
+        return self
+
+    def _update_storage_mount_from_dsc_model(
+        self, dsc_job: oci.data_science.models.Job, overwrite: bool = True
+    ) -> DataScienceJob:
+        """Update the mount storage properties from an OCI data science job model.
+
+        Parameters
+        ----------
+        dsc_job: oci.data_science.models.Job
+            An OCI data science job model.
+
+        overwrite: bool
+            Whether to overwrite the existing values.
+            If this is set to False, only the empty/None properties will be updated.
+
+        Returns
+        -------
+        DataScienceJob
+            The DataScienceJob instance (self)
+        """
+        storage_mount_list = get_value(
+            dsc_job, "job_storage_mount_configuration_details_list"
+        )
+        if storage_mount_list:
+            storage_mount = [
+                self.storage_mount_type_dict[
+                    file_system.storage_type
+                ].update_from_dsc_model(file_system)
+                for file_system in storage_mount_list
+                if file_system.storage_type in self.storage_mount_type_dict
+            ]
+            if overwrite or not self.get_spec(self.CONST_STORAGE_MOUNT):
+                self.set_spec(self.CONST_STORAGE_MOUNT, storage_mount)
         return self
 
     def _update_job_infra(self, dsc_job: DSCJob) -> DataScienceJob:
@@ -1321,6 +1416,17 @@ def _update_job_infra(self, dsc_job: DSCJob) -> DataScienceJob:
             dsc_job.job_infrastructure_configuration_details[
                 "jobInfrastructureType"
             ] = JobInfrastructureConfigurationDetails.JOB_INFRASTRUCTURE_TYPE_STANDALONE
+
+        if self.storage_mount:
+            if not hasattr(
+                oci.data_science.models, "JobStorageMountConfigurationDetails"
+            ):
+                raise EnvironmentError(
+                    "Storage mount hasn't been supported in the current OCI SDK installed."
+                )
+            dsc_job.job_storage_mount_configuration_details_list = [
+                DSCFileSystemManager.initialize(file_system) for file_system in self.storage_mount
+            ]
         return self
 
     def build(self) -> DataScienceJob:
diff --git a/tests/unitary/default_setup/jobs/test_jobs_mount_file_system.py b/tests/unitary/default_setup/jobs/test_jobs_mount_file_system.py
