Enforce compute cluster version

llmfoundry/command_utils/data_prep/convert_delta_to_json.py

@@ -8,6 +8,7 @@
 import urllib.parse
 from collections import namedtuple
 from concurrent.futures import ProcessPoolExecutor
+from importlib.metadata import version as importlib_version
 from typing import TYPE_CHECKING, Iterable, Optional, Union
 from uuid import uuid4
 
@@ -49,7 +50,7 @@
 except ImportError:
     data_frame_installed = False
 
-MINIMUM_DB_CONNECT_DBR_VERSION = '14.1'
+MINIMUM_DB_CONNECT_DBR_VERSION = importlib_version('databricks-connect')
 MINIMUM_SQ_CONNECT_DBR_VERSION = '12.2'
 
 TABLENAME_PATTERN = re.compile(r'(\S+)\.(\S+)\.(\S+)')
@@ -676,6 +677,29 @@ def fetch_DT(
         if e.code(
         ) == grpc.StatusCode.INTERNAL and 'Job aborted due to stage failure' in e.details(
         ):
+            from databricks.sdk import WorkspaceClient
+            w = WorkspaceClient()
+            if cluster_id:
+                res = w.clusters.get(cluster_id=cluster_id)
+                cluster_spark_version_key = res.spark_version
+                assert cluster_spark_version_key
+                cluster_major_version_match = re.search(
+                    r'^(\d+)\.',
+                    cluster_spark_version_key,
+                )
+
+                assert cluster_major_version_match
+                cluster_major_version = int(cluster_major_version_match.group(1))
+                databricks_connect_major_version = int(
+                    MINIMUM_DB_CONNECT_DBR_VERSION.split('.')[0],
+                )
+
+                if cluster_major_version != databricks_connect_major_version:
+                    raise FaultyDataPrepCluster(
+                        message=
+                        f'Data prep cluster version is compatible with current databricks-connect version {MINIMUM_DB_CONNECT_DBR_VERSION}. If you keep encountering this issue, please consider downgrading the compute Databricks Runtime Version to {MINIMUM_DB_CONNECT_DBR_VERSION}.',
+                    ) from e
+
             raise FaultyDataPrepCluster(
                 message=
                 f'Faulty data prep cluster, please try swapping data prep cluster: {e.details()}',