Fix Jobs API independent authentication issue for different objects/threads. (#742)

Author: qiuosier

ads/jobs/ads_job.py

@@ -10,6 +10,7 @@
 
 import fsspec
 import oci
+import yaml
 from ads.common.auth import default_signer
 from ads.common.decorator.utils import class_or_instance_method
 from ads.jobs.builders.base import Builder
@@ -263,6 +264,9 @@ def __init__(
             Job runtime, by default None.
 
         """
+        # Saves a copy of the auth object from the class to the instance.
+        # Future changes to the class level Job.auth will not affect the auth of existing instances.
+        self.auth = self.auth.copy()
         for key in ["config", "signer", "client_kwargs"]:
             if kwargs.get(key):
                 self.auth[key] = kwargs.pop(key)
@@ -545,6 +549,26 @@ def to_dict(self, **kwargs: Dict) -> Dict:
             "spec": spec,
         }
 
+    @class_or_instance_method
+    def from_yaml(
+        cls,
+        yaml_string: str = None,
+        uri: str = None,
+        loader: callable = yaml.SafeLoader,
+        **kwargs,
+    ):
+        if inspect.isclass(cls):
+            job = cls(**cls.auth)
+        else:
+            job = cls.__class__(**cls.auth)
+
+        if yaml_string:
+            return job.from_dict(yaml.load(yaml_string, Loader=loader))
+        if uri:
+            yaml_dict = yaml.load(cls._read_from_file(uri=uri, **kwargs), Loader=loader)
+            return job.from_dict(yaml_dict)
+        raise ValueError("Must provide either YAML string or URI location")
+
     @class_or_instance_method
     def from_dict(cls, config: dict) -> "Job":
         """Initializes a job from a dictionary containing the configurations.
@@ -573,9 +597,9 @@ def from_dict(cls, config: dict) -> "Job":
             "runtime": cls._RUNTIME_MAPPING,
         }
         if inspect.isclass(cls):
-            job = cls()
+            job = cls(**cls.auth)
         else:
-            job = cls.__class__()
+            job = cls.__class__(**cls.auth)
 
         for key, value in spec.items():
             if key in mappings:

ads/jobs/builders/infrastructure/dsc_job.py

@@ -6,8 +6,8 @@
 from __future__ import annotations
 
 import datetime
+import inspect
 import logging
-import oci
 import os
 import time
 import traceback
@@ -17,11 +17,12 @@
 from typing import Any, Dict, List, Optional, Union
 
 import fsspec
+import oci
 import oci.data_science
 import oci.util as oci_util
-import yaml
 from oci.data_science.models import JobInfrastructureConfigurationDetails
 from oci.exceptions import ServiceError
+import yaml
 from ads.common import utils
 from ads.common.oci_datascience import DSCNotebookSession, OCIDataScienceMixin
 from ads.common.oci_logging import OCILog
@@ -782,7 +783,7 @@ def to_yaml(self) -> str:
         # Update runtime from job run
         from ads.jobs import Job
 
-        job = Job.from_dict(job_dict)
+        job = Job(**self.auth).from_dict(job_dict)
         envs = job.runtime.envs
         run_config_override = run_dict.get("jobConfigurationOverrideDetails", {})
         envs.update(run_config_override.get("environmentVariables", {}))
@@ -811,7 +812,7 @@ def job(self):
         """
         from ads.jobs import Job
 
-        return Job.from_datascience_job(self.job_id)
+        return Job(**self.auth).from_datascience_job(self.job_id)
 
     def download(self, to_dir):
         """Downloads files from job run output URI to local.
@@ -953,9 +954,9 @@ def standardize_spec(spec):
             if key not in attribute_map and key.lower() in snake_to_camel_map:
                 value = spec.pop(key)
                 if isinstance(value, dict):
-                    spec[
-                        snake_to_camel_map[key.lower()]
-                    ] = DataScienceJob.standardize_spec(value)
+                    spec[snake_to_camel_map[key.lower()]] = (
+                        DataScienceJob.standardize_spec(value)
+                    )
                 else:
                     spec[snake_to_camel_map[key.lower()]] = value
         return spec
@@ -971,6 +972,9 @@ def __init__(self, spec: Dict = None, **kwargs) -> None:
             Specification as keyword arguments.
             If spec contains the same key as the one in kwargs, the value from kwargs will be used.
         """
+        # Saves a copy of the auth object from the class to the instance.
+        # Future changes to the class level Job.auth will not affect the auth of existing instances.
+        self.auth = self.auth.copy()
         for key in ["config", "signer", "client_kwargs"]:
             if kwargs.get(key):
                 self.auth[key] = kwargs.pop(key)
@@ -1710,6 +1714,15 @@ def from_id(cls, job_id: str) -> DataScienceJob:
         """
         return cls.from_dsc_job(DSCJob(**cls.auth).from_ocid(job_id))
 
+    @class_or_instance_method
+    def from_dict(cls, obj_dict: dict):
+        """Initialize the object from a Python dictionary"""
+        if inspect.isclass(cls):
+            job_cls = cls
+        else:
+            job_cls = cls.__class__
+        return job_cls(spec=obj_dict.get("spec"), **cls.auth)
+
     @class_or_instance_method
     def list_jobs(cls, compartment_id: str = None, **kwargs) -> List[DataScienceJob]:
         """Lists all jobs in a compartment.

ads/pipeline/ads_pipeline_step.py

@@ -1,14 +1,15 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2022 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 copy
 from typing import List
 
 from ads.jobs import Job
 from ads.jobs.builders.infrastructure.dsc_job import DataScienceJob
 from ads.jobs.builders.runtimes.base import Runtime
+from ads.common.utils import get_random_name_for_resource
 
 PIPELINE_STEP_KIND_TO_OCI_MAP = {
     "dataScienceJob": "ML_JOB",
@@ -43,7 +44,7 @@ class PipelineStep(Job):
 
     def __init__(
         self,
-        name: str,
+        name: str = None,
         job_id: str = None,
         infrastructure=None,
         runtime=None,
@@ -174,7 +175,7 @@ def __init__(
 
         super().__init__()
         if not name:
-            raise ValueError("PipelineStep name must be specified.")
+            name = get_random_name_for_resource()
         elif any(char in PIPELINE_STEP_RESTRICTED_CHAR_SET for char in name):
             raise ValueError(
                 "PipelineStep name can not include any of the "
@@ -521,17 +522,15 @@ def to_dict(self) -> dict:
             dict_details["spec"][self.CONST_DESCRIPTION] = self.description
         if self.kind == "ML_JOB":
             if self.environment_variable:
-                dict_details["spec"][self.CONST_ENVIRONMENT_VARIABLES] = (
-                    self.environment_variable
-                )
+                dict_details["spec"][
+                    self.CONST_ENVIRONMENT_VARIABLES
+                ] = self.environment_variable
             if self.argument:
-                dict_details["spec"][self.CONST_COMMAND_LINE_ARGUMENTS] = (
-                    self.argument
-                )
+                dict_details["spec"][self.CONST_COMMAND_LINE_ARGUMENTS] = self.argument
             if self.maximum_runtime_in_minutes:
-                dict_details["spec"][self.CONST_MAXIMUM_RUNTIME_IN_MINUTES] = (
-                    self.maximum_runtime_in_minutes
-                )
+                dict_details["spec"][
+                    self.CONST_MAXIMUM_RUNTIME_IN_MINUTES
+                ] = self.maximum_runtime_in_minutes
 
         dict_details["spec"].pop(self.CONST_DEPENDS_ON, None)
 

tests/unitary/default_setup/auth/test_jobs_auth.py

@@ -0,0 +1,50 @@
+"""Contains tests for Jobs API authentication."""
+
+from unittest import TestCase
+from ads.jobs import Job
+
+JOB_YAML = """
+kind: job
+apiVersion: v1.0
+spec:
+  name: llama2
+  infrastructure:
+    kind: infrastructure
+    spec:
+      blockStorageSize: 256
+      compartmentId: "ocid1.compartment.oc1..aaa"
+      logGroupId: "ocid1.loggroup.oc1.iad.aaa"
+      logId: "ocid1.log.oc1.iad.aaa"
+      projectId: "ocid1.datascienceproject.oc1.iad.aaa"
+      subnetId: "ocid1.subnet.oc1.iad.aaa"
+      shapeName: VM.GPU.A10.2
+    type: dataScienceJob
+  runtime:
+    kind: runtime
+    type: pyTorchDistributed
+    spec:
+      replicas: 2
+      conda:
+        type: service
+        slug: pytorch20_p39_gpu_v2
+      command: >-
+        torchrun examples/finetuning.py
+"""
+
+
+class JobsAuthTest(TestCase):
+    """Contains tests for Jobs API authentication."""
+
+    def test_auth_from_yaml(self):
+        """Test using different endpoints for different jobs."""
+        auth1 = {"client_kwargs": {"endpoint": "endpoint1.com"}}
+        auth2 = {"client_kwargs": {"endpoint": "endpoint2.com"}}
+        job1 = Job(**auth1).from_yaml(JOB_YAML)
+        job2 = Job(**auth2).from_yaml(JOB_YAML)
+        job3 = Job.from_yaml(JOB_YAML)
+        self.assertEqual(job1.auth, auth1)
+        self.assertEqual(job1.infrastructure.auth, auth1)
+        self.assertEqual(job2.auth, auth2)
+        self.assertEqual(job2.infrastructure.auth, auth2)
+        self.assertEqual(job3.auth, {})
+        self.assertEqual(job3.infrastructure.auth, {})