update error as user error

llmfoundry/command_utils/data_prep/convert_delta_to_json.py

@@ -23,8 +23,8 @@
     ClusterInvalidAccessMode,
     FailedToConnectToDatabricksError,
     FailedToCreateSQLConnectionError,
+    FaultyDataPrepCluster,
     InsufficientPermissionsError,
-    InternalError,
 )
 
 if TYPE_CHECKING:
@@ -676,9 +676,9 @@ def fetch_DT(
         if e.code(
         ) == grpc.StatusCode.INTERNAL and 'Job aborted due to stage failure' in e.details(
         ):
-            raise InternalError(
+            raise FaultyDataPrepCluster(
                 message=
-                f'Possible Hardware Failure, try swapping data prep cluster: {e.details()}',
+                f'Faulty data prep cluster, please try swapping data prep cluster: {e.details()}',
             ) from e
         raise e
 

llmfoundry/utils/exceptions.py

@@ -466,3 +466,18 @@ def __reduce__(self):
 
     def __str__(self):
         return self.message
+
+
+class FaultyDataPrepCluster(UserError):
+    """Error thrown when the user uses faulty data prep cluster."""
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+        super().__init__(message)
+
+    def __reduce__(self):
+        # Return a tuple of class, a tuple of arguments, and optionally state
+        return (FaultyDataPrepCluster, (self.message,))
+
+    def __str__(self):
+        return self.message

tests/a_scripts/data_prep/test_convert_delta_to_json.py

@@ -10,8 +10,8 @@
 import grpc
 
 from llmfoundry.command_utils.data_prep.convert_delta_to_json import (
+    FaultyDataPrepCluster,
     InsufficientPermissionsError,
-    InternalError,
     download,
     fetch,
     fetch_DT,
@@ -559,7 +559,7 @@ def test_fetch_DT_grpc_error_handling(
         DATABRICKS_TOKEN = 'test-token'
 
         # Act & Assert
-        with self.assertRaises(InternalError) as context:
+        with self.assertRaises(FaultyDataPrepCluster) as context:
             fetch_DT(
                 delta_table_name=delta_table_name,
                 json_output_folder=json_output_folder,
@@ -570,8 +570,11 @@ def test_fetch_DT_grpc_error_handling(
                 DATABRICKS_TOKEN=DATABRICKS_TOKEN,
             )
 
-        # Verify that the InternalError contains the expected message
-        self.assertIn('Possible Hardware Failure', str(context.exception))
+        # Verify that the FaultyDataPrepCluster contains the expected message
+        self.assertIn(
+            'Faulty data prep cluster, please try swapping data prep cluster: ',
+            str(context.exception),
+        )
         self.assertIn(
             'Job aborted due to stage failure',
             str(context.exception),