mlrun 1.10.0rc24__py3-none-any.whl → 1.10.0rc26__py3-none-any.whl

mlrun/db/base.py

@@ -722,7 +722,7 @@ class RunDBInterface(ABC):
         tsdb_metrics: bool = False,
         metric_list: Optional[list[str]] = None,
         top_level: bool = False,
-        mode: Optional[mlrun.common.schemas.EndpointMode] = None,
+        modes: Optional[list[mm_constants.EndpointMode]] = None,
         uids: Optional[list[str]] = None,
         latest_only: bool = False,
     ) -> mlrun.common.schemas.ModelEndpointList:
@@ -792,6 +792,7 @@ class RunDBInterface(ABC):
         version: Optional[str] = None,
         tag: Optional[str] = None,
         force_refresh: bool = False,
+        object_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 
@@ -803,6 +804,19 @@ class RunDBInterface(ABC):
         version: Optional[str] = None,
         tag: str = "latest",
         force_refresh: bool = False,
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
+    ):
+        pass
+
+    @abstractmethod
+    def get_hub_asset(
+        self,
+        source_name: str,
+        item_name: str,
+        asset_name: str,
+        version: Optional[str] = None,
+        tag: str = "latest",
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 

mlrun/db/httpdb.py

@@ -45,6 +45,7 @@ import mlrun.runtimes.nuclio.api_gateway
 import mlrun.runtimes.nuclio.function
 import mlrun.utils
 from mlrun.alerts.alert import AlertConfig
+from mlrun.common.schemas.hub import HubSourceType
 from mlrun.db.auth_utils import OAuthClientIDTokenProvider, StaticTokenProvider
 from mlrun.errors import MLRunInvalidArgumentError, err_to_str
 from mlrun.secrets import get_secret_or_env
@@ -3770,7 +3771,7 @@ class HTTPRunDB(RunDBInterface):
         tsdb_metrics: bool = False,
         metric_list: Optional[list[str]] = None,
         top_level: bool = False,
-        mode: mm_constants.EndpointMode = None,
+        modes: Optional[list[mm_constants.EndpointMode]] = None,
         uids: Optional[list[str]] = None,
         latest_only: bool = False,
     ) -> mlrun.common.schemas.ModelEndpointList:
@@ -3791,8 +3792,8 @@ class HTTPRunDB(RunDBInterface):
                                 If tsdb_metrics=False, this parameter will be ignored and no tsdb metrics
                                 will be included.
         :param top_level:       Whether to return only top level model endpoints.
-        :param mode:            Specifies the mode of the model endpoint. Can be "real-time" (0), "batch" (1), or
-                                both if set to None.
+        :param modes:           Specifies the modes of the model endpoints. Can be "real-time" (0), "batch" (1),
+                                "batch_legacy" (2). If set to None, all are included.
         :param uids:            A list of unique ids to filter by.
         :param latest_only:     Whether to return only the latest model endpoint version.
         :return:                A list of model endpoints.
@@ -3801,6 +3802,8 @@ class HTTPRunDB(RunDBInterface):
         labels = self._parse_labels(labels)
         if names and isinstance(names, str):
             names = [names]
+        if isinstance(modes, mm_constants.EndpointMode):
+            modes = [modes]
         response = self.api_call(
             method=mlrun.common.types.HTTPMethod.GET,
             path=path,
@@ -3816,7 +3819,7 @@ class HTTPRunDB(RunDBInterface):
                 "tsdb-metrics": tsdb_metrics,
                 "metric": metric_list,
                 "top-level": top_level,
-                "mode": mode,
+                "mode": modes,
                 "uid": uids,
                 "latest-only": latest_only,
             },
@@ -4079,7 +4082,7 @@ class HTTPRunDB(RunDBInterface):
         response = self.api_call(
             method=mlrun.common.types.HTTPMethod.DELETE,
             path=f"projects/{project}/model-monitoring/functions",
-            params={"functions": functions},
+            params={"function": functions},
         )
         deletion_failed = False
         if response.status_code == http.HTTPStatus.ACCEPTED:
@@ -4359,6 +4362,7 @@ class HTTPRunDB(RunDBInterface):
         version: Optional[str] = None,
         tag: Optional[str] = None,
         force_refresh: bool = False,
+        object_type: HubSourceType = HubSourceType.functions,
     ):
         """
         Retrieve the item catalog for a specified hub source.
@@ -4371,6 +4375,7 @@ class HTTPRunDB(RunDBInterface):
             rather than rely on cached information which may exist from previous get requests. For example,
             if the source was re-built,
             this will make the server get the updated information. Default is ``False``.
+        :param object_type: Type of object to retrieve from the hub source (e.g: functions, modules).
         :returns: :py:class:`~mlrun.common.schemas.hub.HubCatalog` object, which is essentially a list
             of :py:class:`~mlrun.common.schemas.hub.HubItem` entries.
         """
@@ -4379,6 +4384,7 @@ class HTTPRunDB(RunDBInterface):
             "version": version,
             "tag": tag,
             "force-refresh": force_refresh,
+            "object_type": object_type,
         }
         response = self.api_call(method="GET", path=path, params=params)
         return mlrun.common.schemas.HubCatalog(**response.json())
@@ -4390,6 +4396,7 @@ class HTTPRunDB(RunDBInterface):
         version: Optional[str] = None,
         tag: str = "latest",
         force_refresh: bool = False,
+        item_type: HubSourceType = HubSourceType.functions,
     ):
         """
         Retrieve a specific hub item.
@@ -4401,6 +4408,7 @@ class HTTPRunDB(RunDBInterface):
         :param force_refresh: Make the server fetch the information from the actual hub
             source, rather than
             rely on cached information. Default is ``False``.
+        :param item_type: The type of item to retrieve from the hub source (e.g: functions, modules).
         :returns: :py:class:`~mlrun.common.schemas.hub.HubItem`.
         """
         path = (f"hub/sources/{source_name}/items/{item_name}",)
@@ -4408,6 +4416,7 @@ class HTTPRunDB(RunDBInterface):
             "version": version,
             "tag": tag,
             "force-refresh": force_refresh,
+            "item_type": item_type,
         }
         response = self.api_call(method="GET", path=path, params=params)
         return mlrun.common.schemas.HubItem(**response.json())
@@ -4419,6 +4428,7 @@ class HTTPRunDB(RunDBInterface):
         asset_name: str,
         version: Optional[str] = None,
         tag: str = "latest",
+        item_type: HubSourceType = HubSourceType.functions,
     ):
         """
         Get hub asset from item.
@@ -4428,13 +4438,14 @@ class HTTPRunDB(RunDBInterface):
         :param asset_name:  Name of the asset to retrieve.
         :param version: Get a specific version of the item. Default is ``None``.
         :param tag: Get a specific version of the item identified by tag. Default is ``latest``.
-
+        :param item_type: The type of item to retrieve from the hub source (e.g: functions, modules).
         :returns: http response with the asset in the content attribute
         """
         path = f"hub/sources/{source_name}/items/{item_name}/assets/{asset_name}"
         params = {
             "version": version,
             "tag": tag,
+            "item_type": item_type,
         }
         response = self.api_call(method="GET", path=path, params=params)
         return response

mlrun/db/nopdb.py

@@ -626,6 +626,7 @@ class NopDB(RunDBInterface):
         tsdb_metrics: bool = False,
         metric_list: Optional[list[str]] = None,
         top_level: bool = False,
+        modes: Optional[list[mm_constants.EndpointMode]] = None,
         uids: Optional[list[str]] = None,
         latest_only: bool = False,
     ) -> mlrun.common.schemas.ModelEndpointList:
@@ -688,6 +689,7 @@ class NopDB(RunDBInterface):
         version: Optional[str] = None,
         tag: Optional[str] = None,
         force_refresh: bool = False,
+        object_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 
@@ -699,6 +701,18 @@ class NopDB(RunDBInterface):
         version: Optional[str] = None,
         tag: str = "latest",
         force_refresh: bool = False,
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
+    ):
+        pass
+
+    def get_hub_asset(
+        self,
+        source_name: str,
+        item_name: str,
+        asset_name: str,
+        version: Optional[str] = None,
+        tag: str = "latest",
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 

mlrun/k8s_utils.py

@@ -470,20 +470,6 @@ def _handle_allow_mode(
     list[kubernetes.client.V1Toleration],
     typing.Optional[kubernetes.client.V1Affinity],
 ]:
-    for op in [
-        mlrun.common.schemas.NodeSelectorOperator.node_selector_op_not_in.value,
-        mlrun.common.schemas.NodeSelectorOperator.node_selector_op_in.value,
-    ]:
-        affinity = _prune_affinity_node_selector_requirement(
-            generate_preemptible_node_selector_requirements(op),
-            affinity=affinity,
-        )
-
-    node_selector = _prune_node_selector(
-        mlconfig.get_preemptible_node_selector(),
-        enriched_node_selector=node_selector,
-    )
-
     tolerations = _merge_tolerations(tolerations, preemptible_tolerations)
     return node_selector, tolerations, affinity
 

mlrun/model_monitoring/api.py

@@ -350,8 +350,8 @@ def _generate_model_endpoint(
             project=project,
             name=model_endpoint_name,
             endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.BATCH_EP,
-            # Due to backwards compatibility, old batch model endpoint will be analyzed as real time endpoint
-            mode=mlrun.common.schemas.model_monitoring.EndpointMode.REAL_TIME,
+            # Due to backwards compatibility, this endpoint will be created as a legacy batch endpoint.
+            mode=mlrun.common.schemas.model_monitoring.EndpointMode.BATCH_LEGACY,
         ),
         spec=mlrun.common.schemas.ModelEndpointSpec(
             function_name=function_name or "function",

mlrun/model_monitoring/applications/base.py

@@ -18,7 +18,7 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Iterator
 from contextlib import contextmanager, nullcontext
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Any, Literal, Optional, Union, cast
 
 import pandas as pd
@@ -347,6 +347,21 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                         feature_stats=feature_stats,
                     )
                 )
