ODSC-50353: Allow to override the config values in the operator specification (#456)

Author: mrDzurb

ads/opctl/config/merger.py

@@ -124,7 +124,10 @@ def _fill_config_with_defaults(self, ads_config_path: str) -> None:
                     exec_config.get("auth") or AuthType.API_KEY
                 )
         # determine profile
-        if self.config["execution"]["auth"] != AuthType.API_KEY:
+        if self.config["execution"]["auth"] in (
+            AuthType.RESOURCE_PRINCIPAL,
+            AuthType.INSTANCE_PRINCIPAL,
+        ):
             profile = self.config["execution"]["auth"].upper()
             exec_config.pop("oci_profile", None)
             self.config["execution"]["oci_profile"] = None

ads/opctl/constants.py

@@ -32,6 +32,7 @@
 # OPERATOR
 OPERATOR_MODULE_PATH = "ads.opctl.operator.lowcode"
 OPERATOR_IMAGE_WORK_DIR = "/etc/operator"
+OVERRIDE_KWARGS = "override_kwargs"
 
 
 class RUNTIME_TYPE(ExtendedEnum):

ads/opctl/decorator/common.py

@@ -7,10 +7,13 @@
 from functools import wraps
 from typing import Callable, Dict, List
 
+import click
+
 from ads.common.auth import AuthContext
 from ads.opctl import logger
 from ads.opctl.config.base import ConfigProcessor
 from ads.opctl.config.merger import ConfigMerger
+from ads.opctl.constants import OVERRIDE_KWARGS
 
 RUN_ID_FIELD = "run_id"
 
@@ -101,3 +104,26 @@ def wrapper(*args, **kwargs) -> Dict:
             return func(*args, **kwargs)
 
     return wrapper
+
+
+def with_click_unknown_args(func: Callable) -> Callable:
+    """The decorator to parse the click unknown arguments and put them into kwargs."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs) -> Dict:
+        kwargs[OVERRIDE_KWARGS] = {}
+        try:
+            click_context = next(
+                item for item in args if isinstance(item, click.core.Context)
+            )
+            kwargs[OVERRIDE_KWARGS] = {
+                key[2:]: value
+                for key, value in zip(click_context.args[::2], click_context.args[1::2])
+            }
+        except Exception as ex:
+            logger.debug(ex)
+            pass
+
+        return func(*args, **kwargs)
+
+    return wrapper

ads/opctl/operator/cli.py

@@ -13,7 +13,7 @@
 from ads.common.auth import AuthType
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.opctl.constants import BACKEND_NAME, RUNTIME_TYPE
-from ads.opctl.decorator.common import click_options, with_auth
+from ads.opctl.decorator.common import click_options, with_auth, with_click_unknown_args
 from ads.opctl.utils import suppress_traceback
 
 from .__init__ import __operators__
@@ -271,7 +271,12 @@ def publish_conda(debug: bool, **kwargs: Dict[str, Any]) -> None:
     suppress_traceback(debug)(cmd_publish_conda)(**kwargs)
 
 
-@commands.command()
+@commands.command(
+    context_settings=dict(
+        ignore_unknown_options=True,
+        allow_extra_args=True,
+    )
+)
 @click_options(DEBUG_OPTION + ADS_CONFIG_OPTION + AUTH_TYPE_OPTION)
 @click.option(
     "--file",
@@ -303,8 +308,10 @@ def publish_conda(debug: bool, **kwargs: Dict[str, Any]) -> None:
     is_flag=True,
     help="During dry run, the actual operation is not performed, only the steps are enumerated.",
 )
+@click.pass_context
+@with_click_unknown_args
 @with_auth
-def run(debug: bool, **kwargs: Dict[str, Any]) -> None:
+def run(ctx: click.core.Context, debug: bool, **kwargs: Dict[str, Any]) -> None:
     """
     Runs the operator with the given specification on the targeted backend.
     """

ads/opctl/operator/cmd.py

@@ -565,7 +565,9 @@ def create(
     raise NotImplementedError()
 
 
-def run(config: Dict, backend: Union[Dict, str] = None, **kwargs) -> None:
+def run(
+    config: Dict, backend: Union[Dict, str] = None, **kwargs: Dict[str, Any]
+) -> None:
     """
     Runs the operator with the given specification on the targeted backend.
 
@@ -575,7 +577,7 @@ def run(config: Dict, backend: Union[Dict, str] = None, **kwargs) -> None:
         The operator's config.
     backend: (Union[Dict, str], optional)
         The backend config or backend name to run the operator.
-    kwargs: (Dict, optional)
+    kwargs: (Dict[str, Any], optional)
         Optional key value arguments to run the operator.
     """
     BackendFactory.backend(

ads/opctl/operator/common/backend_factory.py

@@ -24,10 +24,12 @@
 from ads.opctl.constants import (
     BACKEND_NAME,
     DEFAULT_ADS_CONFIG_FOLDER,
+    OVERRIDE_KWARGS,
     RESOURCE_TYPE,
     RUNTIME_TYPE,
 )
 from ads.opctl.operator.common.const import PACK_TYPE
+from ads.opctl.operator.common.dictionary_merger import DictionaryMerger
 from ads.opctl.operator.common.operator_loader import OperatorInfo, OperatorLoader
 
 
@@ -202,7 +204,10 @@ def backend(
             {**backend, **{"execution": {"backend": backend_kind}}}
         ).step(ConfigMerger, **kwargs)
 
-        config.config["runtime"] = backend
+        # merge backend with the override parameters
+        config.config["runtime"] = DictionaryMerger(
+            updates=kwargs.get(OVERRIDE_KWARGS)
+        ).merge(backend)
         config.config["infrastructure"] = p_backend.config["infrastructure"]
         config.config["execution"] = p_backend.config["execution"]
 

ads/opctl/operator/common/dictionary_merger.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*--
+
+# Copyright (c) 2023 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 Any, Dict, List
+
+
+class DictionaryMerger:
+    """
+    A class to update dictionary values for specified keys and
+    then merge these updates back into the original dictionary.
+
+    Example
+    -------
+    >>> updates = {
+    ...     "infrastructure.blockStorageSize": "20",
+    ...     "infrastructure.projectId": "my_new_project_id"
+    ...     "runtime.conda": "my_conda"
+    ... }
+    >>> updater = DictionaryMerger(updates)
+    >>> source_data = {
+    ...     "infrastructure": {
+    ...         "blockStorageSize": "10",
+    ...         "projectId": "old_project_id",
+    ...     },
+    ...     "runtime": {
+    ...         "conda": "conda",
+    ...     },
+    ... }
+    >>> result = updater.dispatch(source_data)
+    ... {
+    ...     "infrastructure": {
+    ...         "blockStorageSize": "20",
+    ...         "projectId": "my_new_project_id",
+    ...     },
+    ...     "runtime": {
+    ...         "conda": "my_conda",
+    ...     },
+    ... }
+
+    Attributes
+    ----------
+    updates: Dict[str, Any]
+        A dictionary containing the keys with their new values for the update.
+    """
+
+    _SYSTEM_KEYS = set(("kind", "type", "spec", "infrastructure", "runtime"))
+
+    def __init__(self, updates: Dict[str, Any], system_keys: List[str] = None):
+        """
+        Initializes the DictionaryMerger with a dictionary of updates.
+
+        Parameters
+        ----------
+        updates Dict[str, Any]
+            A dictionary with keys that need to be updated and their new values.
+        system_keys: List[str]
+            The list of keys that cannot be replaced in the source dictionary.
+        """
+        self.updates = updates
+        self.system_keys = set(system_keys or []).union(self._SYSTEM_KEYS)
+
+    def _update_keys(
+        self, dict_to_update: Dict[str, Any], parent_key: str = ""
+    ) -> None:
+        """
+        Recursively updates the values of given keys in a dictionary.
+
+        Parameters
+        ----------
+        dict_to_update: Dict[str, Any]
+            The dictionary whose values are to be updated.
+        parent_key: (str, optional)
+            The current path in the dictionary being processed, used for nested dictionaries.
+
+        Returns
+        -------
+        None
+            The method updates the dict_to_update in place.
+        """
+        for key, value in dict_to_update.items():
+            new_key = f"{parent_key}.{key}" if parent_key else key
+            if isinstance(value, dict):
+                self._update_keys(value, new_key)
+            elif new_key in self.updates and key not in self.system_keys:
+                dict_to_update[key] = self.updates[new_key]
+
+    def _merge_updates(
+        self,
+        original_dict: Dict[str, Any],
+        updated_dict: Dict[str, Any],
+        parent_key: str = "",
+    ) -> None:
+        """
+        Merges updated values from the updated_dict into the original_dict based on the provided keys.
+
+        Parameters
+        ----------
+        original_dict: Dict[str, Any]
+            The original dictionary to merge updates into.
+        updated_dict: Dict[str, Any]
+            The updated dictionary with new values.
+        parent_key: str
+            The base key path for recursive merging.
+
+        Returns
+        -------
+        None
+            The method updates the original_dict in place.
+        """
+        for key, value in updated_dict.items():
+            new_key = f"{parent_key}.{key}" if parent_key else key
+            if isinstance(value, dict) and key in original_dict:
+                self._merge_updates(original_dict[key], value, new_key)
+            elif new_key in self.updates:
+                original_dict[key] = value
+
+    def merge(self, src_dict: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Updates the dictionary with new values for specified keys and merges
+        these changes back into the original dictionary.
+
+        Parameters
+        ----------
+        src_dict: Dict[str, Any]
+            The dictionary to be updated and merged.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The updated and merged dictionary.
+        """
+        if not self.updates:
+            return src_dict
+
+        original_dict = copy.deepcopy(src_dict)
+        updated_dict = copy.deepcopy(src_dict)
+
+        # Update the dictionary with the new values
+        self._update_keys(updated_dict)
+
+        # Merge the updates back into the original dictionary
+        self._merge_updates(original_dict, updated_dict)
+
+        return original_dict