diff --git a/llmfoundry/command_utils/data_prep/convert_delta_to_json.py b/llmfoundry/command_utils/data_prep/convert_delta_to_json.py index dac57dcd3d..d68718e2ee 100644 --- a/llmfoundry/command_utils/data_prep/convert_delta_to_json.py +++ b/llmfoundry/command_utils/data_prep/convert_delta_to_json.py @@ -677,7 +677,7 @@ def fetch_DT( ) == grpc.StatusCode.INTERNAL and 'Job aborted due to stage failure' in e.details( ): raise InternalError( - message=f'Possible Hardware Failure: {e.details()}' + message=f'Possible Hardware Failure: {e.details()}', ) from e raise e diff --git a/tests/a_scripts/data_prep/test_convert_delta_to_json.py b/tests/a_scripts/data_prep/test_convert_delta_to_json.py index ce6274974d..8038e6624d 100644 --- a/tests/a_scripts/data_prep/test_convert_delta_to_json.py +++ b/tests/a_scripts/data_prep/test_convert_delta_to_json.py @@ -530,10 +530,12 @@ def test_format_tablename(self): @patch('llmfoundry.command_utils.data_prep.convert_delta_to_json.fetch') @patch( - 'llmfoundry.command_utils.data_prep.convert_delta_to_json.validate_and_get_cluster_info' + 'llmfoundry.command_utils.data_prep.convert_delta_to_json.validate_and_get_cluster_info', ) def test_fetch_DT_grpc_error_handling( - self, mock_validate_cluster_info: MagicMock, mock_fetch: MagicMock + self, + mock_validate_cluster_info: MagicMock, + mock_fetch: MagicMock, ): # Arrange # Mock the validate_and_get_cluster_info to return test values @@ -571,7 +573,8 @@ def test_fetch_DT_grpc_error_handling( # Verify that the InternalError contains the expected message self.assertIn('Possible Hardware Failure', str(context.exception)) self.assertIn( - 'Job aborted due to stage failure', str(context.exception) + 'Job aborted due to stage failure', + str(context.exception), ) # Verify that fetch was called