mlrun 1.7.0rc37__py3-none-any.whl → 1.7.0rc39__py3-none-any.whl

mlrun/alerts/alert.py

@@ -29,6 +29,7 @@ class AlertConfig(ModelObj):
         "reset_policy",
         "state",
         "count",
+        "created",
     ]
     _fields_to_serialize = ModelObj._fields_to_serialize + [
         "entities",
@@ -55,12 +56,13 @@ class AlertConfig(ModelObj):
         created: str = None,
         count: int = None,
     ):
-        """
-        Alert config object
+        """Alert config object
 
         Example::
+
             # create an alert on endpoint_id, which will be triggered to slack if there is a "data_drift_detected" event
-            3 times in the next hour.
+            # 3 times in the next hour.
+
             from mlrun.alerts import AlertConfig
             import mlrun.common.schemas.alert as alert_objects
 
@@ -93,29 +95,29 @@ class AlertConfig(ModelObj):
             )
             project.store_alert_config(alert_data)
 
-        :param project:        name of the project to associate the alert with
-        :param name:           name of the alert
-        :param template:       optional parameter that allows to create an alert based on a predefined template.
-                               you can pass either an AlertTemplate object or a string (the template name).
-                               if a template is used, many fields of the alert will be auto-generated based on the
-                               template. however, you still need to provide the following fields:
+        :param project:        Name of the project to associate the alert with
+        :param name:           Name of the alert
+        :param template:       Optional parameter that allows creating an alert based on a predefined template.
+                               You can pass either an AlertTemplate object or a string (the template name).
+                               If a template is used, many fields of the alert will be auto-generated based on the
+                               template.However, you still need to provide the following fields:
                                `name`, `project`, `entity`, `notifications`
-        :param description:    description of the alert
-        :param summary:        summary of the alert, will be sent in the generated notifications
-        :param severity:       severity of the alert
-        :param trigger:        the events that will trigger this alert, may be a simple trigger based on events or
+        :param description:    Description of the alert
+        :param summary:        Summary of the alert, will be sent in the generated notifications
+        :param severity:       Severity of the alert
+        :param trigger:        The events that will trigger this alert, may be a simple trigger based on events or
                                complex trigger which is based on a prometheus alert
-        :param criteria:       when the alert will be triggered based on the specified number of events within the
+        :param criteria:       When the alert will be triggered based on the specified number of events within the
                                defined time period.
-        :param reset_policy:   when to clear the alert. May be "manual" for manual reset of the alert, or
+        :param reset_policy:   When to clear the alert. May be "manual" for manual reset of the alert, or
                                "auto" if the criteria contains a time period
