mlrun 1.3.3rc1__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/{builder.py → api/utils/builder.py}

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,38 +11,36 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import os.path
 import pathlib
 import re
-import tarfile
 import tempfile
+import typing
 from base64 import b64decode, b64encode
 from os import path
 from urllib.parse import urlparse
 
 from kubernetes import client
 
-import mlrun.api.schemas
+import mlrun.api.utils.singletons.k8s
+import mlrun.common.constants
+import mlrun.common.schemas
 import mlrun.errors
 import mlrun.runtimes.utils
-
-from .config import config
-from .datastore import store_manager
-from .k8s_utils import BasePod, get_k8s_helper
-from .utils import enrich_image_url, get_parsed_docker_registry, logger, normalize_name
-
-IMAGE_NAME_ENRICH_REGISTRY_PREFIX = "."
+import mlrun.utils
+from mlrun.config import config
+from mlrun.utils.helpers import remove_image_protocol_prefix
 
 
 def make_dockerfile(
-    base_image,
-    commands=None,
-    source=None,
-    requirements=None,
-    workdir="/mlrun",
-    extra="",
-    user_unix_id=None,
-    enriched_group_id=None,
+    base_image: str,
+    commands: list = None,
+    source: str = None,
+    requirements_path: str = None,
+    workdir: str = "/mlrun",
+    extra: str = "",
+    user_unix_id: int = None,
+    enriched_group_id: int = None,
 ):
     dock = f"FROM {base_image}\n"
 
