diff --git a/config.yaml b/config.yaml index cf6bcf8..f930980 100644 --- a/config.yaml +++ b/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 - index: testindexcode + - id: 8 + index: 12345 name: Entrance params: roi_points: diff --git a/main_start.py b/main_start.py index 034b85b..e9a8d97 100644 --- a/main_start.py +++ b/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}")