mlrun 1.7.1rc10__py3-none-any.whl → 1.8.0rc8__py3-none-any.whl

mlrun/execution.py

@@ -15,8 +15,9 @@
 import logging
 import os
 import uuid
+import warnings
 from copy import deepcopy
-from typing import Union
+from typing import Optional, Union, cast
 
 import numpy as np
 import yaml
@@ -25,17 +26,23 @@ from dateutil import parser
 import mlrun
 import mlrun.common.constants as mlrun_constants
 import mlrun.common.formatters
-from mlrun.artifacts import ModelArtifact
+from mlrun.artifacts import (
+    Artifact,
+    DatasetArtifact,
+    DocumentArtifact,
+    DocumentLoaderSpec,
+    ModelArtifact,
+)
 from mlrun.datastore.store_resources import get_store_resource
 from mlrun.errors import MLRunInvalidArgumentError
 
-from .artifacts import DatasetArtifact
 from .artifacts.manager import ArtifactManager, dict_to_artifact, extend_artifact_path
 from .datastore import store_manager
 from .features import Feature
 from .model import HyperParamOptions
 from .secrets import SecretsStore
 from .utils import (
+    Logger,
     RunKeys,
     dict_to_json,
     dict_to_yaml,
@@ -152,7 +159,7 @@ class MLClientCtx:
         return self._project
 
     @property
-    def logger(self):
+    def logger(self) -> Logger:
         """Built-in logger interface
 
         Example::
@@ -194,6 +201,11 @@ class MLClientCtx:
         """Dictionary of artifacts (read-only)"""
         return deepcopy(self._artifacts_manager.artifact_list())
 
+    @property
+    def artifact_uris(self):
+        """Dictionary of artifact URIs (read-only)"""
+        return deepcopy(self._artifacts_manager.artifact_uris)
+
     @property
     def in_path(self):
         """Default input path for data objects"""
@@ -296,7 +308,7 @@ class MLClientCtx:
             )
         self._parent.log_iteration_results(self._iteration, None, self.to_dict())
 
-    def get_store_resource(self, url, secrets: dict = None):
+    def get_store_resource(self, url, secrets: Optional[dict] = None):
         """Get mlrun data resource (feature set/vector, artifact, item) from url.
 
         Example::
@@ -317,7 +329,7 @@ class MLClientCtx:
             data_store_secrets=secrets,
         )
 
-    def get_dataitem(self, url, secrets: dict = None):
+    def get_dataitem(self, url, secrets: Optional[dict] = None):
         """Get mlrun dataitem from url
 
         Example::
@@ -425,8 +437,11 @@ class MLClientCtx:
             self._results = status.get("results", self._results)
             for artifact in status.get("artifacts", []):
                 artifact_obj = dict_to_artifact(artifact)
-                key = artifact_obj.key
-                self._artifacts_manager.artifacts[key] = artifact_obj
+                self._artifacts_manager.artifact_uris[artifact_obj.key] = (
+                    artifact_obj.uri
+                )
+            for key, uri in status.get("artifact_uris", {}).items():
+                self._artifacts_manager.artifact_uris[key] = uri
             self._state = status.get("state", self._state)
 
         # No need to store the run for every worker
@@ -486,11 +501,11 @@ class MLClientCtx:
             return default
         return self._parameters[key]
 
-    def get_project_object(self):
+    def get_project_object(self) -> Optional["mlrun.MlrunProject"]:
         """
         Get the MLRun project object by the project name set in the context.
 
-        :return: The project object or None if it couldn't be retrieved.
+        :returns: The project object or None if it couldn't be retrieved.
         """
         return self._load_project_object()
 
@@ -573,22 +588,25 @@ class MLClientCtx:
         """Reserved for internal use"""
 
         if best:
+            # Recreate the best iteration context for the interface of getting its artifacts
+            best_context = MLClientCtx.from_dict(
+                task, store_run=False, include_status=True
+            )
             self._results["best_iteration"] = best
-            for k, v in get_in(task, ["status", "results"], {}).items():
-                self._results[k] = v
-            for artifact in get_in(task, ["status", RunKeys.artifacts], []):
-                self._artifacts_manager.artifacts[artifact["metadata"]["key"]] = (
-                    artifact
-                )
+            for key, result in best_context.results.items():
+                self._results[key] = result
+            for key, artifact_uri in best_context.artifact_uris.items():
+                self._artifacts_manager.artifact_uris[key] = artifact_uri
+                artifact = best_context.get_artifact(key)
                 self._artifacts_manager.link_artifact(
                     self.project,
                     self.name,
                     self.tag,
-                    artifact["metadata"]["key"],
+                    key,
                     self.iteration,
-                    artifact["spec"]["target_path"],
+                    artifact.target_path,
+                    db_key=artifact.db_key,
                     link_iteration=best,
-                    db_key=artifact["spec"]["db_key"],
                 )
 
         if summary is not None:
@@ -611,7 +629,7 @@ class MLClientCtx:
         format=None,
         db_key=None,
         **kwargs,
-    ):
+    ) -> Artifact:
         """Log an output artifact and optionally upload it to datastore
 
         Example::
