完成显卡核心分配

config.yaml

@@ -34,6 +34,7 @@ database:
 # debug_mode: True=前台运行（适合PyCharm调试），False=后台运行（适合生产部署）
 debug_mode: true
 downloader_debug_mode: false    # 下载进程调试模式
+total_devices: 1               # 显卡核心数量，用于自动分配设备
 alert_push_url: "http://123.57.151.210:10000/picenter/websocket/test/process"
 hls_root_path: "D:/ProjectDoc/Police/data/hls"
 hls_downloader_daily_rotate_hour: 3        # 凌晨轮换时间
@@ -45,10 +46,10 @@ service_groups:
   video_source_type: "hls"
   ws_host: "0.0.0.0"             # WebSocket 服务地址
   ws_port: 8765                  # WebSocket 服务端口
-  algorithm: "supervision_room"            # 算法类型
+  algorithm: "checkpoint"            # 算法类型
   cameras:                       # 该组下的摄像头列表
-  - id: 1
+  - id: 8
-    index: testindexcode
+    index: 12345
     name: Entrance
     params:
       roi_points:

main_start.py

@@ -50,6 +50,17 @@ def load_downloader_debug_mode(config_path: str = "config.yaml") -> bool:
 DEBUG_MODE = load_debug_mode()
 DOWNLOADER_DEBUG_MODE = load_downloader_debug_mode()
 
+# 设备分配计数器
+device_counter = 0
+
+
+def get_next_device(total_devices: int) -> int:
+    """获取下一个分配的设备ID"""
+    global device_counter
+    device_id = device_counter % total_devices
+    device_counter += 1
+    return device_id
+
 
 def load_service_groups_from_yaml(config_path: str = "config.yaml") -> List[dict]:
     """从 YAML 文件加载服务组配置"""
@@ -67,7 +78,8 @@ def load_hls_config(config_path: str = "config.yaml") -> dict:
             "hls_root_path": cfg.get("hls_root_path", ""),
             "daily_rotate_hour": cfg.get("hls_downloader_daily_rotate_hour", 3),
             "retention_days": cfg.get("hls_downloader_retention_days", 3),
-            "retry_interval": cfg.get("hls_downloader_retry_interval_seconds", 10)
+            "retry_interval": cfg.get("hls_downloader_retry_interval_seconds", 10),
+            "total_devices": cfg.get("total_devices", 1)
         }
     except Exception as e:
         logger.error(f"[ERROR] Failed to load HLS config: {e}")
@@ -75,7 +87,8 @@ def load_hls_config(config_path: str = "config.yaml") -> dict:
             "hls_root_path": "",
             "daily_rotate_hour": 3,
             "retention_days": 3,
-            "retry_interval": 10
+            "retry_interval": 10,
+            "total_devices": 1
         }
 
 
@@ -168,7 +181,7 @@ def get_script_path(video_source_type: str) -> str:
         return "rtsp_service_ws_kadian.py"
 
 
-def start_service_group(group: dict, hls_config: dict = None):
+def start_service_group(group: dict, hls_config: dict = None, device_id: int = None):
     """启动服务子进程（后台运行）"""
     cameras = group.get("cameras", [])
     ws_host = group.get("ws_host", "0.0.0.0")
@@ -194,13 +207,19 @@ def start_service_group(group: dict, hls_config: dict = None):
     if video_source_type == "hls" and hls_config:
         cmd.extend(["--hls-root-path", hls_config.get("hls_root_path", "")])
 
-    logger.info(f"[INFO] Starting service group '{group_name}': ws={ws_host}:{ws_port}, algorithm={algorithm}, source={video_source_type}")
+    # 设置环境变量（包括设备分配）
+    env = os.environ.copy()
+    if device_id is not None:
+        env["ASCEND_RT_VISIBLE_DEVICES"] = str(device_id)
+        logger.info(f"[INFO] Starting service group '{group_name}': ws={ws_host}:{ws_port}, algorithm={algorithm}, source={video_source_type}, device={device_id}")
+    else:
+        logger.info(f"[INFO] Starting service group '{group_name}': ws={ws_host}:{ws_port}, algorithm={algorithm}, source={video_source_type}")
 
     # DEBUG_MODE=True 时前台运行，方便调试；False 时后台运行，适合部署
     if DEBUG_MODE:
-        process = subprocess.Popen(cmd)
+        process = subprocess.Popen(cmd, env=env)
     else:
-        process = subprocess.Popen(cmd, start_new_session=True)
+        process = subprocess.Popen(cmd, env=env, start_new_session=True)
     
     return process, group_name
 
@@ -250,8 +269,9 @@ def start_service():
     # 1. 读取配置
     service_groups = load_service_groups_from_yaml(config_path)
     hls_config = load_hls_config(config_path)
+    total_devices = hls_config.get("total_devices", 1)
     logger.info(f"[INFO] Loaded {len(service_groups)} service groups from {config_path}")
-    logger.info(f"[INFO] HLS config: root_path={hls_config.get('hls_root_path')}, rotate_hour={hls_config.get('daily_rotate_hour')}")
+    logger.info(f"[INFO] HLS config: root_path={hls_config.get('hls_root_path')}, rotate_hour={hls_config.get('daily_rotate_hour')}, total_devices={total_devices}")
     
     if not service_groups:
         logger.error("[ERROR] No service groups found in config, exiting...")
@@ -274,7 +294,9 @@ def start_service():
             logger.warning(f"[WARN] Service group '{group_name}' is already running with PID {pid}")
         else:
             try:
-                process, name = start_service_group(group, hls_config)
+                # 获取下一个设备ID
+                device_id = get_next_device(total_devices) if total_devices > 1 else None
+                process, name = start_service_group(group, hls_config, device_id=device_id)
                 time.sleep(0.5)
                 save_pid(process.pid, name)
                 logger.info(f"[INFO] Service group '{name}' started with PID {process.pid}")