mlrun 1.9.0rc8__py3-none-any.whl → 1.9.0rc9__py3-none-any.whl

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

@@ -11,8 +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 traceback
+import time
 from collections.abc import Callable
 from enum import Enum
 from typing import Any, Final, Optional, Union
@@ -20,6 +19,9 @@ from typing import Any, Final, Optional, Union
 import taosws
 from taosws import TaosStmt
 
+import mlrun
+from mlrun.utils import logger
+
 
 class _StrEnum(str, Enum):
     pass
@@ -137,40 +139,99 @@ class Statement:
         return statement
 
 
-def _run(connection_string, prefix_statements, q, statements, query):
-    try:
-        conn = taosws.connect(connection_string)
-
-        for statement in prefix_statements + statements:
-            if isinstance(statement, Statement):
-                prepared_statement = statement.prepare(conn.statement())
-                prepared_statement.execute()
-            else:
-                conn.execute(statement)
-
-        if not query:
-            q.put(None)
-            return
-
-        res = conn.query(query)
-
-        # taosws.TaosField is not serializable
-        fields = [
-            Field(field.name(), field.type(), field.bytes()) for field in res.fields
-        ]
-
-        q.put(QueryResult(list(res), fields))
-    except Exception as e:
-        tb = traceback.format_exc()
-        q.put(ErrorResult(tb, e))
-
-
 class TDEngineConnection:
-    def __init__(self, connection_string):
+    def __init__(self, connection_string, max_retries=3, retry_delay=0.5):
         self._connection_string = connection_string
         self.prefix_statements = []
+        self._max_retries = max_retries
+        self._retry_delay = retry_delay
 