@@ -679,9 +697,9 @@ class MLClientCtx:
         db_key=None,
         target_path="",
         extra_data=None,
-        label_column: str = None,
+        label_column: Optional[str] = None,
         **kwargs,
-    ):
+    ) -> DatasetArtifact:
         """Log a dataset artifact and optionally upload it to datastore
 
         If the dataset exists with the same key and tag, it will be overwritten.
@@ -719,7 +737,7 @@ class MLClientCtx:
         :param db_key:        The key to use in the artifact DB table, by default its run name + '_' + key
                               db_key=False will not register it in the artifacts table
 
-        :returns: Artifact object
+        :returns: Dataset artifact object
         """
         ds = DatasetArtifact(
             key,
@@ -732,16 +750,19 @@ class MLClientCtx:
             **kwargs,
         )
 
-        item = self._artifacts_manager.log_artifact(
-            self,
-            ds,
-            local_path=local_path,
-            artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
-            target_path=target_path,
-            tag=tag,
-            upload=upload,
-            db_key=db_key,
-            labels=labels,
+        item = cast(
+            DatasetArtifact,
+            self._artifacts_manager.log_artifact(
+                self,
+                ds,
+                local_path=local_path,
+                artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
+                target_path=target_path,
+                tag=tag,
+                upload=upload,
+                db_key=db_key,
+                labels=labels,
+            ),
         )
         self._update_run()
         return item
@@ -760,16 +781,16 @@ class MLClientCtx:
         artifact_path=None,
         upload=True,
         labels=None,
-        inputs: list[Feature] = None,
-        outputs: list[Feature] = None,
-        feature_vector: str = None,
-        feature_weights: list = None,
+        inputs: Optional[list[Feature]] = None,
+        outputs: Optional[list[Feature]] = None,
+        feature_vector: Optional[str] = None,
+        feature_weights: Optional[list] = None,
         training_set=None,
-        label_column: Union[str, list] = None,
+        label_column: Optional[Union[str, list]] = None,
         extra_data=None,
         db_key=None,
         **kwargs,
-    ):
+    ) -> ModelArtifact:
         """Log a model artifact and optionally upload it to datastore
 
         Example::
@@ -788,7 +809,7 @@ class MLClientCtx:
         :param key:             Artifact key or artifact class ()
         :param body:            Will use the body as the artifact content
         :param model_file:      Path to the local model file we upload (see also model_dir)
-                                or to a model file data url (e.g. http://host/path/model.pkl)
+                                or to a model file data url (e.g. `http://host/path/model.pkl`)
         :param model_dir:       Path to the local dir holding the model file and extra files
         :param artifact_path:   Target artifact path (when not using the default)
                                 to define a subpath under the default location use:
@@ -811,7 +832,7 @@ class MLClientCtx:
         :param db_key:          The key to use in the artifact DB table, by default its run name + '_' + key
                                 db_key=False will not register it in the artifacts table
 
-        :returns: Artifact object
+        :returns: Model artifact object
         """
 
         if training_set is not None and inputs:
@@ -838,13 +859,63 @@ class MLClientCtx:
         if training_set is not None:
             model.infer_from_df(training_set, label_column)
 
