Odsc 44980: fix unittests (#259)

Author: KshitizLohia

.gitleaks.toml

@@ -9,11 +9,13 @@ useDefault = true
 # Paths listed in allowlist will not be scanned.
 [allowlist]
     description = "Global allow list"
-    stopwords = ["test_password", "sample_key"]
     regexes = [
         '''example-password''',
         '''this-is-not-the-secret''',
-        '''<redacted>'''
+        '''<redacted>''',
+        # NVIDIA_GPGKEY_SUM from public documentation:
+        # https://gitlab.com/nvidia/container-images/cuda/-/blob/master/dist/10.1/centos7/base/Dockerfile
+        '''d0664fbbdb8c32356d45de36c5984617217b2d0bef41b93ccecd326ba3b80c87'''
     ]
     paths = [
         '''tests/integration/tests_configs.yaml'''

.pre-commit-config.yaml

@@ -40,7 +40,7 @@ repos:
     -   id: gitleaks
 # Oracle copyright checker
 -   repo: https://github.com/oracle-samples/oci-data-science-ai-samples/
-    rev: cbe0136
+    rev: cbe0136f7aaffe463b31ddf3f34b0e16b4b124ff
     hooks:
     -   id: check-copyright
         name: check-copyright

ads/common/oci_client.py

@@ -17,8 +17,6 @@
 from oci.resource_search import ResourceSearchClient
 from oci.secrets import SecretsClient
 from oci.vault import VaultsClient
-from oci.feature_store import FeatureStoreClient
-
 logger = logging.getLogger(__name__)
 
 
@@ -65,10 +63,15 @@ def _client_impl(self, client):
             "ai_language": AIServiceLanguageClient,
             "data_labeling_dp": DataLabelingClient,
             "data_labeling_cp": DataLabelingManagementClient,
-            "feature_store": FeatureStoreClient,
             "resource_search": ResourceSearchClient,
             "data_catalog": DataCatalogClient
         }
+        try:
+            from oci.feature_store import FeatureStoreClient
+            client_map["feature_store"] = FeatureStoreClient
+        except ImportError:
+            logger.warning("OCI SDK with feature store support is not installed")
+            pass
 
         assert (
             client in client_map

ads/common/oci_mixin.py

@@ -230,7 +230,7 @@ def _parse_kwargs(attribute_map: dict, **kwargs):
 
         return parsed_kwargs
 
-    @classmethod
+    @class_or_instance_method
     def deserialize(cls, data, to_cls):
         """De-serialize data from dictionary to an OCI model"""
         if cls.type_mappings is None:
@@ -549,7 +549,7 @@ def from_dict(cls, data):
         """
         return cls.create_instance(**data)
 
-    @classmethod
+    @class_or_instance_method
     def deserialize(cls, data: dict, to_cls: str = None):
         """Deserialize data
 
@@ -726,7 +726,7 @@ def update_from_oci_model(
         for attr in self.swagger_types.keys():
             if (
                 hasattr(oci_model_instance, attr)
-                and getattr(oci_model_instance, attr) is not None
+                and getattr(oci_model_instance, attr)
                 and (
                     not hasattr(self, attr)
                     or not getattr(self, attr)

ads/common/serializer.py

@@ -15,6 +15,7 @@
 import yaml
 
 from ads.common import logger
+from ads.common.auth import default_signer
 
 try:
     from yaml import CSafeDumper as dumper
@@ -134,6 +135,14 @@ def _read_from_file(uri: str, **kwargs) -> str:
         -------
         string: Contents in file specified by URI
         """
+        # Add default signer if the uri is an object storage uri, and
+        # the user does not specify config or signer.
+        if (
+            uri.startswith("oci://")
+            and "config" not in kwargs
+            and "signer" not in kwargs
+        ):
+            kwargs.update(default_signer())
         with fsspec.open(uri, "r", **kwargs) as f:
             return f.read()
 

ads/jobs/__init__.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2022 Oracle and/or its affiliates.
+# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 
@@ -17,6 +17,7 @@
         DataFlowRuntime,
         DataFlowNotebookRuntime,
     )
+    from ads.jobs.builders.runtimes.pytorch_runtime import PyTorchDistributedRuntime
     from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
     from ads.jobs.ads_job import Job
     from ads.jobs.builders import infrastructure
@@ -44,6 +45,7 @@
     "NotebookRuntime",
     "ScriptRuntime",
     "ContainerRuntime",
+    "PyTorchDistributedRuntime",
     "DataFlow",
     "DataFlowRun",
     "DataFlowRuntime",

ads/jobs/ads_job.py

@@ -11,6 +11,7 @@
 from ads.jobs.builders.base import Builder
 from ads.jobs.builders.infrastructure.dataflow import DataFlow, DataFlowRun
 from ads.jobs.builders.infrastructure.dsc_job import DataScienceJob, DataScienceJobRun
+from ads.jobs.builders.runtimes.pytorch_runtime import PyTorchDistributedRuntime
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
 from ads.jobs.builders.runtimes.python_runtime import (
     DataFlowRuntime,
@@ -131,6 +132,7 @@ class Job(Builder):
             ContainerRuntime,
             ScriptRuntime,
             NotebookRuntime,
+            PyTorchDistributedRuntime,
             DataFlowRuntime,
         ]
     }

ads/jobs/builders/infrastructure/dsc_job.py

@@ -35,7 +35,11 @@
 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, OCIObjectStorage
+from ads.common.dsc_file_system import (
+    OCIFileStorage,
+    DSCFileSystemManager,
+    OCIObjectStorage,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -1454,11 +1458,14 @@ def _update_job_infra(self, dsc_job: DSCJob) -> DataScienceJob:
             if value:
                 dsc_job.job_infrastructure_configuration_details[camel_attr] = value
 
-        if (
-            not dsc_job.job_infrastructure_configuration_details.get("shapeName", "").endswith("Flex")
-            and dsc_job.job_infrastructure_configuration_details.get("jobShapeConfigDetails")
+        if not dsc_job.job_infrastructure_configuration_details.get(
+            "shapeName", ""
+        ).endswith("Flex") and dsc_job.job_infrastructure_configuration_details.get(
+            "jobShapeConfigDetails"
         ):
-            raise ValueError("Shape config is not required for non flex shape from user end.")
+            raise ValueError(
+                "Shape config is not required for non flex shape from user end."
+            )
 
         if dsc_job.job_infrastructure_configuration_details.get("subnetId"):
             dsc_job.job_infrastructure_configuration_details[
@@ -1495,7 +1502,10 @@ def init(self) -> DataScienceJob:
             self.build()
             .with_compartment_id(self.compartment_id or "{Provide a compartment OCID}")
             .with_project_id(self.project_id or "{Provide a project OCID}")
-            .with_subnet_id(self.subnet_id or "{Provide a subnet OCID or remove this field if you use a default networking}")
+            .with_subnet_id(
+                self.subnet_id
+                or "{Provide a subnet OCID or remove this field if you use a default networking}"
+            )
         )
 
     def create(self, runtime, **kwargs) -> DataScienceJob:
@@ -1552,7 +1562,7 @@ def run(
         freeform_tags=None,
         defined_tags=None,
         wait=False,
-        **kwargs
+        **kwargs,
     ) -> DataScienceJobRun:
         """Runs a job on OCI Data Science job
 
@@ -1603,15 +1613,21 @@ def run(
                 envs.update(env_var)
             name = Template(name).safe_substitute(envs)
 
-        return self.dsc_job.run(
+        kwargs = dict(
             display_name=name,
             command_line_arguments=args,
             environment_variables=env_var,
             freeform_tags=freeform_tags,
             defined_tags=defined_tags,
             wait=wait,
-            **kwargs
+            **kwargs,
         )
+        # A Runtime class may define customized run() method.
+        # Use the customized method if the run() method is defined by the runtime.
+        # Otherwise, use the default run() method defined in this class.
+        if hasattr(self.runtime, "run"):
+            return self.runtime.run(self.dsc_job, **kwargs)
+        return self.dsc_job.run(**kwargs)
 
     def delete(self) -> None:
         """Deletes a job"""

ads/jobs/builders/infrastructure/dsc_job_runtime.py

@@ -29,13 +29,19 @@
     GitPythonRuntime,
 )
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
+from ads.jobs.builders.runtimes.pytorch_runtime import (
+    PyTorchDistributedRuntime,
+    PyTorchDistributedArtifact,
+)
 from ads.jobs.builders.runtimes.artifact import (
     ScriptArtifact,
     NotebookArtifact,
     PythonArtifact,
     GitPythonArtifact,
 )
+from ads.opctl.distributed.common import cluster_config_helper
 from ads.jobs.builders.infrastructure.utils import get_value
+from ads.jobs.templates import driver_utils
 
 
 class IncompatibleRuntime(Exception):  # pragma: no cover
@@ -184,7 +190,7 @@ def _translate_config(self, runtime: Runtime) -> dict:
         if runtime.args:
             # shlex.join() is not available until python 3.8
             job_configuration_details["command_line_arguments"] = " ".join(
-                shlex.quote(arg) for arg in runtime.get_spec(runtime.CONST_ARGS)
+                shlex.quote(str(arg)) for arg in runtime.get_spec(runtime.CONST_ARGS)
             )
         return job_configuration_details
 
@@ -653,7 +659,7 @@ def _translate_env(self, runtime: PythonRuntime) -> dict:
 
         if runtime.entrypoint:
             envs[self.CONST_CODE_ENTRYPOINT] = runtime.entrypoint
-        else:
+        elif runtime.script_uri:
             envs[self.CONST_CODE_ENTRYPOINT] = os.path.basename(runtime.script_uri)
 
         envs[self.CONST_JOB_ENTRYPOINT] = PythonArtifact.CONST_DRIVER_SCRIPT
@@ -674,9 +680,13 @@ def _extract_envs(self, dsc_job) -> dict:
         """
         spec = super()._extract_envs(dsc_job)
         envs = spec.pop(PythonRuntime.CONST_ENV_VAR, {})
-        if self.CONST_CODE_ENTRYPOINT not in envs:
+        if (
+            self.__class__ == PythonRuntimeHandler
+            and self.CONST_CODE_ENTRYPOINT not in envs
+        ):
             raise IncompatibleRuntime()
-        envs.pop(PythonRuntimeHandler.CONST_JOB_ENTRYPOINT)
+        # PyTorchDistributedRuntime does not require entrypoint.
+        envs.pop(PythonRuntimeHandler.CONST_JOB_ENTRYPOINT, None)
         spec.update(self._extract_specs(envs, self.SPEC_MAPPINGS))
         if PythonRuntime.CONST_PYTHON_PATH in spec:
             spec[PythonRuntime.CONST_PYTHON_PATH] = spec[
@@ -1035,6 +1045,98 @@ def _extract_envs(self, dsc_job):
         return spec
 
 
+class PyTorchDistributedRuntimeHandler(PythonRuntimeHandler):
+    RUNTIME_CLASS = PyTorchDistributedRuntime
+    CONST_WORKER_COUNT = "OCI__WORKER_COUNT"
+    CONST_COMMAND = "OCI__LAUNCH_CMD"
+    CONST_DEEPSPEED = "OCI__DEEPSPEED"
+
+    GIT_SPEC_MAPPINGS = {
+        cluster_config_helper.OCI__RUNTIME_URI: GitPythonRuntime.CONST_GIT_URL,
+        cluster_config_helper.OCI__RUNTIME_GIT_BRANCH: GitPythonRuntime.CONST_BRANCH,
+        cluster_config_helper.OCI__RUNTIME_GIT_COMMIT: GitPythonRuntime.CONST_COMMIT,
+        cluster_config_helper.OCI__RUNTIME_GIT_SECRET_ID: GitPythonRuntime.CONST_GIT_SSH_SECRET_ID,
+    }
+
+    SPEC_MAPPINGS = PythonRuntimeHandler.SPEC_MAPPINGS
+    SPEC_MAPPINGS.update(
+        {
+            PyTorchDistributedRuntime.CONST_COMMAND: CONST_COMMAND,
+        }
+    )
+
+    def _translate_artifact(self, runtime: PyTorchDistributedRuntime):
+        return PyTorchDistributedArtifact(runtime.source_uri, runtime)
+
+    def _translate_env(self, runtime: PyTorchDistributedRuntime) -> dict:
+        envs = super()._translate_env(runtime)
+        replica = runtime.replica if runtime.replica else 1
+        # WORKER_COUNT = REPLICA - 1 so that it will be same as distributed training
+        envs[self.CONST_WORKER_COUNT] = str(replica - 1)
+        envs[self.CONST_JOB_ENTRYPOINT] = PyTorchDistributedArtifact.CONST_DRIVER_SCRIPT
+        if runtime.inputs:
+            envs[driver_utils.CONST_ENV_INPUT_MAPPINGS] = json.dumps(runtime.inputs)
+        if runtime.git:
+            for env_key, spec_key in self.GIT_SPEC_MAPPINGS.items():
+                if not runtime.git.get(spec_key):
+                    continue
+                envs[env_key] = runtime.git[spec_key]
+        if runtime.dependencies:
+            if PyTorchDistributedRuntime.CONST_PIP_PKG in runtime.dependencies:
+                envs[driver_utils.CONST_ENV_PIP_PKG] = runtime.dependencies[
+                    PyTorchDistributedRuntime.CONST_PIP_PKG
+                ]
+            if PyTorchDistributedRuntime.CONST_PIP_REQ in runtime.dependencies:
+                envs[driver_utils.CONST_ENV_PIP_REQ] = runtime.dependencies[
+                    PyTorchDistributedRuntime.CONST_PIP_REQ
+                ]
+        if runtime.use_deepspeed:
+            envs[self.CONST_DEEPSPEED] = "1"
+        return envs
+
+    def _extract_envs(self, dsc_job) -> dict:
+        spec = super()._extract_envs(dsc_job)
+        envs = spec.pop(PythonRuntime.CONST_ENV_VAR, {})
+        if self.CONST_WORKER_COUNT not in envs:
+            raise IncompatibleRuntime()
+        # Replicas
+        spec[PyTorchDistributedRuntime.CONST_REPLICA] = (
+            int(envs.pop(self.CONST_WORKER_COUNT)) + 1
+        )
+        # Git
+        if cluster_config_helper.OCI__RUNTIME_URI in envs:
+            git_spec = {}
+            for env_key, spec_key in self.GIT_SPEC_MAPPINGS.items():
+                if env_key in envs:
+                    git_spec[spec_key] = envs.pop(env_key)
+            spec[PyTorchDistributedRuntime.CONST_GIT] = git_spec
+        # Inputs
+        input_mappings = envs.pop(driver_utils.CONST_ENV_INPUT_MAPPINGS, None)
+        if input_mappings:
+            try:
+                spec[PyTorchDistributedRuntime.CONST_INPUT] = json.loads(input_mappings)
+            except ValueError:
+                spec[PyTorchDistributedRuntime.CONST_INPUT] = input_mappings
+        # Dependencies
+        dep = {}
+        if driver_utils.CONST_ENV_PIP_PKG in envs:
+            dep[PyTorchDistributedRuntime.CONST_PIP_PKG] = envs.pop(
+                driver_utils.CONST_ENV_PIP_PKG
+            )
+        if driver_utils.CONST_ENV_PIP_REQ in envs:
+            dep[PyTorchDistributedRuntime.CONST_PIP_REQ] = envs.pop(
+                driver_utils.CONST_ENV_PIP_REQ
+            )
+        if dep:
+            spec[PyTorchDistributedRuntime.CONST_DEP] = dep
+        if envs.pop(self.CONST_DEEPSPEED, None):
+            spec[PyTorchDistributedRuntime.CONST_DEEPSPEED] = True
+        # Envs
+        if envs:
+            spec[PythonRuntime.CONST_ENV_VAR] = envs
+        return spec
+
+
 class DataScienceJobRuntimeManager(RuntimeHandler):
     """This class is used by the DataScienceJob infrastructure to handle the runtime conversion.
     The translate() method determines the actual runtime handler by matching the RUNTIME_CLASS.
@@ -1046,6 +1148,7 @@ class DataScienceJobRuntimeManager(RuntimeHandler):
 
     runtime_handlers = [
         ContainerRuntimeHandler,
+        PyTorchDistributedRuntimeHandler,
         GitPythonRuntimeHandler,
         NotebookRuntimeHandler,
         PythonRuntimeHandler,