jiangcang_vj/main_tr.py

import warnings
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import joblib
import gc
import pandas as pd
import numpy as np
import redis
import time
import pickle
import shutil
from datetime import datetime, timedelta
from data_loader import chunk_list, mongo_con_parse, load_train_data
from data_preprocess import preprocess_data
from config import mongodb_config, vj_flight_route_list, vj_flight_route_list_hot, vj_flight_route_list_nothot, \
    CLEAN_VJ_HOT_NEAR_INFO_TAB, CLEAN_VJ_HOT_FAR_INFO_TAB, CLEAN_VJ_NOTHOT_NEAR_INFO_TAB, CLEAN_VJ_NOTHOT_FAR_INFO_TAB

warnings.filterwarnings('ignore')


# 根据环境变量的存在设置分布式开关
if 'LOCAL_RANK' in os.environ:
    FLAG_Distributed = True
else:
    FLAG_Distributed = False


# 定义特征和参数
categorical_features = ['city_pair', 'flight_number_1', 'flight_number_2']
other_features = []
features = []

target_vars = ['target_min_to_price']   # 最低会降到的价格

# 分布式环境初始化
def init_distributed_backend():
    if FLAG_Distributed:
        local_rank = int(os.environ['LOCAL_RANK'])
        # 关键：绑定设备必须在初始化进程组之前
        torch.cuda.set_device(local_rank)            # 显式设置当前进程使用的 GPU
        try:
            dist.init_process_group(
                backend='nccl',
                init_method='env://',
                world_size=int(os.environ['WORLD_SIZE']),
                rank=int(os.environ['RANK']),
                timeout=timedelta(minutes=30)   
            )
            print(f"Process group initialized for rank {dist.get_rank()}")  # 添加日志
        except Exception as e:
            print(f"Failed to initialize process group: {e}")  # 捕获异常
            raise
        device = torch.device("cuda", local_rank)
    else:
        # 如果不在分布式环境中, 使用默认设备
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("use common environment")
    return device

# 初始化模型和相关参数
def initialize_model(device):
    return None

def continue_before_process(redis_client, lock_key):
    # rank0 跳出循环前的处理
    redis_client.set(lock_key, 2)               # 设置 Redis 锁 key 的值为 2
    print("rank0 已将 Redis 锁 key 值设置为 2")
    time.sleep(5)
    print("rank0 5秒等待结束")

def start_train():
    device = init_distributed_backend()

    model = initialize_model(device)

    if FLAG_Distributed:
        rank = dist.get_rank()
        local_rank = int(os.environ.get('LOCAL_RANK'))
        world_size = dist.get_world_size()
    else:
        rank = 0
        local_rank = 0
        world_size = 1

    output_dir = "./data_shards" 
    photo_dir = "./photo"

    date_end = datetime.today().strftime("%Y-%m-%d")
    date_begin = (datetime.today() - timedelta(days=10)).strftime("%Y-%m-%d")

    # 仅在 rank == 0 时要做的
    if rank == 0:
        # 如果处理中断, 注释掉以下代码
        batch_dir = os.path.join(output_dir, "batches")
        try:
            shutil.rmtree(batch_dir)
        except FileNotFoundError:
            print(f"rank:{rank}, {batch_dir} not found")

        # 如果处理中断, 注释掉以下代码
        csv_file_list = ['evaluate_results.csv']
        for csv_file in csv_file_list:
            try:
                csv_path = os.path.join(output_dir, csv_file)
                os.remove(csv_path)
            except Exception as e:
                print(f"remove {csv_path}: {str(e)}")

        # 确保目录存在
        os.makedirs(output_dir, exist_ok=True) 
        os.makedirs(photo_dir, exist_ok=True)

        print(f"最终特征列表：{features}")

    # 定义优化器和损失函数(只回归)
    # criterion = RegressionLoss(loss_func_flag="Quantile", quantile=0.5)
    # optimizer = torch.optim.Adam(model.parameters(), lr=0.0005, weight_decay=1e-5)

    group_size = 1
    num_epochs_per_batch = 200  # 每个批次训练的轮数，可以根据需要调整

    # 初始化 Redis 客户端（请根据实际情况修改 host、port、db）
    redis_client = redis.Redis(host='192.168.20.222', port=6379, db=0)
    lock_key = "data_loading_lock_11"
    barrier_key = 'distributed_barrier_11'

    batch_idx = -1

    # 主干代码
    flight_route_list = vj_flight_route_list_hot + vj_flight_route_list_nothot
    flight_route_list_len = len(flight_route_list)
    route_len_hot = len(vj_flight_route_list_hot)
    route_len_nothot = len(vj_flight_route_list_nothot)

    # 调试代码
    # s = 38   # 菲律宾2025-12-08是节假日 s=38 选到马尼拉 
    # flight_route_list = vj_flight_route_list_hot[:0] + vj_flight_route_list_nothot[s:]
    # flight_route_list_len = len(flight_route_list)
    # route_len_hot = len(vj_flight_route_list_hot[:0])
    # route_len_nothot = len(vj_flight_route_list_nothot[s:])
    
    if local_rank == 0:
        print(f"flight_route_list_len:{flight_route_list_len}")
        print(f"route_len_hot:{route_len_hot}")
        print(f"route_len_nothot:{route_len_nothot}")
    
    chunks = chunk_list(flight_route_list, group_size)

    for idx, group_route_list in enumerate(chunks, start=0):
        # 特殊处理，跳过不好的批次
        pass
        redis_client.set(lock_key, 0)
        redis_client.set(barrier_key, 0)
        # 所有 Rank 同步的标志变量
        valid_batch = torch.tensor([1], dtype=torch.int, device=device)  # 1表示有效批次

        # 仅在 rank == 0 时要做的
        if rank == 0:
            # Rank0 设置 Redis 锁 key 的初始值为 0，表示数据加载尚未完成
            redis_client.set(lock_key, 0)
            print("rank0 开始数据加载...")
            # 使用默认配置
            client, db = mongo_con_parse()
            print(f"第 {idx} 组 :", group_route_list)

            # 根据索引位置决定是 热门 还是 冷门
            if 0 <= idx < route_len_hot:
                is_hot = 1
                table_name = CLEAN_VJ_HOT_NEAR_INFO_TAB
            elif route_len_hot <= idx < route_len_hot + route_len_nothot:
                is_hot = 0
                table_name = CLEAN_VJ_NOTHOT_NEAR_INFO_TAB
            else:
                print(f"无法确定热门还是冷门, 跳过此批次。")
                continue_before_process(redis_client, lock_key)
                continue
            
            # 加载训练数据
            start_time = time.time()
            df_train = load_train_data(db, group_route_list, table_name, date_begin, date_end, output_dir, is_hot)
            end_time = time.time()
            run_time = round(end_time - start_time, 3)
            print(f"用时: {run_time} 秒")

            client.close()

            if df_train.empty:
                print(f"训练数据为空，跳过此批次。")
                continue_before_process(redis_client, lock_key)
                continue
            
            # 数据预处理
            df_train_inputs = preprocess_data(df_train, features, categorical_features, is_training=True)
            pass

        else:
            pass



if __name__ == "__main__":
    start_train()