ADD AutoDL调度基本功能完成

AutoDL/autodl_scheduling/server.py

@@ -0,0 +1,110 @@
+import time
+import traceback
+import uuid
+from concurrent.futures.thread import ThreadPoolExecutor
+
+import loguru
+import uvicorn
+from fastapi import FastAPI, Depends, HTTPException
+from fastapi.routing import APIRoute
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from starlette import status
+
+from AutoDL.autodl_scheduling.entity.instance_pool import InstancePool, Instance
+from AutoDL.autodl_scheduling.entity.result_map import ResultMap
+from AutoDL.autodl_scheduling.entity.running_pool import RunningPool
+from AutoDL.autodl_scheduling.entity.waiting_queue import WaitingQueue
+
+
+class Server:
+    def __init__(self):
+        self.app = FastAPI()
+        self.waiting_queue = WaitingQueue()
+        self.running_pool = RunningPool()
+        self.instance_pool = InstancePool(max_instance=29)
+        self.result_map = ResultMap()
+        self.executor = ThreadPoolExecutor(max_workers=2)
+        self.worker_1 = self.executor.submit(self.scaling_worker)
+        self.worker_2 = self.executor.submit(self.introspect_instance)
+        self.app.routes.append(APIRoute("/", endpoint=self.submit, methods=['POST']))
+        self.app.routes.append(APIRoute("/{uid}", endpoint=self.get_result, methods=['GET']))
+
+    # async def submit(self, video_url, audio_url, token: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
+    async def submit(self, video_url, audio_url):
+        # if token.credentials != "7468B79CA2CF53436C06B29A6F16451C":
+        #     raise HTTPException(
+        #         status_code=status.HTTP_401_UNAUTHORIZED,
+        #         detail="Incorrect bearer token",
+        #         headers={"WWW-Authenticate": "Bearer"},
+        #     )
+        uid = str(uuid.uuid4())
+        try:
+            self.waiting_queue.enqueue(uid, video_url, audio_url)
+            return {"uid": uid}
+        except:
+            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
+    async def get_result(self, uid):
+        try:
+            return self.result_map.get(uid)
+        except:
+            if uid in  self.running_pool.tasks:
+                return {"status": "running", "msg": ""}
+            return {"status": "queuing", "msg":""}
+
+    def introspect_instance(self):
+        loguru.logger.info("Introspecting worker started")
+        while True:
+            self.instance_pool.introspection()
+            time.sleep(0.5)
+
+    def scaling_worker(self):
+        loguru.logger.info("Scaling worker started")
+        try:
+            last_target = 0
+            while True:
+                # 提交任务
+                if last_target != self.waiting_queue.get_size()+self.running_pool.get_running_size():
+                    last_target = self.waiting_queue.get_size()+self.running_pool.get_running_size()
+                    self.instance_pool.scale_instance(last_target)
+                for instance in self.instance_pool.instances:
+                    if not instance.active:
+                        # 从等待队列取出任务
+                        uid, video_path, audio_path = self.waiting_queue.dequeue()
+                        # 提交任务到运行池
+                        if self.running_pool.run(instance.uuid, uid, instance.domain, video_path, audio_path):
+                            # 更新实例池状态
+                            instance.active = True
+                            instance.last_active_time = time.time()
+                            loguru.logger.info("Task Submitted")
+                        else:
+                            loguru.logger.error("Submit Task Failed")
+
+                #查询结果放到结果缓存里
+                for uid, task in list(self.running_pool.tasks.items()):
+                    result = self.running_pool.result(uid)
+                    # 任务运行完成或失败
+                    if result:
+                        self.result_map.set(uid, result)
+                        self.running_pool.tasks.pop(uid)
+                        if result["status"] == "fail":
+                            #删除失败任务的实例
+                            self.instance_pool.remove_instance(Instance(uuid=task["instance_id"]))
+                        else:
+                            #更新成功任务实例的状态
+                            for instance in self.instance_pool.instances:
+                                if instance.uuid == task["instance_id"]:
+                                    instance.active = False
+                                    instance.last_active_time = time.time()
+                time.sleep(0.5)
+        except:
+            traceback.print_exc()
+
+
+if __name__=="__main__":
+    server = Server()
+    uvicorn.run(server.app, host="127.0.0.1", port=8888)
+
+
+
+