Skip to content

Commit

Permalink
Private link error handling (#3689)
Browse files Browse the repository at this point in the history
Co-authored-by: Matthew Ding <[email protected]>
  • Loading branch information
nancyhung and mattyding authored Oct 30, 2024
1 parent e33c9e3 commit 785ad4e
Showing 1 changed file with 13 additions and 1 deletion.
14 changes: 13 additions & 1 deletion composer/utils/object_store/mlflow_object_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 785ad4e

Please sign in to comment.