Simplify router clean up when pools or clusters ends (#3086)

Author: wjsi

MANIFEST.in

@@ -13,3 +13,4 @@ include mars/services/web/static/*
 global-exclude .DS_Store
 include versioneer.py
 include mars/_version.py
+global-exclude conftest.py

benchmarks/tpch/gen_data.py

@@ -26,7 +26,6 @@
 import shutil
 import subprocess
 from multiprocessing import Pool, set_start_method
-import pandas as pd
 import pyarrow.parquet as pq
 
 
@@ -44,6 +43,7 @@
 # Change location of tpch-dbgen if not in same place as this script
 tpch_dbgen_location = "./tpch-dbgen"
 
+
 # First element is the table single character short-hand understood by dbgen
 # Second element is the number of pieces we want the parquet dataset to have for that table
 # Third element is the function that reads generated CSV to a pandas dataframe

benchmarks/tpch/run_queries.py

@@ -301,7 +301,6 @@ def q03(lineitem, orders, customer):
 
 @tpc_query
 def q04(lineitem, orders):
-    t1 = time.time()
     date1 = md.Timestamp("1993-11-01")
     date2 = md.Timestamp("1993-08-01")
     lsel = lineitem.L_COMMITDATE < lineitem.L_RECEIPTDATE
@@ -617,7 +616,7 @@ def g2(x):
 def q13(customer, orders):
     customer_filtered = customer.loc[:, ["C_CUSTKEY"]]
     orders_filtered = orders[
-        ~orders["O_COMMENT"].str.contains("special[\S|\s]*requests")
+        ~orders["O_COMMENT"].str.contains(r"special[\S|\s]*requests")
     ]
     orders_filtered = orders_filtered.loc[:, ["O_ORDERKEY", "O_CUSTKEY"]]
     c_o_merged = customer_filtered.merge(
@@ -696,7 +695,7 @@ def q16(part, partsupp, supplier):
     )
     total = total.loc[:, ["P_BRAND", "P_TYPE", "P_SIZE", "PS_SUPPKEY"]]
     supplier_filtered = supplier[
-        supplier["S_COMMENT"].str.contains("Customer(\S|\s)*Complaints")
+        supplier["S_COMMENT"].str.contains(r"Customer(\S|\s)*Complaints")
     ]
     supplier_filtered = supplier_filtered.loc[:, ["S_SUPPKEY"]].drop_duplicates()
     # left merge to select only PS_SUPPKEY values not in supplier_filtered

mars/conftest.py

@@ -36,6 +36,16 @@ def auto_cleanup(request):
     request.addfinalizer(clear_all_alru_caches)
 
 
+@pytest.fixture(scope="module", autouse=True)
+def check_router_cleaned(request):
+    def route_checker():
+        if Router.get_instance() is not None:
+            assert len(Router.get_instance()._mapping) == 0
+            assert len(Router.get_instance()._local_mapping) == 0
+
+    request.addfinalizer(route_checker)
+
+
 @pytest.fixture(scope="module")
 def ray_start_regular_shared(request):  # pragma: no cover
     yield from _ray_start_regular(request)
@@ -144,7 +154,7 @@ def stop_ray(request):  # pragma: no cover
 
 
 @pytest.fixture
-async def ray_create_mars_cluster(request):
+async def ray_create_mars_cluster(request, check_router_cleaned):
     from mars.deploy.oscar.ray import new_cluster, _load_config
 
     ray_config = _load_config()
@@ -177,11 +187,10 @@ def stop_mars():
         import mars
 
         mars.stop_server()
-        Router.set_instance(None)
 
 
 @pytest.fixture(scope="module")
-def _new_test_session():
+def _new_test_session(check_router_cleaned):
     from .deploy.oscar.tests.session import new_test_session
 
     sess = new_test_session(
@@ -200,7 +209,7 @@ def _new_test_session():
 
 
 @pytest.fixture(scope="module")
-def _new_integrated_test_session():
+def _new_integrated_test_session(check_router_cleaned):
     from .deploy.oscar.tests.session import new_test_session
 
     sess = new_test_session(
@@ -234,7 +243,7 @@ def _new_integrated_test_session():
 
 
 @pytest.fixture(scope="module")
-def _new_gpu_test_session():  # pragma: no cover
+def _new_gpu_test_session(check_router_cleaned):  # pragma: no cover
     from .deploy.oscar.tests.session import new_test_session
     from .resource import cuda_count
 

mars/dataframe/contrib/raydataset/tests/test_mldataset.py

@@ -20,7 +20,6 @@
 from ..... import dataframe as md
 from .....deploy.oscar.ray import new_cluster
 from .....deploy.oscar.session import new_session
-from .....oscar.backends.router import Router
 from .....tests.core import require_ray
 from .....utils import lazy_import
 from ....contrib import raydataset as mdd
@@ -47,11 +46,8 @@ async def create_cluster(request):
         worker_cpu=1,
         worker_mem=256 * 1024**2,
     )
-    try:
-        async with client:
-            yield client
-    finally:
-        Router.set_instance(None)
+    async with client:
+        yield client
 
 
 @require_ray

mars/dataframe/contrib/raydataset/tests/test_raydataset.py

@@ -20,7 +20,6 @@
 from ..... import dataframe as md
 from .....deploy.oscar.ray import new_cluster
 from .....deploy.oscar.session import new_session
-from .....oscar.backends.router import Router
 from .....tests.core import require_ray
 from .....utils import lazy_import
 from ....contrib import raydataset as mdd
@@ -45,11 +44,8 @@ async def create_cluster(request):
         worker_cpu=1,
         worker_mem=256 * 1024**2,
     )
-    try:
-        async with client:
-            yield client
-    finally:
-        Router.set_instance(None)
+    async with client:
+        yield client
 
 
 @require_ray

mars/deploy/oscar/local.py

@@ -25,6 +25,7 @@
 from ... import oscar as mo
 from ...core.entrypoints import init_extension_entrypoints
 from ...lib.aio import get_isolation, stop_isolation
+from ...oscar.backends.router import Router
 from ...resource import cpu_count, cuda_count, mem_total
 from ...services import NodeRole
 from ...services.task.execution.api import ExecutionConfig
@@ -128,6 +129,7 @@ async def stop_cluster(cluster: ClusterType):
     isolation = get_isolation()
     coro = cluster.stop()
     await asyncio.wrap_future(asyncio.run_coroutine_threadsafe(coro, isolation.loop))
+    Router.set_instance(None)
 
 
 class LocalCluster:
@@ -287,6 +289,7 @@ async def stop(self):
         await self._supervisor_pool.stop()
         AbstractSession.reset_default()
         self._exiting_check_task.cancel()
+        Router.set_instance(None)
 
 
 class LocalClient:

mars/deploy/oscar/ray.py

@@ -29,6 +29,7 @@
     process_address_to_placement,
 )
 from ...oscar.backends.ray.pool import RayPoolState
+from ...oscar.backends.router import Router
 from ...oscar.errors import ReconstructWorkerError
 from ...resource import Resource
 from ...services.cluster.backends.base import (
@@ -569,6 +570,7 @@ async def stop(self):
             finally:
                 AbstractSession.reset_default()
                 RayActorDriver.stop_cluster()
+                Router.set_instance(None)
             self._stopped = True
 
 

mars/deploy/oscar/tests/test_fault_injection.py

@@ -20,7 +20,6 @@
 
 from .... import dataframe as md
 from .... import tensor as mt
-from ....oscar.backends.router import Router
 from ....oscar.errors import ServerClosed
 from ....remote import spawn
 from ....services.tests.fault_injection_manager import (
@@ -51,11 +50,8 @@ async def fault_cluster(request):
         n_worker=2,
         n_cpu=2,
     )
-    try:
-        async with client:
-            yield client
-    finally:
-        Router.set_instance(None)
+    async with client:
+        yield client
 
 
 async def create_fault_injection_manager(

mars/deploy/oscar/tests/test_local.py

@@ -35,7 +35,6 @@
 from ....config import option_context
 from ....core.context import get_context
 from ....lib.aio import new_isolation
-from ....oscar.backends.router import Router
 from ....storage import StorageLevel
 from ....services.storage import StorageAPI
 from ....tensor.arithmetic.add import TensorAdd
@@ -120,13 +119,10 @@ async def create_cluster(request):
         n_cpu=2,
         use_uvloop=False,
     )
-    try:
-        async with client:
-            if request.param == "default":
-                assert client.session.client is not None
-            yield client, request.param
-    finally:
-        Router.set_instance(None)
+    async with client:
+        if request.param == "default":
+            assert client.session.client is not None
+        yield client, request.param
 
 
 def _assert_storage_cleaned(session_id: str, addr: str, level: StorageLevel):
@@ -649,7 +645,6 @@ def setup_session(request):
             yield session
     finally:
         session.stop_server()
-        Router.set_instance(None)
 
 
 def test_decref(setup_session):
@@ -892,11 +887,8 @@ async def speculative_cluster():
         n_cpu=10,
         use_uvloop=False,
     )
-    try:
-        async with client:
-            yield client
-    finally:
-        Router.set_instance(None)
+    async with client:
+        yield client
 
 
 @pytest.mark.timeout(timeout=500)