Added support for mounting object storage service for job. (#204)

Author: lu-ohai

ads/common/dsc_file_system.py

@@ -10,6 +10,7 @@
 from dataclasses import dataclass
 
 FILE_STORAGE_TYPE = "FILE_STORAGE"
+OBJECT_STORAGE_TYPE = "OBJECT_STORAGE"
 
 
 @dataclass
@@ -182,9 +183,53 @@ def update_from_dsc_model(cls, dsc_model) -> dict:
             "dest" : dsc_model.destination_directory_name
         }
 
+@dataclass
+class OCIObjectStorage(DSCFileSystem):
+
+    storage_type: str = OBJECT_STORAGE_TYPE
+
+    def update_to_dsc_model(self) -> dict:
+        arguments = {
+            "destinationDirectoryName" : self.dest,
+            "storageType" : self.storage_type
+        }
+        src_list = self.src.split("@")
+        bucket_segment = src_list[0]
+        namespace_segment = src_list[1].strip("/")
+        arguments["bucket"] = bucket_segment[6:]
+        if "/" in namespace_segment:
+            first_slash_index = namespace_segment.index("/")
+            arguments["namespace"] = namespace_segment[:first_slash_index]
+            arguments["prefix"] = namespace_segment[first_slash_index+1:]
+        else:
+            arguments["namespace"] = namespace_segment
+        return arguments
+
+    @classmethod
+    def update_from_dsc_model(cls, dsc_model) -> dict:
+        if not dsc_model.namespace:
+            raise ValueError(
+                "Missing parameter `namespace` from service. Check service log to see the error."
+            )
+        if not dsc_model.bucket:
+            raise ValueError(
+                "Missing parameter `bucket` 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"oci://{dsc_model.bucket}@{dsc_model.namespace}/{dsc_model.prefix or ''}",
+            "dest" : dsc_model.destination_directory_name
+        }
+
 
 class DSCFileSystemManager:
 
+    storage_mount_dest = set()
+
     @classmethod
     def initialize(cls, arguments: dict) -> dict:
         """Initialize and update arguments to dsc model.
@@ -204,6 +249,16 @@ def initialize(cls, arguments: dict) -> dict:
                 "Parameter `dest` is required for mounting file storage system."
             )
 
+        if arguments["dest"] in cls.storage_mount_dest:
+            raise ValueError(
+                "Duplicate `dest` found. Please specify different `dest` for each file system to be mounted."
+            )
+        cls.storage_mount_dest.add(arguments["dest"])
+
+        # case oci://bucket@namespace/prefix
+        if arguments["src"].startswith("oci://") and "@" in arguments["src"]:
+            return OCIObjectStorage(**arguments).update_to_dsc_model()
+
         first_segment = arguments["src"].split(":")[0]
         # case <mount_target_id>:<export_id>
         if first_segment.startswith("ocid"):

ads/jobs/builders/infrastructure/dsc_job.py

@@ -35,14 +35,15 @@
 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
+from ads.common.dsc_file_system import OCIFileStorage, DSCFileSystemManager, OCIObjectStorage
 
 logger = logging.getLogger(__name__)
 
 SLEEP_INTERVAL = 3
 WAIT_SECONDS_AFTER_FINISHED = 90
 MAXIMUM_MOUNT_COUNT = 5
 FILE_STORAGE_TYPE = "FILE_STORAGE"
+OBJECT_STORAGE_TYPE = "OBJECT_STORAGE"
 
 
 class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
@@ -911,7 +912,10 @@ 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}
+    storage_mount_type_dict = {
+        FILE_STORAGE_TYPE: OCIFileStorage,
+        OBJECT_STORAGE_TYPE: OCIObjectStorage,
+    }
 
     @staticmethod
     def standardize_spec(spec):