DatangDataPy/scada/scada_collation.py

import logging
import os
import shutil
import time
from datetime import datetime
from multiprocessing import Pool  # 新增
from pathlib import Path

import pandas as pd

from labeler import apply_labels

FILE_MODIFY_MIN_SECONDS = 60
# ROLLING_DAYS_30 = 30
# ROLLING_DAYS_90 = 90

# ====================== 配置项 ======================
# base_dir = "/data/wind-turbine"
base_dir = "/data/wind_files/headquarter/scada"
BASE_TMP_DIR = Path(f"{base_dir}/tmp")
BASE_LABEL_DIR = Path(f"{base_dir}/label")

# BASE_ROLLING_DAYS_DIR = Path(f"{base_dir}/{ROLLING_DAYS}days")

ERROR_DATA_DIR = Path(f"{base_dir}/error_data")

LOG_DIR = Path("/data/logs/scada_data_py")
LOG_FILE = LOG_DIR / f"process_{datetime.now().strftime('%Y%m%d')}.log"
LOG_LEVEL = logging.INFO

# 多进程配置（根据服务器CPU核心数调整）
MAX_WORKERS = 8


# ====================== 算法 ======================
def wind_data_labeling(df: pd.DataFrame, model_name: str = '') -> pd.DataFrame:
    df = apply_labels(df, model_name)
    return df


# ====================== 日志 ======================
def init_logger():
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    logging.basicConfig(
        level=LOG_LEVEL,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[
            logging.FileHandler(LOG_FILE, encoding="utf-8"),
            logging.StreamHandler()
        ]
    )
    return logging.getLogger(__name__)


logger = init_logger()


# ====================== 工具方法 ======================
def is_file_valid(file_path: Path) -> bool:
    try:
        pd.read_parquet(file_path)
        return True
    except Exception as e:
        logger.error(f"文件损坏，跳过：{file_path}，异常：{str(e)}")
        return False


def get_target_files(base_dir: Path) -> list:
    target_files = []
    now = time.time()
    for root, _, files in os.walk(base_dir):
        for file in files:
            if not file.endswith(".parquet"):
                continue
            file_path = Path(root) / file
            if now - file_path.stat().st_mtime <= FILE_MODIFY_MIN_SECONDS:
                continue
            target_files.append(file_path)
    return target_files


from functools import wraps


def retry_on_exception(max_retries=3, delay=1, backoff=2):
    """重试装饰器"""

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            current_delay = delay
            while retries < max_retries:
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    retries += 1
                    if retries == max_retries:
                        raise
                    logger.warning(f"操作失败，{retries}/{max_retries} 次重试，错误: {e}")
                    time.sleep(current_delay)
                    current_delay *= backoff
            return None

        return wrapper

    return decorator


@retry_on_exception(max_retries=3, delay=0.5)
def safe_read_parquet(file_path):
    """带重试机制的 Parquet 读取"""
    return pd.read_parquet(file_path)


def merge_rolling_days_data(model: str, wind_farm_code: str, wind_turbine_code: str,
                            df_oneday_label: pd.DataFrame, ROLLING_DAYS: int) -> None:
    BASE_ROLLING_DAYS_DIR = Path(f"{base_dir}/{ROLLING_DAYS}days")
    output_path = BASE_ROLLING_DAYS_DIR / model / wind_farm_code / f"{wind_turbine_code}.parquet"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    df_list = []
    if output_path.exists():
        df_list.append(safe_read_parquet(output_path))
    df_list.append(df_oneday_label)

    merged_df = pd.concat(df_list, ignore_index=True)
    merged_df['localtime'] = pd.to_datetime(merged_df['localtime'], errors='coerce')
    merged_df.sort_values(by=['localtime'], inplace=True)

    # 按自然日 00:00 保留最近30天
    latest_time = merged_df['localtime'].max()
    latest_day = latest_time.floor('D')
    thirty_days_ago = latest_day - pd.Timedelta(days=ROLLING_DAYS - 1)
    merged_df = merged_df[merged_df['localtime'] > thirty_days_ago]

    merged_df.to_parquet(output_path, engine="pyarrow", index=False)
    logger.info(f"{ROLLING_DAYS}天滚动文件生成成功：{output_path}")


# ====================== 单个文件处理逻辑（抽成独立函数）======================
def process_single_file(tmp_file):
    # 20251020/CCWE1500-82.DF/EmlMGcty/zKCujSuK.parquet
    parts = tmp_file.parts
    data_date = parts[-4]
    model = parts[-3]
    wind_farm_code = parts[-2]
    wind_turbine_code = tmp_file.stem

    try:
        if not is_file_valid(tmp_file):
            target_path = ERROR_DATA_DIR / data_date / model / wind_farm_code / f"{wind_turbine_code}.parquet"
            target_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(tmp_file, target_path)
            return "fail"

        # 读取 + 打标签
        df = safe_read_parquet(tmp_file)
        for col in df.columns:
            if col != "localtime":
                df[col] = pd.to_numeric(df[col], errors="coerce")

        # 补全数据,查询前一天最后一条数据,向后补全

        df_oneday_label = wind_data_labeling(df, model)

        # 保存label
        label_file = BASE_LABEL_DIR / data_date / model / wind_farm_code / f"{wind_turbine_code}.parquet"
        label_file.parent.mkdir(parents=True, exist_ok=True)
        df_oneday_label.to_parquet(label_file, engine="pyarrow", index=False)
        logger.info(f"{wind_turbine_code} 打标签完成：{label_file}")

        # 30天数据合并
        merge_rolling_days_data(model, wind_farm_code, wind_turbine_code, df_oneday_label, 30)
        # 90天数据合并
        merge_rolling_days_data(model, wind_farm_code, wind_turbine_code, df_oneday_label, 90)
        # 删除临时文件
        tmp_file.unlink()
        logger.info(f"临时文件已删除：{tmp_file}")
        return "success"

    except Exception as e:
        logger.error(f"处理失败：{tmp_file}，异常：{str(e)}", exc_info=True)
        target_path = ERROR_DATA_DIR / data_date / model / wind_farm_code / f"{wind_turbine_code}.parquet"
        target_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(tmp_file, target_path)
        return "fail"


# ====================== 主流程（多进程）======================
def process_wind_data():
    logger.info(f"临时目录：{BASE_TMP_DIR}")
    tmp_files = get_target_files(BASE_TMP_DIR)

    if not tmp_files:
        logger.info("无待处理文件，任务结束")
        return

    logger.info(f"待处理文件总数：{len(tmp_files)}")
    logger.info(f"启动多进程，进程数：{MAX_WORKERS}")

    # ========== 多进程核心 ==========
    with Pool(processes=MAX_WORKERS) as pool:
        results = pool.map(process_single_file, tmp_files)

    success_count = sum(1 for res in results if res == "success")
    fail_count = sum(1 for res in results if res == "fail")

    logger.info(f"任务完成 | 成功：{success_count} | 失败：{fail_count}")
    logger.info("=" * 50)


if __name__ == "__main__":
    process_wind_data()