save code

worker.py

@@ -0,0 +1,321 @@
+import time
+import datetime
+import uuid
+import redis
+import logging
+import requests
+import threading
+import json
+from datetime import datetime, timedelta
+
+from config import REDIS_HOST
+from config import REDIS_PORT
+from config import REDIS_DB
+from config import QUEUE_NAME
+from config import GPU_SERVER_LIST
+from config import SERVER_LOG_EVENT_LEN
+from config import SERVER_LOG_
+from config import GPU_SERVER_TIME_OUT
+from config import KEY_QUEUE_LOCK_NAME, acquire_lock, SERVER_LIST_LOCK_NAME, SERVER_LIST_LOG_LOCK_NAME
+import copy
+
+import logging
+import csv
+
+# 创建 logger
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)  # 设置 logger 的级别
+
+
+# 获取当前日期和时间
+now = datetime.now()
+
+# 格式化为字符串（例如：2023-10-25 14:30:45）
+date_time_str = now.strftime("%Y-%m-%d_%H:%M:%S")
+
+# 创建文件 handler
+file_handler = logging.FileHandler(f'/var/log/fuyan/worker_{date_time_str}.log')
+file_handler.setLevel(logging.INFO)
+file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+file_handler.setFormatter(file_formatter)
+
+# 创建控制台 handler
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.INFO)
+console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+console_handler.setFormatter(console_formatter)
+
+# 添加 handlers 到 logger
+logger.addHandler(file_handler)
+logger.addHandler(console_handler)
+
+# 创建 Redis 连接池
+redis_pool = redis.ConnectionPool(
+    host=REDIS_HOST, 
+    port=REDIS_PORT, 
+    db=REDIS_DB,
+    max_connections=2000  # 根据实际情况调整
+)
+
+def get_redis_conn():
+    """获取 Redis 连接"""
+    return redis.Redis(connection_pool=redis_pool)
+
+redis_conn = get_redis_conn()
+
+def get_time(timestamp):
+    # 转换为UTC时间
+    utc_time = datetime.utcfromtimestamp(timestamp)  
+
+    # 转换为北京时间 (UTC+8)
+    beijing_time = utc_time + timedelta(hours=8)  # 正确
+
+    # 格式化输出
+    formatted_time = beijing_time.strftime("%Y-%m-%d %H:%M:%S")
+    return formatted_time
+
+def server_log_event(name, action, in_data, success=True):
+    
+    log_data = copy.deepcopy(in_data)
+    server_log_name = f"{SERVER_LOG_}{name}"
+
+    lock = acquire_lock(redis_conn, f"server_log_lock_{name}")
+    if lock:
+        try:
+            logger.info(f"server_log_event name:{name} action:{action} success:{success}")
+            server_log_str = redis_conn.get(server_log_name).decode("utf-8")
+            server_log = json.loads(server_log_str)
+            event = {}
+            t = time.time()
+            event['time'] = t
+            event['time_str'] = get_time(t)
+            event['action'] = action
+            
+            server_log['last_call'] = "success"
+            if 'state' in log_data:
+                if log_data['state'] != 0:
+                    server_log['last_call'] = "failure"
+            
+            if not success:
+                server_log['last_call'] = "failure"
+
+            event['data'] = log_data
+            if(len(server_log['events']) > SERVER_LOG_EVENT_LEN):
+                server_log['events'].pop(0)
+            server_log['events'].append(event)
+            server_log['last_event'] = event
+            server_log_str = json.dumps(server_log)
+            redis_conn.set(f"{SERVER_LOG_}{name}", server_log_str)
+
+        finally:
+            lock.release()
+    else:
+        logger.error(f"get lock none {server_log_name}")
+
+def registerGpuServer(name, url, can_use):
+    try:
+        server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8")
+        server_list = json.loads(server_list_str)
+        if name not in server_list:
+            server_list.append(name)
+        server = {}
+        server['name'] = name
+        server['url'] = url
+        server['can_use'] = can_use
+        server_str = json.dumps(server)
+        logger.info(f"registerGpuServer, {server_str}")
+        redis_conn.set(name, server_str)
+
+        server_log = {}
+        server_log['name'] = name
+        event = {}
+        t = time.time()
+        event['time'] = t 
+        event['time_str'] = get_time(t)
+        event['action'] = "register"
+        event['data'] = f"url {url}"
+        server_log['last_event'] = event
+        server_log['events'] = []
+        server_log['last_call'] = "success"
+        if(len(server_log['events']) > SERVER_LOG_EVENT_LEN):
+            server_log['events'].pop(0)
+        server_log['events'].append(event)
+        server_log_str = json.dumps(server_log)
+        redis_conn.set(f"{SERVER_LOG_}{name}", server_log_str)
+
+        redis_conn.set(GPU_SERVER_LIST, json.dumps(server_list))
+    except Exception as e:
+        logger.error(f"registerGpuServer{e.message}")
+
+
+def get_remote_gpu_server():
+    start_time = datetime.now()
+    while (datetime.now() - start_time).seconds < GPU_SERVER_TIME_OUT:
+        lock = acquire_lock(redis_conn, SERVER_LIST_LOCK_NAME)
+        try:
+            server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8")
+            server_list = json.loads(server_list_str)
+            for name in server_list:
+                server_str = redis_conn.get(name).decode("utf-8")
+                server = json.loads(server_str)
+                if server['can_use']:
+                    return server
+        finally:
+            lock.release()
+        time.sleep(0.2)
+    return None
+
+def call_remote_gpu_server(task_data_str, server=None):
+    task_data = json.loads(task_data_str)
+    key = task_data['key']
+    if server == None:
+        server = get_remote_gpu_server()
+
+    result_str = ''
+    result_status_code = '200'
+    if server:
+        headers = {
+            'Content-Type': 'application/json'
+        }
+
+        if 'do_change_cloth' in task_data['api']:
+            data = {
+                "human_url": task_data['request']['human_url'],
+                "cloth_url": task_data['request']["cloth_url"],
+                "output_format":task_data['request']['output_format']
+            }
+        
+        name = server['name']
+
+        # set server busy
+        server['can_use'] = False
+        server_str = json.dumps(server)
+        redis_conn.set(name, server_str)
+
+        server_log_event(server['name'], "Request_GpuServer", data, True)
+        url = server['url'] + task_data['api']
+        try:
+            logger.info(f"call get_remote_gpu_server {url}")
+            start_ms = int(time.time() * 1000)  # 转换为毫秒级整数 
+            base64 = False
+            if data['output_format'] == 'base64':
+                base64 = True
+
+            response = requests.post(url, headers=headers, json=data, timeout=30)
+            
+            logger.info(f"call finshed {response.status_code} {response.headers} {response.text[:128]}")
+            if response.status_code == 200:
+                result = response.json()
+                end_ms = int(time.time() * 1000)
+                if 'do_change_cloth' in task_data['api']:
+                    if base64:
+                        result['result'] = f"data:image/jpeg;base64,{result['result']}"
+
+                result['process_time_ms'] = (end_ms - start_ms)
+                server_log_event(server['name'], "GpuServer_Response", result, True)
+                result_str = json.dumps(result)
+                result_status_code = '200'
+                if 'state' in result:
+                    if result['state'] != 0:
+                        result_status_code = '500'
+            else:
+                try:
+                    result = response.json()
+                except json.JSONDecodeError:
+                    result = {
+                        "state":-1,
+                        "ret":-1,
+                        "data":"",
+                        'msg': f'status_code error{response.status_code} call_remote_gpu_server{url} {response.text}'
+                    }
+                
+                result_str = json.dumps(result)
+                result_status_code = str(response.status_code)
+                server_log_event(server['name'], f"error GpuServer_Response ", result, False)
+        except Exception as e:
+            logger.info(f"End call Exception {url} {type(e).__name__}, 错误信息: {e}")
+            result = {
+                "state":-1,
+                "ret":-1,
+                "data":"",
+                'msg': f'Exception call call_remote_gpu_server{url}'
+            }
+            result_str = json.dumps(result)
+            result_status_code = '500'
+            server_log_event(server['name'], f"error GpuServer_Response ", result, False)
+
+        logger.info(f"End call {url}")
+        # release server
+        server['can_use'] = True
+        server_str = json.dumps(server)
+        redis_conn.set(name, server_str)
+
+    else:
+        result_str = json.dumps({
+            "state":-1,
+            "ret":-1,
+            "data":"",
+            'msg': f'Timeout after all gpu server is busy'
+        })
+        result_status_code = '500'
+    
+    result_key = f"result_{key}"
+    redis_conn.set(result_key, result_str, ex=60)
+    result_status_key = f"result_status_code_{key}"
+    redis_conn.set(result_status_key, result_status_code, ex=60)
+    
+
+def get_task_queue_key():
+    redis_conn = get_redis_conn()
+    lock = acquire_lock(redis_conn, KEY_QUEUE_LOCK_NAME)
+    if not lock:
+        logger.error("get_task_queue_key get lock error")
+        return
+    try:
+        key_queue_str = redis_conn.get(QUEUE_NAME).decode("utf-8")
+        key_queue = json.loads(key_queue_str)
+        queue_len = len(key_queue)
+        if(queue_len > 0):
+            logger.info(f"get_task_queue_key queue len:{queue_len}")
+            key = key_queue.pop(0)
+            key_queue_str = json.dumps(key_queue)
+            logger.info(f"get_task_queue_key queue len:{len(key_queue)}")
+            redis_conn.set(QUEUE_NAME, key_queue_str)
+            return key
+        else:
+            return None
+    finally:
+        lock.release()
+
+def main_worker():
+    while True:
+        try:
+            key = get_task_queue_key()
+            if key:
+                task_data_str = redis_conn.get(key)
+                if task_data_str:
+                    threading.Thread(target=call_remote_gpu_server, args=(task_data_str,)).start()
+                else:
+                    logging.error(f"main_worker get task_data_str error with key{key}")
+                continue
+        except Exception as e:
+            logger.info(f"main_worker  {e}")
+        # check_server_work()
+        time.sleep(0.1)
+
+serverindex = 0
+def regServer(instance):
+    global serverindex
+    registerGpuServer(f"s{serverindex}_1_{instance}", f"http://localhost:8888", True)
+    serverindex += 1
+
+
+if __name__ == '__main__':
+    import os
+    print(f"[Worker] Starting with PID: {os.getpid()}")
+    redis_conn.set(GPU_SERVER_LIST, "[]")
+
+    servers = regServer("localhost")
+    for s in servers:
+        regServer(s)
+    main_worker()