Adding PyTorchDistributedRuntime to support distributed (multi-node) training (#209)

Author: qiuosier

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
 

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,

ads/jobs/builders/runtimes/artifact.py

@@ -8,7 +8,6 @@
 import os
 import shutil
 import tempfile
-import zipfile
 from io import DEFAULT_BUFFER_SIZE
 from urllib import request
 from urllib.parse import urlparse
@@ -41,7 +40,7 @@ class Artifact:
 
     def __init__(self, source, runtime=None) -> None:
         # Get the full path of source file if it is local file.
-        if not urlparse(source).scheme:
+        if source and not urlparse(source).scheme:
             self.source = os.path.abspath(os.path.expanduser(source))
         else:
             self.source = source
@@ -203,6 +202,7 @@ class PythonArtifact(Artifact):
     """Represents a PythonRuntime job artifact"""
 
     CONST_DRIVER_SCRIPT = "driver_python.py"
+    DEFAULT_BASENAME = "artifact"
     # The directory to store user code
     # This directory must match the USER_CODE_DIR in driver_python.py
     USER_CODE_DIR = "code"
@@ -217,7 +217,12 @@ def _copy_artifacts(self, drivers=None):
         """Copies the drivers and artifacts to the temp artifact dir."""
         # The basename of the job artifact,
         # this will be the name of the zip file uploading to OCI
-        self.basename = os.path.basename(str(self.source).rstrip("/")).split(".", 1)[0]
+        if self.source:
+            self.basename = os.path.basename(str(self.source).rstrip("/")).split(
+                ".", 1
+            )[0]
+        else:
+            self.basename = self.DEFAULT_BASENAME
         # The temp dir path for storing the artifacts, including drivers and user code
         self.artifact_dir = os.path.join(self.temp_dir.name, self.basename)
         # The temp dir path for storing the user code
@@ -236,8 +241,9 @@ def _copy_artifacts(self, drivers=None):
             shutil.copy(file_path, os.path.join(self.artifact_dir, filename))
 
         # Copy user code
-        os.makedirs(self.code_dir, exist_ok=True)
-        Artifact.copy_from_uri(self.source, self.code_dir, unpack=True)
+        if self.source:
+            os.makedirs(self.code_dir, exist_ok=True)
+            Artifact.copy_from_uri(self.source, self.code_dir, unpack=True)
 
     def _zip_artifacts(self):
         """Create a zip file from the temp artifact dir."""
@@ -312,7 +318,6 @@ def build(self):
 
 
 class GitPythonArtifact(Artifact):
-
     CONST_DRIVER_SCRIPT = "driver_oci.py"
 
     def __init__(self) -> None:

ads/jobs/builders/runtimes/python_runtime.py

@@ -149,7 +149,7 @@ class ScriptRuntime(CondaRuntime):
             # Environment variable
             .with_environment_variable(NAME="Welcome to OCI Data Science.")
             # Command line argument
-            .with_argument("100 linux \"hi there\"")
+            .with_argument("100 linux 'hi there'")
             # The entrypoint is applicable only to directory or zip file as source
             # The entrypoint should be a path relative to the working dir.
             # Here my_script.sh is a file in the code_dir/my_package directory