32.3.2 分布式计算框架：Ray分布式计算

import timeimport randomfrom dataclasses import dataclass, fieldfrom typing import List, Dict, Optional, Any, Callablefrom enum import Enumfrom collections import defaultdict class TaskState(Enum):    PENDING = 'pending'    RUNNING = 'running'    DONE = 'done'    FAILED = 'failed' @dataclassclass RayTask:    task_id: str    func: Callable    args: tuple = ()    kwargs: dict = field(default_factory=dict)    dependencies: List[str] = field(default_factory=list)    state: TaskState = TaskState.PENDING    result: Any = None    num_cpus: int = 1    num_gpus: int = 0 @dataclassclass ObjectRef:    object_id: str    task_id: str    consumed: bool = False class RayObjectStore:    def __init__(self):        self._store: Dict[str, Any] = {}        self._ref_count: Dict[str, int] = defaultdict(int)     def put(self, obj_id: str, value: Any):        self._store[obj_id] = value        self._ref_count[obj_id] += 1     def get(self, obj_id: str) -> Any:        return self._store.get(obj_id)     def remove(self, obj_id: str):        self._ref_count[obj_id] -= 1        if self._ref_count[obj_id] <= 0:            del self._store[obj_id] class RayActor:    def __init__(self, actor_id: str, cls: type, *args, **kwargs):        self.actor_id = actor_id        self._instance = cls(*args, **kwargs)        self._method_queue: List[Dict] = []        self._results: Dict[str, Any] = {}     def call_method(self, method_name: str, *args, **kwargs) -> str:        call_id = f'{self.actor_id}_{len(self._method_queue)}'        self._method_queue.append({            'call_id': call_id, 'method': method_name,            'args': args, 'kwargs': kwargs,        })        return call_id     def execute(self, call_id: str) -> Any:        for item in self._method_queue:            if item['call_id'] == call_id:                method = getattr(self._instance, item['method'])                result = method(*item['args'], **item['kwargs'])                self._results[call_id] = result                return result        return None class RayScheduler:    def __init__(self, num_workers: int = 4, total_cpus: int = 8, total_gpus: int = 2):        self.num_workers = num_workers        self.total_cpus = total_cpus        self.total_gpus = total_gpus        self.available_cpus = total_cpus        self.available_gpus = total_gpus        self.tasks: Dict[str, RayTask] = {}        self.object_store = RayObjectStore()        self.actors: Dict[str, RayActor] = {}        self.completed_count = 0        self.failed_count = 0     def submit(self, task: RayTask) -> ObjectRef:        self.tasks[task.task_id] = task        obj_id = f'obj_{task.task_id}'        return ObjectRef(object_id=obj_id, task_id=task.task_id)     def get_ready_tasks(self) -> List[RayTask]:        ready = []        for task in self.tasks.values():            if task.state != TaskState.PENDING:                continue            deps_met = all(                self.tasks.get(dep, RayTask('', None)).state == TaskState.DONE                for dep in task.dependencies            )            if deps_met:                ready.append(task)        return ready     def can_schedule(self, task: RayTask) -> bool:        return (self.available_cpus >= task.num_cpus and                self.available_gpus >= task.num_gpus)     def execute_task(self, task: RayTask):        task.state = TaskState.RUNNING        self.available_cpus -= task.num_cpus        self.available_gpus -= task.num_gpus        resolved_args = []        for arg in task.args:            if isinstance(arg, ObjectRef):                resolved_args.append(self.object_store.get(arg.object_id))            else:                resolved_args.append(arg)        try:            result = task.func(*resolved_args, **task.kwargs)            task.result = result            task.state = TaskState.DONE            obj_id = f'obj_{task.task_id}'            self.object_store.put(obj_id, result)            self.completed_count += 1        except Exception as e:            task.state = TaskState.FAILED            self.failed_count += 1        self.available_cpus += task.num_cpus        self.available_gpus += task.num_gpus     def run(self) -> Dict[str, Any]:        total = len(self.tasks)        while self.completed_count + self.failed_count < total:            ready = self.get_ready_tasks()            if not ready:                break            for task in ready:                if self.can_schedule(task):                    self.execute_task(task)        return {tid: t.result for tid, t in self.tasks.items()                if t.state == TaskState.DONE}     def create_actor(self, cls: type, *args, **kwargs) -> str:        actor_id = f'actor_{len(self.actors)}'        self.actors[actor_id] = RayActor(actor_id, cls, *args, **kwargs)        return actor_id def remote_task(x):    return x * 2 + 1 def remote_reduce(a, b):    return a + b class Counter:    def __init__(self):        self.value = 0    def increment(self, n=1):        self.value += n        return self.value    def get_value(self):        return self.value def main():    scheduler = RayScheduler(num_workers=4, total_cpus=8, total_gpus=2)    t1 = RayTask('task_1', remote_task, args=(10,), num_cpus=1)    t2 = RayTask('task_2', remote_task, args=(20,), num_cpus=1)    ref1 = scheduler.submit(t1)    ref2 = scheduler.submit(t2)    t3 = RayTask('task_3', remote_reduce, args=(ref1, ref2),                 dependencies=['task_1', 'task_2'], num_cpus=1)    scheduler.submit(t3)    results = scheduler.run()    print(f'Task results: {results}')    print(f'Completed: {scheduler.completed_count}, Failed: {scheduler.failed_count}')    actor_id = scheduler.create_actor(Counter)    actor = scheduler.actors[actor_id]    call_id = actor.call_method('increment', 5)    print(f'Actor counter: {actor.execute(call_id)}')    print(f'Final value: {actor.execute(actor.call_method("get_value"))}') if __name__ == '__main__':    main()