adding auth

Author: Ziqun Ye

ads/opctl/backend/base.py

@@ -21,7 +21,7 @@ class Backend:
 
     def __init__(self, config: Dict) -> None:
         self.config = config
-        self.auth_type = config["execution"].get("auth")
+        self.auth_type = config["execution"].get("auth", "api_key")
         self.profile = config["execution"].get("oci_profile", None)
         self.oci_config = config["execution"].get("oci_config", None)
 

ads/opctl/backend/local.py

@@ -219,6 +219,7 @@ def _run_with_conda_pack(
             )
             if os.path.exists(os.path.join(conda_pack_path, "spark-defaults.conf")):
                 env_vars["SPARK_CONF_DIR"] = os.path.join(DEFAULT_IMAGE_CONDA_DIR, slug)
+            logger.info(f"Running with conda pack in a container with command {command}")
             return self._activate_conda_env_and_run(
                 image, slug, command, bind_volumes, env_vars
             )
@@ -701,7 +702,8 @@ def predict(self) -> None:
             )
 
             _download_model(
-                oci_auth=self.oci_auth,
+                auth=self.auth_type,
+                profile=self.profile,
                 ocid=ocid,
                 artifact_directory=artifact_directory,
                 region=region,
@@ -724,32 +726,37 @@ def predict(self) -> None:
         # bind_volumnes
         bind_volumes = {}
         SCRIPT = "script.py"
+        dir_path = os.path.dirname(os.path.realpath(__file__))
         if not is_in_notebook_session():
             bind_volumes = {
                 os.path.expanduser(
                     os.path.dirname(self.config["execution"]["oci_config"])
                 ): {"bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, ".oci")}
             }
-            dir_path = os.path.dirname(os.path.realpath(__file__))
-
+            
             self.config["execution"]["source_folder"] = os.path.abspath(
                 os.path.join(dir_path, "..")
             )
             self.config["execution"]["entrypoint"] = SCRIPT
             bind_volumes[artifact_directory] = {"bind": DEFAULT_MODEL_DEPLOYMENT_FOLDER}
 
-        # payload
+        # extra cmd
         data = self.config["execution"].get("payload")
+        extra_cmd = f"--payload '{data}' " + f"--auth {self.auth_type} "
+        if self.auth_type != "resource_principal":
+            extra_cmd += f"--profile {self.profile}"
 
         if is_in_notebook_session() or NO_CONTAINER:
-            script_path = os.path.join(self.config['execution']['source_folder'], SCRIPT)
-            run_command(cmd=f"python {script_path} " + f"{artifact_directory} "+ f"'{data}'", shell=True)
+            # _run_with_conda_pack has code to handle notebook session case,
+            # however, it activate the conda pack and then run the script. 
+            # For the deployment, we just take the current conda env and run it.
+            # Hence we just handle the notebook case directly here.
+            script_path = os.path.join(os.path.join(dir_path, ".."), SCRIPT)
+            cmd = f"python {script_path} " + f"--artifact-directory {artifact_directory} " + extra_cmd
+            logger.info(f"Running in a notebook or NO_CONTAINER with command {cmd}")
+            run_command(cmd=cmd, shell=True)
         else:
-            extra_cmd = (
-            DEFAULT_MODEL_DEPLOYMENT_FOLDER
-            + " "
-            + data
-        )
+            extra_cmd = f"--artifact-directory {DEFAULT_MODEL_DEPLOYMENT_FOLDER} "+ extra_cmd
             exit_code = self._run_with_conda_pack(
                     bind_volumes, extra_cmd, install=True, conda_uri=conda_path
                 )
@@ -758,7 +765,7 @@ def predict(self) -> None:
                     f"`predict` did not complete successfully. Exit code: {exit_code}. "
                     f"Run with the --debug argument to view container logs."
                 )
-
+                
     def _get_conda_info_from_custom_metadata(self, ocid):
         """
         Get conda env info from custom metadata from model catalog.

ads/opctl/cli.py

@@ -622,6 +622,11 @@ def init(debug: bool, **kwargs: Dict[str, Any]) -> None:
     type=click.Choice(AuthType.values()),
     default=None,
 )
+@click.option(
+    "--oci-profile",
+    help="oci profile",
+    default=None,
+)
 @click.option("--debug", "-d", help="set debug mode", is_flag=True, default=False)
 def predict(**kwargs):
     """

ads/opctl/model/cmds.py

@@ -1,12 +1,12 @@
 import os
 import shutil
 
-from ads.common.auth import create_signer
+from ads.common.auth import AuthContext
 from ads.model.datascience_model import DataScienceModel
 from ads.opctl import logger
-from ads.opctl.constants import DEFAULT_MODEL_FOLDER
 from ads.opctl.config.base import ConfigProcessor
 from ads.opctl.config.merger import ConfigMerger
+from ads.opctl.constants import DEFAULT_MODEL_FOLDER
 
 
 def download_model(**kwargs):
@@ -15,12 +15,6 @@ def download_model(**kwargs):
 
     auth_type = p.config["execution"].get("auth")
     profile = p.config["execution"].get("oci_profile", None)
-    oci_config = p.config["execution"].get("oci_config", None)
-    oci_auth = create_signer(
-        auth_type,
-        oci_config,
-        profile,
-    )
     model_folder = os.path.expanduser(
         p.config["execution"].get("model_save_folder", DEFAULT_MODEL_FOLDER)
     )
@@ -44,7 +38,8 @@ def download_model(**kwargs):
             bucket_uri=bucket_uri,
             timeout=timeout,
             force_overwrite=force_overwrite,
-            oci_auth=oci_auth,
+            auth=auth_type,
+            profile=profile
         )
     else:
         logger.error(f"Model already exists. Set `force_overwrite=True` to overwrite.")