+
+                if (
+                    monitoring_context.endpoint_id
+                    and monitoring_context.sample_df.empty
+                ):
+                    # The current sample is empty
+                    context.logger.warning(
+                        "No sample data available for tracking",
+                        application_name=application_name,
+                        endpoint_id=monitoring_context.endpoint_id,
+                        start_time=monitoring_context.start_infer_time,
+                        end_time=monitoring_context.end_infer_time,
+                    )
+                    return
+
                 result = self.do_tracking(monitoring_context)
                 endpoints_output[monitoring_context.endpoint_id].append(
                     (monitoring_context, result)
@@ -591,6 +606,16 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         start_dt = datetime.fromisoformat(start)
         end_dt = datetime.fromisoformat(end)
 
+        # If `start_dt` and `end_dt` do not include time zone information - change them to UTC
+        if (start_dt.tzinfo is None) and (end_dt.tzinfo is None):
+            start_dt = start_dt.replace(tzinfo=timezone.utc)
+            end_dt = end_dt.replace(tzinfo=timezone.utc)
+        elif (start_dt.tzinfo is None) or (end_dt.tzinfo is None):
+            raise mlrun.errors.MLRunValueError(
+                "The start and end times must either both include time zone information or both be naive (no time "
+                f"zone). Asserting the above failed, aborting the evaluate request: start={start}, end={end}."
+            )
+
         if existing_data_handling != ExistingDataHandling.delete_all:
             start_dt = cls._validate_monotonically_increasing_data(
                 application_schedules=application_schedules,
@@ -841,7 +866,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         :py:meth:`~mlrun.model_monitoring.applications.ModelMonitoringApplicationBase.do_tracking`
         model monitoring logic as a :py:class:`~mlrun.runtimes.KubejobRuntime`, which is an MLRun function.
 
-        This function has default values for all of its arguments. You should be change them when you want to pass
+        This function has default values for all of its arguments. You should change them when you want to pass
         data to the application.
 
         :param func_path:         The path to the function. If ``None``, the current notebook is used.
@@ -858,6 +883,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         :param reference_data:    Pandas data-frame or :py:class:`~mlrun.artifacts.dataset.DatasetArtifact` URI as
                                   the reference dataset.
                                   When set, its statistics override the model endpoint's feature statistics.
+                                  You do not need to have a model endpoint to use this option.
         :param image:             Docker image to run the job on (when running remotely).
         :param with_repo:         Whether to clone the current repo to the build source.
         :param class_handler:     The relative path to the class, useful when using Git sources or code from images.
@@ -878,8 +904,9 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         :param start:             The start time of the endpoint's data, not included.
                                   If you want the model endpoint's data at ``start`` included, you need to subtract a
                                   small ``datetime.timedelta`` from it.
-                                  Make sure to include the time zone when constructing `datetime.datetime` objects
-                                  manually.
+                                  Make sure to include the time zone when constructing ``datetime.datetime`` objects
+                                  manually. When both ``start`` and ``end`` times do not include a time zone, they will
+                                  be treated as UTC.
         :param end:               The end time of the endpoint's data, included.
                                   Please note: when ``start`` and ``end`` are set, they create a left-open time interval
                                   ("window") :math:`(\\operatorname{start}, \\operatorname{end}]` that excludes the
@@ -902,13 +929,13 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                                   if ``endpoints`` are passed.
                                   Note: the model monitoring infrastructure must be up for the writing to work.
         :param existing_data_handling:
-                                  How to handle the existing application data for the model endpoints when writing the
-                                  new data. Relevant only when ``write_output=True``. The default is
-                                  ``"fail_on_overlap"``. The options are:
+                                  How to handle the existing application data for the model endpoints when writing
+                                  new data whose requested ``start`` time precedes the ``end`` time of a previous run
+                                  that also wrote to the database. Relevant only when ``write_output=True``.
+                                  The options are:
 
-                                  - ``"fail_on_overlap"``: when the requested ``start`` time precedes the
-                                    ``end`` time of a previous run that also wrote to the database - an error is raised.
-                                  - ``"skip_overlap"``: when the previously described situation occurs, the relevant
+                                  - ``"fail_on_overlap"``: Default. An error is raised.
+                                  - ``"skip_overlap"``:  the overlapping data is ignored and the
                                     time window is cut so that it starts at the earliest possible time after ``start``.
                                   - ``"delete_all"``: delete all the data that was written by the application to the
                                     model endpoints, regardless of the time window, and write the new data.

mlrun/model_monitoring/applications/context.py

@@ -24,15 +24,12 @@ import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.errors
 import mlrun.feature_store as fstore
 import mlrun.feature_store.feature_set as fs
-import mlrun.features
 import mlrun.serving
 import mlrun.utils
 from mlrun.artifacts import Artifact, DatasetArtifact, ModelArtifact, get_model
 from mlrun.common.model_monitoring.helpers import FeatureStats
 from mlrun.common.schemas import ModelEndpoint
-from mlrun.model_monitoring.helpers import (
-    calculate_inputs_statistics,
-)
+from mlrun.model_monitoring.helpers import calculate_inputs_statistics
 
 
 class _ArtifactsLogger(Protocol):

mlrun/model_monitoring/controller.py

@@ -801,8 +801,17 @@ class MonitoringApplicationController:
         logger.info("Starting monitoring controller chief")
         applications_names = []
         endpoints = self.project_obj.list_model_endpoints(
-            tsdb_metrics=False, mode=mm_constants.EndpointMode.REAL_TIME
+            tsdb_metrics=False,
+            modes=[
+                mm_constants.EndpointMode.REAL_TIME,
+                mm_constants.EndpointMode.BATCH_LEGACY,
+            ],
         ).endpoints
+
+        if not endpoints:
+            logger.info("No model endpoints found", project=self.project)
+            return
+
         last_request_dict = self.tsdb_connector.get_last_request(
             endpoint_ids=[mep.metadata.uid for mep in endpoints]
         )
@@ -811,9 +820,6 @@ class MonitoringApplicationController:
                 mm_constants.EventFieldType.ENDPOINT_ID
             )[mm_constants.ModelEndpointSchema.LAST_REQUEST].to_dict()
 
-        if not endpoints:
-            logger.info("No model endpoints found", project=self.project)
-            return
         monitoring_functions = self.project_obj.list_model_monitoring_functions()
         if monitoring_functions:
             # if monitoring_functions: - TODO : ML-7700
@@ -859,7 +865,11 @@ class MonitoringApplicationController:
                 for endpoint in endpoints:
                     last_request = last_request_dict.get(endpoint.metadata.uid, None)
                     if isinstance(last_request, float):
-                        last_request = pd.to_datetime(last_request, unit="ms", utc=True)
+                        last_request = datetime.datetime.fromtimestamp(
+                            last_request, tz=datetime.timezone.utc
+                        )
+                    elif isinstance(last_request, pd.Timestamp):
+                        last_request = last_request.to_pydatetime()
                     endpoint.status.last_request = (
                         last_request or endpoint.status.last_request
                     )

mlrun/model_monitoring/db/_schedules.py

@@ -16,7 +16,7 @@ import json
 import sys
 from abc import ABC, abstractmethod
 from contextlib import AbstractContextManager
-from datetime import datetime, timezone
+from datetime import datetime
 from types import TracebackType
 from typing import TYPE_CHECKING, Final, Optional
 
@@ -281,9 +281,7 @@ class ModelMonitoringSchedulesFileApplication(ModelMonitoringSchedulesFileBase):
         self, endpoint_uid: str, last_analyzed: datetime
     ) -> None:
         self._check_open_schedules()
