mlrun 1.7.0rc50__py3-none-any.whl → 1.7.0rc51__py3-none-any.whl

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

@@ -57,8 +57,25 @@ class TDEngineConnector(TSDBConnector):
             self._connection = self._create_connection()
         return self._connection
 
+    def with_retry_on_closed_connection(self, fn, **kwargs):
+        try:
+            return fn(self.connection, **kwargs)
+        except (taosws.QueryError, taosws.FetchError) as err:
+            logger.warn(f"TDEngine error: {err}")
+            if "Internal error:" in str(err):
+                logger.info("Retrying TDEngine query with a new connection")
+                try:
+                    self._connection.close()
+                except Exception:
+                    pass
+                self._connection = None
+                return fn(self.connection, **kwargs)
+            else:
+                raise err
+
     def _create_connection(self) -> taosws.Connection:
         """Establish a connection to the TSDB server."""
+        logger.debug("Creating a new connection to TDEngine", project=self.project)
         conn = taosws.connect(self._tdengine_connection_string)
         try:
             conn.execute(f"CREATE DATABASE {self.database}")
@@ -71,6 +88,7 @@ class TDEngineConnector(TSDBConnector):
             raise mlrun.errors.MLRunTSDBConnectionFailureError(
                 f"Failed to use TDEngine database {self.database}, {mlrun.errors.err_to_str(e)}"
             )
+        logger.debug("Connected to TDEngine", project=self.project)
         return conn
 
     def _init_super_tables(self):
@@ -91,7 +109,9 @@ class TDEngineConnector(TSDBConnector):
         """Create TDEngine supertables."""
         for table in self.tables:
             create_table_query = self.tables[table]._create_super_table_query()
-            self.connection.execute(create_table_query)
+            self.with_retry_on_closed_connection(
+                lambda conn: conn.execute(create_table_query)
+            )
 
     def write_application_event(
         self,
@@ -137,10 +157,14 @@ class TDEngineConnector(TSDBConnector):
         )
 
         create_table_sql = table._create_subtable_sql(subtable=table_name, values=event)
-        self.connection.execute(create_table_sql)
+        self.with_retry_on_closed_connection(
+            lambda conn: conn.execute(create_table_sql)
+        )
 
