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

ads/config.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2020, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2020, 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 contextlib
@@ -33,8 +33,10 @@ CONDA_BUCKET_NS = os.environ.get("CONDA_BUCKET_NS", "id19sfcrra6z")
 OCI_RESOURCE_PRINCIPAL_RPT_ENDPOINT = os.environ.get(
     "OCI_RESOURCE_PRINCIPAL_RPT_ENDPOINT"
 )
+PROJECT_COMPARTMENT_OCID = os.environ.get("PROJECT_COMPARTMENT_OCID")
 COMPARTMENT_OCID = (
     NB_SESSION_COMPARTMENT_OCID
+    or PROJECT_COMPARTMENT_OCID
     or JOB_RUN_COMPARTMENT_OCID
     or PIPELINE_RUN_COMPARTMENT_OCID
 )
@@ -47,6 +49,43 @@ RESOURCE_OCID = (
 NO_CONTAINER = os.environ.get("NO_CONTAINER")
 TMPDIR = os.environ.get("TMPDIR")
 
+ODSC_MODEL_COMPARTMENT_OCID = os.environ.get("ODSC_MODEL_COMPARTMENT_OCID")
+AQUA_MODEL_DEPLOYMENT_IMAGE = os.environ.get("AQUA_MODEL_DEPLOYMENT_IMAGE")
+AQUA_MODEL_DEPLOYMENT_CONFIG = os.environ.get(
+    "AQUA_DEPLOYMENT_CONFIG", "deployment_config.json"
+)
+AQUA_MODEL_FINETUNING_CONFIG = os.environ.get(
+    "AQUA_MODEL_FINETUNING_CONFIG", "ft_config.json"
+)
+AQUA_CONTAINER_INDEX_CONFIG = os.environ.get(
+    "AQUA_CONTAINER_INDEX_CONFIG", "container_index.json"
+)
+AQUA_MODEL_DEPLOYMENT_CONFIG_DEFAULTS = os.environ.get(
+    "AQUA_MODEL_DEPLOYMENT_CONFIG_DEFAULTS", "deployment_config_defaults.json"
+)
+AQUA_RESOURCE_LIMIT_NAMES_CONFIG = os.environ.get(
+    "AQUA_RESOURCE_LIMIT_NAMES_CONFIG", "resource_limit_names.json"
+)
+AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME = "deployment-container"
+AQUA_FINETUNING_CONTAINER_METADATA_NAME = "finetune-container"
+AQUA_EVALUATION_CONTAINER_METADATA_NAME = "evaluation-container"
+AQUA_MODEL_DEPLOYMENT_FOLDER = "/opt/ds/model/deployed_model/"
+AQUA_SERVED_MODEL_NAME = "odsc-llm"
+AQUA_CONFIG_FOLDER = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "aqua/config/"
+)
+AQUA_JOB_SUBNET_ID = os.environ.get("AQUA_JOB_SUBNET_ID", None)
+AQUA_SERVICE_MODELS_BUCKET = os.environ.get(
+    "AQUA_SERVICE_MODELS_BUCKET", "service-managed-models"
+)
+AQUA_TELEMETRY_BUCKET = os.environ.get(
+    "AQUA_TELEMETRY_BUCKET", "service-managed-models"
+)
+AQUA_TELEMETRY_BUCKET_NS = os.environ.get("AQUA_TELEMETRY_BUCKET_NS", CONDA_BUCKET_NS)
+DEBUG_TELEMETRY = os.environ.get("DEBUG_TELEMETRY", None)
+AQUA_SERVICE_NAME = "aqua"
+DATA_SCIENCE_SERVICE_NAME = "data-science"
+
 
 def export(
     uri: Optional[str] = DEFAULT_CONFIG_PATH,

ads/jobs/ads_job.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 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 inspect
 import time
 from typing import List, Union, Dict
 from urllib.parse import urlparse
@@ -10,12 +11,13 @@ from urllib.parse import urlparse
 import fsspec
 import oci
 from ads.common.auth import default_signer
+from ads.common.decorator.utils import class_or_instance_method
 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, 
-    SLEEP_INTERVAL
+    DataScienceJob,
+    DataScienceJobRun,
+    SLEEP_INTERVAL,
 )
 from ads.jobs.builders.runtimes.pytorch_runtime import PyTorchDistributedRuntime
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
@@ -143,8 +145,10 @@ class Job(Builder):
         ]
     }
 