-        self._conn = taosws.connect(self._connection_string)
+        self._conn = self._create_connection()
+
+    def _create_connection(self):
+        """Create a new TDEngine connection."""
+        return taosws.connect(self._connection_string)
+
+    def _reconnect(self):
+        """Close current connection and create a new one."""
+        try:
+            if hasattr(self, "_conn") and self._conn:
+                self._conn.close()
+        except Exception as e:
+            logger.warning(f"Error closing connection during reconnect: {e}")
+
+        self._conn = self._create_connection()
+        logger.info("Successfully reconnected to TDEngine")
+
+    def _execute_with_retry(self, operation, operation_name, *args, **kwargs):
+        """
+        Execute an operation with retry logic for connection failures.
+
+        :param operation: The function to execute
+        :param operation_name: Name of the operation for logging
+        :param args: Arguments to pass to the operation
+        :param kwargs: Keyword arguments to pass to the operation
+        :return: Result of the operation
+        """
+        last_exception = None
+
+        for attempt in range(self._max_retries + 1):  # +1 for initial attempt
+            try:
+                return operation(*args, **kwargs)
+
+            except taosws.Error as e:
+                last_exception = e
+
+                if attempt < self._max_retries:
+                    logger.warning(
+                        f"Connection error during {operation_name} "
+                        f"(attempt {attempt + 1}/{self._max_retries + 1}): {e}. "
+                        f"Retrying in {self._retry_delay} seconds..."
+                    )
+
+                    # Wait before retrying
+                    time.sleep(self._retry_delay)
+
+                    # Reconnect
+                    try:
+                        self._reconnect()
+                    except Exception as reconnect_error:
+                        logger.error(f"Failed to reconnect: {reconnect_error}")
+                        if attempt == self._max_retries - 1:
+                            # Last attempt, raise the reconnection error
+                            raise TDEngineError(
+                                f"Failed to reconnect after {operation_name} failure: {reconnect_error}"
+                            ) from reconnect_error
+                        continue
+                else:
+                    # Max retries exceeded
+                    logger.error(
+                        f"Max retries ({self._max_retries}) exceeded for {operation_name}"
+                    )
+                    break
+
+            except Exception as e:
+                # Non-TDEngine error, don't retry
+                raise TDEngineError(
+                    f"Unexpected error during {operation_name}: {e}"
+                ) from e
+
+        # If we get here, all retries failed
+        raise TDEngineError(
+            f"Failed to {operation_name} after {self._max_retries} retries: {last_exception}"
+        ) from last_exception
+
+    def _execute_statement(self, statement):
+        """Execute a single statement (string or Statement object)."""
+        if isinstance(statement, Statement):
+            prepared_statement = statement.prepare(self._conn.statement())
+            prepared_statement.execute()
+        else:
+            self._conn.execute(statement)
+
+    def _execute_query(self, query):
+        """Execute a query and return the result."""
+        return self._conn.query(query)
 
     def run(
         self,
@@ -181,33 +242,40 @@ class TDEngineConnection:
         if not isinstance(statements, list):
             statements = [statements]
 
-        for statement in self.prefix_statements + statements:
+        # Execute all statements with retry logic
+        all_statements = self.prefix_statements + statements
+        for i, statement in enumerate(all_statements):
+            operation_name = f"execute statement {i + 1}/{len(all_statements)}"
             if isinstance(statement, Statement):
-                try:
-                    prepared_statement = statement.prepare(self._conn.statement())
-                    prepared_statement.execute()
-                except taosws.Error as e:
-                    raise TDEngineError(
-                        f"Failed to run prepared statement `{self._conn.statement()}`: {e}"
-                    ) from e
+                operation_name += " (prepared)"
             else:
-                try:
-                    self._conn.execute(statement)
-                except taosws.Error as e:
-                    raise TDEngineError(
-                        f"Failed to run statement `{statement}`: {e}"
-                    ) from e
+                operation_name += f" `{statement}`"
+
+            self._execute_with_retry(self._execute_statement, operation_name, statement)
 
         if not query:
             return None
 
-        try:
-            res = self._conn.query(query)
-        except taosws.Error as e:
-            raise TDEngineError(f"Failed to run query `{query}`: {e}") from e
+        # Execute query with retry logic
+        res = self._execute_with_retry(
+            self._execute_query, f"execute query `{query}`", query
+        )
 
+        # Process results
         fields = [
             Field(field.name(), field.type(), field.bytes()) for field in res.fields
         ]
 
         return QueryResult(list(res), fields)
+
+    def close(self):
+        """Close the connection."""
+        try:
+            if self._conn:
+                self._conn.close()
+                logger.debug("TDEngine connection closed")
+                self._conn = None
+        except Exception as e:
+            logger.warning(
+                f"Error closing TDEngine connection: {mlrun.errors.err_to_str(e)}"
+            )

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

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import threading
 from datetime import datetime, timedelta
-from threading import Lock
 from typing import Callable, Final, Literal, Optional, Union
 
 import pandas as pd
@@ -32,8 +32,8 @@ from mlrun.model_monitoring.db.tsdb.tdengine.tdengine_connection import (
 from mlrun.model_monitoring.helpers import get_invocations_fqn
 from mlrun.utils import logger
 
-_connection = None
-_connection_lock = Lock()
+# Thread-local storage for connections
+_thread_local = threading.local()
 
 
 class TDEngineTimestampPrecision(mlrun.common.types.StrEnum):
@@ -76,16 +76,15 @@ class TDEngineConnector(TSDBConnector):
 
     @property
     def connection(self) -> TDEngineConnection:
-        global _connection
-
-        if _connection:
-            return _connection
-
-        with _connection_lock:
-            if not _connection:
-                _connection = self._create_connection()
-
-        return _connection
+        if not hasattr(_thread_local, "connection"):
+            _thread_local.connection = self._create_connection()
+            logger.debug(
+                "Created new TDEngine connection for thread",
+                project=self.project,
+                thread_name=threading.current_thread().name,
+                thread_id=threading.get_ident(),
+            )
+        return _thread_local.connection
 
     def _create_connection(self) -> TDEngineConnection:
         """Establish a connection to the TSDB server."""

mlrun/projects/project.py

@@ -1408,7 +1408,10 @@ class MlrunProject(ModelObj):
                               https://apscheduler.readthedocs.io/en/3.x/modules/triggers/cron.html#module-apscheduler.triggers.cron
                               Note that "local" engine does not support this argument
         :param ttl:           Pipeline ttl in secs (after that the pods will be removed)
-        :param image:         Image for workflow runner job, only for scheduled and remote workflows
+        :param image:         Image for workflow runner job, only for scheduled and remote workflows.
+                              The image must have mlrun[kfp] installed which requires python 3.9.
+                              Therefore, the project default image will not be used for the workflow,
+                              and the image must be specified explicitly.
         :param args:          Argument values (key=value, ..)
         """
 

mlrun/utils/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "4754fd79dbb4d1b30f70ca03493ee7787da6089b",
-  "version": "1.9.0-rc8"
+  "git_commit": "d71319e6a92279074fcbe69811a8d7dbd7a2b642",
+  "version": "1.9.0-rc9"
 }

{mlrun-1.9.0rc8.dist-info → mlrun-1.9.0rc9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlrun
-Version: 1.9.0rc8
+Version: 1.9.0rc9
 Summary: Tracking and config of machine learning runs
 Home-page: https://github.com/mlrun/mlrun
 Author: Yaron Haviv
@@ -15,6 +15,7 @@ Classifier: Operating System :: Microsoft :: Windows
 Classifier: Operating System :: MacOS
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: Software Development :: Libraries

{mlrun-1.9.0rc8.dist-info → mlrun-1.9.0rc9.dist-info}/RECORD

@@ -242,8 +242,8 @@ mlrun/model_monitoring/db/tsdb/helpers.py,sha256=0oUXc4aUkYtP2SGP6jTb3uPPKImIUsV
 mlrun/model_monitoring/db/tsdb/tdengine/__init__.py,sha256=vgBdsKaXUURKqIf3M0y4sRatmSVA4CQiJs7J5dcVBkQ,620
 mlrun/model_monitoring/db/tsdb/tdengine/schemas.py,sha256=EslhaR65jfeNdD5Ibk-3Hb4e5r5qYPfHb9rTChX3sG0,12689
 mlrun/model_monitoring/db/tsdb/tdengine/stream_graph_steps.py,sha256=Uadj0UvAmln2MxDWod-kAzau1uNlqZh981rPhbUH_5M,2857
-mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connection.py,sha256=8xo2O_yQrJGNDoYYB3Bwtdwwvzs3U9dT3BtPot0zENQ,6449
-mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py,sha256=h0ZrNgOwTlBRd_DaYDc6eeVM9f_8CLJMUPEAIrZpbyU,37803
+mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connection.py,sha256=dtkaHaWKWERPXylEWMECeetwrz3rWl0P43AADcTjlls,9330
+mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py,sha256=f-3bycfwGvu7gX_Jh3vRT-s2k_yZRXLeIkBrt11xTdc,37998
 mlrun/model_monitoring/db/tsdb/v3io/__init__.py,sha256=aL3bfmQsUQ-sbvKGdNihFj8gLCK3mSys0qDcXtYOwgc,616
 mlrun/model_monitoring/db/tsdb/v3io/stream_graph_steps.py,sha256=_-zo9relCDtjGgievxAcAP9gVN9nDWs8BzGtFwTjb9M,6284
 mlrun/model_monitoring/db/tsdb/v3io/v3io_connector.py,sha256=aDdyHWJLOG-3EgPTJgbEzOOQf5NetpFQk_HdQbs5_zI,47160
@@ -271,7 +271,7 @@ mlrun/platforms/iguazio.py,sha256=6VBTq8eQ3mzT96tzjYhAtcMQ2VjF4x8LpIPW5DAcX2Q,13
 mlrun/projects/__init__.py,sha256=0Krf0WIKfnZa71WthYOg0SoaTodGg3sV_hK3f_OlTPI,1220
 mlrun/projects/operations.py,sha256=9ntpM8WnnyRk1iCY0NDKW-3aR4j2QZPSJM0SdMnsQKs,20032
 mlrun/projects/pipelines.py,sha256=wud7ezeEmhIJvfYE_wzQbA4ygEfGXHtbOtoOpan6poY,48556
-mlrun/projects/project.py,sha256=O3oeg5ZZLvsvftN_T6syB1xQJ7e1eFg37mrlIaL4VCY,237075
+mlrun/projects/project.py,sha256=VY9Wd1GfMH8CUUdOPN9UDsmcvsddlB_HYd2abodvSXY,237350
 mlrun/runtimes/__init__.py,sha256=J9Sy2HiyMlztNv6VUurMzF5H2XzttNil8nRsWDsqLyg,8923
 mlrun/runtimes/base.py,sha256=K-38p2HYSpaPcenJgEzQqLlLVFxfwLu2Bwe9oXGWCTM,37853
 mlrun/runtimes/daskjob.py,sha256=JwuGvOiPsxEDHHMMUS4Oie4hLlYYIZwihAl6DjroTY0,19521
@@ -341,11 +341,11 @@ mlrun/utils/notifications/notification/mail.py,sha256=ZyJ3eqd8simxffQmXzqd3bgbAq
 mlrun/utils/notifications/notification/slack.py,sha256=kfhogR5keR7Zjh0VCjJNK3NR5_yXT7Cv-x9GdOUW4Z8,7294
 mlrun/utils/notifications/notification/webhook.py,sha256=zxh8CAlbPnTazsk6r05X5TKwqUZVOH5KBU2fJbzQlG4,5330
 mlrun/utils/version/__init__.py,sha256=7kkrB7hEZ3cLXoWj1kPoDwo4MaswsI2JVOBpbKgPAgc,614
-mlrun/utils/version/version.json,sha256=-dffrM4IxXc-h_dn68qDgfH05FcPVSknGFKhWh-Ltsw,88
+mlrun/utils/version/version.json,sha256=9OoUMhTH3Ii0QEe_aO06dt56NtO2_K4FElMMumdqFx0,88
 mlrun/utils/version/version.py,sha256=eEW0tqIAkU9Xifxv8Z9_qsYnNhn3YH7NRAfM-pPLt1g,1878
-mlrun-1.9.0rc8.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mlrun-1.9.0rc8.dist-info/METADATA,sha256=itS4CUiW4wWgXKfGDT7AUaSKoT32z2Ao09gzQI7_HDw,25708
-mlrun-1.9.0rc8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mlrun-1.9.0rc8.dist-info/entry_points.txt,sha256=1Owd16eAclD5pfRCoJpYC2ZJSyGNTtUr0nCELMioMmU,46
-mlrun-1.9.0rc8.dist-info/top_level.txt,sha256=NObLzw3maSF9wVrgSeYBv-fgnHkAJ1kEkh12DLdd5KM,6
-mlrun-1.9.0rc8.dist-info/RECORD,,
+mlrun-1.9.0rc9.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mlrun-1.9.0rc9.dist-info/METADATA,sha256=Zg7y-d5F7YBeE5O0lnyaiDwrVlCNPU9j2A4hlnXoRSk,25759
+mlrun-1.9.0rc9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mlrun-1.9.0rc9.dist-info/entry_points.txt,sha256=1Owd16eAclD5pfRCoJpYC2ZJSyGNTtUr0nCELMioMmU,46
+mlrun-1.9.0rc9.dist-info/top_level.txt,sha256=NObLzw3maSF9wVrgSeYBv-fgnHkAJ1kEkh12DLdd5KM,6
+mlrun-1.9.0rc9.dist-info/RECORD,,