# data_collector.py """数据采集器：管理 Episode 数据记录与上上传。 纯业务逻辑，无 UI 依赖。通过回调与 UI 层通信。 """ import io import json import pickle import datetime import threading import requests from api import base_url, data_record_url, the_folder class DataCollector: """管理控制过程中的 Episode 数据采集与保存""" def __init__(self): self.episode_data_raw = [] # 所有已完成的 Episode self.current_episode = None # 当前正在记录的 Episode self.last_target_record = None self._on_log = None def set_log_callback(self, callback): """设置日志回调""" self._on_log = callback def log(self, message): if self._on_log: self._on_log(message) def _upload_to_cos(self, data_bytes: bytes, filename: str, folder: str) -> bool: """通过云函数获取直传凭证，再将数据直传到腾讯云 COS。""" try: resp = requests.post(data_record_url, json={ "type": "uploadDataFile", "fileName": filename, "folder": folder, }, timeout=30) result = resp.json() except Exception as e: self.log(f"向云函数申请凭证异常: {e}") return False if not result.get("success"): self.log(f"申请上传凭证失败: {result.get('errMsg', result)}") return False meta = result.get("uploadMetadata") if not meta or "url" not in meta or "authorization" not in meta: self.log("云端未返回有效的上传元数据") return False try: form_data = { "key": meta["cosFileId"], "Signature": meta["authorization"], "x-cos-security-token": meta["token"], "x-cos-meta-fileid": meta["fileId"], } files = {"file": (filename, io.BytesIO(data_bytes))} cos_resp = requests.post(meta["url"], data=form_data, files=files, timeout=60) if cos_resp.status_code in [200, 204]: return True else: self.log(f"COS 直传失败，状态码: {cos_resp.status_code}") return False except Exception as e: self.log(f"COS 直传异常: {e}") return False def reset(self): """重置所有采集状态（控制启动时调用）""" self.episode_data_raw = [] self.current_episode = None self.last_target_record = None def record_step(self, cycle_count: int, current_pressure: float, target_pressure: float, valve_opening: float, kp: float, ki: float, kd: float, q_in: float, v: float): """记录一个控制周期的数据点 Args: cycle_count: 控制周期计数 current_pressure: 当前压力 target_pressure: 目标压力 valve_opening: 阀门开度 kp, ki, kd: PID 参数 q_in: 流量 v: 容积 """ # 目标压力变化时自动切分 Episode if self.current_episode is None or target_pressure != self.last_target_record: if self.current_episode is not None: self.episode_data_raw.append(self.current_episode) self.log(f"Episode 结束，已记录 {len(self.current_episode['pressures'])} 个点") self.current_episode = { 'pid': [float(kp), float(ki), float(kd)], 'target_pressure': target_pressure, 'Q_in': q_in, 'V': v, 'steps': [], 'pressures': [], 'errors': [], 'valves': [] } self.last_target_record = target_pressure # 记录当前步数据 error = -(target_pressure - current_pressure) self.current_episode['steps'].append(cycle_count) self.current_episode['pressures'].append(current_pressure) self.current_episode['errors'].append(error) self.current_episode['valves'].append(float(valve_opening)) def finalize_and_upload(self, flow: float, vol: float): """停止控制时：闭合最后一个 Episode，分片上传到云存储。 单文件超过 5MB 时自动拆分为多个分片， 同时上传一个 manifest.json 记录所有分片信息。 Args: flow: 流量值 (用于文件名/路径) vol: 容积值 (用于文件名/路径) """ # 闭合最后一个 Episode if self.current_episode and len(self.current_episode['pressures']) > 0: self.episode_data_raw.append(self.current_episode) self.current_episode = None if not self.episode_data_raw: return try: timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') base_folder = f"{the_folder}/data_record/data_{flow}SLM_{vol}L" # 拆成每片尽量不超过 5MB 的 episode 分组 MAX_CHUNK_BYTES = 5 * 1024 * 1024 # 5MB chunks = [] # [(chunk_index, episodes_subset)] current_chunk = [] for ep in self.episode_data_raw: current_chunk.append(ep) if pickle.dumps(current_chunk).__len__() >= MAX_CHUNK_BYTES: # 当前片已满，回退一个 episode 后保存 current_chunk.pop() chunks.append(current_chunk) current_chunk = [ep] if current_chunk: chunks.append(current_chunk) total_chunks = len(chunks) self.log(f"控制数据共 {len(self.episode_data_raw)} 个 Episode，" f"拆为 {total_chunks} 个分片上传") def upload_all(): part_files = [] for idx, chunk_eps in enumerate(chunks): data_bytes = pickle.dumps(chunk_eps) size_kb = len(data_bytes) / 1024 part_filename = f'episode_raw_data_{timestamp}_part{idx + 1}of{total_chunks}.pkl' self.log(f" 上传分片 {idx + 1}/{total_chunks} ({size_kb:.0f} KB)...") if self._upload_to_cos(data_bytes, part_filename, base_folder): part_files.append(part_filename) else: self.log(f" 分片 {idx + 1} 上传失败") # 上传 manifest manifest = { "timestamp": timestamp, "total_chunks": total_chunks, "uploaded_chunks": len(part_files), "part_files": part_files, "total_episodes": len(self.episode_data_raw), "flow": flow, "volume": vol, } manifest_str = json.dumps(manifest, indent=2, ensure_ascii=False) manifest_bytes = manifest_str.encode('utf-8') manifest_filename = f'episode_raw_data_{timestamp}_manifest.json' self._upload_to_cos(manifest_bytes, manifest_filename, base_folder) if len(part_files) == total_chunks: self.log(f"控制数据上传成功 ({total_chunks} 个分片)") else: self.log(f"控制数据部分上传失败 ({len(part_files)}/{total_chunks})") threading.Thread(target=upload_all, daemon=True).start() except Exception as e: self.log(f"保存收集数据时发生错误: {e}") finally: self.episode_data_raw = []