-    @staticmethod
-    def from_datascience_job(job_id) -> "Job":
+    auth = {}
+
+    @class_or_instance_method
+    def from_datascience_job(cls, job_id) -> "Job":
         """Loads a data science job from OCI.
 
         Parameters
@@ -158,16 +162,18 @@ class Job(Builder):
             A job instance.
 
         """
-        dsc_infra = DataScienceJob.from_id(job_id)
+        dsc_infra = DataScienceJob(**cls.auth).from_id(job_id)
         job = (
-            Job(name=dsc_infra.name)
+            Job(name=dsc_infra.name, **cls.auth)
             .with_infrastructure(dsc_infra)
             .with_runtime(dsc_infra.runtime)
         )
         return job
 
-    @staticmethod
-    def datascience_job(compartment_id: str = None, **kwargs) -> List["DataScienceJob"]:
+    @class_or_instance_method
+    def datascience_job(
+        cls, compartment_id: str = None, **kwargs
+    ) -> List["DataScienceJob"]:
         """Lists the existing data science jobs in the compartment.
 
         Parameters
@@ -183,10 +189,12 @@ class Job(Builder):
             A list of Job objects.
         """
         return [
-            Job(name=dsc_job.name)
+            Job(name=dsc_job.name, **cls.auth)
             .with_infrastructure(dsc_job)
             .with_runtime(dsc_job.runtime)
-            for dsc_job in DataScienceJob.list_jobs(compartment_id, **kwargs)
+            for dsc_job in DataScienceJob(**cls.auth).list_jobs(
+                compartment_id, **kwargs
+            )
         ]
 
     @staticmethod
@@ -229,7 +237,9 @@ class Job(Builder):
             for df in DataFlow.list_jobs(compartment_id, **kwargs)
         ]
 
-    def __init__(self, name: str = None, infrastructure=None, runtime=None) -> None:
+    def __init__(
+        self, name: str = None, infrastructure=None, runtime=None, **kwargs
+    ) -> None:
         """Initializes a job.
 
         The infrastructure and runtime can be configured when initializing the job,
@@ -253,6 +263,9 @@ class Job(Builder):
             Job runtime, by default None.
 
         """
+        for key in ["config", "signer", "client_kwargs"]:
+            if kwargs.get(key):
+                self.auth[key] = kwargs.pop(key)
         super().__init__()
         if name:
             self.set_spec("name", name)
@@ -398,7 +411,7 @@ class Job(Builder):
         freeform_tags=None,
         defined_tags=None,
         wait=False,
-        **kwargs
+        **kwargs,
     ) -> Union[DataScienceJobRun, DataFlowRun]:
         """Runs the job.
 
@@ -454,7 +467,7 @@ class Job(Builder):
             freeform_tags=freeform_tags,
             defined_tags=defined_tags,
             wait=wait,
-            **kwargs
+            **kwargs,
         )
 
     def run_list(self, **kwargs) -> list:
@@ -466,7 +479,7 @@ class Job(Builder):
             A list of job run instances, the actual object type depends on the infrastructure.
         """
         return self.infrastructure.run_list(**kwargs)
-    
+
     def cancel(self, wait_for_completion: bool = True) -> None:
         """Cancels the runs of the job.
 
@@ -479,16 +492,16 @@ class Job(Builder):
         runs = self.run_list()
         for run in runs:
             run.cancel(wait_for_completion=False)
-        
+
         if wait_for_completion:
             for run in runs:
                 while (
-                    run.lifecycle_state !=
-                    oci.data_science.models.JobRun.LIFECYCLE_STATE_CANCELED
+                    run.lifecycle_state
+                    != oci.data_science.models.JobRun.LIFECYCLE_STATE_CANCELED
                 ):
                     run.sync()
                     time.sleep(SLEEP_INTERVAL)
-        
+
     def delete(self) -> None:
         """Deletes the job from the infrastructure."""
         self.infrastructure.delete()
@@ -532,7 +545,7 @@ class Job(Builder):
             "spec": spec,
         }
 
-    @classmethod
+    @class_or_instance_method
     def from_dict(cls, config: dict) -> "Job":
         """Initializes a job from a dictionary containing the configurations.
 
