Stop importing main module when starting Mars local cluster (#3110)

Author: wjsi

.github/workflows/benchmark-ci.yml

@@ -62,5 +62,5 @@ jobs:
         uses: actions/upload-artifact@v2
         with:
           name: Benchmarks log
-          path: benchmarks/asv_bench/benchmarks.log
+          path: benchmarks/asv_bench/results
         if: failure()

benchmarks/asv_bench/benchmarks/storage.py

@@ -14,6 +14,7 @@
 
 import itertools
 
+import cloudpickle
 import numpy as np
 import pandas as pd
 
@@ -89,8 +90,20 @@ class TransferPackageSuite:
     """
 
     def setup(self):
+        try:
+            # make sure all submodules will serial functions instead of refs
+            cloudpickle.register_pickle_by_value(__import__("benchmarks.storage"))
+        except (AttributeError, ImportError):
+            pass
         mars.new_session(n_worker=2, n_cpu=8)
 
+    def teardown(self):
+        mars.stop_server()
+        try:
+            cloudpickle.unregister_pickle_by_value(__import__("benchmarks.storage"))
+        except (AttributeError, ImportError):
+            pass
+
     def time_1_to_1(self):
         return mr.spawn(send_1_to_1).execute().fetch()
 
@@ -99,3 +112,4 @@ def time_1_to_1(self):
     suite = TransferPackageSuite()
     suite.setup()
     print(suite.time_1_to_1())
+    suite.teardown()

mars/deploy/oscar/tests/test_local.py

@@ -15,13 +15,17 @@
 import asyncio
 import copy
 import os
+import subprocess
+import sys
 import threading
 import tempfile
+import textwrap
 import time
 import uuid
 
 import numpy as np
 import pandas as pd
+import psutil
 import pytest
 
 try:
@@ -917,3 +921,49 @@ def time_consuming(start, x):
         .fetch()
         == pd.Series(list(range(series_size))).apply(lambda x: x * x).sum()
     )
+
+
+def test_naive_code_file():
+    code_file = """
+    import mars
+    import mars.tensor as mt
+    import os
+
+    mars.new_session()
+    try:
+        result_path = os.environ["RESULTPATH"]
+        with open(result_path, "w") as outf:
+            outf.write(str(mt.ones((10, 10)).sum().execute()))
+    finally:
+        mars.stop_server()
+    """
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        try:
+            script_path = os.path.join(temp_dir, "test_file.py")
+            result_path = os.path.join(temp_dir, "result.txt")
+
+            with open(script_path, "w") as file_obj:
+                file_obj.write(textwrap.dedent(code_file))
+
+            env = os.environ.copy()
+            env["PYTHONPATH"] = os.path.pathsep.join(sys.path)
+            env["RESULTPATH"] = result_path
+            proc = subprocess.Popen([sys.executable, script_path], env=env)
+            pid = proc.pid
+            proc.wait(120)
+
+            with open(result_path, "r") as inp_file:
+                assert 100 == int(float(inp_file.read()))
+        except subprocess.TimeoutExpired:
+            try:
+                procs = [psutil.Process(pid)]
+                procs.extend(procs[0].children(True))
+                for proc in reversed(procs):
+                    try:
+                        proc.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+            except psutil.NoSuchProcess:
+                pass
+            raise

mars/oscar/backends/mars/pool.py

@@ -14,12 +14,14 @@
 
 import asyncio
 import concurrent.futures as futures
+import contextlib
 import logging.config
 import multiprocessing
 import os
 import random
 import signal
 import sys
+import threading
 import uuid
 from dataclasses import dataclass
 from types import TracebackType
@@ -63,6 +65,37 @@ def _mp_kill(self):
     BaseProcess.kill = _mp_kill
 
 logger = logging.getLogger(__name__)
+_init_main_suspended_local = threading.local()
+
+
+def _patch_spawn_get_preparation_data():
+    try:
+        from multiprocessing import spawn as mp_spawn
+
+        _raw_get_preparation_data = mp_spawn.get_preparation_data
+
+        def _patched_get_preparation_data(*args, **kw):
+            ret = _raw_get_preparation_data(*args, **kw)
+            if getattr(_init_main_suspended_local, "value", False):
+                # make sure user module is not imported when start Mars cluster
+                ret.pop("init_main_from_name", None)
+                ret.pop("init_main_from_path", None)
+            return ret
+
+        _patched_get_preparation_data._mars_patched = True
+        if not getattr(mp_spawn.get_preparation_data, "_mars_patched", False):
+            mp_spawn.get_preparation_data = _patched_get_preparation_data
+    except (ImportError, AttributeError):  # pragma: no cover
+        pass
+
+
+@contextlib.contextmanager
+def _suspend_init_main():
+    try:
+        _init_main_suspended_local.value = True
+        yield
+    finally:
+        _init_main_suspended_local.value = False
 
 
 @dataslots
@@ -131,21 +164,25 @@ async def start_sub_pool(
         def start_pool_in_process():
             ctx = multiprocessing.get_context(method=start_method)
             status_queue = ctx.Queue()
-            process = ctx.Process(
-                target=cls._start_sub_pool,
-                args=(actor_pool_config, process_index, status_queue),
-                name=f"MarsActorPool{process_index}",
-            )
-            process.daemon = True
-            process.start()
+
+            with _suspend_init_main():
+                process = ctx.Process(
+                    target=cls._start_sub_pool,
+                    args=(actor_pool_config, process_index, status_queue),
+                    name=f"MarsActorPool{process_index}",
+                )
+                process.daemon = True
+                process.start()
+
             # wait for sub actor pool to finish starting
             process_status = status_queue.get()
             return process, process_status
 
+        _patch_spawn_get_preparation_data()
         loop = asyncio.get_running_loop()
-        executor = futures.ThreadPoolExecutor(1)
-        create_pool_task = loop.run_in_executor(executor, start_pool_in_process)
-        return await create_pool_task
+        with futures.ThreadPoolExecutor(1) as executor:
+            create_pool_task = loop.run_in_executor(executor, start_pool_in_process)
+            return await create_pool_task
 
     @classmethod
     async def wait_sub_pools_ready(cls, create_pool_tasks: List[asyncio.Task]):
@@ -240,8 +277,11 @@ async def kill_sub_pool(
                 pass
             process.terminate()
             wait_pool = futures.ThreadPoolExecutor(1)
-            loop = asyncio.get_running_loop()
-            await loop.run_in_executor(wait_pool, process.join, 3)
+            try:
+                loop = asyncio.get_running_loop()
+                await loop.run_in_executor(wait_pool, process.join, 3)
+            finally:
+                wait_pool.shutdown(False)
         process.kill()
         await asyncio.to_thread(process.join, 5)
 

mars/oscar/debug.py

@@ -79,7 +79,7 @@ async def _log_timeout(timeout, msg, *args, **kwargs):
         await asyncio.sleep(timeout * rnd)
         rnd += 1
         logger.warning(
-            msg + "(timeout for %.4f seconds).",
+            msg + " (timeout for %.4f seconds).",
             *args,
             time.time() - start_time,
             **kwargs,

mars/tests/test_utils.py

@@ -407,6 +407,7 @@ def write(self, d):
         print("LINE 2", file=sys.stderr, end="\n")
     finally:
         sys.stdout, sys.stderr = old_stdout, old_stderr
+        executor.shutdown(False)
 
     assert stdout_w.content == "LINE T\nLINE 1\n"
     assert stderr_w.content == "LINE 2\n"