ADD AutoDL调度基本功能完成

AutoDL/autodl_scheduling/entity/instance_pool.py

@@ -0,0 +1,115 @@
+import copy
+import os
+import random
+import time
+from concurrent.futures.thread import ThreadPoolExecutor
+from typing import List
+
+import loguru
+
+from AutoDL.autodl_scheduling.util.audodl_sdk import instance_operate, get_autodl_machines, payg
+
+RETRY_LIMIT = 20 # 创建实例重试次数
+
+
+class Instance:
+    def __init__(self, uuid:str, active:bool=False, last_active_time:float=-1., domain:str=""):
+        self.uuid:str=uuid
+        self.active:bool=active
+        self.last_active_time:float=last_active_time
+        self.domain:str=domain
+
+    def __str__(self):
+        return "uuid:%s active:%s last_active_time:%s domain:%s" % (self.uuid, self.active, self.last_active_time, self.domain)
+
+
+class InstancePool:
+    def __init__(self, min_instance=0, max_instance=100, scaledown_window=120, buffer_instance=0, timeout=1200):
+        self.min_instance = min_instance
+        self.max_instance = max_instance
+        self.scaledown_window = scaledown_window
+        self.buffer_instance = buffer_instance
+        self.timeout = timeout
+        self.instances:List[Instance] = []
+        self.executor = ThreadPoolExecutor(max_workers=os.cpu_count()*2)
+        self.threads = []
+
+    def scale_instance(self, target_instance):
+        if target_instance + self.buffer_instance < self.min_instance:
+            return self._scale(self.min_instance)
+        if target_instance + self.buffer_instance > self.max_instance:
+            return self._scale(self.max_instance)
+        if target_instance + self.buffer_instance == len(self.instances):
+            return True
+        return self._scale(target_instance + self.buffer_instance)
+
+    def remove_instance(self, instance:Instance):
+        if instance_operate(instance.uuid, "power_off"):
+            if instance_operate(instance.uuid, "release"):
+                for i in self.instances:
+                    if i.uuid == instance.uuid:
+                        self.instances.remove(i)
+            else:
+                loguru.logger.error("Instance {} failed to release".format(instance.uuid))
+        else:
+            loguru.logger.error("Instance {} failed to power off".format(instance.uuid))
+
+    def _add_instance(self):
+        lim = RETRY_LIMIT
+        while lim > 0:
+            machines = get_autodl_machines()
+            if len(machines) > 0:
+                m = random.choice(machines)
+                result = payg(m["region_name"], m["machine_id"])
+                if result:
+                    self.instances.append(
+                        Instance(uuid=result[0], active=False, last_active_time=time.time(), domain="https://"+result[1]))
+                    break
+                else:
+                    time.sleep(1)
+                    lim -= 1
+        if lim <= 0:
+            loguru.logger.error("Fail to Scale[Add] Instance")
+
+    def introspection(self):
+        # 停止超时实例(运行超时和无任务超时)
+        instance_copy = copy.deepcopy(self.instances)
+        for instance in instance_copy:
+            if instance.active:
+                if (time.time() - instance.last_active_time) > self.timeout:
+                    self.threads.append(self.executor.submit(self.remove_instance, instance=instance))
+            else:
+                if (time.time() - instance.last_active_time) > self.scaledown_window:
+                    self.threads.append(self.executor.submit(self.remove_instance, instance=instance))
+
+    def _scale(self, target_instance:int):
+        loguru.logger.info("Instance Num Before Scaling %d ; Target %d" % (len(self.instances), target_instance))
+        self.introspection()
+        # 调整实例数量
+        instance_copy = copy.deepcopy(self.instances)
+        dest = target_instance - len(instance_copy)
+        if dest < 0:
+            dest = abs(dest)
+            for instance in instance_copy:
+                if not instance.active and dest > 0:
+                    self.threads.append(self.executor.submit(self.remove_instance, instance=instance))
+                    dest -= 1
+        elif dest > 0:
+            for i in range(dest):
+                self.threads.append(self.executor.submit(self._add_instance))
+        while len(self.threads) > 0:
+            for t in self.threads:
+                t.result(timeout=self.timeout//2)
+                self.threads.remove(t)
+        loguru.logger.info("Instance Num After Scaling %d ; Target %d" % (len(self.instances), target_instance))
+        if len(self.instances) == target_instance:
+            return True
+        else:
+            return False
+
+
+if __name__ == "__main__":
+    ip = InstancePool()
+    print(ip.scale_instance(5))
+    time.sleep(5)
+    print(ip.scale_instance(0))

AutoDL/autodl_scheduling/entity/result_map.py

@@ -0,0 +1,36 @@
+from typing import Union
+
+
+class ResultMap:
+    def __init__(self, max_size:int=1000):
+        self.map = dict()
+        self.max_size = max_size
+
+    def get(self, uid:str) -> Union[dict,None]:
+        if uid in self.map:
+            if self.map[uid]:
+                r = self.map[uid]
+                # self.remove(uid)
+                return r
+            else:
+                return None
+        else:
+            raise KeyError(uid)
+
+    def set(self, uid:str, result:dict=None) -> None:
+        if uid in self.map:
+            raise RuntimeError(f"Key {uid} already exists")
+        else:
+            if len(self.map) + 1 > self.max_size:
+                raise RuntimeError(f"ResultMap exceeds max size {self.max_size}")
+            self.map[uid] = result
+
+    def update(self, uid:str, result:dict) -> None:
+        if uid in self.map:
+            self.map[uid] = result
+        else:
+            raise KeyError(uid)
+
+    def remove(self, uid:str) -> None:
+        if uid in self.map:
+            del self.map[uid]

AutoDL/autodl_scheduling/entity/running_pool.py

@@ -0,0 +1,80 @@
+import time
+import uuid
+from concurrent.futures.thread import ThreadPoolExecutor
+from typing import Union
+
+import loguru
+import requests
+from fastapi import UploadFile
+SERVER_TIME_OUT = 60 * 35
+
+class RunningPool:
+    def __init__(self):
+        self.tasks = {}
+        self._headers = {
+            "Authorization": "Bearer 7468B79CA2CF53436C06B29A6F16451C"
+        }
+
+    def run(self, instance_id, uid, base_url:str, video_file:Union[UploadFile,str], audio_file:Union[UploadFile,str]) -> Union[bool, None]:
+        try:
+            code = self._req(base_url, video_file, audio_file)
+        except Exception as e:
+            loguru.logger.error("%s Task Submit Failed: %s" % (uid, e))
+            return None
+        self.tasks[uid] = {"instance_id": instance_id, "base_url":base_url, "code": code, "status": False, "start_time": time.time()}
+        return True
+
+    def result(self, uid:str) -> Union[dict, None]:
+        if uid in self.tasks:
+            try:
+                resp = requests.get(self.tasks[uid]["base_url"] + "/" + self.tasks[uid]["code"], headers=self._headers)
+                if resp.status_code != 200:
+                    raise Exception("%s Get Result Failed: %s" % (uid, resp.status_code))
+                else:
+                    if resp.json()["status"] == "waiting":
+                        if self.tasks[uid]["start_time"] + SERVER_TIME_OUT + 10 < time.time():
+                            self.tasks[uid]["status"] = True
+                            return {"status": "fail", "msg": "Instance Timeout", "instance_id": self.tasks[uid]["instance_id"]}
+                        return None
+                    else:
+                        self.tasks[uid]["status"] = True
+                        return {"status": resp.json()['status'], "msg": resp.json()['msg'],
+                                    "instance_id": self.tasks[uid]["instance_id"]}
+            except Exception as e:
+                loguru.logger.error("%s Get Result Failed: %s" % (uid, e))
+                return None
+        else:
+            raise KeyError(uid)
+
+    def get_running_size(self):
+        return len([i for i in self.tasks.values() if not i["status"]])
+
+    def _req(self, base_url, video_file, audio_file):
+        data = {
+            "video_file": video_file,
+            "audio_file": audio_file
+        }
+        resp = requests.post(base_url, data, headers=self._headers, allow_redirects=True, stream=True)
+        if resp.status_code == 200:
+            if resp.json()["status"] == "success":
+                code = resp.json()["code"]
+                return code
+            else:
+                return None
+        else:
+            return None
+
+if __name__ == "__main__":
+    rp = RunningPool()
+    id = rp.run("https://u336391-b31a-07132bc4.westx.seetacloud.com:8443",
+        "https://sucai-1324682537.cos.ap-shanghai.myqcloud.com/tiktok/video/1111.mp4",
+        "https://sucai-1324682537.cos.ap-shanghai.myqcloud.com/tiktok/video/XBR037ruAZsA.mp3")
+    while True:
+        r = rp.result(id)
+        if r:
+            loguru.logger.success(r)
+            break
+        else:
+            loguru.logger.info("waiting...")
+            time.sleep(5)
+

AutoDL/autodl_scheduling/entity/waiting_queue.py

@@ -0,0 +1,21 @@
+from queue import Queue
+
+
+class WaitingQueue:
+    def __init__(self):
+        self.queue = Queue(maxsize=500)
+
+    def enqueue(self,uid, video_path,audio_path):
+        data = {
+            "uid": uid,
+            "video_path": video_path,
+            "audio_path": audio_path
+        }
+        self.queue.put(data)
+
+    def dequeue(self):
+        data = self.queue.get()
+        return data["uid"], data["video_path"], data["audio_path"]
+
+    def get_size(self):
+        return self.queue.qsize()