mlrun 1.4.0rc25__py3-none-any.whl → 1.5.0rc2__py3-none-any.whl

mlrun/api/runtime_handlers/daskjob.py

@@ -0,0 +1,209 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 typing
+from typing import Dict, List, Optional, Union
+
+from kubernetes.client.rest import ApiException
+from sqlalchemy.orm import Session
+
+import mlrun.common.schemas
+import mlrun.errors
+import mlrun.k8s_utils
+import mlrun.utils
+import mlrun.utils.regex
+from mlrun.api.db.base import DBInterface
+from mlrun.api.runtime_handlers.base import BaseRuntimeHandler
+from mlrun.config import config
+from mlrun.runtimes.base import RuntimeClassMode
+from mlrun.runtimes.utils import get_k8s
+from mlrun.utils import logger
+
+
+class DaskRuntimeHandler(BaseRuntimeHandler):
+    kind = "dask"
+    class_modes = {RuntimeClassMode.run: "dask"}
+
+    # Dask runtime resources are per function (and not per run).
+    # It means that monitoring runtime resources state doesn't say anything about the run state.
+    # Therefore dask run monitoring is done completely by the SDK, so overriding the monitoring method with no logic
+    def monitor_runs(
+        self, db: DBInterface, db_session: Session, leader_session: Optional[str] = None
+    ):
+        return
+
+    @staticmethod
+    def _get_object_label_selector(object_id: str) -> str:
+        return f"mlrun/function={object_id}"
+
+    @staticmethod
+    def resolve_object_id(
+        run: dict,
+    ) -> typing.Optional[str]:
+        """
+        Resolves the object ID from the run object.
+        In dask runtime, the object ID is the function name.
+        :param run: run object
+        :return: function name
+        """
+
+        function = run.get("spec", {}).get("function", None)
+        if function:
+
+            # a dask run's function field is in the format <project-name>/<function-name>@<run-uid>
+            # we only want the function name
+            project_and_function = function.split("@")[0]
+            return project_and_function.split("/")[-1]
+
+        return None
+
+    def _enrich_list_resources_response(
+        self,
+        response: Union[
+            mlrun.common.schemas.RuntimeResources,
+            mlrun.common.schemas.GroupedByJobRuntimeResourcesOutput,
+            mlrun.common.schemas.GroupedByProjectRuntimeResourcesOutput,
+        ],
+        namespace: str,
+        label_selector: str = None,
+        group_by: Optional[
+            mlrun.common.schemas.ListRuntimeResourcesGroupByField
+        ] = None,
+    ) -> Union[
+        mlrun.common.schemas.RuntimeResources,
+        mlrun.common.schemas.GroupedByJobRuntimeResourcesOutput,
+        mlrun.common.schemas.GroupedByProjectRuntimeResourcesOutput,
+    ]:
+        """
+        Handling listing service resources
+        """
+        enrich_needed = self._validate_if_enrich_is_needed_by_group_by(group_by)
+        if not enrich_needed:
+            return response
+        services = get_k8s().v1api.list_namespaced_service(
+            namespace, label_selector=label_selector
+        )
+        service_resources = []
+        for service in services.items:
+            service_resources.append(
+                mlrun.common.schemas.RuntimeResource(
+                    name=service.metadata.name, labels=service.metadata.labels
+                )
+            )
+        return self._enrich_service_resources_in_response(
+            response, service_resources, group_by
+        )
+
+    def _build_output_from_runtime_resources(
+        self,
+        response: Union[
+            mlrun.common.schemas.RuntimeResources,
+            mlrun.common.schemas.GroupedByJobRuntimeResourcesOutput,
+            mlrun.common.schemas.GroupedByProjectRuntimeResourcesOutput,
+        ],
+        runtime_resources_list: List[mlrun.common.schemas.RuntimeResources],
+        group_by: Optional[
+            mlrun.common.schemas.ListRuntimeResourcesGroupByField
+        ] = None,
+    ):
+        enrich_needed = self._validate_if_enrich_is_needed_by_group_by(group_by)
+        if not enrich_needed:
+            return response
+        service_resources = []
+        for runtime_resources in runtime_resources_list:
+            if runtime_resources.service_resources:
+                service_resources += runtime_resources.service_resources
+        return self._enrich_service_resources_in_response(
+            response, service_resources, group_by
+        )
+
+    def _validate_if_enrich_is_needed_by_group_by(
+        self,
+        group_by: Optional[
+            mlrun.common.schemas.ListRuntimeResourcesGroupByField
+        ] = None,
+    ) -> bool:
+        # Dask runtime resources are per function (and not per job) therefore, when grouping by job we're simply
+        # omitting the dask runtime resources
+        if group_by == mlrun.common.schemas.ListRuntimeResourcesGroupByField.job:
+            return False
+        elif group_by == mlrun.common.schemas.ListRuntimeResourcesGroupByField.project:
+            return True
+        elif group_by is not None:
+            raise NotImplementedError(
+                f"Provided group by field is not supported. group_by={group_by}"
+            )
+        return True
+
+    def _enrich_service_resources_in_response(
+        self,
+        response: Union[
+            mlrun.common.schemas.RuntimeResources,
+            mlrun.common.schemas.GroupedByJobRuntimeResourcesOutput,
+            mlrun.common.schemas.GroupedByProjectRuntimeResourcesOutput,
+        ],
+        service_resources: List[mlrun.common.schemas.RuntimeResource],
+        group_by: Optional[
+            mlrun.common.schemas.ListRuntimeResourcesGroupByField
+        ] = None,
+    ):
+        if group_by == mlrun.common.schemas.ListRuntimeResourcesGroupByField.project:
+            for service_resource in service_resources:
+                self._add_resource_to_grouped_by_project_resources_response(
+                    response, "service_resources", service_resource
+                )
+        else:
+            response.service_resources = service_resources
+        return response
+
+    def _delete_extra_resources(
+        self,
+        db: DBInterface,
+        db_session: Session,
+        namespace: str,
+        deleted_resources: List[Dict],
+        label_selector: str = None,
+        force: bool = False,
+        grace_period: int = None,
+    ):
+        """
+        Handling services deletion
+        """
+        if grace_period is None:
+            grace_period = config.runtime_resources_deletion_grace_period
+        service_names = []
+        for pod_dict in deleted_resources:
+            dask_component = (
+                pod_dict["metadata"].get("labels", {}).get("dask.org/component")
+            )
+            cluster_name = (
+                pod_dict["metadata"].get("labels", {}).get("dask.org/cluster-name")
+            )
+            if dask_component == "scheduler" and cluster_name:
+                service_names.append(cluster_name)
+
+        services = get_k8s().v1api.list_namespaced_service(
+            namespace, label_selector=label_selector
+        )
+        for service in services.items:
+            try:
+                if force or service.metadata.name in service_names:
+                    get_k8s().v1api.delete_namespaced_service(
+                        service.metadata.name, namespace
+                    )
+                    logger.info(f"Deleted service: {service.metadata.name}")
+            except ApiException as exc:
+                # ignore error if service is already removed
+                if exc.status != 404:
+                    raise

mlrun/api/runtime_handlers/kubejob.py

@@ -0,0 +1,37 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from mlrun.api.runtime_handlers.base import BaseRuntimeHandler
+from mlrun.runtimes.base import RuntimeClassMode
+
+
+class KubeRuntimeHandler(BaseRuntimeHandler):
+    kind = "job"
+    class_modes = {RuntimeClassMode.run: "job", RuntimeClassMode.build: "build"}
+
+    @staticmethod
+    def _expect_pods_without_uid() -> bool:
+        """
+        builder pods are handled as part of this runtime handler - they are not coupled to run object, therefore they
+        don't have the uid in their labels
+        """
+        return True
+
+    @staticmethod
+    def _are_resources_coupled_to_run_object() -> bool:
+        return True
+
+    @staticmethod
+    def _get_object_label_selector(object_id: str) -> str:
+        return f"mlrun/uid={object_id}"

mlrun/api/runtime_handlers/mpijob.py

@@ -0,0 +1,147 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 typing
+from datetime import datetime
+
+from sqlalchemy.orm import Session
+
+from mlrun.api.db.base import DBInterface
+from mlrun.api.runtime_handlers import BaseRuntimeHandler
+from mlrun.runtimes.base import RuntimeClassMode
+from mlrun.runtimes.constants import MPIJobV1Alpha1States, RunStates
+from mlrun.runtimes.mpijob import MpiRuntimeV1, MpiRuntimeV1Alpha1
+
+
+class MpiV1Alpha1RuntimeHandler(BaseRuntimeHandler):
+    kind = "mpijob"
+    class_modes = {
+        RuntimeClassMode.run: "mpijob",
+    }
+
+    def _resolve_crd_object_status_info(
+        self, db: DBInterface, db_session: Session, crd_object
+    ) -> typing.Tuple[bool, typing.Optional[datetime], typing.Optional[str]]:
+        """
+        https://github.com/kubeflow/mpi-operator/blob/master/pkg/apis/kubeflow/v1alpha1/types.go#L115
+        """
+        launcher_status = crd_object.get("status", {}).get("launcherStatus", "")
+        in_terminal_state = launcher_status in MPIJobV1Alpha1States.terminal_states()
+        desired_run_state = MPIJobV1Alpha1States.mpijob_state_to_run_state(
+            launcher_status
+        )
+        completion_time = None
+        if in_terminal_state:
+            completion_time = datetime.fromisoformat(
+                crd_object.get("status", {})
+                .get("completionTime")
+                .replace("Z", "+00:00")
+            )
+            desired_run_state = {
+                "Succeeded": RunStates.completed,
+                "Failed": RunStates.error,
+            }[launcher_status]
+        return in_terminal_state, completion_time, desired_run_state
+
+    @staticmethod
+    def _are_resources_coupled_to_run_object() -> bool:
+        return True
+
+    @staticmethod
+    def _get_object_label_selector(object_id: str) -> str:
+        return f"mlrun/uid={object_id}"
+
+    @staticmethod
+    def _get_main_runtime_resource_label_selector() -> str:
+        """
+        There are some runtimes which might have multiple k8s resources attached to a one runtime, in this case
+        we don't want to pull logs from all but rather only for the "driver"/"launcher" etc
+        :return: the label selector
+        """
+        return "mpi_role_type=launcher"
+
+    @staticmethod
+    def _get_crd_info() -> typing.Tuple[str, str, str]:
+        return (
+            MpiRuntimeV1Alpha1.crd_group,
+            MpiRuntimeV1Alpha1.crd_version,
+            MpiRuntimeV1Alpha1.crd_plural,
+        )
+
+    @staticmethod
+    def _get_crd_object_status(crd_object) -> str:
+        return crd_object.get("status", {}).get("launcherStatus", "")
+
+
+class MpiV1RuntimeHandler(BaseRuntimeHandler):
+    kind = "mpijob"
+    class_modes = {
+        RuntimeClassMode.run: "mpijob",
+    }
+
+    def _resolve_crd_object_status_info(
+        self, db: DBInterface, db_session: Session, crd_object
+    ) -> typing.Tuple[bool, typing.Optional[datetime], typing.Optional[str]]:
+        """
+        https://github.com/kubeflow/mpi-operator/blob/master/pkg/apis/kubeflow/v1/types.go#L29
+        https://github.com/kubeflow/common/blob/master/pkg/apis/common/v1/types.go#L55
+        """
+        launcher_status = (
+            crd_object.get("status", {}).get("replicaStatuses", {}).get("Launcher", {})
+        )
+        # the launcher status also has running property, but it's empty for
+        # short period after the creation, so we're
+        # checking terminal state by the completion time existence
+        in_terminal_state = (
+            crd_object.get("status", {}).get("completionTime", None) is not None
+        )
+        desired_run_state = RunStates.running
+        completion_time = None
+        if in_terminal_state:
+            completion_time = datetime.fromisoformat(
+                crd_object.get("status", {})
+                .get("completionTime")
+                .replace("Z", "+00:00")
+            )
+            desired_run_state = (
+                RunStates.error
+                if launcher_status.get("failed", 0) > 0
+                else RunStates.completed
+            )
+        return in_terminal_state, completion_time, desired_run_state
+
+    @staticmethod
+    def _are_resources_coupled_to_run_object() -> bool:
+        return True
+
+    @staticmethod
+    def _get_object_label_selector(object_id: str) -> str:
+        return f"mlrun/uid={object_id}"
+
+    @staticmethod
+    def _get_main_runtime_resource_label_selector() -> str:
+        """
+        There are some runtimes which might have multiple k8s resources attached to a one runtime, in this case
+        we don't want to pull logs from all but rather only for the "driver"/"launcher" etc
+        :return: the label selector
+        """
+        return "mpi-job-role=launcher"
+
+    @staticmethod
+    def _get_crd_info() -> typing.Tuple[str, str, str]:
+        return (
+            MpiRuntimeV1.crd_group,
+            MpiRuntimeV1.crd_version,
+            MpiRuntimeV1.crd_plural,
+        )