+        item = cast(
+            ModelArtifact,
+            self._artifacts_manager.log_artifact(
+                self,
+                model,
+                artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
+                tag=tag,
+                upload=upload,
+                db_key=db_key,
+                labels=labels,
+            ),
+        )
+        self._update_run()
+        return item
+
+    def log_document(
+        self,
+        key: str,
+        tag: str = "",
+        local_path: str = "",
+        artifact_path: Optional[str] = None,
+        document_loader: DocumentLoaderSpec = DocumentLoaderSpec(),
+        upload: Optional[bool] = False,
+        labels: Optional[dict[str, str]] = None,
+        target_path: Optional[str] = None,
+        **kwargs,
+    ) -> DocumentArtifact:
+        """
+        Log a document as an artifact.
+
+        :param key: Artifact key
+        :param tag: Version tag
+        :param local_path:    path to the local file we upload, will also be use
+                              as the destination subpath (under "artifact_path")
+        :param artifact_path:   Target artifact path (when not using the default)
+                                to define a subpath under the default location use:
+                                `artifact_path=context.artifact_subpath('data')`
+        :param document_loader: Spec to use to load the artifact as langchain document
+        :param upload: Whether to upload the artifact
+        :param labels: Key-value labels
+        :param target_path: Path to the local file
+        :param kwargs: Additional keyword arguments
+        :return: DocumentArtifact object
+        """
+        doc_artifact = DocumentArtifact(
+            key=key,
+            original_source=local_path or target_path,
+            document_loader=document_loader,
+            **kwargs,
+        )
+
         item = self._artifacts_manager.log_artifact(
             self,
-            model,
+            doc_artifact,
             artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
             tag=tag,
             upload=upload,
-            db_key=db_key,
             labels=labels,
         )
         self._update_run()
@@ -852,10 +923,18 @@ class MLClientCtx:
 
     def get_cached_artifact(self, key):
         """Return a logged artifact from cache (for potential updates)"""
-        return self._artifacts_manager.artifacts[key]
+        warnings.warn(
+            "get_cached_artifact is deprecated in 1.8.0 and will be removed in 1.10.0. Use get_artifact instead.",
+            FutureWarning,
+        )
+        return self.get_artifact(key)
+
+    def get_artifact(self, key: str) -> Artifact:
+        artifact_uri = self._artifacts_manager.artifact_uris[key]
+        return self.get_store_resource(artifact_uri)
 
-    def update_artifact(self, artifact_object):
-        """Update an artifact object in the cache and the DB"""
+    def update_artifact(self, artifact_object: Artifact):
+        """Update an artifact object in the DB and the cached uri"""
         self._artifacts_manager.update_artifact(self, artifact_object)
 
     def commit(self, message: str = "", completed=False):
@@ -885,7 +964,12 @@ class MLClientCtx:
         if completed and not self.iteration:
             mlrun.runtimes.utils.global_context.set(None)
 
-    def set_state(self, execution_state: str = None, error: str = None, commit=True):
+    def set_state(
+        self,
+        execution_state: Optional[str] = None,
+        error: Optional[str] = None,
+        commit=True,
+    ):
         """
         Modify and store the execution state or mark an error and update the run state accordingly.
         This method allows to set the run state to 'completed' in the DB which is discouraged.
@@ -1013,7 +1097,7 @@ class MLClientCtx:
         set_if_not_none(struct["status"], "commit", self._commit)
         set_if_not_none(struct["status"], "iterations", self._iteration_results)
 
-        struct["status"][RunKeys.artifacts] = self._artifacts_manager.artifact_list()
+        struct["status"][RunKeys.artifact_uris] = self._artifacts_manager.artifact_uris
         self._data_stores.to_dict(struct["spec"])
         return struct
 
@@ -1107,7 +1191,9 @@ class MLClientCtx:
         set_if_not_none(struct, "status.commit", self._commit)
         set_if_not_none(struct, "status.iterations", self._iteration_results)
 
-        struct[f"status.{RunKeys.artifacts}"] = self._artifacts_manager.artifact_list()
+        struct[f"status.{RunKeys.artifact_uris}"] = (
+            self._artifacts_manager.artifact_uris
+        )
         return struct
 
     def _init_dbs(self, rundb):
@@ -1121,7 +1207,7 @@ class MLClientCtx:
         self._data_stores = store_manager.set(self._secrets_manager, db=self._rundb)
         self._artifacts_manager = ArtifactManager(db=self._rundb)
 
-    def _load_project_object(self):
+    def _load_project_object(self) -> Optional["mlrun.MlrunProject"]:
         if not self._project_object:
             if not self._project:
                 self.logger.warning(

mlrun/feature_store/__init__.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# flake8: noqa
-
 __all__ = [
     "get_offline_features",
     "get_online_feature_service",

mlrun/feature_store/api.py

@@ -11,6 +11,7 @@
 # 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 copy
 import importlib.util
 import pathlib
@@ -99,21 +100,21 @@ def _features_to_vector_and_check_permissions(features, update_stats):
 def get_offline_features(
     feature_vector: Union[str, FeatureVector],
     entity_rows=None,
-    entity_timestamp_column: str = None,
+    entity_timestamp_column: Optional[str] = None,
     target: DataTargetBase = None,
     run_config: RunConfig = None,
-    drop_columns: list[str] = None,
-    start_time: Union[str, datetime] = None,
-    end_time: Union[str, datetime] = None,
+    drop_columns: Optional[list[str]] = None,
+    start_time: Optional[Union[str, datetime]] = None,
+    end_time: Optional[Union[str, datetime]] = None,
     with_indexes: bool = False,
     update_stats: bool = False,
-    engine: str = None,
-    engine_args: dict = None,
-    query: str = None,
-    order_by: Union[str, list[str]] = None,
-    spark_service: str = None,
-    timestamp_for_filtering: Union[str, dict[str, str]] = None,
-    additional_filters: list = None,
+    engine: Optional[str] = None,
+    engine_args: Optional[dict] = None,
+    query: Optional[str] = None,
+    order_by: Optional[Union[str, list[str]]] = None,
+    spark_service: Optional[str] = None,
+    timestamp_for_filtering: Optional[Union[str, dict[str, str]]] = None,
+    additional_filters: Optional[list] = None,
 ):
     """retrieve offline feature vector results
 