-        :param notifications:  list of notifications to invoke once the alert is triggered
-        :param entities:       entities that the event relates to. The entity object will contain fields that uniquely
-                               identify a given entity in the system
-        :param id:             internal id of the alert (user should not supply it)
-        :param state:          state of the alert, may be active/inactive (user should not supply it)
-        :param created:        when the alert is created (user should not supply it)
-        :param count:          internal counter of the alert (user should not supply it)
+        :param notifications:  List of notifications to invoke once the alert is triggered
+        :param entities:       Entities that the event relates to. The entity object will contain fields that
+                               uniquely identify a given entity in the system
+        :param id:             Internal id of the alert (user should not supply it)
+        :param state:          State of the alert, may be active/inactive (user should not supply it)
+        :param created:        When the alert is created (user should not supply it)
+        :param count:          Internal counter of the alert (user should not supply it)
         """
         self.project = project
         self.name = name
@@ -136,8 +138,8 @@ class AlertConfig(ModelObj):
             self._apply_template(template)
 
     def validate_required_fields(self):
-        if not self.project or not self.name:
-            raise mlrun.errors.MLRunBadRequestError("Project and name must be provided")
+        if not self.name:
+            raise mlrun.errors.MLRunInvalidArgumentError("Alert name must be provided")
 
     def _serialize_field(
         self, struct: dict, field_name: str = None, strip: bool = False
@@ -236,9 +238,11 @@ class AlertConfig(ModelObj):
             db = mlrun.get_run_db()
             template = db.get_alert_template(template)
 
-        # Extract parameters from the template and apply them to the AlertConfig object
-        self.summary = template.summary
-        self.severity = template.severity
-        self.criteria = template.criteria
-        self.trigger = template.trigger
-        self.reset_policy = template.reset_policy
+        # Apply parameters from the template to the AlertConfig object only if they are not already specified by the
+        # user in the current configuration.
+        # User-provided parameters will take precedence over corresponding template values
+        self.summary = self.summary or template.summary
+        self.severity = self.severity or template.severity
+        self.criteria = self.criteria or template.criteria
+        self.trigger = self.trigger or template.trigger
+        self.reset_policy = self.reset_policy or template.reset_policy

mlrun/common/schemas/alert.py

@@ -23,6 +23,7 @@ from mlrun.common.types import StrEnum
 
 class EventEntityKind(StrEnum):
     MODEL_ENDPOINT_RESULT = "model-endpoint-result"
+    MODEL_MONITORING_APPLICATION = "model-monitoring-application"
     JOB = "job"
 
 
@@ -43,6 +44,7 @@ class EventKind(StrEnum):
     SYSTEM_PERFORMANCE_SUSPECTED = "system_performance_suspected"
     MM_APP_ANOMALY_DETECTED = "mm_app_anomaly_detected"
     MM_APP_ANOMALY_SUSPECTED = "mm_app_anomaly_suspected"
+    MM_APP_FAILED = "mm_app_failed"
     FAILED = "failed"
 
 
@@ -57,6 +59,7 @@ _event_kind_entity_map = {
     EventKind.SYSTEM_PERFORMANCE_SUSPECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
     EventKind.MM_APP_ANOMALY_DETECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
     EventKind.MM_APP_ANOMALY_SUSPECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
+    EventKind.MM_APP_FAILED: [EventEntityKind.MODEL_MONITORING_APPLICATION],
     EventKind.FAILED: [EventEntityKind.JOB],
 }
 

mlrun/common/schemas/model_monitoring/constants.py

@@ -53,9 +53,11 @@ class EventFieldType:
     PREDICTIONS = "predictions"
     NAMED_PREDICTIONS = "named_predictions"
     ERROR_COUNT = "error_count"
+    MODEL_ERROR = "model_error"
     ENTITIES = "entities"
     FIRST_REQUEST = "first_request"
     LAST_REQUEST = "last_request"
+    LAST_REQUEST_TIMESTAMP = "last_request_timestamp"
     METRIC = "metric"
     METRICS = "metrics"
     BATCH_INTERVALS_DICT = "batch_intervals_dict"
@@ -217,6 +219,7 @@ class FileTargetKind:
     APP_METRICS = "app_metrics"
     MONITORING_SCHEDULES = "monitoring_schedules"
     MONITORING_APPLICATION = "monitoring_application"
+    ERRORS = "errors"
 
 
 class ModelMonitoringMode(str, Enum):
@@ -240,6 +243,7 @@ class V3IOTSDBTables(MonitoringStrEnum):
     APP_RESULTS = "app-results"
     METRICS = "metrics"
     EVENTS = "events"
+    ERRORS = "errors"
 
 
 class TDEngineSuperTables(MonitoringStrEnum):

mlrun/common/schemas/notification.py

@@ -52,6 +52,7 @@ class NotificationLimits(enum.Enum):
 class Notification(pydantic.BaseModel):
     """
     Notification object schema
+
     :param kind: notification implementation kind - slack, webhook, etc.
     :param name: for logging and identification
     :param message: message content in the notification
@@ -71,9 +72,9 @@ class Notification(pydantic.BaseModel):
 
     kind: NotificationKind
     name: str
-    message: str
-    severity: NotificationSeverity
-    when: list[str]
+    message: typing.Optional[str] = None
+    severity: typing.Optional[NotificationSeverity] = None
+    when: typing.Optional[list[str]] = None
     condition: typing.Optional[str] = None
     params: typing.Optional[dict[str, typing.Any]] = None
     status: typing.Optional[NotificationStatus] = None

mlrun/datastore/alibaba_oss.py

@@ -22,7 +22,7 @@ from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
-from .base import DataStore, FileStats, makeDatastoreSchemaSanitizer
+from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 
 
 class OSSStore(DataStore):
@@ -53,7 +53,7 @@ class OSSStore(DataStore):
         except ImportError as exc:
             raise ImportError("ALIBABA ossfs not installed") from exc
         filesystem_class = get_filesystem_class(protocol=self.kind)