mlrun/api/runtime_handlers/remotesparkjob.py

@@ -0,0 +1,29 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from mlrun.api.runtime_handlers.kubejob import KubeRuntimeHandler
+from mlrun.runtimes.base import RuntimeClassMode
+
+
+class RemoteSparkRuntimeHandler(KubeRuntimeHandler):
+    kind = "remote-spark"
+    class_modes = {RuntimeClassMode.run: "remote-spark"}
+
+    @staticmethod
+    def _are_resources_coupled_to_run_object() -> bool:
+        return True
+
+    @staticmethod
+    def _get_object_label_selector(object_id: str) -> str:
+        return f"mlrun/uid={object_id}"

mlrun/api/runtime_handlers/sparkjob.py

@@ -0,0 +1,148 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 typing
+from datetime import datetime
+from typing import Dict, Optional, Tuple
+
+from kubernetes.client.rest import ApiException
+from sqlalchemy.orm import Session
+
+from mlrun.api.db.base import DBInterface
+from mlrun.api.runtime_handlers.base import BaseRuntimeHandler
+from mlrun.runtimes.base import RuntimeClassMode
+from mlrun.runtimes.constants import RunStates, SparkApplicationStates
+from mlrun.runtimes.sparkjob.abstract import AbstractSparkRuntime
+from mlrun.runtimes.utils import get_k8s
+from mlrun.utils import logger
+
+
+class SparkRuntimeHandler(BaseRuntimeHandler):
+    kind = "spark"
+    class_modes = {
+        RuntimeClassMode.run: "spark",
+    }
+
+    def _resolve_crd_object_status_info(
+        self, db: DBInterface, db_session: Session, crd_object
+    ) -> Tuple[bool, Optional[datetime], Optional[str]]:
+        state = crd_object.get("status", {}).get("applicationState", {}).get("state")
+        in_terminal_state = state in SparkApplicationStates.terminal_states()
+        desired_run_state = SparkApplicationStates.spark_application_state_to_run_state(
+            state
+        )
+        completion_time = None
+        if in_terminal_state:
+            if crd_object.get("status", {}).get("terminationTime"):
+                completion_time = datetime.fromisoformat(
+                    crd_object.get("status", {})
+                    .get("terminationTime")
+                    .replace("Z", "+00:00")
+                )
+            else:
+                last_submission_attempt_time = crd_object.get("status", {}).get(
+                    "lastSubmissionAttemptTime"
+                )
+                if last_submission_attempt_time:
+                    last_submission_attempt_time = last_submission_attempt_time.replace(
+                        "Z", "+00:00"
+                    )
+                    completion_time = datetime.fromisoformat(
+                        last_submission_attempt_time
+                    )
+        return in_terminal_state, completion_time, desired_run_state
+
+    def _update_ui_url(
+        self,
+        db: DBInterface,
+        db_session: Session,
+        project: str,
+        uid: str,
+        crd_object,
+        run: Dict = None,
+    ):
+        app_state = (
+            crd_object.get("status", {}).get("applicationState", {}).get("state")
+        )
+        state = SparkApplicationStates.spark_application_state_to_run_state(app_state)
+        ui_url = None
+        if state == RunStates.running:
+            ui_url = (
+                crd_object.get("status", {})
+                .get("driverInfo", {})
+                .get("webUIIngressAddress")
+            )
+        db_ui_url = run.get("status", {}).get("ui_url")
+        if db_ui_url == ui_url:
+            return
+        run.setdefault("status", {})["ui_url"] = ui_url
+        db.store_run(db_session, run, uid, project)
+
+    @staticmethod
+    def _are_resources_coupled_to_run_object() -> bool:
+        return True
+
+    @staticmethod
+    def _get_object_label_selector(object_id: str) -> str:
+        return f"mlrun/uid={object_id}"
+
+    @staticmethod
+    def _get_main_runtime_resource_label_selector() -> str:
+        """
+        There are some runtimes which might have multiple k8s resources attached to a one runtime, in this case
+        we don't want to pull logs from all but rather only for the "driver"/"launcher" etc
+        :return: the label selector
+        """
+        return "spark-role=driver"
+
+    @staticmethod
+    def _get_crd_info() -> Tuple[str, str, str]:
+        return (
+            AbstractSparkRuntime.group,
+            AbstractSparkRuntime.version,
+            AbstractSparkRuntime.plural,
+        )
+
+    def _delete_extra_resources(
+        self,
+        db: DBInterface,
+        db_session: Session,
+        namespace: str,
+        deleted_resources: typing.List[Dict],
+        label_selector: str = None,
+        force: bool = False,
+        grace_period: int = None,
+    ):
+        """
+        Handling config maps deletion
+        """
+        uids = []
+        for crd_dict in deleted_resources:
+            uid = crd_dict["metadata"].get("labels", {}).get("mlrun/uid", None)
+            uids.append(uid)
+
+        config_maps = get_k8s().v1api.list_namespaced_config_map(
+            namespace, label_selector=label_selector
+        )
+        for config_map in config_maps.items:
+            try:
+                uid = config_map.metadata.labels.get("mlrun/uid", None)
+                if force or uid in uids:
+                    get_k8s().v1api.delete_namespaced_config_map(
+                        config_map.metadata.name, namespace
+                    )
+                    logger.info(f"Deleted config map: {config_map.metadata.name}")
+            except ApiException as exc:
+                # ignore error if config map is already removed
+                if exc.status != 404:
+                    raise

