Add Mlflow 403 PL UserError

llmfoundry/callbacks/hf_checkpointer.py

@@ -42,6 +42,7 @@
     PreTrainedTokenizerBase,
 )
 
+from llmfoundry.utils.exceptions import ProbablyNetworkingError
 from llmfoundry.models.mpt import MPTConfig, MPTForCausalLM
 from llmfoundry.models.utils import init_empty_weights
 from llmfoundry.utils.huggingface_hub_utils import \
@@ -296,7 +297,13 @@ def run_event(self, event: Event, state: State, logger: Logger) -> None:
                     + f'Got {type(state.model)} instead.',
                 )
             if self.remote_ud is not None:
-                self.remote_ud.init(state, logger)
+                try:
+                    self.remote_ud.init(state, logger)
+                except PermissionError as e:
+                    log.error(
+                        f'Error initializing remote uploader/downloader. This is likely private link related. {e}',
+                    )
+                    raise ProbablyNetworkingError(is_contd_pretrain=state.run_name.startswith('contd-pretrain'))
                 state.callbacks.append(self.remote_ud)
 
             if self.mlflow_registered_model_name is not None:

llmfoundry/utils/exceptions.py

@@ -30,6 +30,7 @@
     'MisconfiguredHfDatasetError',
     'DatasetTooSmallError',
     'RunTimeoutError',
+    'ProbablyNetworkingError',
 ]
 
 ALLOWED_RESPONSE_KEYS = {'response', 'completion'}
@@ -524,3 +525,11 @@ def __reduce__(self):
 
     def __str__(self):
         return self.message
+
+
+class ProbablyNetworkingError(UserError):
+    """Error thrown due to a networking restriction."""
+
+    def __init__(self, is_contd_pretrain: bool) -> None:
+        message = f"An error occurred while trying to access a remote URL. This is likely due to some networking restriction such as private link or SEG."
+        super().__init__(message, is_contd_pretrain=is_contd_pretrain)