spawn to start vertex uploader process (apple#1187)

Co-authored-by: Jinhao Lei <jlei2@apple.com>

Author: JinhaoLei

axlearn/cloud/gcp/vertexai_tensorboard.py

@@ -137,7 +137,13 @@ def upload(self):
                 resource_name=self._resource_name,
                 logdir=cfg.summary_dir,
             )
-            self._uploader_proc = multiprocessing.Process(
+            # Uses spawn method to start the process as it was observed that using fork method may
+            # cause training job fail to exit correctly after completion. It is also suggested by
+            # the warning message: "RuntimeWarning: os.fork() was called. os.fork() is incompatible
+            # with multithreaded code, and JAX is multithreaded, so this will likely lead to a
+            # deadlock.".
+            ctx = multiprocessing.get_context("spawn")
+            self._uploader_proc = ctx.Process(
                 target=_start_vertexai_tensorboard, kwargs=kwargs, daemon=True
             )
             self._uploader_proc.start()

axlearn/cloud/gcp/vertexai_tensorboard_test.py

@@ -30,7 +30,7 @@ def start(self):  # pylint: disable=no-self-use
 class VertexAITensorboardUploaderTest(absltest.TestCase):
     """Tests VertexAITensorboardUploader."""
 
-    @mock.patch("multiprocessing.Process", side_effect=fake_process)
+    @mock.patch("multiprocessing.get_context")
     @mock.patch(f"{uploader.TensorBoardUploader.__module__}.TensorBoardUploader", autospec=True)
     @mock.patch(
         f"{initializer.global_config.__module__}.global_config.create_client",
@@ -45,9 +45,15 @@ class VertexAITensorboardUploaderTest(absltest.TestCase):
         return_value=("fake_bucket", "fake_folder"),
     )
     def test_uploader_calls(
-        self, bucket_folder_fn, create_client_fn, tb_uploader_class, unused
+        self,
+        bucket_folder_fn,
+        create_client_fn,
+        tb_uploader_class,
+        mock_get_context,
     ):  # pylint: disable=no-self-use
-        del unused
+        mock_context = mock.MagicMock()
+        mock_context.Process.side_effect = fake_process
+        mock_get_context.return_value = mock_context
 
         mock_settings = {
             "vertexai_tensorboard": "fake_tb_instance",
@@ -61,6 +67,7 @@ def test_uploader_calls(
             )
         tb_uploader = cfg.instantiate()
         tb_uploader.upload()
+        mock_get_context.assert_called_once_with("spawn")
         create_client_fn.assert_called_once()
         bucket_folder_fn.assert_called_once()
         tb_uploader_class.return_value.create_experiment.assert_called_once()