mlrun/api/schemas/__init__.py

@@ -36,6 +36,7 @@ import mlrun.common.schemas.function as old_function
 import mlrun.common.schemas.http as old_http
 import mlrun.common.schemas.k8s as old_k8s
 import mlrun.common.schemas.memory_reports as old_memory_reports
+import mlrun.common.schemas.model_monitoring.grafana
 import mlrun.common.schemas.object as old_object
 import mlrun.common.schemas.pipeline as old_pipeline
 import mlrun.common.schemas.project as old_project
@@ -163,13 +164,23 @@ MostCommonObjectTypesReport = DeprecationHelper(
 ObjectTypeReport = DeprecationHelper(mlrun.common.schemas.ObjectTypeReport)
 Features = DeprecationHelper(mlrun.common.schemas.Features)
 FeatureValues = DeprecationHelper(mlrun.common.schemas.FeatureValues)
-GrafanaColumn = DeprecationHelper(mlrun.common.schemas.GrafanaColumn)
-GrafanaDataPoint = DeprecationHelper(mlrun.common.schemas.GrafanaDataPoint)
-GrafanaNumberColumn = DeprecationHelper(mlrun.common.schemas.GrafanaNumberColumn)
-GrafanaStringColumn = DeprecationHelper(mlrun.common.schemas.GrafanaStringColumn)
-GrafanaTable = DeprecationHelper(mlrun.common.schemas.GrafanaTable)
+GrafanaColumn = DeprecationHelper(
+    mlrun.common.schemas.model_monitoring.grafana.GrafanaColumn
+)
+GrafanaDataPoint = DeprecationHelper(
+    mlrun.common.schemas.model_monitoring.grafana.GrafanaDataPoint
+)
+GrafanaNumberColumn = DeprecationHelper(
+    mlrun.common.schemas.model_monitoring.grafana.GrafanaNumberColumn
+)
+GrafanaStringColumn = DeprecationHelper(
+    mlrun.common.schemas.model_monitoring.grafana.GrafanaStringColumn
+)
+GrafanaTable = DeprecationHelper(
+    mlrun.common.schemas.model_monitoring.grafana.GrafanaTable
+)
 GrafanaTimeSeriesTarget = DeprecationHelper(
-    mlrun.common.schemas.GrafanaTimeSeriesTarget
+    mlrun.common.schemas.model_monitoring.grafana.GrafanaTimeSeriesTarget
 )
 ModelEndpoint = DeprecationHelper(mlrun.common.schemas.ModelEndpoint)
 ModelEndpointList = DeprecationHelper(mlrun.common.schemas.ModelEndpointList)

mlrun/api/utils/builder.py

@@ -420,10 +420,7 @@ def build_image(
         # 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:]
+        relative_workdir = relative_workdir.removeprefix("./")
 
         runtime.spec.clone_target_dir = path.join(tmpdir, "mlrun", relative_workdir)
 

mlrun/api/utils/clients/chief.py

@@ -98,6 +98,20 @@ class Client(
             "DELETE", f"projects/{project}/schedules/{name}", request
         )
 
+    async def submit_workflow(
+        self,
+        project: str,
+        name: str,
+        request: fastapi.Request,
+        json: dict,
+    ) -> fastapi.Response:
+        """
+        Workflow schedules are running only on chief
+        """
+        return await self._proxy_request_to_chief(
+            "POST", f"projects/{project}/workflows/{name}/submit", request, json
+        )
+
     async def delete_schedules(
         self, project: str, request: fastapi.Request
     ) -> fastapi.Response: