From 785ad4e5409d3b3f5e6ae4c2ab2cef7b50460f8e Mon Sep 17 00:00:00 2001 From: Nancy Hung Date: Wed, 30 Oct 2024 15:38:26 -0700 Subject: [PATCH] Private link error handling (#3689) Co-authored-by: Matthew Ding --- composer/utils/object_store/mlflow_object_store.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/composer/utils/object_store/mlflow_object_store.py b/composer/utils/object_store/mlflow_object_store.py index 26b4074b8d..6a6655fba8 100644 --- a/composer/utils/object_store/mlflow_object_store.py +++ b/composer/utils/object_store/mlflow_object_store.py @@ -62,9 +62,21 @@ def _wrap_mlflow_exceptions(uri: str, e: Exception): retryable_client_codes = [ErrorCode.Name(code) for code in [ABORTED, REQUEST_LIMIT_EXCEEDED, RESOURCE_EXHAUSTED]] not_found_codes = [ErrorCode.Name(code) for code in [RESOURCE_DOES_NOT_EXIST, NOT_FOUND, ENDPOINT_NOT_FOUND]] + # MLflow wraps Azure data exceptions as INTERNAL_ERROR. Need to unwrap and check msg for the specific error. + permission_error_codes = [ + '401', + '403', + ] + if isinstance(e, MlflowException): error_code = e.error_code # pyright: ignore - if error_code in retryable_server_codes or error_code in retryable_client_codes: + if error_code == ErrorCode.Name(INTERNAL_ERROR): + error_message = e.message # pyright: ignore + if any(f'{code} Client Error' in error_message for code in permission_error_codes): + raise PermissionError( + f'Permission denied for object {uri} from the data provider. Details: {error_message}', + ) from e + elif error_code in retryable_server_codes or error_code in retryable_client_codes: raise ObjectStoreTransientError(error_code) from e elif error_code in not_found_codes: raise FileNotFoundError(f'Object {uri} not found') from e