oracle-ads 2.10.0__py3-none-any.whl → 2.11.0__py3-none-any.whl

ads/jobs/builders/runtimes/pytorch_runtime.py

@@ -1,16 +1,20 @@
-import json
-import time
-import traceback
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2023, 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
 from ads.jobs.builders.runtimes.artifact import PythonArtifact, GitPythonArtifact
+from ads.jobs.builders.runtimes.base import MultiNodeRuntime
 from ads.jobs.builders.runtimes.python_runtime import (
     PythonRuntime,
     GitPythonRuntime,
 )
 
 
-class PyTorchDistributedRuntime(PythonRuntime):
+class PyTorchDistributedRuntime(PythonRuntime, MultiNodeRuntime):
+    """Represents runtime supporting PyTorch Distributed training."""
     CONST_GIT = "git"
-    CONST_REPLICA = "replicas"
     CONST_INPUT = "inputs"
     CONST_DEP = "dependencies"
     CONST_PIP_REQ = "pipRequirements"
@@ -79,26 +83,6 @@ class PyTorchDistributedRuntime(PythonRuntime):
         """The input files to be copied into the job run."""
         return self.get_spec(self.CONST_INPUT)
 
-    def with_replica(self, count: int):
-        """Specifies the number of nodes (job runs) for the job.
-
-        Parameters
-        ----------
-        count : int
-            Number of nodes (job runs)
-
-        Returns
-        -------
-        self
-            The runtime instance.
-        """
-        return self.set_spec(self.CONST_REPLICA, count)
-
-    @property
-    def replica(self) -> int:
-        """The number of nodes (job runs)."""
-        return self.get_spec(self.CONST_REPLICA)
-
     def with_dependency(self, pip_req=None, pip_pkg=None):
         """Specifies additional dependencies to be installed using pip.
 
@@ -185,7 +169,7 @@ class PyTorchDistributedRuntime(PythonRuntime):
     def command(self):
         """The command for launching the workload."""
         return self.get_spec(self.CONST_COMMAND)
-
+    
     @property
     def use_deepspeed(self):
         """Indicate whether whether to configure deepspeed for multi-node workload"""
@@ -193,43 +177,6 @@ class PyTorchDistributedRuntime(PythonRuntime):
             return True
         return False
 
-    def run(self, dsc_job, **kwargs):
-        """Starts the job runs"""
-        replicas = self.replica if self.replica else 1
-        main_run = None
-        job_runs = []
-        try:
-            for i in range(replicas):
-                replica_kwargs = kwargs.copy()
-                envs = replica_kwargs.get("environment_variables")
-                if not envs:
-                    envs = {}
-                # Huggingface accelerate requires machine rank
-                # Here we use NODE_RANK to store the machine rank
-                envs["NODE_RANK"] = str(i)
-                envs["WORLD_SIZE"] = str(replicas)
-                if main_run:
-                    envs["MAIN_JOB_RUN_OCID"] = main_run.id
-                name = replica_kwargs.get("display_name")
-                if not name:
-                    name = dsc_job.display_name
-
-                replica_kwargs["display_name"] = f"{name}-{str(i)}"
-                replica_kwargs["environment_variables"] = envs
-                run = dsc_job.run(**replica_kwargs)
-                job_runs.append(run)
-                if i == 0:
-                    main_run = run
-        except Exception:
-            traceback.print_exc()
-            # Wait a few second to avoid the job run being in a transient state.
-            time.sleep(2)
-            # If there is any error when creating the job runs
-            # cancel all the job runs.
-            for run in job_runs:
-                run.cancel()
-        return main_run
-
 
 class PyTorchDistributedArtifact(PythonArtifact):
     CONST_DRIVER_SCRIPT = "driver_pytorch.py"

ads/jobs/templates/driver_pytorch.py

@@ -5,6 +5,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 """This module requires oracle-ads>=2.6.8
 """
+import getpass
 import ipaddress
 import logging
 import multiprocessing
@@ -40,17 +41,32 @@ logger = logging.getLogger(__name__)
 logger = driver_utils.set_log_level(logger)
 
 
+# Envs provisioned by the service
 CONST_ENV_HOST_JOB_RUN_OCID = "MAIN_JOB_RUN_OCID"
 CONST_ENV_JOB_RUN_OCID = "JOB_RUN_OCID"
-CONST_ENV_LD_PRELOAD = "LD_PRELOAD"
+# Envs set by the ADS API
+OCI__WORKER_COUNT = "OCI__WORKER_COUNT"
+CONST_ENV_NODE_RANK = "NODE_RANK"
+CONST_ENV_NODE_COUNT = "NODE_COUNT"
 CONST_ENV_LAUNCH_CMD = "OCI__LAUNCH_CMD"
 CONST_ENV_DEEPSPEED = "OCI__DEEPSPEED"
+# Envs set by this module
+CONST_ENV_WORLD_SIZE = "WORLD_SIZE"
+CONST_ENV_LD_PRELOAD = "LD_PRELOAD"
+# Envs for debugging only
+# OCI_ODSC_SERVICE_ENDPOINT is used for all processes in the job run
+CONST_ENV_ODSC_SERVICE_ENDPOINT = "OCI_ODSC_SERVICE_ENDPOINT"
+# OCI_DS_SERVICE_ENDPOINT is used only by the training process
+CONST_ENV_DS_SERVICE_ENDPOINT = "OCI_DS_SERVICE_ENDPOINT"
+
+# Constants used in logs
 LOG_PREFIX_HOST_IP = "Distributed Training HOST IP: "
 LOG_PREFIX_NODE_IP = "Node IP: "
 LOG_PREFIX_PUBLIC_KEY = "HOST PUBLIC KEY: "
-SSH_DIR = "/home/datascience/.ssh"
-# Working count is the number of node - 1
-OCI__WORKER_COUNT = "OCI__WORKER_COUNT"
+# Other constants used within this script
+# Other constants used within this script
+USER_HOME = os.environ.get("HOME", f"/home/{getpass.getuser()}")
+SSH_DIR = os.environ.get("OCI__SSH_DIR", os.path.join(USER_HOME, ".ssh"))
 DEFAULT_LAUNCHER = "torchrun"
 
 # Set authentication method to resource principal
@@ -131,8 +147,11 @@ class Runner(driver_utils.JobRunner):
 
         self.host_job_run = DataScienceJobRun.from_ocid(self.host_ocid)
         self.entrypoint_env = PythonRuntimeHandler.CONST_CODE_ENTRYPOINT
-        # The total number of node is OCI__WORKER_COUNT + 1
-        self.node_count = int(os.environ.get(OCI__WORKER_COUNT, 0)) + 1
+        # The total number of nodes is OCI__WORKER_COUNT + 1
+        if CONST_ENV_NODE_COUNT in os.environ:
+            self.node_count = int(os.environ[CONST_ENV_NODE_COUNT])
+        else:
+            self.node_count = int(os.environ.get(OCI__WORKER_COUNT, 0)) + 1
         logger.debug("Node count: %s", self.node_count)
         self.gpu_count = torch.cuda.device_count()
         logger.debug("GPU count on this node: %s", self.gpu_count)
@@ -343,9 +362,7 @@ class Runner(driver_utils.JobRunner):
         if self.launch_cmd:
             if self.LAUNCHER:
                 if not self.launch_cmd.startswith(self.LAUNCHER):
-                    raise ValueError(
-                        f"Command not supported: '{self.launch_cmd}'. "
-                    )
+                    raise ValueError(f"Command not supported: '{self.launch_cmd}'. ")
 
                 launch_args.append(self.launch_cmd[len(self.LAUNCHER) + 1 :])
             else:
@@ -689,7 +706,7 @@ class GenericRunner(TorchRunner, DeepSpeedRunner):
     def set_env_var(self):
         """Set default environment variables."""
         defaults = {
-            "WORLD_SIZE": self.node_count,
+            "WORLD_SIZE": self.node_count * self.gpu_count,
             "MASTER_ADDR": self.host_ip,
             "MASTER_PORT": self.RDZV_PORT,
         }

ads/model/artifact_downloader.py

@@ -1,12 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2022, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import os
 import shutil
-import tempfile
 import uuid
 from abc import ABC, abstractmethod
 from typing import Dict, Optional
@@ -15,6 +14,7 @@ from zipfile import ZipFile
 from ads.common import utils
 from ads.common.utils import extract_region
 from ads.model.service.oci_datascience_model import OCIDataScienceModel
+from ads.common.object_storage_details import ObjectStorageDetails
 
 
 class ArtifactDownloader(ABC):
@@ -85,17 +85,32 @@ class SmallArtifactDownloader(ArtifactDownloader):
     def _download(self):
         """Downloads model artifacts."""
         self.progress.update("Importing model artifacts from catalog")
-        zip_content = self.dsc_model.get_model_artifact_content()
+
+        artifact_info = self.dsc_model.get_artifact_info()
+        artifact_name = artifact_info["Content-Disposition"].replace(
+            "attachment; filename=", ""
+        )
+        _, file_extension = os.path.splitext(artifact_name)
+        file_extension = file_extension.lower() if file_extension else ".zip"
+
+        file_content = self.dsc_model.get_model_artifact_content()
         self.progress.update("Copying model artifacts to the artifact directory")
-        
-        zip_file_path = os.path.join(self.target_dir, f"{str(uuid.uuid4())}.zip")
-        with open(zip_file_path, "wb") as zip_file:
-            zip_file.write(zip_content)
-        self.progress.update("Extracting model artifacts")
-        with ZipFile(zip_file_path) as zip_file:
-            zip_file.extractall(self.target_dir)
-        
-        utils.remove_file(zip_file_path)
+
+        file_name = (
+            "model_description" if file_extension == ".json" else str(uuid.uuid4())
+        )
+        artifact_file_path = os.path.join(
+            self.target_dir, f"{file_name}{file_extension}"
+        )
+        with open(artifact_file_path, "wb") as _file:
+            _file.write(file_content)
+
+        if file_extension == ".zip":
+            self.progress.update("Extracting model artifacts")
+            with ZipFile(artifact_file_path) as _file:
+                _file.extractall(self.target_dir)
+            utils.remove_file(artifact_file_path)
+
 
 class LargeArtifactDownloader(ArtifactDownloader):
     PROGRESS_STEPS_COUNT = 4
@@ -110,6 +125,7 @@ class LargeArtifactDownloader(ArtifactDownloader):
         bucket_uri: Optional[str] = None,
         overwrite_existing_artifact: Optional[bool] = True,
         remove_existing_artifact: Optional[bool] = True,
+        model_file_description: Optional[dict] = None,
     ):
         """Initializes `LargeArtifactDownloader` instance.
 