-        self._filesystem = makeDatastoreSchemaSanitizer(
+        self._filesystem = make_datastore_schema_sanitizer(
             filesystem_class,
             using_bucket=self.using_bucket,
             **self.get_storage_options(),

mlrun/datastore/azure_blob.py

@@ -16,12 +16,13 @@ import time
 from pathlib import Path
 from urllib.parse import urlparse
 
+from azure.storage.blob import BlobServiceClient
 from azure.storage.blob._shared.base_client import parse_connection_str
 from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
-from .base import DataStore, FileStats, makeDatastoreSchemaSanitizer
+from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 
 # Azure blobs will be represented with the following URL: az://<container name>. The storage account is already
 # pointed to by the connection string, so the user is not expected to specify it in any way.
@@ -29,47 +30,131 @@ from .base import DataStore, FileStats, makeDatastoreSchemaSanitizer
 
 class AzureBlobStore(DataStore):
     using_bucket = True
+    max_concurrency = 100
+    max_blocksize = 1024 * 1024 * 4
+    max_single_put_size = (
+        1024 * 1024 * 8
+    )  # for service_client property only, does not affect filesystem
 
     def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
+        self._service_client = None
+        self._storage_options = None
+
+    def get_storage_options(self):
+        return self.storage_options
+
+    @property
+    def storage_options(self):
+        if not self._storage_options:
+            res = dict(
+                account_name=self._get_secret_or_env("account_name")
+                or self._get_secret_or_env("AZURE_STORAGE_ACCOUNT_NAME"),
+                account_key=self._get_secret_or_env("account_key")
+                or self._get_secret_or_env("AZURE_STORAGE_ACCOUNT_KEY"),
+                connection_string=self._get_secret_or_env("connection_string")
+                or self._get_secret_or_env("AZURE_STORAGE_CONNECTION_STRING"),
+                tenant_id=self._get_secret_or_env("tenant_id")
+                or self._get_secret_or_env("AZURE_STORAGE_TENANT_ID"),
+                client_id=self._get_secret_or_env("client_id")
+                or self._get_secret_or_env("AZURE_STORAGE_CLIENT_ID"),
+                client_secret=self._get_secret_or_env("client_secret")
+                or self._get_secret_or_env("AZURE_STORAGE_CLIENT_SECRET"),
+                sas_token=self._get_secret_or_env("sas_token")
+                or self._get_secret_or_env("AZURE_STORAGE_SAS_TOKEN"),
+                credential=self._get_secret_or_env("credential"),
+            )
+            self._storage_options = self._sanitize_storage_options(res)
+        return self._storage_options
 
     @property
     def filesystem(self):
         """return fsspec file system object, if supported"""
-        if self._filesystem:
-            return self._filesystem
         try:
             import adlfs  # noqa
         except ImportError as exc:
             raise ImportError("Azure adlfs not installed") from exc
-        # in order to support az and wasbs kinds.
-        filesystem_class = get_filesystem_class(protocol=self.kind)
-        self._filesystem = makeDatastoreSchemaSanitizer(
-            filesystem_class,
-            using_bucket=self.using_bucket,
-            **self.get_storage_options(),
-        )
+
+        if not self._filesystem:
+            # in order to support az and wasbs kinds
+            filesystem_class = get_filesystem_class(protocol=self.kind)
+            self._filesystem = make_datastore_schema_sanitizer(
+                filesystem_class,
+                using_bucket=self.using_bucket,
+                blocksize=self.max_blocksize,
+                **self.storage_options,
+            )
         return self._filesystem
 
-    def get_storage_options(self):
-        res = dict(
-            account_name=self._get_secret_or_env("account_name")
-            or self._get_secret_or_env("AZURE_STORAGE_ACCOUNT_NAME"),
-            account_key=self._get_secret_or_env("account_key")
-            or self._get_secret_or_env("AZURE_STORAGE_KEY"),
-            connection_string=self._get_secret_or_env("connection_string")
-            or self._get_secret_or_env("AZURE_STORAGE_CONNECTION_STRING"),
-            tenant_id=self._get_secret_or_env("tenant_id")
-            or self._get_secret_or_env("AZURE_STORAGE_TENANT_ID"),
-            client_id=self._get_secret_or_env("client_id")
-            or self._get_secret_or_env("AZURE_STORAGE_CLIENT_ID"),
-            client_secret=self._get_secret_or_env("client_secret")
-            or self._get_secret_or_env("AZURE_STORAGE_CLIENT_SECRET"),
-            sas_token=self._get_secret_or_env("sas_token")
-            or self._get_secret_or_env("AZURE_STORAGE_SAS_TOKEN"),
-            credential=self._get_secret_or_env("credential"),
-        )
-        return self._sanitize_storage_options(res)
+    @property
+    def service_client(self):
+        try:
+            import azure  # noqa
+        except ImportError as exc:
+            raise ImportError("Azure not installed") from exc
+
+        if not self._service_client:
+            self._do_connect()
+        return self._service_client
+
+    def _do_connect(self):
+        """
+
+        Creates a client for azure.
+        Raises MLRunInvalidArgumentError if none of the connection details are available
+        based on do_connect in AzureBlobFileSystem:
+        https://github.com/fsspec/adlfs/blob/2023.9.0/adlfs/spec.py#L422
+        """
+        from azure.identity import ClientSecretCredential
+
+        storage_options = self.storage_options
+        connection_string = storage_options.get("connection_string")
+        client_name = storage_options.get("account_name")
+        account_key = storage_options.get("account_key")
+        sas_token = storage_options.get("sas_token")
+        client_id = storage_options.get("client_id")
+        credential = storage_options.get("credential")
+
+        credential_from_client_id = None
+        if (
+            credential is None
+            and account_key is None
+            and sas_token is None
+            and client_id is not None
+        ):
+            credential_from_client_id = ClientSecretCredential(
+                tenant_id=storage_options.get("tenant_id"),
+                client_id=client_id,
+                client_secret=storage_options.get("client_secret"),
+            )
+        try:
+            if connection_string is not None:
+                self._service_client = BlobServiceClient.from_connection_string(
+                    conn_str=connection_string,
+                    max_block_size=self.max_blocksize,
+                    max_single_put_size=self.max_single_put_size,
+                )
+            elif client_name is not None:
+                account_url = f"https://{client_name}.blob.core.windows.net"
+                cred = credential_from_client_id or credential or account_key
+                if not cred and sas_token is not None:
+                    if not sas_token.startswith("?"):
+                        sas_token = f"?{sas_token}"
+                    account_url = account_url + sas_token
+                self._service_client = BlobServiceClient(
+                    account_url=account_url,
+                    credential=cred,
+                    max_block_size=self.max_blocksize,
+                    max_single_put_size=self.max_single_put_size,
+                )
+            else:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Must provide either a connection_string or account_name with credentials"
+                )
+        except Exception as e:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"unable to connect to account for {e}"
+            )
 
     def _convert_key_to_remote_path(self, key):
         key = key.strip("/")
@@ -82,7 +167,15 @@ class AzureBlobStore(DataStore):
 
     def upload(self, key, src_path):
         remote_path = self._convert_key_to_remote_path(key)
-        self.filesystem.put_file(src_path, remote_path, overwrite=True)
+        container, remote_path = remote_path.split("/", 1)
+        container_client = self.service_client.get_container_client(container=container)
+        with open(file=src_path, mode="rb") as data:
+            container_client.upload_blob(
+                name=remote_path,
+                data=data,
+                overwrite=True,
+                max_concurrency=self.max_concurrency,
+            )
 
     def get(self, key, size=None, offset=0):
         remote_path = self._convert_key_to_remote_path(key)
@@ -135,7 +228,7 @@ class AzureBlobStore(DataStore):
 
     def get_spark_options(self):
         res = {}
-        st = self.get_storage_options()
+        st = self.storage_options
         service = "blob"
         primary_url = None
         if st.get("connection_string"):

mlrun/datastore/base.py

@@ -748,7 +748,7 @@ class HttpStore(DataStore):
 # As an example, it converts an S3 URL 's3://s3bucket/path' to just 's3bucket/path'.
 # Since 'ds' schemas are not inherently processed by fsspec, we have adapted the _strip_protocol()
 # method specifically to strip away the 'ds' schema as required.
-def makeDatastoreSchemaSanitizer(cls, using_bucket=False, *args, **kwargs):
+def make_datastore_schema_sanitizer(cls, using_bucket=False, *args, **kwargs):
     if not issubclass(cls, fsspec.AbstractFileSystem):
         raise ValueError("Class must be a subclass of fsspec.AbstractFileSystem")
 

mlrun/datastore/dbfs_store.py

@@ -19,7 +19,7 @@ from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
-from .base import DataStore, FileStats, makeDatastoreSchemaSanitizer
+from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 
 
 class DatabricksFileBugFixed(DatabricksFile):
@@ -89,7 +89,7 @@ class DBFSStore(DataStore):
         """return fsspec file system object, if supported"""
         filesystem_class = get_filesystem_class(protocol=self.kind)
         if not self._filesystem:
-            self._filesystem = makeDatastoreSchemaSanitizer(
+            self._filesystem = make_datastore_schema_sanitizer(
                 cls=filesystem_class,
                 using_bucket=False,
                 **self.get_storage_options(),

mlrun/datastore/google_cloud_storage.py

@@ -12,44 +12,82 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import os
 from pathlib import Path
 
 from fsspec.registry import get_filesystem_class
+from google.auth.credentials import Credentials
+from google.cloud.storage import Client, transfer_manager
+from google.oauth2 import service_account
 
 import mlrun.errors
 from mlrun.utils import logger
 
-from .base import DataStore, FileStats, makeDatastoreSchemaSanitizer
+from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 
 # Google storage objects will be represented with the following URL: gcs://<bucket name>/<path> or gs://...
 
 
 class GoogleCloudStorageStore(DataStore):
     using_bucket = True
+    workers = 8
+    chunk_size = 32 * 1024 * 1024
 
     def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
+        self._storage_client = None
+        self._storage_options = None
+
+    @property
+    def storage_client(self):
+        if self._storage_client:
+            return self._storage_client
+
+        token = self._get_credentials().get("token")
+        access = "https://www.googleapis.com/auth/devstorage.full_control"
+        if isinstance(token, str):
+            if os.path.exists(token):
+                credentials = service_account.Credentials.from_service_account_file(
+                    token, scopes=[access]
+                )
+            else:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "gcsfs authentication file not found!"
+                )
+        elif isinstance(token, dict):
+            credentials = service_account.Credentials.from_service_account_info(
+                token, scopes=[access]
+            )
+        elif isinstance(token, Credentials):
+            credentials = token
+        else:
+            raise ValueError(f"Unsupported token type: {type(token)}")
+        self._storage_client = Client(credentials=credentials)
+        return self._storage_client
 
     @property
     def filesystem(self):
         """return fsspec file system object, if supported"""
-        if self._filesystem:
-            return self._filesystem
-        try:
-            import gcsfs  # noqa
-        except ImportError as exc:
-            raise ImportError(
-                "Google gcsfs not installed, run pip install gcsfs"
-            ) from exc
-        filesystem_class = get_filesystem_class(protocol=self.kind)
-        self._filesystem = makeDatastoreSchemaSanitizer(
-            filesystem_class,
-            using_bucket=self.using_bucket,
-            **self.get_storage_options(),
-        )
+        if not self._filesystem:
+            filesystem_class = get_filesystem_class(protocol=self.kind)
+            self._filesystem = make_datastore_schema_sanitizer(
+                filesystem_class,
+                using_bucket=self.using_bucket,
+                **self.storage_options,
+            )
         return self._filesystem
 
-    def get_storage_options(self):
+    @property
+    def storage_options(self):
+        if self._storage_options:
+            return self._storage_options
+        credentials = self._get_credentials()
+        # due to caching problem introduced in gcsfs 2024.3.1 (ML-7636)
+        credentials["use_listings_cache"] = False
+        self._storage_options = credentials
+        return self._storage_options
+
+    def _get_credentials(self):
         credentials = self._get_secret_or_env(
             "GCP_CREDENTIALS"
         ) or self._get_secret_or_env("GOOGLE_APPLICATION_CREDENTIALS")
@@ -71,6 +109,9 @@ class GoogleCloudStorageStore(DataStore):
             )
             return self._sanitize_storage_options(None)
 
+    def get_storage_options(self):
+        return self.storage_options
+
     def _make_path(self, key):
         key = key.strip("/")
         path = Path(self.endpoint, key).as_posix()
@@ -103,8 +144,29 @@ class GoogleCloudStorageStore(DataStore):
             f.write(data)
 
     def upload(self, key, src_path):
-        path = self._make_path(key)
-        self.filesystem.put_file(src_path, path, overwrite=True)
+        file_size = os.path.getsize(src_path)
+        united_path = self._make_path(key)
+
+        # Multiple upload limitation recommendations as described in
+        # https://cloud.google.com/storage/docs/multipart-uploads#storage-upload-object-chunks-python
+
+        if file_size <= self.chunk_size:
+            self.filesystem.put_file(src_path, united_path, overwrite=True)
+            return
+
+        bucket = self.storage_client.bucket(self.endpoint)
+        blob = bucket.blob(key.strip("/"))
+
+        try:
+            transfer_manager.upload_chunks_concurrently(
+                src_path, blob, chunk_size=self.chunk_size, max_workers=self.workers
+            )
+        except Exception as upload_chunks_concurrently_exception:
+            logger.warning(
+                f"gcs: failed to concurrently upload {src_path},"
+                f" exception: {upload_chunks_concurrently_exception}. Retrying with single part upload."
+            )
+            self.filesystem.put_file(src_path, united_path, overwrite=True)
 
     def stat(self, key):
         path = self._make_path(key)
