chore: import upstream snapshot with attribution

swift/ray_utils/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+from .arguments import RayArguments
+from .base import RayHelper
+
+
+def try_init_ray():
+    import argparse
+    import json
+    from transformers.utils import strtobool
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--use_ray', type=str, default='0')
+    parser.add_argument('--device_groups', type=str, default=None)
+    args, _ = parser.parse_known_args()
+    args.use_ray = strtobool(args.use_ray)
+    if args.use_ray:
+        RayHelper.initialize(json.loads(args.device_groups))

swift/ray_utils/arguments.py

@@ -0,0 +1,30 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class RayArguments:
+    """A dataclass that holds the configuration and usage for Ray.
+
+    Args:
+        use_ray (bool): Whether to use Ray for distributed operations. Defaults to False.
+        ray_exp_name (Optional[str]): The name of the Ray experiment. This is used as a prefix for cluster and worker
+            names. This argument is optional. Defaults to None.
+        device_groups (Optional[str]): A JSON string that defines the device groups for Ray. This field is mandatory
+            when `use_ray` is True. Defaults to None. For the specific format and details, please refer to the
+            [Ray documentation](https://swift.readthedocs.io/zh-cn/latest/Instruction/Ray.html)
+    """
+    use_ray: bool = False
+
+    ray_exp_name: Optional[str] = None
+
+    device_groups: Optional[str] = None
+
+    def __post_init__(self):
+        if isinstance(self.device_groups, str):
+            self.device_groups = json.loads(self.device_groups)
+        if self.ray_exp_name:
+            os.environ['RAY_SWIFT_EXP_NAME'] = self.ray_exp_name.strip()

swift/ray_utils/base.py

