Catch bad data prep

llmfoundry/command_utils/data_prep/convert_delta_to_json.py

@@ -1,6 +1,7 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+import json
 import logging
 import os
 import re
@@ -26,6 +27,8 @@
     FailedToCreateSQLConnectionError,
     FaultyDataPrepCluster,
     InsufficientPermissionsError,
+    MisconfiguredHfDatasetError,
+    StoragePermissionError,
     UCNotEnabledError,
 )
 
@@ -681,7 +684,7 @@ def fetch_DT(
 
     log.info(f'Directory {json_output_folder} created.')
 
-    # validate_and_get_cluster_info allows cluster_id to be None if use_serverless is True
+    # Validate_and_get_cluster_info allows cluster_id to be None if use_serverless is True.
     method, dbsql, sparkSession = validate_and_get_cluster_info(
         cluster_id=cluster_id,
         databricks_host=DATABRICKS_HOST,
@@ -732,12 +735,41 @@ def fetch_DT(
     if dbsql is not None:
         dbsql.close()
 
-    # combine downloaded jsonl into one big jsonl for IFT
+    # Combine downloaded jsonl into one big jsonl for IFT.
     iterative_combine_jsons(
         json_output_folder,
         os.path.join(json_output_folder, json_output_filename),
     )
 
+    _validate_written_file(
+        json_output_folder,
+        json_output_filename,
+        delta_table_name,
+    )
+
+
+def _validate_written_file(
+    json_output_folder: str,
+    json_output_filename: str,
+    delta_table_name: str,
+):
+    # Validate downloaded dataset is actually downloaded.
+    with open(os.path.join(json_output_folder, json_output_filename)) as f:
+        is_empty = True
+        for line in f.readlines():
+            is_empty = False
+            try:
+                json.loads(line)
+            except:
+                raise MisconfiguredHfDatasetError(
+                    delta_table_name,
+                    split=json_output_folder,
+                ) from ValueError('line')
+        if is_empty:
+            raise StoragePermissionError(
+                f'Unable to download {delta_table_name}, check network permissions.',
+            )
+
 
 def _check_imports():
     try:

tests/a_scripts/data_prep/test_convert_delta_to_json.py

@@ -1,9 +1,11 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+import os
 import sys
 import unittest
 from argparse import Namespace
+from tempfile import NamedTemporaryFile
 from typing import Any
 from unittest.mock import MagicMock, mock_open, patch
 
@@ -14,14 +16,18 @@
 from llmfoundry.command_utils.data_prep.convert_delta_to_json import (
     FaultyDataPrepCluster,
     InsufficientPermissionsError,
+    _validate_written_file,
     download,
     fetch,
     fetch_DT,
     format_tablename,
     iterative_combine_jsons,
     run_query,
 )
-from llmfoundry.utils.exceptions import DeltaTableNotFoundError
+from llmfoundry.utils.exceptions import (
+    DeltaTableNotFoundError,
+    StoragePermissionError,
+)
 
 
 class TestConvertDeltaToJsonl(unittest.TestCase):
@@ -102,17 +108,20 @@ def test_stream_delta_to_json(
         )
         mock_workspace_client.return_value.clusters.get = mock_cluster_get
 
-        fetch_DT(
-            delta_table_name=delta_table_name,
-            json_output_folder=json_output_folder,
-            http_path=http_path,
-            cluster_id=cluster_id,
-            DATABRICKS_HOST=DATABRICKS_HOST,
-            DATABRICKS_TOKEN=DATABRICKS_TOKEN,
-            use_serverless=use_serverless,
-            batch_size=batch_size,
-            json_output_filename=json_output_filename,
-        )
+        try:
+            fetch_DT(
+                delta_table_name=delta_table_name,
+                json_output_folder=json_output_folder,
+                http_path=http_path,
+                cluster_id=cluster_id,
+                DATABRICKS_HOST=DATABRICKS_HOST,
+                DATABRICKS_TOKEN=DATABRICKS_TOKEN,
+                use_serverless=use_serverless,
+                batch_size=batch_size,
+                json_output_filename=json_output_filename,
+            )
+        except FileNotFoundError:
+            pass
         mock_sql_connect.assert_called_once_with(
             server_hostname='test_host',
             http_path='test_path',
@@ -287,15 +296,18 @@ def test_dbconnect_called(
         )  # Mock return value for getOrCreate
         mock_databricks_session.builder.remote.return_value = mock_remote
 
-        fetch_DT(
-            delta_table_name=delta_table_name,
-            json_output_folder=json_output_folder,
-            http_path=http_path,
-            cluster_id=cluster_id,
-            DATABRICKS_HOST=DATABRICKS_HOST,
-            DATABRICKS_TOKEN=DATABRICKS_TOKEN,
-            use_serverless=use_serverless,
-        )
+        try:
+            fetch_DT(
+                delta_table_name=delta_table_name,
+                json_output_folder=json_output_folder,
+                http_path=http_path,
+                cluster_id=cluster_id,
+                DATABRICKS_HOST=DATABRICKS_HOST,
+                DATABRICKS_TOKEN=DATABRICKS_TOKEN,
+                use_serverless=use_serverless,
+            )
+        except FileNotFoundError:
+            pass
         mock_databricks_session.builder.remote.assert_called_once_with(
             host=DATABRICKS_HOST,
             token=DATABRICKS_TOKEN,
@@ -342,15 +354,19 @@ def test_sqlconnect_called_dbr13(
         )
         mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response
 
-        fetch_DT(
-            delta_table_name=delta_table_name,
-            json_output_folder=json_output_folder,
-            http_path=http_path,
-            cluster_id=cluster_id,
-            DATABRICKS_HOST=DATABRICKS_HOST,
-            DATABRICKS_TOKEN=DATABRICKS_TOKEN,
-            use_serverless=use_serverless,
-        )
+        try:
+            fetch_DT(
+                delta_table_name=delta_table_name,
+                json_output_folder=json_output_folder,
+                http_path=http_path,
+                cluster_id=cluster_id,
+                DATABRICKS_HOST=DATABRICKS_HOST,
+                DATABRICKS_TOKEN=DATABRICKS_TOKEN,
+                use_serverless=use_serverless,
+            )
+        except FileNotFoundError:
+            pass
+
         mock_sql_connect.assert_called_once_with(
             server_hostname=DATABRICKS_HOST,
             http_path=http_path,
@@ -397,15 +413,19 @@ def test_sqlconnect_called_dbr14(
         )
         mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response
 
-        fetch_DT(
-            delta_table_name=delta_table_name,
-            json_output_folder=json_output_folder,
-            http_path=http_path,
-            cluster_id=cluster_id,
-            DATABRICKS_HOST=DATABRICKS_HOST,
-            DATABRICKS_TOKEN=DATABRICKS_TOKEN,
-            use_serverless=use_serverless,
-        )
+        try:
+            fetch_DT(
+                delta_table_name=delta_table_name,
+                json_output_folder=json_output_folder,
+                http_path=http_path,
+                cluster_id=cluster_id,
+                DATABRICKS_HOST=DATABRICKS_HOST,
+                DATABRICKS_TOKEN=DATABRICKS_TOKEN,
+                use_serverless=use_serverless,
+            )
+        except FileNotFoundError:
+            pass
+
         mock_sql_connect.assert_called_once_with(
             server_hostname=DATABRICKS_HOST,
             http_path=http_path,
@@ -452,15 +472,18 @@ def test_sqlconnect_called_https(
         )
         mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response
 
-        fetch_DT(
-            delta_table_name=delta_table_name,
-            json_output_folder=json_output_folder,
-            http_path=http_path,
-            cluster_id=cluster_id,
-            DATABRICKS_HOST=DATABRICKS_HOST,
-            DATABRICKS_TOKEN=DATABRICKS_TOKEN,
-            use_serverless=use_serverless,
-        )
+        try:
+            fetch_DT(
+                delta_table_name=delta_table_name,
+                json_output_folder=json_output_folder,
+                http_path=http_path,
+                cluster_id=cluster_id,
+                DATABRICKS_HOST=DATABRICKS_HOST,
+                DATABRICKS_TOKEN=DATABRICKS_TOKEN,
+                use_serverless=use_serverless,
+            )
+        except FileNotFoundError:
+            pass
         mock_sql_connect.assert_called_once_with(
             server_hostname='test-host',
             http_path=http_path,
@@ -504,15 +527,19 @@ def test_serverless(
         mock_cluster_response = Namespace(spark_version='14.2.0-scala2.12')
         mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response
 
-        fetch_DT(
-            delta_table_name=delta_table_name,
-            json_output_folder=json_output_folder,
-            http_path=http_path,
-            cluster_id=cluster_id,
-            DATABRICKS_HOST=DATABRICKS_HOST,
-            DATABRICKS_TOKEN=DATABRICKS_TOKEN,
-            use_serverless=use_serverless,
-        )
+        try:
+            fetch_DT(
+                delta_table_name=delta_table_name,
+                json_output_folder=json_output_folder,
+                http_path=http_path,
+                cluster_id=cluster_id,
+                DATABRICKS_HOST=DATABRICKS_HOST,
+                DATABRICKS_TOKEN=DATABRICKS_TOKEN,
+                use_serverless=use_serverless,
+            )
+        except FileNotFoundError:
+            pass
+
         assert not mock_sql_connect.called
         assert not mock_databricks_session.builder.remote.called
 
@@ -644,3 +671,14 @@ def test_fetch_nonexistent_table_error(
 
         # Verify that get_total_rows was called
         mock_gtr.assert_called_once()
+
+    def test_fetch_DT_catches_bad_download(self):
+        with NamedTemporaryFile() as tf:
+            file_name = tf.name
+            file_folder, file_name = os.path.split(file_name)
+            with self.assertRaises(StoragePermissionError):
+                _validate_written_file(
+                    file_folder,
+                    file_name,
+                    'test_delta_table',
+                )