-        insert_statement = table._insert_subtable_stmt(
-            self.connection, subtable=table_name, values=event
+        insert_statement = self.with_retry_on_closed_connection(
+            lambda conn: table._insert_subtable_stmt(
+                conn, subtable=table_name, values=event
+            )
         )
         insert_statement.add_batch()
         insert_statement.execute()
@@ -200,18 +224,25 @@ class TDEngineConnector(TSDBConnector):
         """
         Delete all project resources in the TSDB connector, such as model endpoints data and drift results.
         """
+        logger.debug(
+            "Deleting all project resources using the TDEngine connector",
+            project=self.project,
+        )
         for table in self.tables:
             get_subtable_names_query = self.tables[table]._get_subtables_query(
                 values={mm_schemas.EventFieldType.PROJECT: self.project}
             )
-            subtables = self.connection.query(get_subtable_names_query)
+            subtables = self.with_retry_on_closed_connection(
+                lambda conn: conn.query(get_subtable_names_query)
+            )
             for subtable in subtables:
                 drop_query = self.tables[table]._drop_subtable_query(
                     subtable=subtable[0]
                 )
                 self.connection.execute(drop_query)
-        logger.info(
-            f"Deleted all project resources in the TSDB connector for project {self.project}"
+        logger.debug(
+            "Deleted all project resources using the TDEngine connector",
+            project=self.project,
         )
 
     def get_model_endpoint_real_time_metrics(
@@ -282,7 +313,9 @@ class TDEngineConnector(TSDBConnector):
         )
         logger.debug("Querying TDEngine", query=full_query)
         try:
-            query_result = self.connection.query(full_query)
+            query_result = self.with_retry_on_closed_connection(
+                lambda conn: conn.query(full_query)
+            )
         except taosws.QueryError as e:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"Failed to query table {table} in database {self.database}, {str(e)}"

mlrun/utils/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "fe89f8ec27cecdb512643519d74f5debf285aacb",
-  "version": "1.7.0-rc50"
+  "git_commit": "a8919d5f6a549f9434544e023887b886d432d5c9",
+  "version": "1.7.0-rc51"
 }

{mlrun-1.7.0rc50.dist-info → mlrun-1.7.0rc51.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlrun
-Version: 1.7.0rc50
+Version: 1.7.0rc51
 Summary: Tracking and config of machine learning runs
 Home-page: https://github.com/mlrun/mlrun
 Author: Yaron Haviv
@@ -84,7 +84,7 @@ Requires-Dist: redis ~=4.3 ; extra == 'all'
 Requires-Dist: s3fs <2024.7,>=2023.9.2 ; extra == 'all'
 Requires-Dist: snowflake-connector-python ~=3.7 ; extra == 'all'
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'all'
-Requires-Dist: taos-ws-py ~=0.3.3 ; extra == 'all'
+Requires-Dist: taos-ws-py ==0.3.2 ; extra == 'all'
 Provides-Extra: api
 Requires-Dist: uvicorn ~=0.27.1 ; extra == 'api'
 Requires-Dist: dask-kubernetes ~=0.11.0 ; extra == 'api'
@@ -137,7 +137,7 @@ Requires-Dist: redis ~=4.3 ; extra == 'complete'
 Requires-Dist: s3fs <2024.7,>=2023.9.2 ; extra == 'complete'
 Requires-Dist: snowflake-connector-python ~=3.7 ; extra == 'complete'
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'complete'
-Requires-Dist: taos-ws-py ~=0.3.3 ; extra == 'complete'
+Requires-Dist: taos-ws-py ==0.3.2 ; extra == 'complete'
 Provides-Extra: complete-api
 Requires-Dist: adlfs ==2023.9.0 ; extra == 'complete-api'
 Requires-Dist: aiobotocore <2.16,>=2.5.0 ; extra == 'complete-api'
@@ -174,7 +174,7 @@ Requires-Dist: redis ~=4.3 ; extra == 'complete-api'
 Requires-Dist: s3fs <2024.7,>=2023.9.2 ; extra == 'complete-api'
 Requires-Dist: snowflake-connector-python ~=3.7 ; extra == 'complete-api'
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'complete-api'
-Requires-Dist: taos-ws-py ~=0.3.3 ; extra == 'complete-api'
+Requires-Dist: taos-ws-py ==0.3.2 ; extra == 'complete-api'
 Requires-Dist: timelength ~=1.1 ; extra == 'complete-api'
 Requires-Dist: uvicorn ~=0.27.1 ; extra == 'complete-api'
 Requires-Dist: memray ~=1.12 ; (sys_platform != "win32") and extra == 'complete-api'
@@ -209,7 +209,7 @@ Requires-Dist: snowflake-connector-python ~=3.7 ; extra == 'snowflake'
 Provides-Extra: sqlalchemy
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'sqlalchemy'
 Provides-Extra: tdengine
-Requires-Dist: taos-ws-py ~=0.3.3 ; extra == 'tdengine'
+Requires-Dist: taos-ws-py ==0.3.2 ; extra == 'tdengine'
 
 <a id="top"></a>
 [![Build Status](https://github.com/mlrun/mlrun/actions/workflows/build.yaml/badge.svg?branch=development)](https://github.com/mlrun/mlrun/actions/workflows/build.yaml?query=branch%3Adevelopment)

{mlrun-1.7.0rc50.dist-info → mlrun-1.7.0rc51.dist-info}/RECORD

@@ -245,7 +245,7 @@ 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=7yZFn42sF597TBumVM-xhh1bjIQCbIo6qIvMK5WpWO0,10503
 mlrun/model_monitoring/db/tsdb/tdengine/stream_graph_steps.py,sha256=Hb0vcCBP-o0ET78mU4P32fnhUL65QZv-pMuv2lnCby4,1586
-mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py,sha256=L4cDFfuGOVyF_bnPbUJ_xhMEt_DGwY6FWwoO4VEXSW4,18671
+mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py,sha256=w9BzsqMWGvc3knMm6J0B6jeDAo4iWLKm-66QKIf3JaQ,19942
 mlrun/model_monitoring/db/tsdb/v3io/__init__.py,sha256=aL3bfmQsUQ-sbvKGdNihFj8gLCK3mSys0qDcXtYOwgc,616
 mlrun/model_monitoring/db/tsdb/v3io/stream_graph_steps.py,sha256=mbmhN4f_F58ptVjhwoMF6ifZSdnZWhK7x8eNsWS39IA,6217
 mlrun/model_monitoring/db/tsdb/v3io/v3io_connector.py,sha256=1H-IBXPNJPRAaxDMGWpUU25QqfR87LpZbJ03vaJkICs,32858
@@ -341,11 +341,11 @@ mlrun/utils/notifications/notification/ipython.py,sha256=ZtVL30B_Ha0VGoo4LxO-voT
 mlrun/utils/notifications/notification/slack.py,sha256=wqpFGr5BTvFO5KuUSzFfxsgmyU1Ohq7fbrGeNe9TXOk,7006
 mlrun/utils/notifications/notification/webhook.py,sha256=cb9w1Mc8ENfJBdgan7iiVHK9eVls4-R3tUxmXM-P-8I,4746
 mlrun/utils/version/__init__.py,sha256=7kkrB7hEZ3cLXoWj1kPoDwo4MaswsI2JVOBpbKgPAgc,614
-mlrun/utils/version/version.json,sha256=slRKoCCm9V9GXOEcyhpf1FNyOfO-3lk6SydAW7BkmVg,89
+mlrun/utils/version/version.json,sha256=mIhH_8cnM4P-kdwi7BT1WdUxal7WaIr2QqHt5HbvdtI,89
 mlrun/utils/version/version.py,sha256=eEW0tqIAkU9Xifxv8Z9_qsYnNhn3YH7NRAfM-pPLt1g,1878
-mlrun-1.7.0rc50.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mlrun-1.7.0rc50.dist-info/METADATA,sha256=0FpVzUdgaF0PueIufozDnrdMySc9QYVTznJOkgnEEjw,24262
-mlrun-1.7.0rc50.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-mlrun-1.7.0rc50.dist-info/entry_points.txt,sha256=1Owd16eAclD5pfRCoJpYC2ZJSyGNTtUr0nCELMioMmU,46
-mlrun-1.7.0rc50.dist-info/top_level.txt,sha256=NObLzw3maSF9wVrgSeYBv-fgnHkAJ1kEkh12DLdd5KM,6
-mlrun-1.7.0rc50.dist-info/RECORD,,
+mlrun-1.7.0rc51.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mlrun-1.7.0rc51.dist-info/METADATA,sha256=gB94QVVHmvDU2L8d0NTi70uF2KFrdBxG5n5qMPUuTyA,24262
+mlrun-1.7.0rc51.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+mlrun-1.7.0rc51.dist-info/entry_points.txt,sha256=1Owd16eAclD5pfRCoJpYC2ZJSyGNTtUr0nCELMioMmU,46
+mlrun-1.7.0rc51.dist-info/top_level.txt,sha256=NObLzw3maSF9wVrgSeYBv-fgnHkAJ1kEkh12DLdd5KM,6
+mlrun-1.7.0rc51.dist-info/RECORD,,