@@ -133,12 +195,13 @@ class GoogleCloudStorageStore(DataStore):
 
     def rm(self, path, recursive=False, maxdepth=None):
         path = self._make_path(path)
+        # in order to raise an error in case of a connection error (ML-7056)
         self.filesystem.exists(path)
-        self.filesystem.rm(path=path, recursive=recursive, maxdepth=maxdepth)
+        super().rm(path, recursive=recursive, maxdepth=maxdepth)
 
     def get_spark_options(self):
         res = {}
-        st = self.get_storage_options()
+        st = self._get_credentials()
         if "token" in st:
             res = {"spark.hadoop.google.cloud.auth.service.account.enable": "true"}
             if isinstance(st["token"], str):

mlrun/datastore/s3.py

@@ -20,7 +20,7 @@ from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
-from .base import DataStore, FileStats, get_range, makeDatastoreSchemaSanitizer
+from .base import DataStore, FileStats, get_range, make_datastore_schema_sanitizer
 
 
 class S3Store(DataStore):
@@ -126,7 +126,7 @@ class S3Store(DataStore):
         except ImportError as exc:
             raise ImportError("AWS s3fs not installed") from exc
         filesystem_class = get_filesystem_class(protocol=self.kind)
-        self._filesystem = makeDatastoreSchemaSanitizer(
+        self._filesystem = make_datastore_schema_sanitizer(
             filesystem_class,
             using_bucket=self.using_bucket,
             **self.get_storage_options(),

mlrun/datastore/sources.py

@@ -32,6 +32,7 @@ from mlrun.config import config
 from mlrun.datastore.snowflake_utils import get_snowflake_spark_options
 from mlrun.datastore.utils import transform_list_filters_to_tuple
 from mlrun.secrets import SecretsStore
+from mlrun.utils import logger
 
 from ..model import DataSource
 from ..platforms.iguazio import parse_path
@@ -1163,6 +1164,59 @@ class KafkaSource(OnlineSource):
             "to a Spark dataframe is not possible, as this operation is not supported by Spark"
         )
 
+    def create_topics(
+        self,
+        num_partitions: int = 4,
+        replication_factor: int = 1,
+        topics: list[str] = None,
+    ):
+        """
+        Create Kafka topics with the specified number of partitions and replication factor.
+
+        :param num_partitions:      number of partitions for the topics
+        :param replication_factor:  replication factor for the topics
+        :param topics:              list of topic names to create, if None,
+                                    the topics will be taken from the source attributes
+        """
+        from kafka.admin import KafkaAdminClient, NewTopic
+
+        brokers = self.attributes.get("brokers")
+        if not brokers:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "brokers must be specified in the KafkaSource attributes"
+            )
+        topics = topics or self.attributes.get("topics")
+        if not topics:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "topics must be specified in the KafkaSource attributes"
+            )
+        new_topics = [
+            NewTopic(topic, num_partitions, replication_factor) for topic in topics
+        ]
+        kafka_admin = KafkaAdminClient(
+            bootstrap_servers=brokers,
+            sasl_mechanism=self.attributes.get("sasl", {}).get("sasl_mechanism"),
+            sasl_plain_username=self.attributes.get("sasl", {}).get("username"),
+            sasl_plain_password=self.attributes.get("sasl", {}).get("password"),
+            sasl_kerberos_service_name=self.attributes.get("sasl", {}).get(
+                "sasl_kerberos_service_name", "kafka"
+            ),
+            sasl_kerberos_domain_name=self.attributes.get("sasl", {}).get(
+                "sasl_kerberos_domain_name"
+            ),
+            sasl_oauth_token_provider=self.attributes.get("sasl", {}).get("mechanism"),
+        )
+        try:
+            kafka_admin.create_topics(new_topics)
+        finally:
+            kafka_admin.close()
+        logger.info(
+            "Kafka topics created successfully",
+            topics=topics,
+            num_partitions=num_partitions,
+            replication_factor=replication_factor,
+        )
+
 
 class SQLSource(BaseSourceDriver):
     kind = "sqldb"