@@ -0,0 +1,368 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import argparse
+import functools
+import inspect
+import json
+import numpy as np
+import os
+from typing import Any, Callable, Dict, List, Literal, Optional, TypeVar, Union
+
+from swift.utils import find_free_port, find_node_ip
+from .arguments import RayArguments
+from .resource_manager import ResourceManager
+
+T = TypeVar('T')
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args()
+    return json.dumps(unknown)
+
+
+class RayHelper:
+    resource_manager: Optional[ResourceManager] = None
+
+    worker_cls: Dict = {}
+
+    args: RayArguments = None
+
+    worker_instance: Dict = {}
+
+    initialized = False
+
+    device_groups: Dict[str, Any] = None
+
+    @staticmethod
+    def initialize(device_groups: Dict[str, Any]):
+        """Initialize RayHelper.
+
+        Args:
+            device_groups: The device groups to initialize.
+
+        Returns:
+            None
+        """
+        if RayHelper.ray_inited():
+            return
+        import ray
+        RayHelper.device_groups = device_groups
+        ray.init()
+        if RayHelper.resource_manager is None:
+            # Resource manager initialize only once in the pipeline process.
+            RayHelper.resource_manager = ResourceManager(device_groups)
+
+    @staticmethod
+    def teardown():
+        if RayHelper.resource_manager is not None:
+            RayHelper.resource_manager.destroy_placement_group()
+            RayHelper.resource_manager = None
+
+    @staticmethod
+    def is_called_from_init():
+        """If some function called from __init__.
+
+        Ray functions perform different behaviors depending on whether they are called from __init__.
+
+        Returns:
+            Boolean.
+        """
+        stack = inspect.stack()
+        for frame_info in stack[1:]:
+            if frame_info.function == '__init__':
+                return True
+        return False
+
+    @staticmethod
+    def ray_inited():
+        try:
+            import ray
+        except ImportError:
+            # not installed, not inited
+            return False
+        return ray.is_initialized()
+
+    @staticmethod
+    def is_worker():
+        import ray
+        return RayHelper.ray_inited() and ray._private.worker.global_worker.mode == ray._private.worker.WORKER_MODE
+
+    @staticmethod
+    def worker(group: Union[str, List[str]]):
+
+        def decorator(cls):
+            if not RayHelper.ray_inited():
+                return cls
+            if RayHelper.is_worker():
+                return cls
+            cls.decorated = True
+            groups = [group] if isinstance(group, str) else group
+            import ray
+            _cls = ray.remote(cls)
+            for g in groups:
+                RayHelper.worker_cls[g] = _cls
+
+            init_method = cls.__init__
+
+            @functools.wraps(init_method)
+            def new_init(self, *args, **kwargs):
+                if not RayHelper.is_worker():
+                    # Create remote workers
+                    RayHelper._create_workers(group, *args, **kwargs)
+                init_method(self, *args, **kwargs)
+
+            cls.__init__ = new_init
+
+            return cls
+
+        return decorator
+
+    @staticmethod
+    def collect_func(method: Union[Literal['none', 'flatten'], Callable], result):
+        if isinstance(result[0], tuple):
+            output = []
+            for i in range(len(result[0])):
+                _single_result = [r[i] for r in result]
+                output.append(RayHelper.collect_func(method, _single_result))
+            return output
+        if method == 'none':
+            return result
+        elif method == 'flatten':
+            flatten = [item for sublist in result for item in sublist]
+            if isinstance(result[0], np.ndarray):
+                return np.array(flatten)
+            return type(result[0])(flatten)
+        elif isinstance(method, Callable):
+            # Callable
+            return method(result)
+        else:
+            raise ValueError(f'Unsupported collect method: {method}')
+
+    @staticmethod
+    def function(group: str,
+                 dispatch: Union[Literal['slice', 'all'], Callable] = 'all',
+                 execute: Literal['first', 'all'] = 'all',
+                 collect: Union[Literal['none', 'flatten'], Callable] = 'none'):
+        """Remote execution function.
+
+        Args:
+            group: The group to execute.
+            dispatch: How to dispatch the arguments.
+                'slice': load balance
+                'all': all processes do the same thing
+            execute: How to execute
+                'first': Only first worker
+                'all': All processes
+            collect: How to collect the results.
+                'none': Return as-is
+                'flatten': Return a flattened list
+        Returns:
+            The execution result.
+        """
+
+        def decorator(func: Callable[..., T]) -> Callable[..., T]:
+
+            @functools.wraps(func)
+            def wrapper(self, *args, **kwargs) -> T:
+                if not RayHelper.ray_inited():
+                    return func(self, *args, **kwargs)
+                if RayHelper.is_worker():
+                    if not hasattr(self, 'group'):
+                        # pass through env
+                        self.group = os.environ['RAY_SWIFT_GROUP'].split(',')
+                    if group not in self.group:
+                        if RayHelper.is_called_from_init():
+                            # Functions in init of different group, do nothing
+                            return None
+                        else:
+                            # Should not happen
+                            raise ValueError()
+                    else:
+                        return func(self, *args, **kwargs)
+                else:
+                    if RayHelper.is_called_from_init():
+                        # each worker do its own init
+                        return None
+                    result = RayHelper.execute_all_sync(group, dispatch, execute, func.__name__, *args, **kwargs)
+                    return RayHelper.collect_func(collect, result)
+
+            return wrapper
+
+        return decorator
+
+    @staticmethod
+    def execute_all_sync(group, dispatch, execute, method_name: str, *args, **kwargs):
+        import ray
+        return ray.get(RayHelper.execute_all_async(group, dispatch, execute, method_name, *args, **kwargs))
+
+    @staticmethod
+    def execute_all_async(group, dispatch, execute, method_name: str, *args, **kwargs):
+        workers = RayHelper.worker_instance[group]
+        length = len(workers)
+        if execute == 'first':
+            return getattr(workers[0], method_name).remote(*args, **kwargs)
+        elif dispatch == 'all':
+            return [getattr(worker, method_name).remote(*args, **kwargs) for worker in workers]
+        elif dispatch == 'slice':
+            result = []
+
+            def dispatch_func(arg, n):
+                if isinstance(arg, list):
+                    k, m = divmod(len(arg), n)
+                    return [arg[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n)]
+                else:
+                    return [arg] * n
+
+            args = [dispatch_func(arg, length) for arg in args]
+            kwargs = {k: dispatch_func(v, length) for k, v in kwargs.items()}
+            for i in range(length):
+                sliced_args = tuple(arg[i] for arg in args)
+                sliced_kwargs = {k: v[i] for k, v in kwargs.items()}
+                if (sliced_args and sliced_args[0]) or (kwargs and list(kwargs.values())):
+                    # skip empty input
+                    remote_call = getattr(workers[i], method_name)
+                    result.append(remote_call.remote(*sliced_args, **sliced_kwargs))
+            return result
+        elif isinstance(dispatch, Callable):
+            # dispatch is Callable
+            result = []
+            for i in range(length):
+                sliced_args, sliced_kwargs = dispatch(length, i, *args, **kwargs)
+                remote_call = getattr(workers[i], method_name)
+                result.append(remote_call.remote(*sliced_args, **sliced_kwargs))
+            return result
+        else:
+            raise ValueError(f'Invalid dispatch method: {dispatch}')
+
+    @staticmethod
+    def _create_workers(group: Union[str, List[str]], *args, **kwargs):
+        import ray
+        from ray.runtime_env import RuntimeEnv
+        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+        exp_name = os.environ.get('RAY_SWIFT_EXP_NAME')
+        if not exp_name:
+            exp_name = ''
+        else:
+            exp_name += '-'
+
+        if isinstance(group, str):
+            group = [group]
+
+        for _group in group:
+            if _group in RayHelper.worker_instance:
+                continue
+
+            worker_cls = RayHelper.worker_cls[_group]
+
+            _config = None
+            for name, config in RayHelper.device_groups.items():
+                if name in RayHelper.resource_manager.possible_keys:
+                    continue
+
+                if _group in config['workers']:
+                    _config = config
+                    break
+
+            assert _config is not None
+            local_groups = _config['workers']
+
+            VISIBLE_ENV_MAPPING = {
+                'GPU': 'CUDA_VISIBLE_DEVICES',
+                'NPU': 'ASCEND_VISIBLE_DEVICES',
+            }
+
+            if _config['device'].upper() != 'CPU':
+                world_size = len(_config['ranks'])
+                placement_groups: List[List[Dict]] = RayHelper.resource_manager.resource(_group)
+                workers = []
+                ip, port = None, None
+                for rank, (deploy_pg, gpu) in enumerate(zip(placement_groups, _config['ranks'])):
+                    deploy_pg: Dict
+                    cluster_name = exp_name + '-'.join(local_groups)
+                    worker_name = cluster_name + '-' + str(rank)
+                    env_vars = os.environ.copy()
+                    env_vars.update({
+                        'WORLD_SIZE':
+                        str(world_size),
+                        'RANK':
+                        str(rank),
+                        'LOCAL_RANK':
+                        str(0),
+                        'CLUSTER_NAME':
+                        cluster_name,
+                        'WORKER_NAME':
+                        worker_name,
+                        VISIBLE_ENV_MAPPING[_config['device'].upper()]:
+                        ','.join([str(r) for r in deploy_pg['gpu_rank']]),  # TODO npu
+                        'RAY_SWIFT_ARGS':
+                        get_args(),  # pass through env
+                    })
+
+                    @ray.remote
+                    def get_node_address():
+                        return find_node_ip(), find_free_port()
+
+                    if rank == 0:
+                        ip, port = ray.get(
+                            get_node_address.options(placement_group=deploy_pg['placement_group']).remote())
+
+                    env_vars['MASTER_ADDR'] = ip
+                    env_vars['MASTER_PORT'] = str(port)
+                    env_vars['RAY_SWIFT_GROUP'] = ','.join(local_groups)
+
+                    runtime_env = RuntimeEnv(env_vars=env_vars)
+
+                    worker_options = {
+                        'scheduling_strategy':
+                        PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']),
+                        'name':
+                        worker_name,
+                        'namespace':
+                        'default',
+                        'runtime_env':
+                        runtime_env,
+                        'num_cpus':
+                        0.01,
+                        'num_gpus':
+                        0.01,
+                    }
+
+                    worker = worker_cls.options(**worker_options).remote(*args, **kwargs)
+                    workers.append(worker)
+            else:
+                world_size = _config['ranks']
+                placement_groups: List[List[Dict]] = RayHelper.resource_manager.resource(_group)
+                workers = []
+                for deploy_pg, index in zip(placement_groups, list(range(world_size))):
+                    deploy_pg: Dict
+                    cluster_name = '-'.join(local_groups)
+                    worker_name = cluster_name + '-' + str(index)
+                    env_vars = os.environ.copy()
+                    env_vars.update({
+                        'CLUSTER_NAME': cluster_name,
+                        'WORKER_NAME': worker_name,
+                        VISIBLE_ENV_MAPPING[_config['device'].upper()]: '',
+                        'RAY_SWIFT_ARGS': get_args(),  # pass through env
+                    })
+                    env_vars['RAY_SWIFT_GROUP'] = ','.join(local_groups)
+
+                    runtime_env = RuntimeEnv(env_vars=env_vars)
+
+                    worker_options = {
+                        'scheduling_strategy':
+                        PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']),
+                        'name':
+                        worker_name,
+                        'namespace':
+                        'default',
+                        'runtime_env':
+                        runtime_env,
+                        'num_cpus':
+                        0.01,
+                    }
+
+                    worker = worker_cls.options(**worker_options).remote(*args, **kwargs)
+                    workers.append(worker)
+
+            for g in local_groups:
+                RayHelper.worker_instance[g] = workers

swift/ray_utils/resource_manager.py

@@ -0,0 +1,137 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+# Some code borrowed from ROLL: https://github.com/alibaba/ROLL
+import math
+import os
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+
+@dataclass
+class NodeGroup:
+    device_count: int
+    nodes: List[Any] = field(default_factory=list)
+
+
+def get_node_rank():
+    return int(os.environ.get('NODE_RANK', '0'))
+
+
+class ResourceManager:
+
+    possible_keys = ['nproc_per_node', 'nnodes']
+
+    def __init__(self, groups: Dict[str, Any]):
+        import ray
+        from ray.util.placement_group import PlacementGroup
+        nproc_per_node = int(groups['nproc_per_node'])
+        device_types = set([group['device'].upper()
+                            for group in groups.values() if hasattr(group, '__getitem__')]) - {'CPU'}
+        assert len(device_types) == 1
+        device_type = next(iter(device_types))
+        all_ranks = []
+        last_rank = -1
+        cpu_proc_count = 0
+        for group_name, group in groups.items():
+            if group_name in self.possible_keys:
+                continue
+            ranks = group['ranks']
+            device = group['device'].upper()
+            if device == 'CPU':
+                assert isinstance(ranks, int), 'CPU group only supports integer ranks'
+                cpu_proc_count += ranks
+                continue
+            try:
+                ranks = int(ranks)  # int type
+                ranks = list(range(last_rank + 1, last_rank + 1 + ranks))
+            except Exception:  # noqa
+                if isinstance(ranks, str):
+                    ranks = eval(ranks, {'__builtins__': {'list': list, 'range': range}})
+            finally:
+                all_ranks.extend(ranks)
+                group['ranks'] = ranks
+                last_rank = ranks[-1]
+
+        assert len(set(all_ranks)) == len(all_ranks)
+        groups['nnodes'] = math.ceil(len(all_ranks) / nproc_per_node)
+
+        self.nodes = []
+        for node in ray.nodes():
+            resource = node['Resources']
+            node_gpu_num = int(resource.get(device_type, 0))
+            if node_gpu_num >= nproc_per_node:
+                self.nodes.append(node)
+
+        bundles = []
+        cpu_bundles = []
+        for i in range(groups['nnodes']):
+            node = self.nodes[i]
+            node_cpu = int(node['Resources']['CPU'])
+            bundles.append({device_type: nproc_per_node, 'CPU': node_cpu // 2 + 1})
+            cpu_bundles.append({'CPU': node_cpu // 4 + 1})  # TODO dynamic scheduling
+
+        nproc_cpu_per_node = cpu_proc_count // len(cpu_bundles) + 1
+        self.cpu_node_map = {}
+        for i in range(cpu_proc_count):
+            node_idx = i // nproc_cpu_per_node
+            cpu_cnt = cpu_bundles[node_idx]['CPU']
+            self.cpu_node_map[i] = (node_idx, cpu_cnt // nproc_cpu_per_node)
+
+        self.placement_groups = [ray.util.placement_group([bundle]) for bundle in bundles]
+        self.cpu_placement_groups = [ray.util.placement_group([bundle]) for bundle in cpu_bundles]
+        cpu_bundles.sort(key=lambda bundle: bundle['CPU'], reverse=True)
+        ray.get([pg.ready() for pg in self.placement_groups])
+        ray.get([pg.ready() for pg in self.cpu_placement_groups])
+
+        self.node_ranks = ray.get(
+            [ray.remote(get_node_rank).options(placement_group=pg).remote() for pg in self.placement_groups])
+        if self.node_ranks.count(0) > 1:
+            self.node_ranks = list(range(len(self.placement_groups)))
+
+        self.node2pg: Dict[int, PlacementGroup] = {}
+        for node_rank, placement_group in zip(self.node_ranks, self.placement_groups):
+            self.node2pg[node_rank] = placement_group
+
+        self.device_groups = {}
+        ray_address = str(ray.get_runtime_context().gcs_address)
+        for group_name, group in groups.items():
+            if group_name in self.possible_keys:
+                continue
+
+            if group['device'] != 'CPU':
+                ranks = group['ranks']
+                local_device_groups = []
+                for rank in ranks:
+                    node_rank = rank // nproc_per_node
+                    gpu_rank = rank % nproc_per_node
+                    local_device_groups.append(
+                        dict(
+                            node_rank=node_rank,
+                            gpu_rank=[gpu_rank],
+                            placement_group=self.node2pg[node_rank],
+                            ray_address=ray_address))
+                for worker in group['workers']:
+                    self.device_groups[worker] = local_device_groups
+            else:
+                ranks = group['ranks']
+                local_device_groups = []
+                global_cpu_proc_idx = 0
+                for _ in range(ranks):
+                    local_device_groups.append(
+                        dict(
+                            placement_group=self.cpu_placement_groups[self.cpu_node_map[global_cpu_proc_idx][0]],
+                            ray_address=ray_address))
+                    global_cpu_proc_idx += 1
+                for worker in group['workers']:
+                    self.device_groups[worker] = local_device_groups
+
+        self.groups = groups
+
+    def resource(self, worker):
+        return self.device_groups[worker]
+
+    def destroy_placement_group(self):
+        import ray
+        for pg in self.placement_groups:
+            ray.util.remove_placement_group(pg)
+        for pg in self.cpu_placement_groups:
+            ray.util.remove_placement_group(pg)