@@ -137,6 +153,9 @@ class LargeArtifactDownloader(ArtifactDownloader):
             Overwrite target bucket artifact if exists.
         remove_existing_artifact: (bool, optional). Defaults to `True`.
             Wether artifacts uploaded to object storage bucket need to be removed or not.
+        model_file_description: (dict, optional). Defaults to None.
+            Contains object path details for models created by reference.
+
         """
         super().__init__(
             dsc_model=dsc_model, target_dir=target_dir, force_overwrite=force_overwrite
@@ -146,12 +165,19 @@ class LargeArtifactDownloader(ArtifactDownloader):
         self.bucket_uri = bucket_uri
         self.overwrite_existing_artifact = overwrite_existing_artifact
         self.remove_existing_artifact = remove_existing_artifact
+        self.model_file_description = model_file_description
 
     def _download(self):
         """Downloads model artifacts."""
         self.progress.update(f"Importing model artifacts from catalog")
 
+        if self.dsc_model.is_model_by_reference() and self.model_file_description:
+            self.download_from_model_file_description()
+            self.progress.update()
+            return
+
         bucket_uri = self.bucket_uri
+
         if not os.path.basename(bucket_uri):
             bucket_uri = os.path.join(bucket_uri, f"{self.dsc_model.id}.zip")
         elif not bucket_uri.lower().endswith(".zip"):
@@ -159,7 +185,6 @@ class LargeArtifactDownloader(ArtifactDownloader):
 
         self.dsc_model.import_model_artifact(bucket_uri=bucket_uri, region=self.region)
         self.progress.update("Copying model artifacts to the artifact directory")
-        
         zip_file_path = os.path.join(self.target_dir, f"{str(uuid.uuid4())}.zip")
         zip_file_path = utils.copy_file(
             uri_src=bucket_uri,
@@ -172,7 +197,6 @@ class LargeArtifactDownloader(ArtifactDownloader):
             zip_file.extractall(self.target_dir)
 
         utils.remove_file(zip_file_path)
-
         if self.remove_existing_artifact:
             self.progress.update(
                 "Removing temporary artifacts from the Object Storage bucket"
@@ -180,3 +204,49 @@ class LargeArtifactDownloader(ArtifactDownloader):
             utils.remove_file(bucket_uri)
         else:
             self.progress.update()
+
+    def download_from_model_file_description(self):
+        """Helper function to download the objects using model file description content to the target directory."""
+
+        models = self.model_file_description["models"]
+        total_size = 0
+
+        for model in models:
+            namespace, bucket_name, prefix = (
+                model["namespace"],
+                model["bucketName"],
+                model["prefix"],
+            )
+            bucket_uri = f"oci://{bucket_name}@{namespace}/{prefix}"
+
+            message = f"Copying model artifacts by reference from {bucket_uri} to {self.target_dir}"
+            self.progress.update(message)
+
+            objects = model["objects"]
+            os_details_list = list()
+
+            for obj in objects:
+                name = obj["name"]
+                version = None if obj["version"] == "" else obj["version"]
+                size = obj["sizeInBytes"]
+                if size == 0:
+                    continue
+                total_size += size
+                object_uri = f"oci://{bucket_name}@{namespace}/{name}"
+
+                os_details = ObjectStorageDetails.from_path(object_uri)
+                os_details.version = version
+                os_details_list.append(os_details)
+            try:
+                ObjectStorageDetails.from_path(
+                    bucket_uri
+                ).bulk_download_from_object_storage(
+                    paths=os_details_list,
+                    target_dir=self.target_dir,
+                    progress_bar=self.progress,
+                )
+            except Exception as ex:
+                raise RuntimeError(
+                    f"Failed to download model artifact by reference from the given Object Storage path `{bucket_uri}`."
+                    f"See Exception: {ex}"
+                )

ads/model/artifact_uploader.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
+import logging
 
-# Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2022, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import os
@@ -38,7 +39,7 @@ class ArtifactUploader(ABC):
 
         self.dsc_model = dsc_model
         self.artifact_path = artifact_path
-        self.artifact_zip_path = None
+        self.artifact_file_path = None
         self.progress = None
 
     def upload(self):
@@ -48,8 +49,8 @@ class ArtifactUploader(ABC):
                 ArtifactUploader.PROGRESS_STEPS_COUNT + self.PROGRESS_STEPS_COUNT
             ) as progress:
                 self.progress = progress
-                self.progress.update("Preparing model artifacts ZIP archive.")
-                self._prepare_artifact_tmp_zip()
+                self.progress.update("Preparing model artifacts file.")
+                self._prepare_artifact_tmp_file()
                 self.progress.update("Uploading model artifacts.")
                 self._upload()
                 self.progress.update(
@@ -59,35 +60,35 @@ class ArtifactUploader(ABC):
         except Exception:
             raise
         finally:
-            self._remove_artifact_tmp_zip()
+            self._remove_artifact_tmp_file()
 
-    def _prepare_artifact_tmp_zip(self) -> str:
-        """Prepares model artifacts ZIP archive.
+    def _prepare_artifact_tmp_file(self) -> str:
+        """Prepares model artifacts file.
 
         Returns
         -------
         str
-            Path to the model artifact ZIP archive.
+            Path to the model artifact file.
         """
         if ObjectStorageDetails.is_oci_path(self.artifact_path):
