Expand slot scheduler to resource scheduler (mars-project#2846)

Author: zhongchun

asv_bench/benchmarks/graph_assigner.py

@@ -15,6 +15,7 @@
 import mars.tensor as mt
 import mars.dataframe as md
 from mars.core.graph import TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
+from mars.resource import Resource
 from mars.services.task.analyzer import GraphAnalyzer
 from mars.services.task.analyzer.assigner import GraphAssigner
 
@@ -39,8 +40,10 @@ def setup(self):
 
     def time_assigner(self):
         start_ops = list(GraphAnalyzer._iter_start_ops(self.chunk_graph))
-        band_slots = {(f"worker-{i}", "numa-0"): 16 for i in range(50)}
+        band_resource = {
+            (f"worker-{i}", "numa-0"): Resource(num_cpus=16) for i in range(50)
+        }
         current_assign = {}
-        assigner = GraphAssigner(self.chunk_graph, start_ops, band_slots)
+        assigner = GraphAssigner(self.chunk_graph, start_ops, band_resource)
         assigned_result = assigner.assign(current_assign)
         assert len(assigned_result) == len(start_ops)

mars/_resource.pyx

@@ -16,15 +16,15 @@ cdef class Resource:
     cdef readonly:
         float num_cpus
         float num_gpus
-        float num_mem_bytes
+        float mem_bytes
 
-    def __init__(self, float num_cpus=0, float num_gpus=0, float num_mem_bytes=0):
+    def __init__(self, float num_cpus=0, float num_gpus=0, float mem_bytes=0):
         self.num_cpus = num_cpus
         self.num_gpus = num_gpus
-        self.num_mem_bytes = num_mem_bytes
+        self.mem_bytes = mem_bytes
 
     def __eq__(self, Resource other):
-        return self.num_mem_bytes == other.num_mem_bytes and \
+        return self.mem_bytes == other.mem_bytes and \
                self.num_gpus == other.num_gpus and \
                self.num_cpus == other.num_cpus
 
@@ -33,22 +33,22 @@ cdef class Resource:
 
     def __le__(self, Resource other):
         # memory first, then gpu, cpu last
-        return self.num_mem_bytes <= other.num_mem_bytes and \
+        return self.mem_bytes <= other.mem_bytes and \
                self.num_gpus <= other.num_gpus and \
                self.num_cpus <= other.num_cpus
 
     def __add__(self, Resource other):
         return Resource(num_cpus=self.num_cpus + other.num_cpus,
                         num_gpus=self.num_gpus + other.num_gpus,
-                        num_mem_bytes=self.num_mem_bytes + other.num_mem_bytes)
+                        mem_bytes=self.mem_bytes + other.mem_bytes)
     def __sub__(self, Resource other):
         return Resource(num_cpus=self.num_cpus - other.num_cpus,
                         num_gpus=self.num_gpus - other.num_gpus,
-                        num_mem_bytes=self.num_mem_bytes - other.num_mem_bytes)
+                        mem_bytes=self.mem_bytes - other.mem_bytes)
     def __neg__(self):
-        return Resource(num_cpus=-self.num_cpus, num_gpus=-self.num_gpus, num_mem_bytes=-self.num_mem_bytes)
+        return Resource(num_cpus=-self.num_cpus, num_gpus=-self.num_gpus, mem_bytes=-self.mem_bytes)
 
     def __repr__(self):
-        return f"Resource(num_cpus={self.num_cpus}, num_gpus={self.num_gpus}, num_mem_bytes={self.num_mem_bytes})"
+        return f"Resource(num_cpus={self.num_cpus}, num_gpus={self.num_gpus}, mem_bytes={self.mem_bytes})"
 
-ZeroResource = Resource(num_cpus=0, num_gpus=0, num_mem_bytes=0)
+ZeroResource = Resource(num_cpus=0, num_gpus=0, mem_bytes=0)

mars/deploy/kubernetes/worker.py

