mlrun 1.6.0rc11__py3-none-any.whl → 1.6.0rc13__py3-none-any.whl

mlrun/projects/project.py

@@ -1931,7 +1931,7 @@ class MlrunProject(ModelObj):
 
         return resolved_function_name, function_object, func
 
-    def enable_model_monitoring_controller(
+    def enable_model_monitoring(
         self,
         default_controller_image: str = "mlrun/mlrun",
         base_period: int = 10,
@@ -1944,8 +1944,9 @@ class MlrunProject(ModelObj):
         :param default_controller_image: The default image of the model monitoring controller job. Note that the writer
                                          function, which is a real time nuclio functino, will be deployed with the same
                                          image. By default, the image is mlrun/mlrun.
-        :param base_period:              Minutes to determine the frequency in which the model monitoring controller job
-                                         is running. By default, the base period is 5 minutes.
+        :param base_period:              The time period in minutes in which the model monitoring controller job
+                                         runs. By default, the base period is 10 minutes. The schedule for the job
+                                         will be the following cron expression: "*/{base_period} * * * *".
         :return: model monitoring controller job as a dictionary.
         """
         db = mlrun.db.get_run_db(secrets=self._secrets)
@@ -1955,7 +1956,7 @@ class MlrunProject(ModelObj):
             base_period=base_period,
         )
 
-    def disable_model_monitoring_controller(self):
+    def disable_model_monitoring(self):
         db = mlrun.db.get_run_db(secrets=self._secrets)
         db.delete_function(
             project=self.name,
@@ -2981,6 +2982,12 @@ class MlrunProject(ModelObj):
         :param extra_args:  A string containing additional builder arguments in the format of command-line options,
             e.g. extra_args="--skip-tls-verify --build-arg A=val"
         """
+        if not overwrite_build_params:
+            # TODO: change overwrite_build_params default to True in 1.8.0
+            warnings.warn(
+                "The `overwrite_build_params` parameter default will change from 'False' to 'True in 1.8.0.",
+                mlrun.utils.OverwriteBuildParamsWarning,
+            )
         default_image_name = mlrun.mlconf.default_project_image_name.format(
             name=self.name
         )
@@ -3052,35 +3059,48 @@ class MlrunProject(ModelObj):
                 FutureWarning,
             )
 
-        self.build_config(
-            image=image,
-            set_as_default=set_as_default,
-            base_image=base_image,
-            commands=commands,
-            secret_name=secret_name,
-            with_mlrun=with_mlrun,
-            requirements=requirements,
-            requirements_file=requirements_file,
-            overwrite_build_params=overwrite_build_params,
-        )
+        if not overwrite_build_params:
+            # TODO: change overwrite_build_params default to True in 1.8.0
+            warnings.warn(
+                "The `overwrite_build_params` parameter default will change from 'False' to 'True in 1.8.0.",
+                mlrun.utils.OverwriteBuildParamsWarning,
+            )
 
-        function = mlrun.new_function("mlrun--project--image--builder", kind="job")
-
-        build = self.spec.build
-        result = self.build_function(
-            function=function,
-            with_mlrun=build.with_mlrun,
-            image=build.image,
-            base_image=build.base_image,
-            commands=build.commands,
-            secret_name=build.secret,
-            requirements=build.requirements,
-            overwrite_build_params=overwrite_build_params,
-            mlrun_version_specifier=mlrun_version_specifier,
-            builder_env=builder_env,
-            extra_args=extra_args,
-            force_build=force_build,
-        )
+        # TODO: remove filter once overwrite_build_params default is changed to True in 1.8.0
+        with warnings.catch_warnings():
+            warnings.simplefilter(
+                "ignore", category=mlrun.utils.OverwriteBuildParamsWarning
+            )
+
+            self.build_config(
+                image=image,
+                set_as_default=set_as_default,
+                base_image=base_image,
+                commands=commands,
+                secret_name=secret_name,
+                with_mlrun=with_mlrun,
+                requirements=requirements,
+                requirements_file=requirements_file,
+                overwrite_build_params=overwrite_build_params,
+            )
+
+            function = mlrun.new_function("mlrun--project--image--builder", kind="job")
+
+            build = self.spec.build
+            result = self.build_function(
+                function=function,
+                with_mlrun=build.with_mlrun,
+                image=build.image,
+                base_image=build.base_image,
+                commands=build.commands,
+                secret_name=build.secret,
+                requirements=build.requirements,
+                overwrite_build_params=overwrite_build_params,
+                mlrun_version_specifier=mlrun_version_specifier,
+                builder_env=builder_env,
+                extra_args=extra_args,
+                force_build=force_build,
+            )
 
         try:
             mlrun.db.get_run_db(secrets=self._secrets).delete_function(

mlrun/runtimes/databricks_job/databricks_runtime.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from ast import FunctionDef, parse, unparse
 from base64 import b64decode, b64encode
 from typing import Callable, Dict, List, Optional, Union
 
@@ -22,61 +23,92 @@ from mlrun.model import HyperParamOptions, RunObject
 from mlrun.runtimes.kubejob import KubejobRuntime
 
 
+def get_log_artifacts_code(runobj: RunObject, task_parameters: dict):
+    artifact_json_dir = task_parameters.get(
+        "artifact_json_dir",
+        mlrun.mlconf.function.databricks.artifact_directory_path,
+    )
+    artifact_json_path = (
+        f"{artifact_json_dir}/mlrun_artifact_{runobj.metadata.uid}.json"
+    )
+    return (
+        log_artifacts_code_template.format(f"/dbfs{artifact_json_path}"),
+        artifact_json_path,
+    )
+
+
+def replace_log_artifact_function(code: str, log_artifacts_code: str):
+    #  user can use a dummy function in oder to avoid edit his code.
+    #  replace mlrun_log_artifact function if already exist.
+    is_replaced = False
+    parsed_code = parse(code)
+    for node in parsed_code.body:
+        if isinstance(node, FunctionDef) and node.name == "mlrun_log_artifact":
+            new_function_ast = parse(log_artifacts_code)
+            node.args = new_function_ast.body[0].args
+            node.body = new_function_ast.body[0].body
+            is_replaced = True
+            break
+    return unparse(parsed_code), is_replaced
+
+
 class DatabricksRuntime(KubejobRuntime):
     kind = "databricks"
     _is_remote = True
 
-    def _get_log_artifacts_code(self, runobj: RunObject, task_parameters: dict):
-        artifact_json_dir = task_parameters.get(
-            "artifact_json_dir",
-            mlrun.mlconf.function.databricks.artifact_directory_path,
-        )
-        artifact_json_path = (
-            f"{artifact_json_dir}/mlrun_artifact_{runobj.metadata.uid}.json"
+    @staticmethod
+    def _verify_returns(returns):
+        # TODO complete returns feature
+        if returns:
+            raise MLRunInvalidArgumentError(
+                "Databricks function does not support returns."
+            )
+
+    def _get_modified_user_code(self, original_handler: str, log_artifacts_code: str):
+        encoded_code = (
+            self.spec.build.functionSourceCode if hasattr(self.spec, "build") else None
         )
-        return (
-            artifacts_code_template.format(f"/dbfs{artifact_json_path}"),
-            artifact_json_path,
+        if not encoded_code:
+            raise ValueError("Databricks function must be provided with user code")
+
+        decoded_code = b64decode(encoded_code).decode("utf-8")
+        decoded_code, is_replaced = replace_log_artifact_function(
+            code=decoded_code, log_artifacts_code=log_artifacts_code
         )
+        if is_replaced:
+            decoded_code = (
+                logger_and_consts_code + _databricks_script_code + decoded_code
+            )
+        else:
+            decoded_code = (
+                logger_and_consts_code
+                + log_artifacts_code
+                + _databricks_script_code
+                + decoded_code
+            )
+        if original_handler:
+            decoded_code += f"\nresult = {original_handler}(**handler_arguments)\n"
+            decoded_code += _return_artifacts_code
+        return b64encode(decoded_code.encode("utf-8")).decode("utf-8")
 
     def get_internal_parameters(self, runobj: RunObject):
         """
-        Return the internal function code.
+        Return the internal function parameters + code.
         """
         task_parameters = runobj.spec.parameters.get("task_parameters", {})
         if "original_handler" in task_parameters:
             original_handler = task_parameters["original_handler"]
         else:
             original_handler = runobj.spec.handler or ""
-        encoded_code = (
-            self.spec.build.functionSourceCode if hasattr(self.spec, "build") else None
-        )
-        if not encoded_code:
-            raise ValueError("Databricks function must be provided with user code")
-        decoded_code = b64decode(encoded_code).decode("utf-8")
-        artifacts_code, artifact_json_path = self._get_log_artifacts_code(
+        log_artifacts_code, artifact_json_path = get_log_artifacts_code(
             runobj=runobj, task_parameters=task_parameters
         )
-        code = artifacts_code + _databricks_script_code + decoded_code
-        if original_handler:
-            code += f"\nresult = {original_handler}(**handler_arguments)\n"
-            code += """\n
-default_key_template = 'mlrun_return_value_'
-if result:
-    if isinstance(result, dict):
-        for key, path in result.items():
-            mlrun_log_artifact(name=key, path=path)
-    elif isinstance(result, (list, tuple, set)):
-        for index, value in enumerate(result):
-            key = f'{default_key_template}{index+1}'
-            mlrun_log_artifact(name=key, path=value)
-    elif isinstance(result, str):
-        mlrun_log_artifact(name=f'{default_key_template}1', path=result)
-    else:
-        mlrun_logger.warning(f'cannot log artifacts with the result of handler function \
-- result in unsupported type. {type(result)}')
-"""
-        code = b64encode(code.encode("utf-8")).decode("utf-8")
+        returns = runobj.spec.returns or []
+        self._verify_returns(returns=returns)
+        code = self._get_modified_user_code(
+            original_handler=original_handler,
+            log_artifacts_code=log_artifacts_code,
+        )
         updated_task_parameters = {
             "original_handler": original_handler,
             "artifact_json_path": artifact_json_path,
@@ -159,6 +191,22 @@ if result:
         )
 
 
+logger_and_consts_code = """ \n
+import os
+import logging
+mlrun_logger = logging.getLogger('mlrun_logger')
+mlrun_logger.setLevel(logging.DEBUG)
+
+mlrun_console_handler = logging.StreamHandler()
+mlrun_console_handler.setLevel(logging.DEBUG)
+mlrun_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+mlrun_console_handler.setFormatter(mlrun_formatter)
+mlrun_logger.addHandler(mlrun_console_handler)
+
+mlrun_default_artifact_template = 'mlrun_return_value_'
+mlrun_artifact_index = 0
+"""
+
 _databricks_script_code = """
 
 import argparse
@@ -168,16 +216,17 @@ parser.add_argument('handler_arguments')
 handler_arguments = parser.parse_args().handler_arguments
 handler_arguments = json.loads(handler_arguments)
 
-"""
 
-artifacts_code_template = """\n
-import logging
-mlrun_logger = logging.getLogger('mlrun_logger')
-mlrun_logger.setLevel(logging.DEBUG)
+"""
 
-def mlrun_log_artifact(name, path):
-    if not name or not path:
-        mlrun_logger.error(f'name and path required for logging an mlrun artifact - {{name}} : {{path}}')
+log_artifacts_code_template = """\n
+def mlrun_log_artifact(name='', path=''):
+    global mlrun_artifact_index
+    mlrun_artifact_index+=1  #  by how many artifacts we tried to log, not how many succeed.
+    if name is None or name == '':
+        name = f'{{mlrun_default_artifact_template}}{{mlrun_artifact_index}}'
+    if not path:
+        mlrun_logger.error(f'path required for logging an mlrun artifact - {{name}} : {{path}}')
         return
     if not isinstance(name, str) or not isinstance(path, str):
         mlrun_logger.error(f'name and path must be in string type for logging an mlrun artifact - {{name}} : {{path}}')
@@ -186,21 +235,38 @@ def mlrun_log_artifact(name, path):
         mlrun_logger.error(f'path for an mlrun artifact must start with /dbfs or dbfs:/ - {{name}} : {{path}}')
         return
     mlrun_artifacts_path = '{}'
-    import json
-    import os
-    new_data = {{name:path}}
-    if os.path.exists(mlrun_artifacts_path):
-        with open(mlrun_artifacts_path, 'r+') as json_file:
-            existing_data = json.load(json_file)
-            existing_data.update(new_data)
-            json_file.seek(0)
-            json.dump(existing_data, json_file)
-    else:
-        parent_dir = os.path.dirname(mlrun_artifacts_path)
-        if parent_dir != '/dbfs':
-            os.makedirs(parent_dir, exist_ok=True)
-        with open(mlrun_artifacts_path, 'w') as json_file:
-            json.dump(new_data, json_file)
-    mlrun_logger.info(f'successfully wrote artifact details to the artifact JSON file in DBFS - {{name}} : {{path}}')
+    try:
+        new_data = {{name:path}}
+        if os.path.exists(mlrun_artifacts_path):
+            with open(mlrun_artifacts_path, 'r+') as json_file:
+                existing_data = json.load(json_file)
+                existing_data.update(new_data)
+                json_file.seek(0)
+                json.dump(existing_data, json_file)
+        else:
+            parent_dir = os.path.dirname(mlrun_artifacts_path)
+            if parent_dir != '/dbfs':
+                os.makedirs(parent_dir, exist_ok=True)
+            with open(mlrun_artifacts_path, 'w') as json_file:
+                json.dump(new_data, json_file)
+        success_log = f'successfully wrote artifact details to the artifact JSON file in DBFS - {{name}} : {{path}}'
+        mlrun_logger.info(success_log)
+    except Exception as unknown_exception:
+        mlrun_logger.error(f'log mlrun artifact failed - {{name}} : {{path}}. error: {{unknown_exception}}')
 \n
 """
+
+_return_artifacts_code = """\n
+if result:
+    if isinstance(result, dict):
+        for key, path in result.items():
+            mlrun_log_artifact(name=key, path=path)
+    elif isinstance(result, (list, tuple, set)):
+        for artifact_path in result:
+            mlrun_log_artifact(path=artifact_path)
+    elif isinstance(result, str):
+        mlrun_log_artifact(path=result)
+    else:
+        mlrun_logger.warning(f'can not log artifacts with the result of handler function \
+- result in unsupported type. {type(result)}')
+"""

mlrun/runtimes/databricks_job/databricks_wrapper.py

@@ -200,7 +200,6 @@ def run_mlrun_databricks_job(
                 is_finished=True,
             )
         run_output = workspace.jobs.get_run_output(get_task(run).run_id)
-        context.log_result("databricks_runtime_task", run_output.as_dict())
     finally:
         workspace.dbfs.delete(script_path_on_dbfs)
         workspace.dbfs.delete(artifact_json_path)

mlrun/runtimes/function.py

@@ -550,6 +550,13 @@ class RemoteRuntime(KubeResource):
         """
         # todo: verify that the function name is normalized
 
+        old_http_session = getattr(self, "_http_session", None)
+        if old_http_session:
+            # ensure existing http session is terminated prior to (re)deploy to ensure that a connection to an old
+            # replica will not be reused
+            old_http_session.close()
+            self._http_session = None
+
         verbose = verbose or self.verbose
         if verbose:
             self.set_env("MLRUN_LOG_LEVEL", "DEBUG")
@@ -939,7 +946,7 @@ class RemoteRuntime(KubeResource):
                 http_client_kwargs["json"] = body
         try:
             logger.info("invoking function", method=method, path=path)
-            if not hasattr(self, "_http_session"):
+            if not getattr(self, "_http_session", None):
                 self._http_session = requests.Session()
             resp = self._http_session.request(
                 method, path, headers=headers, **http_client_kwargs

mlrun/runtimes/kubejob.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import time
+import warnings
 
 import mlrun.common.schemas
 import mlrun.db
@@ -130,7 +131,12 @@ class KubejobRuntime(KubeResource):
         :param builder_env: Kaniko builder pod env vars dict (for config/credentials)
             e.g. builder_env={"GIT_TOKEN": token}
         """
-
+        if not overwrite:
+            # TODO: change overwrite default to True in 1.8.0
+            warnings.warn(
+                "The `overwrite` parameter default will change from 'False' to 'True in 1.8.0.",
+                mlrun.utils.OverwriteBuildParamsWarning,
+            )
         image = mlrun.utils.helpers.remove_image_protocol_prefix(image)
         self.spec.build.build_config(
             image=image,

mlrun/serving/server.py

@@ -357,7 +357,7 @@ def v2_serving_init(context, namespace=None):
     if hasattr(context, "platform") and hasattr(
         context.platform, "set_termination_callback"
     ):
-        context.logger.debug(
+        context.logger.info(
             "Setting termination callback to terminate graph on worker shutdown"
         )
 
@@ -368,6 +368,23 @@ def v2_serving_init(context, namespace=None):
 
         context.platform.set_termination_callback(termination_callback)
 
+    if hasattr(context, "platform") and hasattr(context.platform, "set_drain_callback"):
+        context.logger.info(
+            "Setting drain callback to terminate and restart the graph on a drain event (such as rebalancing)"
+        )
+
+        def drain_callback():
+            context.logger.info("Drain callback called")
+            server.wait_for_completion()
+            context.logger.info(
+                "Termination of async flow is completed. Rerunning async flow."
+            )
+            # Rerun the flow without reconstructing it
+            server.graph._run_async_flow()
+            context.logger.info("Async flow restarted")
+
+        context.platform.set_drain_callback(drain_callback)
+
 
 def v2_serving_handler(context, event, get_body=False):
     """hook for nuclio handler()"""

mlrun/serving/states.py

@@ -921,6 +921,7 @@ class FlowStep(BaseStep):
 
         if self.engine != "sync":
             self._build_async_flow()
+            self._run_async_flow()
 
     def check_and_process_graph(self, allow_empty=False):
         """validate correct graph layout and initialize the .next links"""
@@ -1075,7 +1076,10 @@ class FlowStep(BaseStep):
                         if next_state.async_object and error_step.async_object:
                             error_step.async_object.to(next_state.async_object)
 
-        self._controller = source.run()
+        self._async_flow = source
+
+    def _run_async_flow(self):
+        self._controller = self._async_flow.run()
 
     def get_queue_links(self):
         """return dict of function and queue its listening on, for building stream triggers"""

mlrun/utils/helpers.py

@@ -65,6 +65,10 @@ DEFAULT_TIME_PARTITIONS = ["year", "month", "day", "hour"]
 DEFAULT_TIME_PARTITIONING_GRANULARITY = "hour"
 
 
+class OverwriteBuildParamsWarning(FutureWarning):
+    pass
+
+
 # TODO: remove in 1.7.0
 @deprecated(
     version="1.5.0",
@@ -183,7 +187,7 @@ def verify_field_regex(
             if mode == mlrun.common.schemas.RegexMatchModes.all:
                 if raise_on_failure:
                     raise mlrun.errors.MLRunInvalidArgumentError(
-                        f"Field '{field_name}' is malformed. {field_value} does not match required pattern: {pattern}"
+                        f"Field '{field_name}' is malformed. '{field_value}' does not match required pattern: {pattern}"
                     )
                 return False
         elif mode == mlrun.common.schemas.RegexMatchModes.any:
@@ -193,7 +197,7 @@ def verify_field_regex(
     elif mode == mlrun.common.schemas.RegexMatchModes.any:
         if raise_on_failure:
             raise mlrun.errors.MLRunInvalidArgumentError(
-                f"Field '{field_name}' is malformed. {field_value} does not match any of the"
+                f"Field '{field_name}' is malformed. '{field_value}' does not match any of the"
                 f" required patterns: {patterns}"
             )
         return False
@@ -333,7 +337,7 @@ def remove_image_protocol_prefix(image: str) -> str:
 def verify_field_of_type(field_name: str, field_value, expected_type: type):
     if not isinstance(field_value, expected_type):
         raise mlrun.errors.MLRunInvalidArgumentError(
-            f"Field '{field_name}' should be of type {expected_type.__name__} "
+            f"Field '{field_name}' should be of type '{expected_type.__name__}' "
             f"(got: {type(field_value).__name__} with value: {field_value})."
         )
 
@@ -357,14 +361,14 @@ def verify_dict_items_type(
     if dictionary:
         if type(dictionary) != dict:
             raise mlrun.errors.MLRunInvalidArgumentTypeError(
-                f"{name} expected to be of type dict, got type : {type(dictionary)}"
+                f"'{name}' expected to be of type dict, got type: {type(dictionary)}"
             )
         try:
             verify_list_items_type(dictionary.keys(), expected_keys_types)
             verify_list_items_type(dictionary.values(), expected_values_types)
         except mlrun.errors.MLRunInvalidArgumentTypeError as exc:
             raise mlrun.errors.MLRunInvalidArgumentTypeError(
-                f"{name} should be of type Dict[{get_pretty_types_names(expected_keys_types)},"
+                f"'{name}' should be of type Dict[{get_pretty_types_names(expected_keys_types)},"
                 f"{get_pretty_types_names(expected_values_types)}]."
             ) from exc
 
@@ -407,7 +411,7 @@ def normalize_name(name: str, verbose: bool = True):
         if verbose:
             warnings.warn(
                 "Names with underscore '_' are about to be deprecated, use dashes '-' instead. "
-                f"Replacing {name} underscores with dashes.",
+                f"Replacing '{name}' underscores with dashes.",
                 FutureWarning,
             )
         name = name.replace("_", "-")
@@ -669,7 +673,7 @@ def parse_artifact_uri(uri, default_project=""):
             iteration = int(iteration)
         except ValueError:
             raise ValueError(
-                f"illegal store path {uri}, iteration must be integer value"
+                f"illegal store path '{uri}', iteration must be integer value"
             )
     return (
         group_dict["project"] or default_project,
@@ -1199,7 +1203,7 @@ def get_function(function, namespace):
         function_object = create_function(function)
     except (ImportError, ValueError) as exc:
         raise ImportError(
-            f"state/function init failed, handler {function} not found"
+            f"state/function init failed, handler '{function}' not found"
         ) from exc
     return function_object
 
@@ -1383,7 +1387,7 @@ def get_in_artifact(artifact: dict, key, default=None, raise_on_missing=False):
 
         if raise_on_missing:
             raise mlrun.errors.MLRunInvalidArgumentError(
-                f"artifact {artifact} is missing metadata/spec/status"
+                f"artifact '{artifact}' is missing metadata/spec/status"
             )
         return default
 
@@ -1419,7 +1423,7 @@ def is_running_in_jupyter_notebook() -> bool:
 
 def as_number(field_name, field_value):
     if isinstance(field_value, str) and not field_value.isnumeric():
-        raise ValueError(f"{field_name} must be numeric (str/int types)")
+        raise ValueError(f"'{field_name}' must be numeric (str/int types)")
     return int(field_value)
 
 

mlrun/utils/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "420a54197d2ce4d81806ce83dbaeacf0e275bf2e",
-  "version": "1.6.0-rc11"
+  "git_commit": "abca3cad961e987ec75d8ecfcdf0b6856dc9b7d1",
+  "version": "1.6.0-rc13"
 }