-        self._schedules[endpoint_uid] = last_analyzed.astimezone(
-            timezone.utc
-        ).isoformat()
+        self._schedules[endpoint_uid] = last_analyzed.isoformat()
 
     def delete_endpoints_last_analyzed(self, endpoint_uids: list[str]) -> None:
         self._check_open_schedules()

mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py

@@ -721,7 +721,9 @@ class TDEngineConnector(TSDBConnector):
         endpoint_ids: Union[str, list[str]],
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
-    ) -> pd.DataFrame:
+    ) -> Union[pd.DataFrame, dict[str, float]]:
+        if not endpoint_ids:
+            return {}
         filter_query = self._generate_filter_query(
             filter_column=mm_schemas.EventFieldType.ENDPOINT_ID,
             filter_values=endpoint_ids,

mlrun/model_monitoring/db/tsdb/v3io/stream_graph_steps.py

@@ -25,10 +25,12 @@ from mlrun.utils import logger
 
 def _normalize_dict_for_v3io_frames(event: dict[str, Any]) -> dict[str, Any]:
     """
-    Normalize user defined keys - input data to a model and its predictions,
-    to a form V3IO frames tolerates.
+    Normalize user-defined keys (e.g., model input data and predictions) to a format V3IO Frames tolerates.
 
-    The dictionary keys should conform to '^[a-zA-Z_:]([a-zA-Z0-9_:])*$'.
+    - Keys must match regex: '^[a-zA-Z_:]([a-zA-Z0-9_:])*$'
+    - Replace invalid characters (e.g., '-') with '_'.
+    - Prefix keys starting with digits with '_'.
+    - Flatten nested dictionaries using dot notation, while normalizing keys recursively.
     """
     prefix = "_"
 
@@ -38,7 +40,18 @@ def _normalize_dict_for_v3io_frames(event: dict[str, Any]) -> dict[str, Any]:
             return prefix + key
         return key
 
-    return {norm_key(k): v for k, v in event.items()}
+    def flatten_dict(d: dict[str, Any], parent_key: str = "") -> dict[str, Any]:
+        items = {}
+        for k, v in d.items():
+            new_key = norm_key(k)
+            full_key = f"{parent_key}.{new_key}" if parent_key else new_key
+            if isinstance(v, dict):
+                items.update(flatten_dict(v, full_key))
+            else:
+                items[full_key] = v
+        return items
+
+    return flatten_dict(event)
 
 
 class ProcessBeforeTSDB(mlrun.feature_store.steps.MapClass):

mlrun/model_monitoring/db/tsdb/v3io/v3io_connector.py

@@ -973,6 +973,9 @@ class V3IOTSDBConnector(TSDBConnector):
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
     ) -> dict[str, float]:
+        if not endpoint_ids:
+            return {}
+
         # Get the last request timestamp for each endpoint from the KV table.
         # The result of the query is a list of dictionaries,
         # each dictionary contains the endpoint id and the last request timestamp.

mlrun/model_monitoring/helpers.py

@@ -143,7 +143,7 @@ def get_stream_path(
         return stream_uri.replace("v3io://", f"ds://{profile.name}")
 
     elif isinstance(
-        profile, mlrun.datastore.datastore_profile.DatastoreProfileKafkaSource
+        profile, mlrun.datastore.datastore_profile.DatastoreProfileKafkaStream
     ):
         topic = mlrun.common.model_monitoring.helpers.get_kafka_topic(
             project=project, function_name=function_name
@@ -152,7 +152,7 @@ def get_stream_path(
     else:
         raise mlrun.errors.MLRunValueError(
             f"Received an unexpected stream profile type: {type(profile)}\n"
-            "Expects `DatastoreProfileV3io` or `DatastoreProfileKafkaSource`."
+            "Expects `DatastoreProfileV3io` or `DatastoreProfileKafkaStream`."
         )
 
 
@@ -300,7 +300,7 @@ def _get_v3io_output_stream(
 
 def _get_kafka_output_stream(
     *,
-    kafka_profile: mlrun.datastore.datastore_profile.DatastoreProfileKafkaSource,
+    kafka_profile: mlrun.datastore.datastore_profile.DatastoreProfileKafkaStream,
     project: str,
     function_name: str,
     mock: bool = False,
@@ -356,7 +356,7 @@ def get_output_stream(
         )
 
     elif isinstance(
-        profile, mlrun.datastore.datastore_profile.DatastoreProfileKafkaSource
+        profile, mlrun.datastore.datastore_profile.DatastoreProfileKafkaStream
     ):
         return _get_kafka_output_stream(
             kafka_profile=profile,
@@ -368,7 +368,7 @@ def get_output_stream(
     else:
         raise mlrun.errors.MLRunValueError(
             f"Received an unexpected stream profile type: {type(profile)}\n"
-            "Expects `DatastoreProfileV3io` or `DatastoreProfileKafkaSource`."
+            "Expects `DatastoreProfileV3io` or `DatastoreProfileKafkaStream`."
         )
 
 

mlrun/platforms/iguazio.py

@@ -96,7 +96,11 @@ class OutputStream:
         if access_key:
             v3io_client_kwargs["access_key"] = access_key
 
-        self._v3io_client = v3io.dataplane.Client(**v3io_client_kwargs)
+        if not mock:
+            self._v3io_client = v3io.dataplane.Client(**v3io_client_kwargs)
+        else:
+            self._v3io_client = None
+
         self._container, self._stream_path = split_path(stream_path)
         self._shards = shards
         self._retention_in_hours = retention_in_hours
@@ -105,7 +109,7 @@ class OutputStream:
         self._mock = mock
         self._mock_queue = []
 
-    def create_stream(self):
+    def create_stream(self) -> None:
         # this import creates an import loop via the utils module, so putting it in execution path
         from mlrun.utils.helpers import logger
 
@@ -210,7 +214,7 @@ class KafkaOutputStream:
         self._initialized = False
 
     def _lazy_init(self):
-        if self._initialized:
+        if self._initialized or self._mock:
             return
 
         import kafka