@@ -33,7 +33,7 @@ async def start_services(self):
         await start_worker(
             self.pool.external_address,
             self.args.supervisors,
-            self.band_to_slot,
+            self.band_to_resource,
             list(self.args.load_modules),
             self.config,
             mark_ready=False,

mars/deploy/oscar/local.py

@@ -24,7 +24,7 @@
 from ... import oscar as mo
 from ...core.entrypoints import init_extension_entrypoints
 from ...lib.aio import get_isolation, stop_isolation
-from ...resource import cpu_count, cuda_count
+from ...resource import cpu_count, cuda_count, mem_total, Resource
 from ...services import NodeRole
 from ...typing import ClusterType, ClientType
 from ..utils import get_third_party_modules_from_config
@@ -51,6 +51,7 @@ async def new_cluster_in_isolation(
     address: str = "0.0.0.0",
     n_worker: int = 1,
     n_cpu: Union[int, str] = "auto",
+    mem_bytes: Union[int, str] = "auto",
     cuda_devices: Union[List[int], str] = "auto",
     subprocess_start_method: str = None,
     backend: str = None,
@@ -65,6 +66,7 @@ async def new_cluster_in_isolation(
         address,
         n_worker,
         n_cpu,
+        mem_bytes,
         cuda_devices,
         subprocess_start_method,
         config,
@@ -79,6 +81,7 @@ async def new_cluster(
     address: str = "0.0.0.0",
     n_worker: int = 1,
     n_cpu: Union[int, str] = "auto",
+    mem_bytes: Union[int, str] = "auto",
     cuda_devices: Union[List[int], str] = "auto",
     subprocess_start_method: str = None,
     config: Union[str, Dict] = None,
@@ -91,6 +94,7 @@ async def new_cluster(
         address,
         n_worker=n_worker,
         n_cpu=n_cpu,
+        mem_bytes=mem_bytes,
         cuda_devices=cuda_devices,
         subprocess_start_method=subprocess_start_method,
         config=config,
@@ -116,6 +120,7 @@ def __init__(
         address: str = "0.0.0.0",
         n_worker: int = 1,
         n_cpu: Union[int, str] = "auto",
+        mem_bytes: Union[int, str] = "auto",
         cuda_devices: Union[List[int], List[List[int]], str] = "auto",
         subprocess_start_method: str = None,
         config: Union[str, Dict] = None,
@@ -132,6 +137,7 @@ def __init__(
         self._subprocess_start_method = subprocess_start_method
         self._config = config
         self._n_cpu = cpu_count() if n_cpu == "auto" else n_cpu
+        self._mem_bytes = mem_total() if mem_bytes == "auto" else mem_bytes
         self._n_supervisor_process = n_supervisor_process
         if cuda_devices == "auto":
             total = cuda_count()
@@ -148,19 +154,22 @@ def __init__(
 
         self._n_worker = n_worker
         self._web = web
-        self._bands_to_slot = bands_to_slot = []
+        self._bands_to_resource = bands_to_resource = []
         worker_cpus = self._n_cpu // n_worker
         if sum(len(devices) for devices in devices_list) == 0:
             assert worker_cpus > 0, (
                 f"{self._n_cpu} cpus are not enough "
                 f"for {n_worker}, try to decrease workers."
             )
+        mem_bytes = self._mem_bytes // n_worker
         for _, devices in zip(range(n_worker), devices_list):
-            worker_band_to_slot = dict()
-            worker_band_to_slot["numa-0"] = worker_cpus
+            worker_band_to_resource = dict()
+            worker_band_to_resource["numa-0"] = Resource(
+                num_cpus=worker_cpus, mem_bytes=mem_bytes
+            )
             for i in devices:  # pragma: no cover
-                worker_band_to_slot[f"gpu-{i}"] = 1
-            bands_to_slot.append(worker_band_to_slot)
+                worker_band_to_resource[f"gpu-{i}"] = Resource(num_gpus=1)
+            bands_to_resource.append(worker_band_to_resource)
         self._supervisor_pool = None
         self._worker_pools = []
 
@@ -211,10 +220,10 @@ async def _start_worker_pools(self):
         worker_modules = get_third_party_modules_from_config(
             self._config, NodeRole.WORKER
         )
-        for band_to_slot in self._bands_to_slot:
+        for band_to_resource in self._bands_to_resource:
             worker_pool = await create_worker_actor_pool(
                 self._address,
-                band_to_slot,
+                band_to_resource,
                 modules=worker_modules,
                 subprocess_start_method=self._subprocess_start_method,
                 metrics=self._config.get("metrics", {}),
@@ -225,11 +234,13 @@ async def _start_service(self):
         self._web = await start_supervisor(
             self.supervisor_address, config=self._config, web=self._web
         )
-        for worker_pool, band_to_slot in zip(self._worker_pools, self._bands_to_slot):
+        for worker_pool, band_to_resource in zip(
+            self._worker_pools, self._bands_to_resource
+        ):
             await start_worker(
                 worker_pool.external_address,
                 self.supervisor_address,
-                band_to_slot,
+                band_to_resource,
                 config=self._config,
             )
 

mars/deploy/oscar/pool.py

@@ -16,7 +16,7 @@
 from typing import Dict, List
 
 from ... import oscar as mo
-from ...resource import cuda_count
+from ...resource import cuda_count, Resource
 
 try:
     from IPython import get_ipython
@@ -46,7 +46,7 @@ async def create_supervisor_actor_pool(
 
 async def create_worker_actor_pool(
     address: str,
-    band_to_slots: Dict[str, int],
+    band_to_resource: Dict[str, Resource],
     n_io_process: int = 1,
     modules: List[str] = None,
     ports: List[int] = None,
@@ -55,7 +55,10 @@ async def create_worker_actor_pool(
     **kwargs,
 ):
     # TODO: support NUMA when ready
-    n_process = sum(slot for slot in band_to_slots.values())
+    n_process = sum(
+        int(resource.num_cpus) or int(resource.num_gpus)
+        for resource in band_to_resource.values()
+    )
     envs = []
     labels = ["main"]
 
@@ -67,15 +70,16 @@ async def create_worker_actor_pool(
             cuda_devices = [int(i) for i in env_devices.split(",")]
 
     i_gpu = iter(sorted(cuda_devices))
-    for band, slot in band_to_slots.items():
+    for band, resource in band_to_resource.items():
         if band.startswith("gpu"):  # pragma: no cover
             idx = str(next(i_gpu))
             envs.append({"CUDA_VISIBLE_DEVICES": idx})
             labels.append(f"gpu-{idx}")
         else:
             assert band.startswith("numa")
-            envs.extend([dict() for _ in range(slot)])
-            labels.extend([band] * slot)
+            num_cpus = int(resource.num_cpus)
+            envs.extend([dict() for _ in range(num_cpus)])
+            labels.extend([band] * num_cpus)
 
     suspend_sigint = get_ipython is not None and get_ipython() is not None
     return await mo.create_actor_pool(

mars/deploy/oscar/ray.py

@@ -30,6 +30,7 @@
 )
 from ...oscar.backends.ray.pool import RayPoolState
 from ...oscar.errors import ReconstructWorkerError
+from ...resource import Resource
 from ...services.cluster.backends.base import (
     register_cluster_backend,
     AbstractClusterBackend,
@@ -130,7 +131,7 @@ def get_cluster_state_ref(self):
 class ClusterStateActor(mo.StatelessActor):
     def __init__(self):
         self._worker_cpu, self._worker_mem, self._config = None, None, None
-        self._pg_name, self._band_to_slot, self._worker_modules = None, None, None
+        self._pg_name, self._band_to_resource, self._worker_modules = None, None, None
         self._pg_counter = itertools.count()
         self._worker_count = 0
         self._workers = {}
@@ -147,7 +148,9 @@ def set_config(self, worker_cpu, worker_mem, config):
             config,
         )
         # TODO(chaokunyang) Support gpu
-        self._band_to_slot = {"numa-0": self._worker_cpu}
+        self._band_to_resource = {
+            "numa-0": Resource(num_cpus=self._worker_cpu, mem_bytes=self._worker_mem)
+        }
         self._worker_modules = get_third_party_modules_from_config(
             self._config, NodeRole.WORKER
         )
@@ -156,11 +159,14 @@ async def request_worker(
         self, worker_cpu: int = None, worker_mem: int = None, timeout: int = None
     ) -> Optional[str]:
         worker_cpu = worker_cpu or self._worker_cpu
+        worker_mem = worker_mem or self._worker_mem
         bundle = {
             "CPU": worker_cpu,
             # "memory": worker_mem or self._worker_mem
         }
-        band_to_slot = {"numa-0": worker_cpu}
+        band_to_resource = {
+            "numa-0": Resource(num_cpus=worker_cpu, mem_bytes=worker_mem)
+        }
         start_time = time.time()
         logger.info("Start to request worker with resource %s.", bundle)
         # TODO rescale ray placement group instead of creating new placement group
@@ -193,7 +199,7 @@ async def request_worker(
         )
         worker_address = process_placement_to_address(pg_name, 0, 0)
         worker_pool = await self.create_worker(worker_address)
-        await self.start_worker(worker_address, band_to_slot=band_to_slot)
+        await self.start_worker(worker_address, band_to_resource=band_to_resource)
         logger.info(
             "Request worker %s succeeds in %.4f seconds",
             worker_address,
@@ -206,7 +212,7 @@ async def create_worker(self, worker_address):
         start_time = time.time()
         worker_pool = await create_worker_actor_pool(
             worker_address,
-            self._band_to_slot,
+            self._band_to_resource,
             modules=self._worker_modules,
             metrics=self._config.get("metrics", {}),
         )
@@ -217,12 +223,12 @@ async def create_worker(self, worker_address):
         )
         return worker_pool
 
-    async def start_worker(self, worker_address, band_to_slot=None):
+    async def start_worker(self, worker_address, band_to_resource=None):
         self._worker_count += 1
         start_time = time.time()
-        band_to_slot = band_to_slot or self._band_to_slot
+        band_to_resource = band_to_resource or self._band_to_resource
         await start_worker(
-            worker_address, self.address, band_to_slot, config=self._config
+            worker_address, self.address, band_to_resource, config=self._config
         )
         worker_pool = ray.get_actor(worker_address)
         await worker_pool.mark_service_ready.remote()
@@ -290,7 +296,7 @@ async def _reconstruct_worker():
 
             start_time = time.time()
             await start_worker(
-                address, self.address, self._band_to_slot, config=self._config
+                address, self.address, self._band_to_resource, config=self._config
             )
             await actor.mark_service_ready.remote()
             logger.info(
@@ -514,7 +520,11 @@ async def start(self):
             asyncio.create_task(
                 create_worker_actor_pool(
                     addr,
-                    {"numa-0": self._worker_cpu},
+                    {
+                        "numa-0": Resource(
+                            num_cpus=self._worker_cpu, mem_bytes=self._worker_mem
+                        )
+                    },
                     modules=get_third_party_modules_from_config(
                         self._config, NodeRole.WORKER
                     ),

mars/deploy/oscar/service.py

@@ -16,6 +16,7 @@
 import os
 from typing import List, Dict, Union
 
+from ...resource import Resource
 from ...services import start_services, stop_services, NodeRole
 from ..utils import load_service_config_file
 
@@ -75,7 +76,7 @@ async def stop_supervisor(address: str, config: Dict = None):
 async def start_worker(
     address: str,
     lookup_address: str,
-    band_to_slots: Dict[str, int],
+    band_to_resource: Dict[str, Resource],
     modules: Union[List, str, None] = None,
     config: Dict = None,
     mark_ready: bool = True,
@@ -87,9 +88,9 @@ async def start_worker(
     if backend == "fixed" and config["cluster"].get("lookup_address") is None:
         config["cluster"]["lookup_address"] = lookup_address
     if config["cluster"].get("resource") is None:
-        config["cluster"]["resource"] = band_to_slots
+        config["cluster"]["resource"] = band_to_resource
     if any(
-        band_name.startswith("gpu-") for band_name in band_to_slots
+        band_name.startswith("gpu-") for band_name in band_to_resource
     ):  # pragma: no cover
         if "cuda" not in config["storage"]["backends"]:
             config["storage"]["backends"].append("cuda")

mars/deploy/oscar/session.py

@@ -1256,10 +1256,10 @@ async def fetch_tileable_op_logs(
     async def get_total_n_cpu(self):
         all_bands = await self._cluster_api.get_all_bands()
         n_cpu = 0
-        for band, size in all_bands.items():
+        for band, resource in all_bands.items():
             _, band_name = band
             if band_name.startswith("numa-"):
-                n_cpu += size
+                n_cpu += resource.num_cpus
         return n_cpu
 
     async def get_cluster_versions(self) -> List[str]: