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 utils import chunk_list_with_index, create_fixed_length_sequences from model import PriceDropClassifiTransModel from data_loader import mongo_con_parse, load_train_data from data_preprocess import preprocess_data, standardization from train import prepare_data_distribute, train_model_distribute from evaluate import printScore_cc 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_day', 'flight_number_1', 'flight_number_2'] # 这个与gid的分组条件一致 common_features = ['hours_until_departure', 'days_to_departure', 'seats_remaining', 'is_cross_country', 'is_transfer', 'fly_duration', 'stop_duration', 'flight_by_hour', 'flight_by_day', 'flight_day_of_month', 'flight_day_of_week', 'flight_day_of_quarter', 'flight_day_is_weekend', 'dep_country_is_holiday', 'arr_country_is_holiday', 'any_country_is_holiday', 'days_to_holiday', ] price_features = ['adult_total_price', 'price_change_times_total', 'price_last_change_hours'] encoded_columns = ['from_city_num', 'to_city_num', 'flight_1_num', 'flight_2_num', 'baggage_level'] features = encoded_columns + price_features + common_features target_vars = ['target_will_price_drop'] # 是否降价 # 分布式环境初始化 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): input_size = len(features) model = PriceDropClassifiTransModel(input_size, num_periods=2, hidden_size=64, num_layers=3, output_size=1, dropout=0.2) model.to(device) if FLAG_Distributed: model = DDP(model, device_ids=[device], find_unused_parameters=True) # 使用DDP包装模型 if FLAG_Distributed: print(f"Rank:{dist.get_rank()}, 模型已初始化，输入尺寸：{input_size}") else: print(f"模型已初始化，输入尺寸：{input_size}") return model 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=41)).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 = None # 后面在训练之前定义 optimizer = torch.optim.Adam(model.parameters(), lr=0.0005, weight_decay=1e-5) group_size = 1 # 每几组作为一个批次 num_epochs_per_batch = 200 # 每个批次训练的轮数，可以根据需要调整 feature_scaler = None # 初始化特征缩放器 target_scaler = None # 初始化目标缩放器 # 初始化 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' assemble_size = 1 # 几个batch作为一个集群assemble batch_idx = -1 batch_flight_routes = None # 占位, 避免其它rank找不到定义 # 主干代码 # 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}") # 如果处理中断，打开注释加载批次顺序 # with open(os.path.join(output_dir, f'order.pkl'), "rb") as f: # flight_route_list = pickle.load(f) if rank == 0: pass # 保存批次顺序, 如果处理临时中断, 将这段代码注释掉 with open(os.path.join(output_dir, f'order.pkl'), "wb") as f: pickle.dump(flight_route_list, f) chunks = chunk_list_with_index(flight_route_list, group_size) # 新增部分：计算总批次数并初始化 scaler 列表 if rank == 0: total_batches = len(chunks) feature_scaler_list = [None] * total_batches # 预分配列表空间 # target_scaler_list = [None] * total_batches # 预分配列表空间 # 中断时，打开下面注释, 临时加载一下 scaler 列表 # if rank == 0: # feature_scaler_path = os.path.join(output_dir, 'feature_scalers.joblib') # target_scaler_path = os.path.join(output_dir, 'target_scalers.joblib') # if os.path.exists(feature_scaler_path): # # 加载旧的scaler列表 # old_feature_scaler_list = joblib.load(feature_scaler_path) # # 计算旧的总批次数 # old_total_batches = len(old_feature_scaler_list) # # 只替换重叠部分 # min_batches = min(old_total_batches, total_batches) # feature_scaler_list[:min_batches] = old_feature_scaler_list[:min_batches] # if os.path.exists(target_scaler_path): # # 加载旧的scaler列表 # old_target_scaler_list = joblib.load(target_scaler_path) # # 计算旧的总批次数 # old_total_batches = len(old_target_scaler_list) # # 只替换重叠部分 # min_batches = min(old_total_batches, total_batches) # target_scaler_list[:min_batches] = old_target_scaler_list[:min_batches] # 如果临时处理中断，从日志里找到 中断的索引 修改它 resume_chunk_idx = 0 chunks = chunks[resume_chunk_idx:] if local_rank == 0: batch_starts = [start_idx for start_idx, _ in chunks] print(f"rank:{rank}, local_rank:{local_rank}, 训练阶段起始索引顺序：{batch_starts}") # 训练阶段 for i, (_, group_route_list) in enumerate(chunks, start=resume_chunk_idx): # 特殊处理，跳过不好的批次 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"第 {i} 组 :", group_route_list) batch_flight_routes = group_route_list # 根据索引位置决定是 热门 还是 冷门 if 0 <= i < route_len_hot: is_hot = 1 table_name = CLEAN_VJ_HOT_NEAR_INFO_TAB elif route_len_hot <= i < 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) print("预处理后数据样本:\n", df_train_inputs.head()) total_rows = df_train_inputs.shape[0] print(f"行数: {total_rows}") if total_rows == 0: print(f"预处理后的训练数据为空，跳过此批次。") continue_before_process(redis_client, lock_key) continue # 标准化与归一化处理 df_train_inputs, feature_scaler, target_scaler = standardization(df_train_inputs, feature_scaler=None, target_scaler=None) # 将 scaler 存入列表 batch_idx = i print("batch_idx:", batch_idx) feature_scaler_list[batch_idx] = feature_scaler # target_scaler_list[batch_idx] = target_scaler # 每个批次保存一下scaler feature_scaler_path = os.path.join(output_dir, f'feature_scalers.joblib') # target_scaler_path = os.path.join(output_dir, f'target_scalers.joblib') joblib.dump(feature_scaler_list, feature_scaler_path) # joblib.dump(target_scaler_list, target_scaler_path) # 生成序列 sequences, targets, group_ids = create_fixed_length_sequences(df_train_inputs, features, target_vars, input_length=452) # 新增有效性检查 if len(sequences) == 0 or len(targets) == 0 or len(group_ids) == 0: valid_batch[0] = 0 print("警告：当前批次数据为空，标记为无效批次") # 数据加载及预处理完成，设置 Redis 锁 key 的值为 1 redis_client.set(lock_key, 1) print("rank0 数据加载完成，已将 Redis 锁 key 值设置为 1") else: val = None # 其它 rank 等待：只有当 lock key 存在且其值为 "1" 时才算数据加载完成 print(f"rank{rank} 正在等待 rank0 完成数据加载...") while True: val = redis_client.get(lock_key) if val is not None and val.decode('utf-8') in ["1", "2"]: break time.sleep(1) if val is not None and val.decode('utf-8') == "2": print(f"rank{rank} 跳过空批次 {i}") time.sleep(3) continue print(f"rank{rank} 检测到数据加载已完成，继续后续处理...") # 同步点：所有 Rank 在此等待 if FLAG_Distributed: # 确保所有 CUDA 操作完成并释放缓存 print(f"rank{rank} ready synchronize ...") torch.cuda.synchronize() print(f"rank{rank} ready empty_cache ...") torch.cuda.empty_cache() print(f"rank{rank} ready barrier ...") dist.barrier() # 移除 device_ids 参数 # dist.barrier(device_ids=[local_rank]) print(f"rank{rank} done barrier ...") # 广播批次有效性标志 if FLAG_Distributed: dist.broadcast(valid_batch, src=0) # 所有 Rank 检查批次有效性 if valid_batch.item() == 0: print(f"Rank {rank} 跳过无效批次 {i}") continue # 所有 Rank 跳过当前循环 # 所有 Rank 同时进入数据分发 if rank == 0: # 分片并分发 my_sequences, my_targets, my_group_ids = distribute_sharded_data(sequences, targets, group_ids, world_size, rank, device, flag_distributed=FLAG_Distributed) else: # 其它 Rank 接收数据 my_sequences, my_targets, my_group_ids = distribute_sharded_data([], [], [], world_size, rank, device, flag_distributed=FLAG_Distributed) # 查看一下各rank是否分到数据 debug_print_shard_info([], my_targets, my_group_ids, rank, local_rank, world_size) pre_flag, train_single, val_single, test_single = prepare_data_distribute(my_sequences, my_targets, my_group_ids, flag_distributed=FLAG_Distributed, rank=rank, local_rank=local_rank, world_size=world_size) del my_sequences del my_targets del my_group_ids gc.collect() if not pre_flag: print(f"Rank {rank} 跳过无效数据批次 {i}") continue train_sequences = train_single['sequences'] train_targets = train_single['targets'] train_group_ids = train_single['group_ids'] val_sequences = val_single['sequences'] val_targets = val_single['targets'] val_group_ids = val_single['group_ids'] # test_sequences = test_single['sequences'] # test_targets = test_single['targets'] # test_group_ids = test_single['group_ids'] if FLAG_Distributed: dist.barrier() # 训练模型 model = train_model_distribute(train_sequences, train_targets, train_group_ids, val_sequences, val_targets, val_group_ids, model, criterion, optimizer, device, num_epochs=num_epochs_per_batch, batch_size=16, target_scaler=target_scaler, flag_distributed=FLAG_Distributed, rank=rank, local_rank=local_rank, world_size=world_size, output_dir=output_dir, photo_dir=photo_dir, batch_idx=batch_idx, batch_flight_routes=batch_flight_routes, patience=40, delta=0.001) del train_single del val_single del test_single gc.collect() # 重置模型参数 if (i + 1) % assemble_size == 0: if FLAG_Distributed: dist.barrier() del model, optimizer torch.cuda.empty_cache() # 清理GPU缓存 model = initialize_model(device) # 重置模型 optimizer = torch.optim.Adam(model.parameters(), lr=0.0005, weight_decay=1e-5) # 重置优化器 print(f"Rank {rank}, Reset Model at batch {i} due to performance drop") ############################################################################################################### # 在整体批次训练结束后 if rank == 0: # pass # torch.save(model.state_dict(), os.path.join(output_dir, 'final_model.pth')) print("模型训练完成并已保存。") csv_file = 'evaluate_results.csv' csv_path = os.path.join(output_dir, csv_file) # 汇总评估结果 try: df = pd.read_csv(csv_path) except Exception as e: print(f"read {csv_path} error: {str(e)}") df = None if df is not None: # 提取真实值和预测值 y_trues_class_labels = df['Actual_Will_Price_Drop'] y_preds_class_labels = df['Predicted_Will_Price_Drop'] printScore_cc(y_trues_class_labels, y_preds_class_labels, batch_fn_str='validate', batch_idx='') if FLAG_Distributed: dist.destroy_process_group() # 显式调用 destroy_process_group 来清理 NCCL 的进程组资源 def distribute_sharded_data(sequences, targets, group_ids, world_size, rank, device, flag_distributed): # --- 非分布式模式：直接返回全量数据 if not flag_distributed: return sequences, targets, group_ids # ================== 第一阶段：元数据广播 ================== if rank == 0: # 将 group_ids 序列化为字节流 group_bytes = pickle.dumps(group_ids) # 转换为张量用于分块传输 group_tensor = torch.frombuffer(bytearray(group_bytes), dtype=torch.uint8).to(device) # 处理其他数据 seq_tensor = torch.stack(sequences, dim=0).to(device) # shape [N, 2, 452, 25] tgt_tensor = torch.stack(targets, dim=0).to(device) # shape [N, 1] meta_data = { # sequences/targets 元数据 'seq_shape': seq_tensor.shape, 'tgt_shape': tgt_tensor.shape, 'seq_dtype': str(seq_tensor.dtype).replace('torch.', ''), # 关键修改点 'tgt_dtype': str(tgt_tensor.dtype).replace('torch.', ''), # group_ids 元数据 'group_shape': group_tensor.shape, 'group_bytes_len': len(group_bytes), 'pickle_protocol': pickle.HIGHEST_PROTOCOL } else: meta_data = None # 广播元数据（所有rank都需要） meta_data = broadcast(meta_data, src=0, rank=rank, device=device) # ================== 第二阶段：分块传输 ================== # 初始化接收缓冲区（所有Rank） if rank == 0: group_tensor = group_tensor seq_tensor = seq_tensor tgt_tensor = tgt_tensor else: seq_dtype = getattr(torch, meta_data['seq_dtype']) # 例如 meta_data['seq_dtype'] = "float32" tgt_dtype = getattr(torch, meta_data['tgt_dtype']) group_tensor = torch.zeros(meta_data['group_shape'], dtype=torch.uint8, device=device) seq_tensor = torch.zeros(meta_data['seq_shape'], dtype=seq_dtype, device=device) tgt_tensor = torch.zeros(meta_data['tgt_shape'], dtype=tgt_dtype, device=device) # 并行传输所有数据（按传输量排序：先大后小） _chunked_broadcast(seq_tensor, src=0, rank=rank) # 最大数据优先 _chunked_broadcast(tgt_tensor, src=0, rank=rank) _chunked_broadcast(group_tensor, src=0, rank=rank) # 最后传输group_ids # ================== 第三阶段：数据重建 ================== # 重建 sequences 和 targets sequences_list = [seq.cpu().clone() for seq in seq_tensor] # 自动按第0维切分 targets_list = [tgt.cpu().clone() for tgt in tgt_tensor] # 重建 group_ids（关键步骤） if rank == 0: # Rank0直接使用原始数据避免重复序列化 group_ids_rebuilt = group_ids else: # 1. 提取有效字节（去除填充) group_bytes = bytes(group_tensor.cpu().numpy().tobytes()[:meta_data['group_bytes_len']]) # 2. 反序列化 try: group_ids_rebuilt = pickle.loads(group_bytes) except pickle.UnpicklingError as e: raise RuntimeError(f"反序列化 group_ids 失败: {str(e)}") # 3. 结构校验 _validate_group_structure(group_ids_rebuilt) return sequences_list, targets_list, group_ids_rebuilt def broadcast(data, src, rank, device): """安全地广播任意数据，确保张量在正确的设备上""" if rank == src: # 序列化数据 data_bytes = pickle.dumps(data) data_size = torch.tensor([len(data_bytes)], dtype=torch.long, device=device) # 创建数据张量并移动到设备 data_tensor = torch.frombuffer(bytearray(data_bytes), dtype=torch.uint8).to(device) # 先广播数据大小 dist.broadcast(data_size, src=src) # 然后广播数据 dist.broadcast(data_tensor, src=src) return data else: # 接收数据大小 data_size = torch.tensor([0], dtype=torch.long, device=device) dist.broadcast(data_size, src=src) # 分配数据张量 data_tensor = torch.empty(data_size.item(), dtype=torch.uint8, device=device) dist.broadcast(data_tensor, src=src) # 反序列化 data = pickle.loads(data_tensor.cpu().numpy().tobytes()) return data def _chunked_broadcast(tensor, src, rank, chunk_size=1024*1024*128): # chunk_size 单位是字节 """分块广播张量优化通信效率""" # Step 1. 准备连续内存缓冲 buffer = tensor.detach().contiguous() # Step 2. 计算字节数 element_size = buffer.element_size() # 每个元素的字节数（如 float32 是 4） total_elements = buffer.numel() # 计算每个块最多包含多少元素（根据字节数换算） elements_per_chunk = chunk_size // element_size # 分块数量 num_chunks = (total_elements + elements_per_chunk - 1) // elements_per_chunk # Step 4. 逐块广播 for chunk_idx in range(num_chunks): # 计算当前块的字节范围 start_element = chunk_idx * elements_per_chunk end_element = min((chunk_idx+1)*elements_per_chunk, total_elements) # Step 5. 从大张量中切出当前块 chunk = buffer.view(-1).narrow(0, start_element, end_element - start_element) # Step 6. 执行广播 dist.broadcast(chunk, src=src) # 说明: 虽然单个chunk是一维的, 但通过其内部的 1.严格的传输顺序 2.接收端的内存预分配 3.最终reshape操作 原始张量的形状得以完美恢复 def _validate_group_structure(group_ids): """校验 group_ids 数据结构完整性""" assert isinstance(group_ids, list), "Group IDs 必须是列表" if len(group_ids) == 0: print("还原的 group_ids 长度为0") return sample = group_ids[0] assert isinstance(sample, tuple), "元素必须是元组" assert len(sample) == 9, "元组长度必须为9" def debug_print_shard_info(sequences, targets, group_ids, rank, local_rank, world_size): """分布式环境下按Rank顺序打印分片前5条样本""" # 同步所有进程 if FLAG_Distributed: dist.barrier(device_ids=[local_rank]) # 按Rank顺序逐个打印（避免输出混杂） for r in range(world_size): if r == rank: print(f"\n=== Rank {rank}/{world_size} Data Shard Samples (showing first 5) ===") # 打印序列数据 # print("[Sequences]") # for i, seq in enumerate(sequences[:5]): # print(f"Sample {i}: {seq[:3]}...") # 只显示前3元素示意 # 打印目标数据 print("\n[Targets]") print(targets[:5]) # 打印Group ID分布 print("\n[Group IDs]") # unique_gids = list(set(group_ids[:50])) # 检查前50条的group分布 print(f"First 5 GIDs: {group_ids[:5]}") # sys.stdout.flush() # 确保立即输出 if FLAG_Distributed: dist.barrier(device_ids=[local_rank]) # 等待当前Rank打印完成 if __name__ == "__main__": start_train()