Merge branch 'main' into objectstore-registry

.github/workflows/pr-gpu.yaml

@@ -1,6 +1,6 @@
 name: PR GPU tests
 on:
-  pull_request:
+  pull_request_target:
   workflow_dispatch:
 # Cancel old runs when a new commit is pushed to the same branch if not on main
 # or dev

composer/callbacks/checkpoint_saver.py

@@ -294,10 +294,15 @@ def __init__(
         num_concurrent_uploads: int = 1,
         upload_timeout_in_seconds: int = 3600,
     ):
+
         backend, _, local_folder = parse_uri(str(folder))
         if local_folder == '':
             local_folder = '.'
 
+        is_remote_folder = backend != ''
+        if is_remote_folder:  # If uploading to a remote path, use a temporary directory to save local checkpoints.
+            local_folder = os.path.join(tempfile.mkdtemp(), local_folder)
+
         filename = str(filename)
         remote_file_name = str(remote_file_name) if remote_file_name is not None else None
         latest_filename = str(latest_filename) if latest_filename is not None else None

composer/loggers/mlflow_logger.py

@@ -301,10 +301,10 @@ def _start_mlflow_run(self, state):
             )
             self.monitor_process.start()
 
-    def _global_exception_handler(self, exc_type, exc_value, exc_traceback):
+    def _global_exception_handler(self, original_excepthook, exc_type, exc_value, exc_traceback):
         """Catch global exception."""
         self._global_exception_occurred += 1
-        sys.__excepthook__(exc_type, exc_value, exc_traceback)
+        original_excepthook(exc_type, exc_value, exc_traceback)
 
     def init(self, state: State, logger: Logger) -> None:
         del logger  # unused
@@ -322,7 +322,13 @@ def init(self, state: State, logger: Logger) -> None:
             self.run_name += f'-rank{dist.get_global_rank()}'
 
         # Register the global exception handler so that uncaught exception is tracked.
-        sys.excepthook = self._global_exception_handler
+        original_excepthook = sys.excepthook
+        sys.excepthook = lambda exc_type, exc_value, exc_traceback: self._global_exception_handler(
+            original_excepthook,
+            exc_type,
+            exc_value,
+            exc_traceback,
+        )
         # Start run
         if self._enabled:
             self._start_mlflow_run(state)

tests/trainer/test_checkpoint.py

@@ -882,7 +882,8 @@ def _get_tmp_dir(self):
 
                     if delete_local:
                         # delete files locally, forcing trainer to look in object store
-                        shutil.rmtree('first')
+                        assert trainer_1._checkpoint_saver is not None
+                        shutil.rmtree(trainer_1._checkpoint_saver.folder)
 
                     trainer_2 = self.get_trainer(
                         latest_filename=latest_filename,