@@ -54,22 +49,22 @@ def download_model(**kwargs):
 
 
 def _download_model(
-    ocid, artifact_directory, oci_auth, region, bucket_uri, timeout, force_overwrite
+    ocid, artifact_directory, region, bucket_uri, timeout, force_overwrite, auth, profile
 ):
     os.makedirs(artifact_directory, exist_ok=True)
 
     try:
-        dsc_model = DataScienceModel.from_id(ocid)
-        dsc_model.download_artifact(
-            target_dir=artifact_directory,
-            force_overwrite=force_overwrite,
-            overwrite_existing_artifact=True,
-            remove_existing_artifact=True,
-            auth=oci_auth,
-            region=region,
-            timeout=timeout,
-            bucket_uri=bucket_uri,
-        )
+        with AuthContext(auth=auth, profile=profile):
+            dsc_model = DataScienceModel.from_id(ocid)
+            dsc_model.download_artifact(
+                target_dir=artifact_directory,
+                force_overwrite=force_overwrite,
+                overwrite_existing_artifact=True,
+                remove_existing_artifact=True,
+                region=region,
+                timeout=timeout,
+                bucket_uri=bucket_uri,
+            )
     except Exception as e:
         print(type(e))
         shutil.rmtree(artifact_directory, ignore_errors=True)

ads/opctl/script.py

@@ -1,28 +1,37 @@
+import argparse
 import json
-import sys
 
+from ads.common.auth import AuthContext
 from ads.model.generic_model import GenericModel
 
 
-def verify(artifact_dir, data): # pragma: no cover
+def verify(artifact_dir, payload, auth, profile): # pragma: no cover
+    kwargs = {"auth": auth}
+    if profile != 'None':
+        kwargs["profile"] = profile
+    with AuthContext(**kwargs):
+        model = GenericModel.from_model_artifact(
+            uri=artifact_dir,
+            artifact_dir=artifact_dir,
+            force_overwrite=True,
+        )
 
-    model = GenericModel.from_model_artifact(
-        uri=artifact_dir,
-        artifact_dir=artifact_dir,
-        force_overwrite=True,
-    )
-
-    try:
-        data = json.loads(data)
-    except:
-        pass
-    print(model.verify(data, auto_serialize_data=False))
+        try:
+            payload = json.loads(payload)
+        except:
+            pass
+        print(model.verify(payload, auto_serialize_data=False))
 
 
 def main(): # pragma: no cover
-    args = sys.argv[1:]
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--payload", type=str, required=True)
+    parser.add_argument("--artifact-directory", type=str, required=True)
+    parser.add_argument("--auth", type=str, required=True)
+    parser.add_argument("--profile", type=str,required=False)
+    args = parser.parse_args()
     verify(
-        artifact_dir=args[0], data=args[1]
+        artifact_dir=args.artifact_directory, payload=args.payload, auth=args.auth, profile=args.profile
     )
     return 0
 

tests/unitary/with_extras/opctl/test_opctl_model.py

@@ -3,7 +3,6 @@
 from unittest.mock import ANY, call, patch
 from ads.model.datascience_model import DataScienceModel
 from unittest.mock import MagicMock, Mock
-from ads.opctl.model.cmds import create_signer
 import os
 
 
@@ -36,12 +35,9 @@ def test_model__download_model_error(mock_from_id):
 
 
 @patch("ads.opctl.model.cmds._download_model")
-@patch("ads.opctl.model.cmds.create_signer")
-def test_download_model(mock_create_signer, mock__download_model):
+def test_download_model( mock__download_model):
     auth_mock = MagicMock()
-    mock_create_signer.return_value = auth_mock
     download_model(ocid="fake_model_id")
-    mock_create_signer.assert_called_once()
     mock__download_model.assert_called_once_with(
         ocid="fake_model_id",
         artifact_directory=os.path.expanduser("~/.ads_ops/models/fake_model_id"),