@@ -559,7 +572,10 @@ class Job(Builder):
             "infrastructure": cls._INFRASTRUCTURE_MAPPING,
             "runtime": cls._RUNTIME_MAPPING,
         }
-        job = cls()
+        if inspect.isclass(cls):
+            job = cls()
+        else:
+            job = cls.__class__()
 
         for key, value in spec.items():
             if key in mappings:
@@ -569,9 +585,11 @@ class Job(Builder):
                     raise NotImplementedError(
                         f"{key.title()} type: {child_config.get('type')} is not supported."
                     )
-                job.set_spec(
-                    key, mapping[child_config.get("type")].from_dict(child_config)
-                )
+                spec_class = mapping[child_config.get("type")]
+                if spec_class == DataScienceJob:
+                    spec_class = DataScienceJob(**job.auth)
+
+                job.set_spec(key, spec_class.from_dict(child_config))
             else:
                 job.set_spec(key, value)
 

ads/jobs/builders/infrastructure/base.py

@@ -1,10 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 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.common import utils as common_utils
+from ads.common.decorator.utils import class_or_instance_method
 from ads.jobs.builders.base import Builder
 from ads.jobs.builders.runtimes.base import Runtime
 import logging
@@ -97,7 +98,8 @@ class Infrastructure(Builder):
         """
         raise NotImplementedError()
 
-    def list_jobs(self, **kwargs) -> list:
+    @class_or_instance_method
+    def list_jobs(cls, **kwargs) -> list:
         """
         List jobs from the infrastructure.
 

ads/jobs/builders/infrastructure/dsc_job.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 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 __future__ import annotations
 
@@ -40,6 +40,7 @@ from ads.common.dsc_file_system import (
     DSCFileSystemManager,
     OCIObjectStorage,
 )
+from ads.common.decorator.utils import class_or_instance_method
 
 logger = logging.getLogger(__name__)
 
@@ -171,10 +172,12 @@ class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
             )
             subnet_id = infra.get(
                 "subnetId",
-                nb_config.subnet_id
-                if infra_type
-                != JobInfrastructureConfigurationDetails.JOB_INFRASTRUCTURE_TYPE_ME_STANDALONE
-                else None,
+                (
+                    nb_config.subnet_id
+                    if infra_type
+                    != JobInfrastructureConfigurationDetails.JOB_INFRASTRUCTURE_TYPE_ME_STANDALONE
+                    else None
+                ),
             )
             job_shape_config_details = infra.get("jobShapeConfigDetails", {})
             memory_in_gbs = job_shape_config_details.get(
@@ -302,9 +305,9 @@ class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
         return self
 
     def _create_with_oci_api(self) -> None:
-        res = self.client.create_job(
-            self.to_oci_model(oci.data_science.models.CreateJobDetails)
-        )
+        oci_model = self.to_oci_model(oci.data_science.models.CreateJobDetails)
+        logger.debug(oci_model)
+        res = self.client.create_job(oci_model)
         self.update_from_oci_model(res.data)
         if self.lifecycle_state == "ACTIVE":
             return
@@ -353,9 +356,16 @@ class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
         """Updates the Data Science Job."""
         raise NotImplementedError("Updating Job is not supported at the moment.")
 
-    def delete(self) -> DSCJob:
+    def delete(self, force_delete: bool = False) -> DSCJob:
         """Deletes the job and the corresponding job runs.
 
+        Parameters
+        ----------
+        force_delete : bool, optional, defaults to False
+            the deletion fails when associated job runs are in progress, but if force_delete to true, then
+            the job run will be canceled, then it will be deleted. In this case, delete job has to wait till
+            job has been canceled.
+
         Returns
         -------
         DSCJob
@@ -364,6 +374,12 @@ class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
         """
         runs = self.run_list()
         for run in runs:
+            if run.lifecycle_state in [
+                DataScienceJobRun.LIFECYCLE_STATE_ACCEPTED,
+                DataScienceJobRun.LIFECYCLE_STATE_IN_PROGRESS,
+                DataScienceJobRun.LIFECYCLE_STATE_NEEDS_ATTENTION,
+            ]:
+                run.cancel(wait_for_completion=True)
             run.delete()
         self.client.delete_job(self.id)
         return self
@@ -430,7 +446,7 @@ class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
         items = oci.pagination.list_call_get_all_results(
             self.client.list_job_runs, self.compartment_id, job_id=self.id, **kwargs
         ).data
-        return [DataScienceJobRun.from_oci_model(item) for item in items]
+        return [DataScienceJobRun(**self.auth).from_oci_model(item) for item in items]
 
     def run(self, **kwargs) -> DataScienceJobRun:
         """Runs the job
@@ -496,28 +512,11 @@ class DSCJob(OCIDataScienceMixin, oci.data_science.models.Job):
             kwargs["job_configuration_override_details"] = config_override
 
         wait = kwargs.pop("wait", False)
-        run = DataScienceJobRun(**kwargs).create()
+        run = DataScienceJobRun(**kwargs, **self.auth).create()
         if wait:
             return run.watch()
         return run
 
-    @classmethod
-    def from_ocid(cls, ocid) -> DSCJob:
-        """Gets a job by OCID
-
-        Parameters
-        ----------
-        ocid : str
-            The OCID of the job.
-
-        Returns
-        -------
-        DSCJob
-            An instance of DSCJob.
-
-        """
-        return super().from_ocid(ocid)
-
 
 class DataScienceJobRun(
     OCIDataScienceMixin, oci.data_science.models.JobRun, RunInstance
@@ -575,7 +574,12 @@ class DataScienceJobRun(
         if not self.log_id:
             raise ValueError("Log OCID is not specified for this job run.")
         # Specifying log group ID when initializing OCILog can reduce the number of API calls.
-        return OCILog(id=self.log_id, log_group_id=self.log_details.log_group_id)
+        auth = self.auth
+        if "client_kwargs" in auth and isinstance(auth["client_kwargs"], dict):
+            auth["client_kwargs"].pop("service_endpoint", None)
+        return OCILog(
+            id=self.log_id, log_group_id=self.log_details.log_group_id, **auth
+        )
 
     @staticmethod
     def _format_log(message: str, date_time: datetime.datetime) -> dict:
@@ -928,6 +932,8 @@ class DataScienceJob(Infrastructure):
         OBJECT_STORAGE_TYPE: OCIObjectStorage,
     }
 
+    auth = {}
+
     @staticmethod
     def standardize_spec(spec):
         if not spec:
@@ -965,12 +971,16 @@ class DataScienceJob(Infrastructure):
             Specification as keyword arguments.
             If spec contains the same key as the one in kwargs, the value from kwargs will be used.
         """
+        for key in ["config", "signer", "client_kwargs"]:
+            if kwargs.get(key):
+                self.auth[key] = kwargs.pop(key)
+
         self.standardize_spec(spec)
         self.standardize_spec(kwargs)
         super().__init__(spec=spec, **kwargs)
         if not self.job_type:
             self.with_job_type("DEFAULT")
-        self.dsc_job = DSCJob()
+        self.dsc_job = DSCJob(**self.auth)
         self.runtime = None
         self._name = None
 
@@ -1557,7 +1567,7 @@ class DataScienceJob(Infrastructure):
         if not payload.get("defined_tags"):
             payload["defined_tags"] = self.defined_tags
 
-        self.dsc_job = DSCJob(**payload)
+        self.dsc_job = DSCJob(**payload, **self.auth)
         # Set Job infra to user values after DSCJob initialized the defaults
         self._update_job_infra(self.dsc_job)
         self.dsc_job.create()
@@ -1682,7 +1692,7 @@ class DataScienceJob(Infrastructure):
         instance.runtime = DataScienceJobRuntimeManager(instance).extract(dsc_job)
         return instance
 
-    @classmethod
+    @class_or_instance_method
     def from_id(cls, job_id: str) -> DataScienceJob:
         """Gets an existing job using Job OCID
 
@@ -1698,9 +1708,9 @@ class DataScienceJob(Infrastructure):
             An instance of DataScienceJob
 
         """
-        return cls.from_dsc_job(DSCJob.from_ocid(job_id))
+        return cls.from_dsc_job(DSCJob(**cls.auth).from_ocid(job_id))
 
-    @classmethod
+    @class_or_instance_method
     def list_jobs(cls, compartment_id: str = None, **kwargs) -> List[DataScienceJob]:
         """Lists all jobs in a compartment.
 
@@ -1721,10 +1731,10 @@ class DataScienceJob(Infrastructure):
         """
         return [
             cls.from_dsc_job(job)
-            for job in DSCJob.list_resource(compartment_id, **kwargs)
+            for job in DSCJob(**cls.auth).list_resource(compartment_id, **kwargs)
         ]
 
-    @classmethod
+    @class_or_instance_method
     def instance_shapes(cls, compartment_id: str = None, **kwargs) -> list:
         """Lists the supported shapes for running jobs in a compartment.
 
@@ -1752,13 +1762,13 @@ class DataScienceJob(Infrastructure):
 
         """
         shapes = oci.pagination.list_call_get_all_results(
-            DSCJob.init_client().list_job_shapes,
+            DSCJob(**cls.auth).init_client().list_job_shapes,
             DSCJob.check_compartment_id(compartment_id),
             **kwargs,
         ).data
         return shapes
 
-    @classmethod
+    @class_or_instance_method
     def fast_launch_shapes(cls, compartment_id: str = None, **kwargs) -> list:
         """Lists the supported fast launch shapes for running jobs in a compartment.
 
@@ -1786,7 +1796,7 @@ class DataScienceJob(Infrastructure):
 
         """
         shapes = oci.pagination.list_call_get_all_results(
-            DSCJob.init_client().list_fast_launch_job_configs,
+            DSCJob(**cls.auth).init_client().list_fast_launch_job_configs,
             DSCJob.check_compartment_id(compartment_id),
             **kwargs,
         ).data

ads/jobs/builders/runtimes/base.py

@@ -1,10 +1,13 @@
 #!/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/
 from __future__ import annotations
 import re
+import time
+import traceback
+
 from typing import Dict, TypeVar
 from ads.jobs.builders.base import Builder
 from ads.jobs import env_var_parser
@@ -253,3 +256,70 @@ class Runtime(Builder):
             )
             .with_argument(**kwargs.get(self.attribute_map[self.CONST_ARGS], {}))
         )
+
+
+class MultiNodeRuntime(Runtime):
+    """Represents runtime supporting multi-node jobs."""
+
+    CONST_REPLICA = "replicas"
+
+    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 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()
+
+                # Only update display name and env vars if replica is specified (not None).
+                if self.replica is not None:
+                    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["NODE_COUNT"] = 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 as ex:
+            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()
+            raise ex
+        return main_run

ads/jobs/builders/runtimes/container_runtime.py

@@ -1,13 +1,13 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 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 typing import Union
-from ads.jobs.builders.runtimes.base import Runtime
+from ads.jobs.builders.runtimes.base import MultiNodeRuntime
 
 
-class ContainerRuntime(Runtime):
+class ContainerRuntime(MultiNodeRuntime):
     """Represents a container job runtime
 
     To define container runtime:
@@ -51,7 +51,7 @@ class ContainerRuntime(Runtime):
         CONST_ENTRYPOINT: CONST_ENTRYPOINT,
         CONST_CMD: CONST_CMD,
     }
-    attribute_map.update(Runtime.attribute_map)
+    attribute_map.update(MultiNodeRuntime.attribute_map)
 
     @property
     def image(self) -> str: