Add initial implementation of scheduling service (#2111)

Author: wjsi

.github/workflows/run-tests.sh

@@ -21,8 +21,10 @@ if [ -z "$NO_COMMON_TESTS" ]; then
   pytest $PYTEST_CONFIG --cov-config .coveragerc-threaded mars/tensor mars/dataframe mars/web \
     mars/learn mars/remote mars/storage mars/lib
   mv .coverage build/.coverage.tensor.file
+  pytest $PYTEST_CONFIG --cov-config .coveragerc-threaded --forked mars/oscar mars/services
+  mv .coverage build/.coverage.oscar.file
   pytest $PYTEST_CONFIG --cov-config .coveragerc --forked --ignore mars/tensor --ignore mars/dataframe \
-    --ignore mars/learn --ignore mars/remote mars
+    --ignore mars/learn --ignore mars/remote --ignore mars/services --ignore mars/oscar mars
   mv .coverage build/.coverage.main.file
   coverage combine build/ && coverage report
 

mars/core/graph/core.pyx

@@ -310,6 +310,9 @@ cdef class DirectedGraph:
             return f'"{val}"'
         return val
 
+    def _extract_operands(self, node):
+        return [node.op]
+
     def to_dot(self, graph_attrs=None, node_attrs=None, trunc_key=5, result_chunk_keys=None):
         sio = StringIO()
         sio.write('digraph {\n')
@@ -328,32 +331,32 @@ cdef class DirectedGraph:
 
         visited = set()
         for node in self.iter_nodes():
-            op = node.op
-            op_name = type(op).__name__
-            if op.stage is not None:
-                op_name = f'{op_name}:{op.stage.name}'
-            if op.key in visited:
-                continue
-            for input_chunk in (op.inputs or []):
-                if input_chunk.key not in visited:
-                    sio.write(f'"Chunk:{input_chunk.key[:trunc_key]}" {chunk_style}\n')
-                    visited.add(input_chunk.key)
-                if op.key not in visited:
-                    sio.write(f'"{op_name}:{op.key[:trunc_key]}" {operand_style}\n')
-                    visited.add(op.key)
-                sio.write(f'"Chunk:{input_chunk.key[:trunc_key]}" -> "{op_name}:{op.key[:5]}"\n')
-
-            for output_chunk in (op.outputs or []):
-                if output_chunk.key not in visited:
-                    tmp_chunk_style = chunk_style
-                    if result_chunk_keys and output_chunk.key in result_chunk_keys:
-                        tmp_chunk_style = '[shape=box,style=filled,fillcolor=cadetblue1]'
-                    sio.write(f'"Chunk:{output_chunk.key[:trunc_key]}" {tmp_chunk_style}\n')
-                    visited.add(output_chunk.key)
-                if op.key not in visited:
-                    sio.write(f'"{op_name}:{op.key[:trunc_key]}" {operand_style}\n')
-                    visited.add(op.key)
-                sio.write(f'"{op_name}:{op.key[:trunc_key]}" -> "Chunk:{output_chunk.key[:5]}"\n')
+            for op in self._extract_operands(node):
+                op_name = type(op).__name__
+                if op.stage is not None:
+                    op_name = f'{op_name}:{op.stage.name}'
+                if op.key in visited:
+                    continue
+                for input_chunk in (op.inputs or []):
+                    if input_chunk.key not in visited:
+                        sio.write(f'"Chunk:{input_chunk.key[:trunc_key]}" {chunk_style}\n')
+                        visited.add(input_chunk.key)
+                    if op.key not in visited:
+                        sio.write(f'"{op_name}:{op.key[:trunc_key]}" {operand_style}\n')
+                        visited.add(op.key)
+                    sio.write(f'"Chunk:{input_chunk.key[:trunc_key]}" -> "{op_name}:{op.key[:5]}"\n')
+
+                for output_chunk in (op.outputs or []):
+                    if output_chunk.key not in visited:
+                        tmp_chunk_style = chunk_style
+                        if result_chunk_keys and output_chunk.key in result_chunk_keys:
+                            tmp_chunk_style = '[shape=box,style=filled,fillcolor=cadetblue1]'
+                        sio.write(f'"Chunk:{output_chunk.key[:trunc_key]}" {tmp_chunk_style}\n')
+                        visited.add(output_chunk.key)
+                    if op.key not in visited:
+                        sio.write(f'"{op_name}:{op.key[:trunc_key]}" {operand_style}\n')
+                        visited.add(op.key)
+                    sio.write(f'"{op_name}:{op.key[:trunc_key]}" -> "Chunk:{output_chunk.key[:5]}"\n')
 
         sio.write('}')
         return sio.getvalue()

mars/deploy/oscar/base_config.yml

@@ -5,6 +5,8 @@ services:
   - meta
   - lifecycle
   - task
+  - scheduling
+  - subtask
   - web
 cluster:
   backend: fixed

mars/deploy/oscar/tests/check_enabled_config.yml

@@ -1,4 +1,5 @@
 inherits: '@default'
 task:
   task_processor_cls: mars.services.task.supervisor.tests.CheckedTaskProcessor
-  subtask_processor_cls: mars.services.task.worker.tests.CheckedSubtaskProcessor
+subtask:
+  subtask_processor_cls: mars.services.subtask.worker.tests.CheckedSubtaskProcessor

mars/deploy/oscar/tests/test_checked_session.py

@@ -23,7 +23,7 @@
 from mars.deploy.oscar.service import _load_config
 from mars.deploy.oscar.tests.session import new_test_session, CONFIG_FILE
 from mars.services.task.supervisor.tests import CheckedTaskProcessor
-from mars.services.task.worker.tests import CheckedSubtaskProcessor
+from mars.services.subtask.worker.tests import CheckedSubtaskProcessor
 
 
 class FakeCheckedTaskProcessor(CheckedTaskProcessor):
@@ -81,7 +81,7 @@ def test_check_task_processor(setup):
 
 def test_check_subtask_processor(setup):
     config = _load_config(CONFIG_FILE)
-    config['task']['subtask_processor_cls'] = \
+    config['subtask']['subtask_processor_cls'] = \
         'mars.deploy.oscar.tests.' \
         'test_checked_session.FakeCheckedSubtaskProcessor'
 

mars/deploy/oscar/tests/test_ray.py

@@ -18,10 +18,10 @@
 import mars.tensor as mt
 from mars.core.session import get_default_session, new_session
 from mars.deploy.oscar.ray import new_cluster, _load_config
+from mars.deploy.oscar.tests import test_local
 from mars.serialization.ray import register_ray_serializers
 from mars.tests.core import require_ray
-from ....utils import lazy_import
-from . import test_local
+from mars.utils import lazy_import
 
 ray = lazy_import('ray')
 

mars/lib/aio/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import contextlib
 import sys
 
 from .file import AioFileObject, AioFilesystem
@@ -28,6 +29,10 @@
     asyncio.get_running_loop = get_running_loop
     asyncio.create_task = asyncio.ensure_future
 
+    # patch async generator
+    from async_generator import asynccontextmanager
+    contextlib.asynccontextmanager = asynccontextmanager
+
 
 def create_lock(loop=None):
     async def _create_lock():

mars/oscar/backends/core.py

@@ -55,7 +55,7 @@ async def _listen(self, client: Client):
                     except (ConnectionError, BrokenPipeError):
                         # close failed, ignore it
                         pass
-                    raise ServerClosed(f'Remote server {client.dest_address} closed')
+                    raise ServerClosed(f'Remote server {client.dest_address} closed') from None
                 future = self._client_to_message_futures[client].pop(message.message_id)
                 future.set_result(message)
             except DeserializeMessageFailed as e:  # pragma: no cover

mars/oscar/backends/pool.py

@@ -828,7 +828,7 @@ async def handle_control_command(self,
                     processor.result = await super().handle_control_command(message)
             elif message.control_message_type == ControlMessageType.stop:
                 timeout, force = message.content if message.content is not None \
-                    else (None, None)
+                    else (None, False)
                 await self.stop_sub_pool(
                     message.address,
                     self.sub_processes[message.address],

mars/services/cluster/api/oscar.py

@@ -147,7 +147,7 @@ async def get_nodes_info(self, nodes: List[str] = None, role: NodeRole = None,
         return await self._node_info_ref.get_nodes_info(
             nodes=nodes, role=role, env=env, resource=resource, state=state)
 
-    async def get_all_bands(self) -> Dict[BandType, int]:
+    async def get_all_bands(self, role: NodeRole = None, watch: bool = False) -> Dict[BandType, int]:
         """
         Get all bands that can be used for computation.
 
@@ -156,7 +156,9 @@ async def get_all_bands(self) -> Dict[BandType, int]:
         band_to_slots : dict
             Band to n_slot.
         """
-        return await self._node_info_ref.get_all_bands()
+        if watch:
+            return await self._node_info_ref.watch_all_bands(role)
+        return await self._node_info_ref.get_all_bands(role)
 
     async def get_bands(self):
         """
@@ -226,6 +228,7 @@ async def create(cls: Type[APIType], address: str, **kw) -> APIType:
                             address=address),
             mo.create_actor(NodeInfoUploaderActor, NodeRole.WORKER,
                             interval=kw.get('upload_interval'),
+                            band_to_slots=kw.get('band_to_slots'),
                             use_gpu=kw.get('use_gpu', False),
                             uid=NodeInfoUploaderActor.default_uid(),
                             address=address),

mars/services/cluster/locator.py

@@ -13,13 +13,16 @@
 # limitations under the License.
 
 import asyncio
+import logging
 from typing import List, Optional
 
 from mars import oscar as mo
 from mars.lib.uhashring import HashRing
 from mars.services.cluster.backends import AbstractClusterBackend, get_cluster_backend
 from mars.utils import extensible
 
+logger = logging.getLogger(__name__)
+
 
 class SupervisorLocatorActor(mo.Actor):
     _backend: Optional[AbstractClusterBackend]

mars/services/cluster/supervisor/node_info.py

@@ -108,8 +108,9 @@ def get_nodes_info(self, nodes: List[str] = None, role: NodeRole = None,
             )
         return ret_infos
 
-    def get_all_bands(self) -> Dict[BandType, int]:
-        nodes = self._role_to_nodes.get(NodeRole.WORKER, list())
+    def get_all_bands(self, role: NodeRole = None) -> Dict[BandType, int]:
+        role = role or NodeRole.WORKER
+        nodes = self._role_to_nodes.get(role, [])
         band_slots = dict()
         for node in nodes:
             node_resource = self._node_infos[node].resource
@@ -136,3 +137,17 @@ async def waiter():
                 self._role_to_events[role].remove(event)
 
         return waiter()
+
+    async def watch_all_bands(self, role: NodeRole = None):
+        role = role or NodeRole.WORKER
+        event = asyncio.Event()
+        self._role_to_events[role].add(event)
+
+        async def waiter():
+            try:
+                await event.wait()
+                return self.get_all_bands(role=role)
+            finally:
+                self._role_to_events[role].remove(event)
+
+        return waiter()

mars/services/cluster/tests/test_api.py

@@ -60,13 +60,18 @@ async def test_api(actor_pool):
     assert pool_addr in nodes_info
     assert 'custom_usage' in nodes_info[pool_addr]['resource']['numa-0']
 
+    bands = await api.get_all_bands()
+    assert (pool_addr, 'numa-0') in bands
+
     with pytest.raises(asyncio.TimeoutError):
         await asyncio.wait_for(api.get_supervisors(watch=True), timeout=0.1)
     with pytest.raises(asyncio.TimeoutError):
         await asyncio.wait_for(api.get_supervisor_refs(
             [TestActor.default_uid()], watch=True), timeout=0.1)
     with pytest.raises(asyncio.TimeoutError):
         await asyncio.wait_for(api.watch_nodes(NodeRole.WORKER), timeout=0.1)
+    with pytest.raises(asyncio.TimeoutError):
+        await asyncio.wait_for(api.get_all_bands(watch=True), timeout=0.1)
 
 
 @pytest.mark.asyncio

mars/services/cluster/uploader.py

@@ -64,6 +64,9 @@ async def mark_node_ready(self):
         # upload info in time to reduce latency
         await self.upload_node_info(False)
 
+    def is_node_ready(self):
+        return self._upload_enabled
+
     async def upload_node_info(self, call_next: bool = True):
         try:
             if not self._info.env:

mars/services/meta/tests/test_api.py

@@ -92,7 +92,7 @@ async def test_meta_web_api():
 
     async with pool:
         config = {
-            "services": ["cluster", "session", "meta", "task", "web"],
+            "services": ["cluster", "session", "meta", "web"],
             "cluster": {
                 "backend": "fixed",
                 "lookup_address": pool.external_address,

mars/services/meta/tests/test_service.py

@@ -26,7 +26,7 @@ async def test_meta_service():
 
     async with pool:
         config = {
-            "services": ["cluster", "session", "meta", "task"],
+            "services": ["cluster", "session", "meta"],
             "cluster": {
                 "backend": "fixed",
                 "lookup_address": pool.external_address,

mars/services/scheduling/__init__.py

@@ -11,3 +11,5 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from .api import MockSchedulingAPI, SchedulingAPI