Update function entry points when updating code (#451)

mlrun/run.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import typing
+import importlib.util as imputil
 import json
 import socket
+import typing
 import uuid
-from ast import literal_eval
 from base64 import b64decode
 from copy import deepcopy
 from os import environ, makedirs, path
@@ -26,14 +26,11 @@
 import yaml
 from kfp import Client
 from nuclio import build_file
-import importlib.util as imputil
 
-from .utils import retry_until_successful
 from .config import config as mlconf
 from .datastore import store_manager
 from .db import get_or_set_dburl, get_run_db
 from .execution import MLClientCtx
-from .funcdoc import find_handlers
 from .k8s_utils import get_k8s_helper
 from .model import RunObject, BaseMetadata, RunTemplate
 from .runtimes import (
@@ -43,7 +40,7 @@
     RuntimeKinds,
     get_runtime_class,
 )
-from .runtimes.base import FunctionEntrypoint
+from .runtimes.funcdoc import update_function_entry_points
 from .runtimes.utils import add_code_metadata, global_context
 from .utils import (
     get_in,
@@ -53,6 +50,7 @@
     new_pipe_meta,
     extend_hub_uri,
 )
+from .utils import retry_until_successful
 
 
 class RunStatuses(object):
@@ -628,8 +626,7 @@ def update_meta(fn):
             update_in(spec, "kind", "Function")
             r.spec.base_spec = spec
             if with_doc:
-                handlers = find_handlers(code)
-                r.spec.entry_points = {h["name"]: as_func(h) for h in handlers}
+                update_function_entry_points(r, code)
         else:
             r.spec.source = filename
             r.spec.function_handler = handler
@@ -680,8 +677,7 @@ def update_meta(fn):
             r.spec.volume_mounts.append(vol.get("volumeMount"))
 
     if with_doc:
-        handlers = find_handlers(code)
-        r.spec.entry_points = {h["name"]: as_func(h) for h in handlers}
+        update_function_entry_points(r, code)
     r.spec.default_handler = handler
     update_meta(r)
     return r
@@ -918,33 +914,6 @@ def list_piplines(
     return resp.total_size, resp.next_page_token, runs
 
 
-def as_func(handler):
-    ret = clean(handler["return"])
-    return FunctionEntrypoint(
-        name=handler["name"],
-        doc=handler["doc"],
-        parameters=[clean(p) for p in handler["params"]],
-        outputs=[ret] if ret else None,
-        lineno=handler["lineno"],
-    ).to_dict()
-
-
-def clean(struct: dict):
-    if not struct:
-        return None
-    if "default" in struct:
-        struct["default"] = py_eval(struct["default"])
-    return {k: v for k, v in struct.items() if v or k == "default"}
-
-
-def py_eval(data):
-    try:
-        value = literal_eval(data)
-        return value
-    except (SyntaxError, ValueError):
-        return data
-
-
 def get_object(url, secrets=None, size=None, offset=0, db=None):
     """get mlrun dataitem body (from path/url)"""
     stores = store_manager.set(secrets, db=db)

mlrun/funcdoc.py → mlrun/runtimes/funcdoc.py

@@ -16,6 +16,8 @@
 import inspect
 import re
 
+from mlrun.runtimes.base import FunctionEntrypoint
+
 
 def type_name(ann):
     if ann is inspect.Signature.empty:
@@ -242,6 +244,40 @@ def find_handlers(code: str, handlers=None):
         return filter_funcs(funcs, markers)
 
 
+def as_func(handler):
+    ret = clean(handler["return"])
+    return FunctionEntrypoint(
+        name=handler["name"],
+        doc=handler["doc"],
+        parameters=[clean(p) for p in handler["params"]],
+        outputs=[ret] if ret else None,
+        lineno=handler["lineno"],
+    ).to_dict()
+
+
+def update_function_entry_points(function, code):
+    handlers = find_handlers(code)
+    function.spec.entry_points = {
+        handler["name"]: as_func(handler) for handler in handlers
+    }
+
+
+def clean(struct: dict):
+    if not struct:
+        return None
+    if "default" in struct:
+        struct["default"] = py_eval(struct["default"])
+    return {k: v for k, v in struct.items() if v or k == "default"}
+
+
+def py_eval(data):
+    try:
+        value = ast.literal_eval(data)
+        return value
+    except (SyntaxError, ValueError):
+        return data
+
+
 def filter_funcs(funcs, markers):
     markers = list(markers)
     if not markers:

mlrun/runtimes/kubejob.py

@@ -20,6 +20,7 @@
 
 from mlrun.runtimes.base import BaseRuntimeHandler
 from .base import RunError
+from .funcdoc import update_function_entry_points
 from .pod import KubeResource
 from .utils import AsyncLogWriter, default_image_name
 from ..builder import build_runtime
@@ -35,22 +36,29 @@ class KubejobRuntime(KubeResource):
 
     _is_remote = True
 
-    def with_code(self, from_file="", body=None):
+    def with_code(self, from_file="", body=None, with_doc=True):
+        """Update the function code
+        This function eliminates the need to build container images every time we edit the code
+
+        :param from_file:   blank for current notebook, or path to .py/.ipynb file
+        :param body:        will use the body as the function code
+        :param with_doc:    update the document of the function parameters
+
+        :return: function object
+        """
         if (not body and not from_file) or (from_file and from_file.endswith(".ipynb")):
             from nuclio import build_file
 
-            name, spec, code = build_file(from_file)
-            self.spec.build.functionSourceCode = get_in(
-                spec, "spec.build.functionSourceCode"
-            )
-            return self
+            _, _, body = build_file(from_file)
 
         if from_file:
             with open(from_file) as fp:
                 body = fp.read()
         self.spec.build.functionSourceCode = b64encode(body.encode("utf-8")).decode(
             "utf-8"
         )
+        if with_doc:
+            update_function_entry_points(self, body)
         return self
 
     @property

tests/test_funcdoc.py → tests/runtimes/test_funcdoc.py

@@ -18,8 +18,8 @@
 import pytest
 import yaml
 
+from mlrun.runtimes import funcdoc
 from tests.conftest import here
-from mlrun import funcdoc
 
 
 def load_rst_cases(name):