diff --git a/.gitignore b/.gitignore index 3563961..9cc405d 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,9 @@ *.swp .vscode/ .vs/ -.idea/ \ No newline at end of file +.idea/ + +# OS / Python cache +.DS_Store +__pycache__/ +*.pyc diff --git a/GPU-Virtual-Service/gpu-remoting/dispatcher.py b/GPU-Virtual-Service/gpu-remoting/dispatcher.py index 1280450..405404b 100644 --- a/GPU-Virtual-Service/gpu-remoting/dispatcher.py +++ b/GPU-Virtual-Service/gpu-remoting/dispatcher.py @@ -1,24 +1,15 @@ import socket import threading import json -import redis from ctypes import Structure, c_char, c_int, c_size_t, sizeof from cffi import FFI import struct import logging +from runtime_config import create_redis_connection, load_runtime_config -logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') -def create_redis_connection(redis_config): - - r = redis.Redis( - host=redis_config['RedisConfig']['host'], - port=redis_config['RedisConfig']['port'], - db=redis_config['RedisConfig']['db'], - password=redis_config['RedisConfig']['password'] - ) - return r +logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') def sort_gpu_info(gpu_Info_properties,conf): if conf == 1: @@ -37,9 +28,8 @@ def sort_gpu_info(gpu_Info_properties,conf): def find_available_gpus(gpuInfoProperties, cnt): - with open('config.json', 'r') as f: - config = json.load(f) - r = create_redis_connection(config) + config = load_runtime_config() + r = create_redis_connection() avail_gpus = list() avail_gpus_properties = list() if config['DispatcherMethod']['method'] == 1: @@ -131,8 +121,7 @@ def find_available_gpus(gpuInfoProperties, cnt): def handle_message(conn, addr, message, gpuInfoProperties): if message.startswith("TypeA:"): - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() if config['DispatcherMethod']['method'] != 1: gpuInfoProperties = handle_redis() logging.info(f"Received message from {addr}: {message}") @@ -170,8 +159,7 @@ def handle_client(conn, addr, gpuInfo): def handle_redis(): gpu_info_properties = list() # gpu_properties = list() - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() conf = config['DispatcherMethod']['method'] if conf == 1:#简单轮询调度 logging.info("RoundRobin scheduling") @@ -181,7 +169,7 @@ def handle_redis(): logging.info("Utilization scheduling") else: logging.info("Free Memory and Utilization scheduling") - r = create_redis_connection(config) + r = create_redis_connection() serv_ip = config['ServerConfig']['serverIp_'] cursor = 0 while True: @@ -195,8 +183,7 @@ def handle_redis(): return gpu_info_properties def main(): - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() HOST = config['DispatcherConfig']['dpcIp_'] PORT = config['DispatcherConfig']['dpcPort_'] @@ -222,4 +209,4 @@ def main(): server_socket.close() if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/GPU-Virtual-Service/gpu-remoting/dispatcher_.py b/GPU-Virtual-Service/gpu-remoting/dispatcher_.py index 4c666fd..8a71282 100644 --- a/GPU-Virtual-Service/gpu-remoting/dispatcher_.py +++ b/GPU-Virtual-Service/gpu-remoting/dispatcher_.py @@ -1,24 +1,15 @@ import socket import threading import json -import redis from ctypes import Structure, c_char, c_int, c_size_t, sizeof from cffi import FFI import struct import logging +from runtime_config import create_redis_connection, load_runtime_config -logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') -def create_redis_connection(redis_config): - - r = redis.Redis( - host=redis_config['RedisConfig']['host'], - port=redis_config['RedisConfig']['port'], - db=redis_config['RedisConfig']['db'], - password=redis_config['RedisConfig']['password'] - ) - return r +logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') def sort_gpu_info(gpu_Info_properties,conf): if conf == 1: @@ -40,9 +31,8 @@ def sort_gpu_info(gpu_Info_properties,conf): def find_available_gpus(gpuInfoProperties, cnt, priority): - with open('config.json', 'r') as f: - config = json.load(f) - r = create_redis_connection(config) + config = load_runtime_config() + r = create_redis_connection() avail_gpus = list() avail_gpus_properties = list() if config['DispatcherMethod']['method'] == 1: @@ -138,8 +128,7 @@ def find_available_gpus(gpuInfoProperties, cnt, priority): def handle_message(conn, addr, message, gpuInfoProperties): if message.startswith("TypeA:"): - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() if config['DispatcherMethod']['method'] != 1: gpuInfoProperties = handle_redis() logging.info(f"Received message from {addr}: {message}") @@ -178,8 +167,7 @@ def handle_client(conn, addr, gpuInfo): def handle_redis(): gpu_info_properties = list() # gpu_properties = list() - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() conf = config['DispatcherMethod']['method'] if conf == 1:#简单轮询调度 logging.info("RoundRobin scheduling") @@ -189,7 +177,7 @@ def handle_redis(): logging.info("Utilization scheduling") else: logging.info("Free Memory and Utilization scheduling") - r = create_redis_connection(config) + r = create_redis_connection() serv_ip = config['ServerConfig']['serverIp_'] cursor = 0 while True: @@ -203,8 +191,7 @@ def handle_redis(): return gpu_info_properties def main(): - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() HOST = config['DispatcherConfig']['dpcIp_'] PORT = config['DispatcherConfig']['dpcPort_'] @@ -230,4 +217,4 @@ def main(): server_socket.close() if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/GPU-Virtual-Service/gpu-remoting/monitor.py b/GPU-Virtual-Service/gpu-remoting/monitor.py index dce1c99..78773ea 100644 --- a/GPU-Virtual-Service/gpu-remoting/monitor.py +++ b/GPU-Virtual-Service/gpu-remoting/monitor.py @@ -1,7 +1,6 @@ import socket import threading import json -import redis import pynvml import os import numpy as np @@ -12,22 +11,14 @@ import struct import logging +from runtime_config import create_redis_connection, load_runtime_config + # 删除 CUDA_VISIBLE_DEVICES 环境变量 if 'CUDA_VISIBLE_DEVICES' in os.environ: del os.environ['CUDA_VISIBLE_DEVICES'] logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') -def create_redis_connection(redis_config): - - r = redis.Redis( - host=redis_config['host'], - port=redis_config['port'], - db=redis_config['db'], - password=redis_config['password'] - ) - return r - # 加载 CUDA 运行时库 cuda = ctypes.CDLL("/usr/local/cuda/lib64/libcudart.so") @@ -48,8 +39,7 @@ def get_device_properties_bytes(device_id=0): #在初始化的时候调用，初始化redis里的GPU信息 def get_gpu_info(): # global r - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() logging.debug("Get GPU info") gpuInfo = list() compressed_datas = list() @@ -60,14 +50,11 @@ def get_gpu_info(): dev_cnt = pynvml.nvmlDeviceGetCount() # torch.cuda.device_count() logging.info(f"Total {dev_cnt} GPU(s) found") - with open('config.json', 'r') as f: - config = json.load(f) - IP_addr = config['ServerConfig']['serverIp_'] Port = config['ServerConfig']['serverPort_'] logging.debug(f"IP_addr: {IP_addr}, Port: {Port}") - r = create_redis_connection(config['RedisConfig']) + r = create_redis_connection() # 打印每个GPU的内存信息和利用率 for i in range(dev_cnt): @@ -117,8 +104,7 @@ def get_gpu_info(): #更新redis中GPU信息 def update_gpu_info(): # global r - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() logging.info("Updating GPU info...") gpuInfo = list() # 初始化NVML，这对于使用pynvml库是必需的 @@ -130,7 +116,7 @@ def update_gpu_info(): # IP_addr = '192.168.0.208' Ip_addr = config['ServerConfig']['serverIp_'] - r = create_redis_connection(config['RedisConfig']) + r = create_redis_connection() # 更新每个GPU的内存信息和利用率 @@ -164,9 +150,8 @@ def update_gpu_info(): def reduce_gpu_job(data): # 解析data字符串，提取GPU ID gpu_ids = data.split(':')[1].split(',') - with open('config.json', 'r') as f: - config = json.load(f) - r = create_redis_connection(config['RedisConfig']) + config = load_runtime_config() + r = create_redis_connection() Ip_addr = config['ServerConfig']['serverIp_'] for gpu_id in gpu_ids: @@ -184,9 +169,8 @@ def reduce_gpu_job(data): def handle_new_job(data): # 解析data字符串，提取GPU ID gpu_ids = data.split(':')[1].split(',') - with open('config.json', 'r') as f: - config = json.load(f) - r = create_redis_connection(config['RedisConfig']) + config = load_runtime_config() + r = create_redis_connection() Ip_addr = config['ServerConfig']['serverIp_'] for gpu_id in gpu_ids: key = f"{Ip_addr}:{gpu_id}" @@ -202,15 +186,14 @@ def handle_new_job(data): logging.info(f"GPU {gpu_id} not found") def collect_gpu_utilization(): - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() # 初始化pynvml pynvml.nvmlInit() dev_cnt = pynvml.nvmlDeviceGetCount() IP_addr = config['ServerConfig']['serverIp_'] Port = config['ServerConfig']['serverPort_'] - r = create_redis_connection(config['RedisConfig']) + r = create_redis_connection() # 创建一个字典来存储每个GPU的utilization数据 gpu_utilizations = {i: [] for i in range(dev_cnt)} @@ -311,8 +294,7 @@ def schedule_update(): def main(): #一开始先初始化redis里的服务器信息，然后启动TCP连接，监听Server请求 get_gpu_info() - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() HOST = config['MonitorConfig']['monitorIp_'] PORT = config['MonitorConfig']['monitorPort_'] @@ -349,4 +331,4 @@ def main(): if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/GPU-Virtual-Service/gpu-remoting/proxy.py b/GPU-Virtual-Service/gpu-remoting/proxy.py index ec74659..0bc0d26 100644 --- a/GPU-Virtual-Service/gpu-remoting/proxy.py +++ b/GPU-Virtual-Service/gpu-remoting/proxy.py @@ -8,6 +8,8 @@ import lz4.frame import struct +from runtime_config import load_runtime_config + logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') clientId2gpuInfos = {} @@ -215,9 +217,7 @@ def handle_client_worker(conn, addr, proxy, proxy_lock): logging.info(f"Client#{clientID} released all GPU resources") def main(): - - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() dpcIp = config['DispatcherConfig']['dpcIp_'] dpcPort = config['DispatcherConfig']['dpcPort_'] @@ -250,4 +250,4 @@ def main(): if __name__ == "__main__": main() - \ No newline at end of file + diff --git a/GPU-Virtual-Service/gpu-remoting/proxy_.py b/GPU-Virtual-Service/gpu-remoting/proxy_.py index 4f1567a..fea70ea 100644 --- a/GPU-Virtual-Service/gpu-remoting/proxy_.py +++ b/GPU-Virtual-Service/gpu-remoting/proxy_.py @@ -8,6 +8,8 @@ import lz4.frame import struct +from runtime_config import load_runtime_config + logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') clientId2gpuInfos = {} @@ -216,9 +218,7 @@ def handle_client_worker(conn, addr, proxy, proxy_lock): logging.info(f"Client#{clientID} released all GPU resources") def main(): - - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() dpcIp = config['DispatcherConfig']['dpcIp_'] dpcPort = config['DispatcherConfig']['dpcPort_'] @@ -251,4 +251,4 @@ def main(): if __name__ == "__main__": main() - \ No newline at end of file + diff --git a/GPU-Virtual-Service/gpu-remoting/proxy_msg.py b/GPU-Virtual-Service/gpu-remoting/proxy_msg.py index 4495535..1764ce8 100644 --- a/GPU-Virtual-Service/gpu-remoting/proxy_msg.py +++ b/GPU-Virtual-Service/gpu-remoting/proxy_msg.py @@ -10,6 +10,7 @@ import time from scheduler.msg_queue import * +from runtime_config import load_runtime_config logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') @@ -321,9 +322,7 @@ def handle_client_worker(conn, addr): def main(): - - with open('config.json', 'r') as f: - config = json.load(f) + config = load_runtime_config() listenIp = config['ClientConfig']['proxyIp_'] listenPort = config['ClientConfig']['proxyPort_'] @@ -347,4 +346,4 @@ def main(): if __name__ == "__main__": main() - \ No newline at end of file + diff --git a/GPU-Virtual-Service/gpu-remoting/runtime_config.py b/GPU-Virtual-Service/gpu-remoting/runtime_config.py new file mode 100644 index 0000000..8bf891e --- /dev/null +++ b/GPU-Virtual-Service/gpu-remoting/runtime_config.py @@ -0,0 +1,32 @@ +import json +from functools import lru_cache +from pathlib import Path + +import redis + + +GPU_REMOTING_ROOT = Path(__file__).resolve().parent +CONFIG_PATH = GPU_REMOTING_ROOT / "config.json" +SCHEDULER_DIR = GPU_REMOTING_ROOT / "scheduler" +JOB_INFO_PATH = SCHEDULER_DIR / "job_info.csv" + + +@lru_cache(maxsize=1) +def load_runtime_config(): + with CONFIG_PATH.open("r", encoding="utf-8") as config_file: + return json.load(config_file) + + +def get_job_info_path(): + return JOB_INFO_PATH + + +def create_redis_connection(redis_config=None, db_key="db"): + config = redis_config or load_runtime_config()["RedisConfig"] + db = config.get(db_key, config["db"]) + return redis.Redis( + host=config["host"], + port=config["port"], + db=db, + password=config["password"], + ) diff --git a/GPU-Virtual-Service/gpu-remoting/scheduler/util.py b/GPU-Virtual-Service/gpu-remoting/scheduler/util.py index 8337d2c..74922ee 100644 --- a/GPU-Virtual-Service/gpu-remoting/scheduler/util.py +++ b/GPU-Virtual-Service/gpu-remoting/scheduler/util.py @@ -16,25 +16,16 @@ from scheduler.gpu_info import * from scheduler.job import * import pulp +from runtime_config import create_redis_connection, get_job_info_path, load_runtime_config logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') def get_config_file(): - # 获取当前文件的目录 - current_dir = os.path.dirname(os.path.abspath(__file__)) - # 获取上一层目录 - parent_dir = os.path.dirname(current_dir) - # 构建 config.json 文件的路径 - config_path = os.path.join(parent_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - return config + return load_runtime_config() def query_job_info(model, batch_size): - # 读取 CSV 文件 - df = pd.read_csv('/home/zss/djh_test/NewVersion15/FlexGV_S/scheduler/job_info.csv') + df = pd.read_csv(get_job_info_path()) # 根据 model 和 batch_size 查询 result = df[(df['model'] == model) & (df['batchsize'] == batch_size)] @@ -43,26 +34,10 @@ def query_job_info(model, batch_size): def redis_connection(): - config = get_config_file() - # 使用 config 中的配置信息创建 Redis 连接 - r = redis.Redis( - host=config['RedisConfig']['host'], - port=config['RedisConfig']['port'], - db=config['RedisConfig']['db'], - password=config['RedisConfig']['password'] - ) - return r + return create_redis_connection() def redis_job_connection(): - config = get_config_file() - # 使用 config 中的配置信息创建 Redis 连接 - r = redis.Redis( - host=config['RedisConfig']['host'], - port=config['RedisConfig']['port'], - db=config['RedisConfig']['jobdb'], - password=config['RedisConfig']['password'] - ) - return r + return create_redis_connection(db_key="jobdb") def allocate_gpus(gpu_free_memory, k, m): """ @@ -170,4 +145,3 @@ def optimize_task_preemption(gpus, tasks, m, k): }) return result -