@@ -209,20 +210,20 @@ def get_offline_features(
 def _get_offline_features(
     feature_vector: Union[str, FeatureVector],
     entity_rows=None,
-    entity_timestamp_column: str = None,
+    entity_timestamp_column: Optional[str] = None,
     target: DataTargetBase = None,
     run_config: RunConfig = None,
-    drop_columns: list[str] = None,
-    start_time: Union[str, datetime] = None,
-    end_time: Union[str, datetime] = None,
+    drop_columns: Optional[list[str]] = None,
+    start_time: Optional[Union[str, datetime]] = None,
+    end_time: Optional[Union[str, datetime]] = None,
     with_indexes: bool = False,
     update_stats: bool = False,
-    engine: str = None,
-    engine_args: dict = None,
-    query: str = None,
-    order_by: Union[str, list[str]] = None,
-    spark_service: str = None,
-    timestamp_for_filtering: Union[str, dict[str, str]] = None,
+    engine: Optional[str] = None,
+    engine_args: Optional[dict] = None,
+    query: Optional[str] = None,
+    order_by: Optional[Union[str, list[str]]] = None,
+    spark_service: Optional[str] = None,
+    timestamp_for_filtering: Optional[Union[str, dict[str, str]]] = None,
     additional_filters=None,
 ) -> Union[OfflineVectorResponse, RemoteVectorResponse]:
     if entity_rows is None and entity_timestamp_column is not None:
@@ -297,9 +298,9 @@ def get_online_feature_service(
     feature_vector: Union[str, FeatureVector],
     run_config: RunConfig = None,
     fixed_window_type: FixedWindowType = FixedWindowType.LastClosedWindow,
-    impute_policy: dict = None,
+    impute_policy: Optional[dict] = None,
     update_stats: bool = False,
-    entity_keys: list[str] = None,
+    entity_keys: Optional[list[str]] = None,
 ):
     """initialize and return online feature vector service api,
     returns :py:class:`~mlrun.feature_store.OnlineVectorService`
@@ -378,9 +379,9 @@ def _get_online_feature_service(
     feature_vector: Union[str, FeatureVector],
     run_config: RunConfig = None,
     fixed_window_type: FixedWindowType = FixedWindowType.LastClosedWindow,
-    impute_policy: dict = None,
+    impute_policy: Optional[dict] = None,
     update_stats: bool = False,
-    entity_keys: list[str] = None,
+    entity_keys: Optional[list[str]] = None,
 ) -> OnlineVectorService:
     if isinstance(feature_vector, FeatureVector):
         update_stats = True
@@ -450,7 +451,7 @@ def _get_namespace(run_config: RunConfig) -> dict[str, Any]:
 def ingest(
     featureset: Union[FeatureSet, str] = None,
     source=None,
-    targets: list[DataTargetBase] = None,
+    targets: Optional[list[DataTargetBase]] = None,
     namespace=None,
     return_df: bool = True,
     infer_options: InferOptions = InferOptions.default(),
@@ -530,7 +531,7 @@ def ingest(
 def _ingest(
     featureset: Union[FeatureSet, str] = None,
     source=None,
-    targets: list[DataTargetBase] = None,
+    targets: Optional[list[DataTargetBase]] = None,
     namespace=None,
     return_df: bool = True,
     infer_options: InferOptions = InferOptions.default(),
@@ -783,11 +784,11 @@ def _ingest(
 def preview(
     featureset: FeatureSet,
     source,
-    entity_columns: list = None,
+    entity_columns: Optional[list] = None,
     namespace=None,
     options: InferOptions = None,
     verbose: bool = False,
-    sample_size: int = None,
+    sample_size: Optional[int] = None,
 ) -> pd.DataFrame:
     """run the ingestion pipeline with local DataFrame/file data and infer features schema and stats
 
@@ -825,11 +826,11 @@ def preview(
 def _preview(
     featureset: FeatureSet,
     source,
-    entity_columns: list = None,
+    entity_columns: Optional[list] = None,
     namespace=None,
     options: InferOptions = None,
     verbose: bool = False,
-    sample_size: int = None,
+    sample_size: Optional[int] = None,
 ) -> pd.DataFrame:
     if isinstance(source, pd.DataFrame):
         source = _rename_source_dataframe_columns(source)
@@ -895,8 +896,8 @@ def _preview(
 def _run_ingestion_job(
     featureset: Union[FeatureSet, str],
     source: DataSource = None,
-    targets: list[DataTargetBase] = None,
-    name: str = None,
+    targets: Optional[list[DataTargetBase]] = None,
+    name: Optional[str] = None,
     infer_options: InferOptions = InferOptions.default(),
     run_config: RunConfig = None,
 ):
@@ -920,8 +921,8 @@ def _run_ingestion_job(
 def deploy_ingestion_service_v2(
     featureset: Union[FeatureSet, str],
     source: DataSource = None,
-    targets: list[DataTargetBase] = None,
-    name: str = None,
+    targets: Optional[list[DataTargetBase]] = None,
+    name: Optional[str] = None,
     run_config: RunConfig = None,
     verbose=False,
 ) -> tuple[str, BaseRuntime]:
@@ -963,8 +964,8 @@ def deploy_ingestion_service_v2(
 def _deploy_ingestion_service_v2(
     featureset: Union[FeatureSet, str],
     source: DataSource = None,
-    targets: list[DataTargetBase] = None,
-    name: str = None,
+    targets: Optional[list[DataTargetBase]] = None,
+    name: Optional[str] = None,
     run_config: RunConfig = None,
     verbose=False,
 ) -> tuple[str, BaseRuntime]:
@@ -1026,7 +1027,7 @@ def _ingest_with_spark(
     spark=None,
     featureset: Union[FeatureSet, str] = None,
     source: BaseSourceDriver = None,
-    targets: list[BaseStoreTarget] = None,
+    targets: Optional[list[BaseStoreTarget]] = None,
     infer_options: InferOptions = InferOptions.default(),
     mlrun_context=None,
     namespace=None,
@@ -1199,8 +1200,8 @@ def _infer_from_static_df(
 def set_task_params(
     featureset: FeatureSet,
     source: DataSource = None,
-    targets: list[DataTargetBase] = None,
-    parameters: dict = None,
+    targets: Optional[list[DataTargetBase]] = None,
+    parameters: Optional[dict] = None,
     infer_options: InferOptions = InferOptions.Null,
     overwrite=None,
 ):

mlrun/feature_store/common.py

@@ -178,17 +178,17 @@ class RunConfig:
     def __init__(
         self,
         function: typing.Union[str, FunctionReference, BaseRuntime] = None,
-        local: bool = None,
-        image: str = None,
-        kind: str = None,
-        handler: str = None,
-        parameters: dict = None,
-        watch: bool = None,
+        local: typing.Optional[bool] = None,
+        image: typing.Optional[str] = None,
+        kind: typing.Optional[str] = None,
+        handler: typing.Optional[str] = None,
+        parameters: typing.Optional[dict] = None,
+        watch: typing.Optional[bool] = None,
         owner=None,
         credentials: typing.Optional[mlrun.model.Credentials] = None,
-        code: str = None,
-        requirements: typing.Union[str, list[str]] = None,
-        extra_spec: dict = None,
+        code: typing.Optional[str] = None,
+        requirements: typing.Optional[typing.Union[str, list[str]]] = None,
+        extra_spec: typing.Optional[dict] = None,
         auth_info=None,
     ):
         """class for holding function and run specs for jobs and serving functions