-            self.artifact_zip_path = self.artifact_path
+            self.artifact_file_path = self.artifact_path
         elif os.path.isfile(self.artifact_path) and self.artifact_path.lower().endswith(
-            ".zip"
+            (".zip", ".json")
         ):
-            self.artifact_zip_path = self.artifact_path
+            self.artifact_file_path = self.artifact_path
         else:
-            self.artifact_zip_path = model_utils.zip_artifact(
+            self.artifact_file_path = model_utils.zip_artifact(
                 artifact_dir=self.artifact_path
             )
-        return self.artifact_zip_path
+        return self.artifact_file_path
 
-    def _remove_artifact_tmp_zip(self):
-        """Removes temporary created artifact zip archive."""
+    def _remove_artifact_tmp_file(self):
+        """Removes temporary created artifact file."""
         if (
-            self.artifact_zip_path
-            and self.artifact_zip_path.lower() != self.artifact_path.lower()
+            self.artifact_file_path
+            and self.artifact_file_path.lower() != self.artifact_path.lower()
         ):
-            shutil.rmtree(self.artifact_zip_path, ignore_errors=True)
+            shutil.rmtree(self.artifact_file_path, ignore_errors=True)
 
     @abstractmethod
     def _upload(self):
@@ -101,9 +102,10 @@ class SmallArtifactUploader(ArtifactUploader):
 
     def _upload(self):
         """Uploads model artifacts to the model catalog."""
+        _, ext = os.path.splitext(self.artifact_file_path)
         self.progress.update("Uploading model artifacts to the catalog")
-        with open(self.artifact_zip_path, "rb") as file_data:
-            self.dsc_model.create_model_artifact(file_data)
+        with open(self.artifact_file_path, "rb") as file_data:
+            self.dsc_model.create_model_artifact(bytes_content=file_data, extension=ext)
 
 
 class LargeArtifactUploader(ArtifactUploader):
@@ -117,7 +119,7 @@ class LargeArtifactUploader(ArtifactUploader):
             - object storage path to zip archive. Example: `oci://<bucket_name>@<namespace>/prefix/mymodel.zip`.
             - local path to zip archive. Example: `./mymodel.zip`.
             - local path to folder with artifacts. Example: `./mymodel`.
-    artifact_zip_path: str
+    artifact_file_path: str
         The uri of the zip of model artifact.
     auth: dict
         The default authetication is set using `ads.set_auth` API.
@@ -222,7 +224,7 @@ class LargeArtifactUploader(ArtifactUploader):
         """Uploads model artifacts to the model catalog."""
         bucket_uri = self.bucket_uri
         self.progress.update("Copying model artifact to the Object Storage bucket")
-        if not bucket_uri == self.artifact_zip_path:
+        if not bucket_uri == self.artifact_file_path:
             bucket_uri_file_name = os.path.basename(bucket_uri)
 
             if not bucket_uri_file_name:
@@ -240,7 +242,7 @@ class LargeArtifactUploader(ArtifactUploader):
 
             try:
                 utils.upload_to_os(
-                    src_uri=self.artifact_zip_path,
+                    src_uri=self.artifact_file_path,
                     dst_uri=bucket_uri,
                     auth=self.auth,
                     parallel_process_count=self._parallel_process_count,