final save code
This commit is contained in:
@@ -0,0 +1,428 @@
|
||||
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/meidaojia/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"
|
||||
|
||||
if 'data' in log_data:
|
||||
imgstr = log_data['data']
|
||||
imglen = len(imgstr)
|
||||
if imglen > 256:
|
||||
log_data['data'] = imgstr[:64]
|
||||
|
||||
if 'img' in log_data:
|
||||
imgstr = log_data['img']
|
||||
imglen = len(imgstr)
|
||||
if imglen > 256:
|
||||
log_data['img'] = imgstr[:64]
|
||||
|
||||
if 'result' in log_data:
|
||||
imgstr = log_data['result']
|
||||
imglen = len(imgstr)
|
||||
if imglen > 256:
|
||||
log_data['result'] = imgstr[:64]
|
||||
|
||||
if 'user_img_path' in log_data:
|
||||
imgstr = log_data['user_img_path']
|
||||
imglen = len(imgstr)
|
||||
if imglen > 256:
|
||||
log_data['user_img_path'] = imgstr[:32]
|
||||
|
||||
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 'hairColor' in task_data['api']:
|
||||
data = {
|
||||
"img": task_data['request']['img'],
|
||||
"rgb": task_data['request']["rgb"],
|
||||
"ratio": task_data['request']['ratio'],
|
||||
"userId": task_data['request']['userId'],
|
||||
"output_format":task_data['request']['output_format']
|
||||
}
|
||||
else:
|
||||
data = {
|
||||
"hair_id": task_data['request']['hair_id'],
|
||||
"task_id": task_data['request']["task_id"],
|
||||
"user_img_path": task_data['request']['user_img_path'],
|
||||
"is_hr": task_data['request']['is_hr'],
|
||||
"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
|
||||
|
||||
# if 'hairColor' in task_data['api']:
|
||||
# if data['img'] == 'base64':
|
||||
# imgstr = redis_conn.get(f"img_{key}").decode("utf-8")
|
||||
# data['img'] = imgstr
|
||||
# else:
|
||||
# if data['user_img_path'] == 'base64':
|
||||
# imgstr = redis_conn.get(f"img_{key}").decode("utf-8")
|
||||
# data['user_img_path'] = imgstr
|
||||
|
||||
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 'hairColor' in task_data['api']:
|
||||
if base64:
|
||||
result['result'] = f"data:image/jpeg;base64,{result['result']}"
|
||||
else:
|
||||
if base64:
|
||||
result['data'] = f"data:image/jpeg;base64,{result['data']}"
|
||||
|
||||
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,
|
||||
"data":"",
|
||||
'msg': f'status_code error{response.status_code} call_remote_gpu_server{url} {response.text}'
|
||||
}
|
||||
time.sleep(1)
|
||||
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}")
|
||||
time.sleep(1)
|
||||
result = {
|
||||
"state":-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,
|
||||
"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 check_server_work_call(server):
|
||||
task_data = {}
|
||||
request = {}
|
||||
request['hair_id'] = "1907651680352395265"
|
||||
request['task_id'] = "1907651680352395265"
|
||||
request['user_img_path'] = "https://cdn.meidaojia.com/ZoeFiles/user2_1_%E5%89%AF%E6%9C%AC.JPG"
|
||||
request['output_format'] = "url"
|
||||
request['is_hr'] = "false"
|
||||
task_data['request'] = request
|
||||
key = str(uuid.uuid4())
|
||||
task_data['key'] = key
|
||||
task_data['api'] = "/api/swapHair/v1"
|
||||
task_data_str = json.dumps(task_data)
|
||||
threading.Thread(target=call_remote_gpu_server, args=(task_data_str, server)).start()
|
||||
|
||||
def check_server_work():
|
||||
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 not server['can_use']:
|
||||
continue
|
||||
|
||||
server_log_name = f"{SERVER_LOG_}{name}"
|
||||
server_log_str = redis_conn.get(server_log_name).decode("utf-8")
|
||||
server_log = json.loads(server_log_str)
|
||||
if time.time() - server_log['last_event']['time'] > (5*60):
|
||||
logger.info("check_server_work and call name")
|
||||
check_server_work_call(server)
|
||||
|
||||
|
||||
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"https://{instance}-http-8801.northwest1.gpugeek.com:8443", True)
|
||||
# registerGpuServer(f"s{serverindex}_2_{instance}", f"https://{instance}-http-8801.northwest1.gpugeek.com:8443", True)
|
||||
serverindex += 1
|
||||
|
||||
def get_server_names(csv_file='servers.csv'):
|
||||
"""
|
||||
从servers.csv文件中提取所有主机名称
|
||||
|
||||
参数:
|
||||
csv_file (str): CSV文件路径，默认为'servers.csv'
|
||||
|
||||
返回:
|
||||
list: 包含所有主机名称的列表
|
||||
"""
|
||||
server_names = []
|
||||
|
||||
try:
|
||||
with open(csv_file, mode='r', encoding='utf-8') as file:
|
||||
csv_reader = csv.reader(file)
|
||||
for row in csv_reader:
|
||||
if row: # 确保不是空行
|
||||
# 取第一列并去除前后空格
|
||||
server_name = row[0].strip()
|
||||
if server_name: # 确保名称不为空
|
||||
server_names.append(server_name)
|
||||
except FileNotFoundError:
|
||||
print(f"错误：文件 {csv_file} 未找到")
|
||||
except Exception as e:
|
||||
print(f"读取文件时出错: {e}")
|
||||
|
||||
return server_names
|
||||
|
||||
if __name__ == '__main__':
|
||||
import os
|
||||
print(f"[Worker] Starting with PID: {os.getpid()}")
|
||||
redis_conn.set(GPU_SERVER_LIST, "[]")
|
||||
|
||||
|
||||
# servers = get_server_names()
|
||||
# for s in servers:
|
||||
# regServer(s)
|
||||
regServer("693570664882181")
|
||||
regServer("693565557084165")
|
||||
|
||||
main_worker()
|
||||
Reference in New Issue
Block a user