@@ -55,20 +53,20 @@ def make_dockerfile(
         dock += f"ARG {build_arg_key}={build_arg_value}\n"
 
     if source:
-        dock += f"RUN mkdir -p {workdir}\n"
         dock += f"WORKDIR {workdir}\n"
         # 'ADD' command does not extract zip files - add extraction stage to the dockerfile
         if source.endswith(".zip"):
+            source_dir = os.path.join(workdir, "source")
             stage1 = f"""
             FROM {base_image} AS extractor
             RUN apt-get update -qqy && apt install --assume-yes unzip
-            RUN mkdir -p /source
-            COPY {source} /source
-            RUN cd /source && unzip {source} && rm {source}
+            RUN mkdir -p {source_dir}
+            COPY {source} {source_dir}
+            RUN cd {source_dir} && unzip {source} && rm {source}
             """
             dock = stage1 + "\n" + dock
 
-            dock += f"COPY --from=extractor /source/ {workdir}\n"
+            dock += f"COPY --from=extractor {source_dir}/ {workdir}\n"
         else:
             dock += f"ADD {source} {workdir}\n"
 
@@ -76,13 +74,16 @@ def make_dockerfile(
             dock += f"RUN chown -R {user_unix_id}:{enriched_group_id} {workdir}\n"
 
         dock += f"ENV PYTHONPATH {workdir}\n"
-    if requirements:
-        dock += f"RUN python -m pip install -r {requirements}\n"
     if commands:
         dock += "".join([f"RUN {command}\n" for command in commands])
+    if requirements_path:
+        dock += (
+            f"RUN echo 'Installing {requirements_path}...'; cat {requirements_path}\n"
+        )
+        dock += f"RUN python -m pip install -r {requirements_path}\n"
     if extra:
         dock += extra
-    logger.debug("Resolved dockerfile", dockfile_contents=dock)
+    mlrun.utils.logger.debug("Resolved dockerfile", dockfile_contents=dock)
     return dock
 
 
@@ -95,6 +96,7 @@ def make_kaniko_pod(
     inline_code=None,
     inline_path=None,
     requirements=None,
+    requirements_path=None,
     secret_name=None,
     name="",
     verbose=False,
@@ -134,7 +136,16 @@ def make_kaniko_pod(
     if dockertext:
         dockerfile = "/empty/Dockerfile"
 
-    args = ["--dockerfile", dockerfile, "--context", context, "--destination", dest]
+    args = [
+        "--dockerfile",
+        dockerfile,
+        "--context",
+        context,
+        "--destination",
+        dest,
+        "--image-fs-extract-retry",
+        config.httpdb.builder.kaniko_image_fs_extraction_retries,
+    ]
     for value, flag in [
         (config.httpdb.builder.insecure_pull_registry_mode, "--insecure-pull"),
         (config.httpdb.builder.insecure_push_registry_mode, "--insecure"),
@@ -158,7 +169,7 @@ def make_kaniko_pod(
             mem=default_requests.get("memory"), cpu=default_requests.get("cpu")
         )
     }
-    kpod = BasePod(
+    kpod = mlrun.api.utils.singletons.k8s.BasePod(
         name or "mlrun-build",
         config.httpdb.builder.kaniko_image,
         args=args,
@@ -193,19 +204,23 @@ def make_kaniko_pod(
         commands = []
         env = {}
         if dockertext:
-            commands.append("echo ${DOCKERFILE} | base64 -d > /empty/Dockerfile")
+            # set and encode docker content to the DOCKERFILE environment variable in the kaniko pod
             env["DOCKERFILE"] = b64encode(dockertext.encode("utf-8")).decode("utf-8")
+            # dump dockerfile content and decode to Dockerfile destination
+            commands.append("echo ${DOCKERFILE} | base64 -d > /empty/Dockerfile")
         if inline_code:
             name = inline_path or "main.py"
-            commands.append("echo ${CODE} | base64 -d > /empty/" + name)
             env["CODE"] = b64encode(inline_code.encode("utf-8")).decode("utf-8")
+            commands.append("echo ${CODE} | base64 -d > /empty/" + name)
         if requirements:
-            commands.append(
-                "echo ${REQUIREMENTS} | base64 -d > /empty/requirements.txt"
-            )
+            # set and encode requirements to the REQUIREMENTS environment variable in the kaniko pod
             env["REQUIREMENTS"] = b64encode(
                 "\n".join(requirements).encode("utf-8")
             ).decode("utf-8")
+            # dump requirement content and decode to the requirement.txt destination
+            commands.append(
+                "echo ${REQUIREMENTS}" + " | " + f"base64 -d > {requirements_path}"
+            )
 
         kpod.append_init_container(
             config.httpdb.builder.kaniko_init_container_image,
@@ -221,7 +236,21 @@ def make_kaniko_pod(
         if end == -1:
             end = len(dest)
         repo = dest[dest.find("/") + 1 : end]
-        configure_kaniko_ecr_init_container(kpod, registry, repo)
+
+        # if no secret is given, assume ec2 instance has attached role which provides read/write access to ECR
+        assume_instance_role = not config.httpdb.builder.docker_registry_secret
+        configure_kaniko_ecr_init_container(kpod, registry, repo, assume_instance_role)
+
+        # project secret might conflict with the attached instance role
+        # ensure "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY" have no values or else kaniko will fail
+        # due to credentials conflict / lack of permission on given credentials
+        if assume_instance_role:
+            kpod.pod.spec.containers[0].env.extend(
+                [
+                    client.V1EnvVar(name="AWS_ACCESS_KEY_ID", value=""),
+                    client.V1EnvVar(name="AWS_SECRET_ACCESS_KEY", value=""),
+                ]
+            )
 
     # mount regular docker config secret
     elif secret_name:
@@ -231,7 +260,9 @@ def make_kaniko_pod(
     return kpod
 
 
-def configure_kaniko_ecr_init_container(kpod, registry, repo):
+def configure_kaniko_ecr_init_container(
+    kpod, registry, repo, assume_instance_role=True
+):
     region = registry.split(".")[3]
 
     # fail silently in order to ignore "repository already exists" errors
@@ -242,12 +273,13 @@ def configure_kaniko_ecr_init_container(kpod, registry, repo):
     )
     init_container_env = {}
 
-    if not config.httpdb.builder.docker_registry_secret:
+    if assume_instance_role:
 
         # assume instance role has permissions to register and store a container image
         # https://github.com/GoogleContainerTools/kaniko#pushing-to-amazon-ecr
         # we only need this in the kaniko container
         kpod.env.append(client.V1EnvVar(name="AWS_SDK_LOAD_CONFIG", value="true"))
+
     else:
         aws_credentials_file_env_key = "AWS_SHARED_CREDENTIALS_FILE"
         aws_credentials_file_env_value = "/tmp/credentials"
@@ -278,24 +310,12 @@ def configure_kaniko_ecr_init_container(kpod, registry, repo):
     )
 
 
-def upload_tarball(source_dir, target, secrets=None):
-
-    # will delete the temp file
-    with tempfile.NamedTemporaryFile(suffix=".tar.gz") as temp_fh:
-        with tarfile.open(mode="w:gz", fileobj=temp_fh) as tar:
-            tar.add(source_dir, arcname="")
-        stores = store_manager.set(secrets)
-        datastore, subpath = stores.get_or_create_store(target)
-        datastore.upload(subpath, temp_fh.name)
-
-
 def build_image(
-    auth_info: mlrun.api.schemas.AuthInfo,
+    auth_info: mlrun.common.schemas.AuthInfo,
     project: str,
     image_target,
     commands=None,
     source="",
-    mounter="v3io",
     base_image=None,
     requirements=None,
     inline_code=None,
@@ -315,43 +335,24 @@ def build_image(
 ):
     runtime_spec = runtime.spec if runtime else None
     builder_env = builder_env or {}
-    image_target, secret_name = _resolve_image_target_and_registry_secret(
+    image_target, secret_name = resolve_image_target_and_registry_secret(
         image_target, registry, secret_name
     )
 
-    if isinstance(requirements, list):
-        requirements_list = requirements
-        requirements_path = "requirements.txt"
-        if source:
-            raise ValueError("requirements list only works with inline code")
-    else:
-        requirements_list = None
-        requirements_path = requirements
-
-    commands = commands or []
-    if with_mlrun:
-        # mlrun prerequisite - upgrade pip
-        upgrade_pip_command = resolve_upgrade_pip_command(commands)
-        if upgrade_pip_command:
-            commands.append(upgrade_pip_command)
-
-        mlrun_command = resolve_mlrun_install_command(
-            mlrun_version_specifier, client_version, commands
-        )
-        if mlrun_command:
-            commands.append(mlrun_command)
+    commands, requirements_list, requirements_path = _resolve_build_requirements(
+        requirements, commands, with_mlrun, mlrun_version_specifier, client_version
+    )
 
-    if not inline_code and not source and not commands:
-        logger.info("skipping build, nothing to add")
+    if not inline_code and not source and not commands and not requirements:
+        mlrun.utils.logger.info("skipping build, nothing to add")
         return "skipped"
 
     context = "/context"
     to_mount = False
-    v3io = (
-        source.startswith("v3io://") or source.startswith("v3ios://")
-        if source
-        else None
-    )
+    is_v3io_source = False
+    if source:
+        is_v3io_source = source.startswith("v3io://") or source.startswith("v3ios://")
+
     access_key = builder_env.get(
         "V3IO_ACCESS_KEY", auth_info.data_session or auth_info.access_key
     )
@@ -365,7 +366,8 @@ def build_image(
     if inline_code or runtime_spec.build.load_source_on_run or not source:
         context = "/empty"
 
-    elif source and "://" in source and not v3io:
+    # source is remote
+    elif source and "://" in source and not is_v3io_source:
         if source.startswith("git://"):
             # if the user provided branch (w/o refs/..) we add the "refs/.."
             fragment = parsed_url.fragment or ""
@@ -376,19 +378,34 @@ def build_image(
         context = source
         source_to_copy = "."
 
+    # source is local / v3io
     else:
-        if v3io:
+        if is_v3io_source:
             source = parsed_url.path
             to_mount = True
             source_dir_to_mount, source_to_copy = path.split(source)
-        else:
+
+        # source is a path without a scheme, we allow to copy absolute paths assuming they are valid paths
+        # in the image, however, it is recommended to use `workdir` instead in such cases
+        # which is set during runtime (mlrun.runtimes.local.LocalRuntime._pre_run).
+        # relative paths are not supported at build time
+        # "." and "./" are considered as 'project context'
+        # TODO: enrich with project context if pulling on build time
+        elif path.isabs(source):
             source_to_copy = source
 
+        else:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Load of relative source ({source}) is not supported at build time "
+                "see 'mlrun.runtimes.kubejob.KubejobRuntime.with_source_archive' or "
+                "'mlrun.projects.project.MlrunProject.set_source' for more details"
+            )
+
     user_unix_id = None
     enriched_group_id = None
     if (
         mlrun.mlconf.function.spec.security_context.enrichment_mode
-        != mlrun.api.schemas.SecurityContextEnrichmentModes.disabled.value
+        != mlrun.common.schemas.SecurityContextEnrichmentModes.disabled.value
     ):
         from mlrun.api.api.utils import ensure_function_security_context
 
@@ -396,14 +413,29 @@ def build_image(
         user_unix_id = runtime.spec.security_context.run_as_user
         enriched_group_id = runtime.spec.security_context.run_as_group
 
+    if source_to_copy and (
+        not runtime.spec.clone_target_dir
+        or not os.path.isabs(runtime.spec.clone_target_dir)
+    ):
+        # use a temp dir for permissions and set it as the workdir
+        tmpdir = tempfile.mkdtemp()
+        relative_workdir = runtime.spec.clone_target_dir or ""
+        if relative_workdir.startswith("./"):
+            # TODO: use 'removeprefix' when we drop python 3.7 support
+            # relative_workdir.removeprefix("./")
+            relative_workdir = relative_workdir[2:]
+
+        runtime.spec.clone_target_dir = path.join(tmpdir, "mlrun", relative_workdir)
+
     dock = make_dockerfile(
         base_image,
         commands,
         source=source_to_copy,
-        requirements=requirements_path,
+        requirements_path=requirements_path,
         extra=extra,
         user_unix_id=user_unix_id,
         enriched_group_id=enriched_group_id,
+        workdir=runtime.spec.clone_target_dir,
     )
 
     kpod = make_kaniko_pod(
@@ -414,6 +446,7 @@ def build_image(
         inline_code=inline_code,
         inline_path=inline_path,
         requirements=requirements_list,
+        requirements_path=requirements_path,
         secret_name=secret_name,
         name=name,
         verbose=verbose,
@@ -430,14 +463,16 @@ def build_image(
             user=username,
         )
 
-    k8s = get_k8s_helper()
+    k8s = mlrun.api.utils.singletons.k8s.get_k8s_helper(silent=False)
     kpod.namespace = k8s.resolve_namespace(namespace)
 
     if interactive:
         return k8s.run_job(kpod)
     else:
         pod, ns = k8s.create_pod(kpod)
-        logger.info(f'started build, to watch build logs use "mlrun watch {pod} {ns}"')
+        mlrun.utils.logger.info(
+            "Build started", pod=pod, namespace=ns, project=project, image=image_target
+        )
         return f"build:{pod}"
 
 
@@ -453,7 +488,7 @@ def get_kaniko_spec_attributes_from_runtime():
     ]
 
 
-def resolve_mlrun_install_command(
+def resolve_mlrun_install_command_version(
     mlrun_version_specifier=None, client_version=None, commands=None
 ):
     commands = commands or []
@@ -483,7 +518,7 @@ def resolve_mlrun_install_command(
             mlrun_version_specifier = (
                 f"{config.package_path}[complete]=={config.version}"
             )
-    return f'python -m pip install "{mlrun_version_specifier}"'
+    return mlrun_version_specifier
 
 
 def resolve_upgrade_pip_command(commands=None):
@@ -497,7 +532,7 @@ def resolve_upgrade_pip_command(commands=None):
 
 
 def build_runtime(
-    auth_info: mlrun.api.schemas.AuthInfo,
+    auth_info: mlrun.common.schemas.AuthInfo,
     runtime,
     with_mlrun=True,
     mlrun_version_specifier=None,
@@ -511,7 +546,7 @@ def build_runtime(
     namespace = runtime.metadata.namespace
     project = runtime.metadata.project
     if skip_deployed and runtime.is_deployed():
-        runtime.status.state = mlrun.api.schemas.FunctionState.ready
+        runtime.status.state = mlrun.common.schemas.FunctionState.ready
         return True
     if build.base_image:
         mlrun_images = [
@@ -523,7 +558,13 @@ def build_runtime(
         # if the base is one of mlrun images - no need to install mlrun
         if any([image in build.base_image for image in mlrun_images]):
             with_mlrun = False
-    if not build.source and not build.commands and not build.extra and not with_mlrun:
+    if (
+        not build.source
+        and not build.commands
+        and not build.requirements
+        and not build.extra
+        and not with_mlrun
+    ):
         if not runtime.spec.image:
             if build.base_image:
                 runtime.spec.image = build.base_image
@@ -536,7 +577,7 @@ def build_runtime(
                 "The deployment was not successful because no image was specified or there are missing build parameters"
                 " (commands/source)"
             )
-        runtime.status.state = mlrun.api.schemas.FunctionState.ready
+        runtime.status.state = mlrun.common.schemas.FunctionState.ready
         return True
 
     build.image = mlrun.runtimes.utils.resolve_function_image_name(runtime, build.image)
@@ -549,17 +590,23 @@ def build_runtime(
         raise mlrun.errors.MLRunInvalidArgumentError(
             "build spec must have a target image, set build.image = <target image>"
         )
-    logger.info(f"building image ({build.image})")
+    name = mlrun.utils.normalize_name(f"mlrun-build-{runtime.metadata.name}")
 
-    name = normalize_name(f"mlrun-build-{runtime.metadata.name}")
     base_image: str = (
         build.base_image or runtime.spec.image or config.default_base_image
     )
-    enriched_base_image = enrich_image_url(
+    enriched_base_image = mlrun.utils.enrich_image_url(
         base_image,
         client_version,
         client_python_version,
     )
+    mlrun.utils.logger.info(
+        "Building runtime image",
+        base_image=enriched_base_image,
+        image=build.image,
+        project=project,
+        name=name,
+    )
 
     status = build_image(
         auth_info,
@@ -567,8 +614,8 @@ def build_runtime(
         image_target=build.image,
         base_image=enriched_base_image,
         commands=build.commands,
+        requirements=build.requirements,
         namespace=namespace,
-        # inline_code=inline,
         source=build.source,
         secret_name=build.secret,
         interactive=interactive,
@@ -586,11 +633,11 @@ def build_runtime(
         # using enriched base image for the runtime spec image, because this will be the image that the function will
         # run with
         runtime.spec.image = enriched_base_image
-        runtime.status.state = mlrun.api.schemas.FunctionState.ready
+        runtime.status.state = mlrun.common.schemas.FunctionState.ready
         return True
 
     if status.startswith("build:"):
-        runtime.status.state = mlrun.api.schemas.FunctionState.deploying
+        runtime.status.state = mlrun.common.schemas.FunctionState.deploying
         runtime.status.build_pod = status[6:]
         # using the base_image, and not the enriched one so we won't have the client version in the image, useful for
         # exports and other cases where we don't want to have the client version in the image, but rather enriched on
@@ -598,48 +645,34 @@ def build_runtime(
         runtime.spec.build.base_image = base_image
         return False
 
-    logger.info(f"build completed with {status}")
+    mlrun.utils.logger.info(f"build completed with {status}")
     if status in ["failed", "error"]:
-        runtime.status.state = mlrun.api.schemas.FunctionState.error
+        runtime.status.state = mlrun.common.schemas.FunctionState.error
         return False
 
     local = "" if build.secret or build.image.startswith(".") else "."
     runtime.spec.image = local + build.image
-    runtime.status.state = mlrun.api.schemas.FunctionState.ready
+    runtime.status.state = mlrun.common.schemas.FunctionState.ready
     return True
 
 
-def _generate_builder_env(project, builder_env):
-    k8s = get_k8s_helper()
-    secret_name = k8s.get_project_secret_name(project)
-    existing_secret_keys = k8s.get_project_secret_keys(project, filter_internal=True)
-
-    # generate env list from builder env and project secrets
-    env = []
-    for key in existing_secret_keys:
-        if key not in builder_env:
-            value_from = client.V1EnvVarSource(
-                secret_key_ref=client.V1SecretKeySelector(name=secret_name, key=key)
-            )
-            env.append(client.V1EnvVar(name=key, value_from=value_from))
-    for key, value in builder_env.items():
-        env.append(client.V1EnvVar(name=key, value=value))
-    return env
-
-
-def _resolve_image_target_and_registry_secret(
+def resolve_image_target_and_registry_secret(
     image_target: str, registry: str = None, secret_name: str = None
 ) -> (str, str):
     if registry:
         return "/".join([registry, image_target]), secret_name
 
     # if dest starts with a dot, we add the configured registry to the start of the dest
-    if image_target.startswith(IMAGE_NAME_ENRICH_REGISTRY_PREFIX):
+    if image_target.startswith(
+        mlrun.common.constants.IMAGE_NAME_ENRICH_REGISTRY_PREFIX
+    ):
 
         # remove prefix from image name
-        image_target = image_target[len(IMAGE_NAME_ENRICH_REGISTRY_PREFIX) :]
+        image_target = image_target[
+            len(mlrun.common.constants.IMAGE_NAME_ENRICH_REGISTRY_PREFIX) :
+        ]
 
-        registry, repository = get_parsed_docker_registry()
+        registry, repository = mlrun.utils.get_parsed_docker_registry()
         secret_name = secret_name or config.httpdb.builder.docker_registry_secret
         if not registry:
             raise ValueError(
@@ -652,4 +685,64 @@ def _resolve_image_target_and_registry_secret(
 
         return "/".join(image_target_components), secret_name
 
+    image_target = remove_image_protocol_prefix(image_target)
+
     return image_target, secret_name
+
+
+def _generate_builder_env(project, builder_env):
+    k8s = mlrun.api.utils.singletons.k8s.get_k8s_helper(silent=False)
+    secret_name = k8s.get_project_secret_name(project)
+    existing_secret_keys = k8s.get_project_secret_keys(project, filter_internal=True)
+
+    # generate env list from builder env and project secrets
+    env = []
+    for key in existing_secret_keys:
+        if key not in builder_env:
+            value_from = client.V1EnvVarSource(
+                secret_key_ref=client.V1SecretKeySelector(name=secret_name, key=key)
+            )
+            env.append(client.V1EnvVar(name=key, value_from=value_from))
+    for key, value in builder_env.items():
+        env.append(client.V1EnvVar(name=key, value=value))
+    return env
+
+
+def _resolve_build_requirements(
+    requirements: typing.Union[typing.List, str],
+    commands: typing.List,
+    with_mlrun: bool,
+    mlrun_version_specifier: typing.Optional[str],
+    client_version: typing.Optional[str],
+):
+    """
+    Resolve build requirements list, requirements path and commands.
+    If mlrun requirement is needed, we add a pip upgrade command to the commands list (prerequisite).
+    """
+    requirements_path = "/empty/requirements.txt"
+    if requirements and isinstance(requirements, list):
+        requirements_list = requirements
+    else:
+        requirements_list = []
+        requirements_path = requirements or requirements_path
+    commands = commands or []
+
+    if with_mlrun:
+        # mlrun prerequisite - upgrade pip
+        upgrade_pip_command = resolve_upgrade_pip_command(commands)
+        if upgrade_pip_command:
+            commands.append(upgrade_pip_command)
+
+        mlrun_version = resolve_mlrun_install_command_version(
+            mlrun_version_specifier, client_version, commands
+        )
+
+        # mlrun must be installed with other python requirements in the same pip command to avoid version conflicts
+        if mlrun_version:
+            requirements_list.insert(0, mlrun_version)
+
+    if not requirements_list:
+        # no requirements, we don't need a requirements file
+        requirements_path = ""
+
+    return commands, requirements_list, requirements_path

mlrun/api/utils/clients/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mlrun/api/utils/clients/chief.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -21,8 +21,8 @@ import urllib.parse
 import aiohttp
 import fastapi
 
-import mlrun.api.schemas
 import mlrun.api.utils.projects.remotes.follower
+import mlrun.common.schemas
 import mlrun.errors
 import mlrun.utils.singleton
 from mlrun.utils import logger
@@ -157,7 +157,7 @@ class Client(
 
     async def get_clusterization_spec(
         self, return_fastapi_response: bool = True, raise_on_failure: bool = False
-    ) -> typing.Union[fastapi.Response, mlrun.api.schemas.ClusterizationSpec]:
+    ) -> typing.Union[fastapi.Response, mlrun.common.schemas.ClusterizationSpec]:
         """
         This method is used both for proxying requests from worker to chief and for aligning the worker state
         with the clusterization spec brought from the chief
@@ -172,7 +172,22 @@ class Client(
                     chief_response
                 )
 
-            return mlrun.api.schemas.ClusterizationSpec(**(await chief_response.json()))
+            return mlrun.common.schemas.ClusterizationSpec(
+                **(await chief_response.json())
+            )
+
+    async def set_schedule_notifications(
+        self, project: str, schedule_name: str, request: fastapi.Request, json: dict
+    ) -> fastapi.Response:
+        """
+        Schedules are running only on chief
+        """
+        return await self._proxy_request_to_chief(
+            "PUT",
+            f"projects/{project}/schedules/{schedule_name}/notifications",
+            request,
+            json,
+        )
 
     async def _proxy_request_to_chief(
         self,