Added support for mounting object storage service for job.

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,6 +183,48 @@ 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:
 
@@ -204,6 +247,10 @@ def initialize(cls, arguments: dict) -> dict:
                 "Parameter `dest` is required for mounting file storage system."
             )
 
+        # 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):

tests/unitary/default_setup/jobs/test_jobs_mount_file_system.py

@@ -83,7 +83,11 @@
             {
                 "src" : "2.2.2.2:test_export_path_two",
                 "dest" : "test_mount_two",
-            },  
+            }, 
+            {
+                "src" : "oci://bucket_name@namespace/synthetic/",
+                "dest" : "test_mount_three",
+            } 
         )
     )
     .with_runtime(
@@ -111,6 +115,8 @@
         dest: test_mount_one
       - src: 2.2.2.2:test_export_path_two
         dest: test_mount_two
+      - src: oci://bucket_name@namespace/synthetic/
+        dest: test_mount_three
       subnetId: ocid1.subnet.oc1.iad.xxxx
     type: dataScienceJob
   name: My Job
@@ -141,6 +147,11 @@ def test_data_science_job_initialize(self):
         assert dsc_file_storage_two["src"] == "2.2.2.2:test_export_path_two"
         assert dsc_file_storage_two["dest"] == "test_mount_two"
 
+        dsc_object_storage = job.infrastructure.storage_mount[2]
+        assert isinstance(dsc_object_storage, dict)
+        assert dsc_object_storage["src"] == "oci://bucket_name@namespace/synthetic/"
+        assert dsc_object_storage["dest"] == "test_mount_three"
+
     def test_data_science_job_from_yaml(self):
         job_from_yaml = Job.from_yaml(job_yaml_string)
 
@@ -155,6 +166,11 @@ def test_data_science_job_from_yaml(self):
         assert dsc_file_storage_two["src"] == "2.2.2.2:test_export_path_two"
         assert dsc_file_storage_two["dest"] == "test_mount_two"
 
+        dsc_object_storage = job.infrastructure.storage_mount[2]
+        assert isinstance(dsc_object_storage, dict)
+        assert dsc_object_storage["src"] == "oci://bucket_name@namespace/synthetic/"
+        assert dsc_object_storage["dest"] == "test_mount_three"
+
     def test_data_science_job_to_dict(self):
         assert job.to_dict() == {
             "kind": "job",
@@ -189,6 +205,10 @@ def test_data_science_job_to_dict(self):
                                 "src" : "2.2.2.2:test_export_path_two",
                                 "dest" : "test_mount_two",
                             },
+                            {
+                                "src" : "oci://bucket_name@namespace/synthetic/",
+                                "dest" : "test_mount_three",
+                            } 
                         ],
                     },
                 },
@@ -265,7 +285,7 @@ def test_update_job_infra(
 
         assert (
             len(dsc_job_payload_copy.job_storage_mount_configuration_details_list)
-            == 2
+            == 3
         )
         assert dsc_job_payload_copy.job_storage_mount_configuration_details_list[
             0