import warnings, os, datetime warnings.filterwarnings("ignore") import json import queue import threading import torch import numpy as np from tqdm import tqdm import pyarrow.parquet as pq from rdkit import Chem import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from model.CSU_EI import CSUEI from model.tokenizer import build_tokenizer from rdkit import RDLogger RDLogger.DisableLog('rdApp.*') def init_environment(): if "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1: dist.init_process_group(backend="nccl", timeout=datetime.timedelta(hours=2)) local_rank = int(os.environ["LOCAL_RANK"]) rank = dist.get_rank() world_size = dist.get_world_size() torch.cuda.set_device(local_rank) device = torch.device(f"cuda:{local_rank}") return True, rank, local_rank, world_size, device else: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") return False, 0, 0, 1, device def distributed_extract_features(parquet_path, config_json_path, read_chunk_size=50000, per_gpu_batch=4096, precision='float16'): is_ddp, rank, local_rank, world_size, device = init_environment() if is_ddp and rank != 0: import sys sys.stdout = open(os.devnull, 'w') sys.stderr = open(os.devnull, 'w') exp_dir = os.path.dirname(config_json_path) time_stamp = '_'.join(exp_dir.split('_')[-2:]) with open(config_json_path, 'r') as f: config_json = json.load(f) config_json["model_config"]["roberta_model_path"] = None model_path = os.path.join(exp_dir, f'checkpoints/trained_model_{time_stamp}.pth') model_dict = torch.load(model_path, map_location=device) model = CSUEI(**config_json["model_config"]) model.load_state_dict(model_dict, strict=False) model.to(device) if is_ddp: model = DDP(model, device_ids=[local_rank]) smiles_encoder = model.module.smiles_encoder else: smiles_encoder = model.smiles_encoder model.eval() tokenizer = build_tokenizer(config_json['roberta_tokenizer_path']) max_len = config_json["dataset_config"]["smiles_maxlen"] if rank == 0: print(f"Initializing Pipelined CPU-GPU Stream processing from {parquet_path}...") pf = pq.ParquetFile(parquet_path) total_len = pf.metadata.num_rows feature_dim = config_json["model_config"]["feature_dim"] parent_folder = os.path.basename(os.path.dirname(parquet_path)) output_dir = f"database/molecular_vector_lib/{parent_folder}_{time_stamp}" os.makedirs(output_dir, exist_ok=True) npy_path = os.path.join(output_dir, f"{parent_folder}_DB_smiles_feature_{precision}_{time_stamp}.bin") if rank == 0: if os.path.exists(npy_path): os.remove(npy_path) print(f"Pre-allocating {precision} memory-mapped file on disk ({total_len} rows)...") dummy = np.memmap(npy_path, dtype=precision, mode='w+', shape=(total_len, feature_dim)) dummy.flush() del dummy if is_ddp: dist.barrier() features_mmap = np.memmap(npy_path, dtype=precision, mode='r+', shape=(total_len, feature_dim)) if rank == 0: print(f"Starting pipelined streaming (Disk Chunk Size: {read_chunk_size}, GPU Batch Size: {per_gpu_batch})...") data_queue = queue.Queue(maxsize=4) def cpu_producer(): try: global_idx = 0 batch_ids, batch_masks, batch_idxs = [], [], [] for batch in pf.iter_batches(batch_size=read_chunk_size, columns=['smiles']): smiles_col = batch.column('smiles').to_pylist() for smi in smiles_col: absolute_idx = global_idx global_idx += 1 if is_ddp and (absolute_idx % world_size != rank): continue if not isinstance(smi, str): continue mol = Chem.MolFromSmiles(smi) if mol is None: continue encode_dict = tokenizer( Chem.MolToSmiles(mol), max_length=max_len, padding='max_length', truncation=True ) batch_ids.append(encode_dict['input_ids']) batch_masks.append(encode_dict['attention_mask']) batch_idxs.append(absolute_idx) if len(batch_ids) == per_gpu_batch: data_queue.put((batch_ids, batch_masks, batch_idxs)) batch_ids, batch_masks, batch_idxs = [], [], [] if batch_ids: data_queue.put((batch_ids, batch_masks, batch_idxs)) except Exception as e: print(f"[Rank {rank}] CPU Producer error: {e}") finally: data_queue.put(None) producer_thread = threading.Thread(target=cpu_producer) producer_thread.daemon = True producer_thread.start() with torch.no_grad(): pbar = tqdm(total=total_len, desc=f"Pipeline Extracting (Rank {rank})", disable=(rank != 0)) while True: item = data_queue.get() if item is None: break batch_ids, batch_masks, batch_idxs = item current_batch_size = len(batch_idxs) _process_and_write(smiles_encoder, device, batch_ids, batch_masks, batch_idxs, features_mmap, precision) pbar.update(current_batch_size * world_size if is_ddp else current_batch_size) pbar.close() producer_thread.join() if rank == 0: print("Flushing final data to disk...") features_mmap.flush() del features_mmap if is_ddp: dist.barrier() dist.destroy_process_group() if rank == 0: print(f"Inference completed. {precision.upper()} binary file saved to {npy_path}") print(f"Done! All assets saved in {output_dir}") def _process_and_write(smiles_encoder, device, input_ids_list, attn_mask_list, indices, features_mmap, precision): input_ids = torch.tensor(input_ids_list, device=device) attn_mask = torch.tensor(attn_mask_list, device=device) with torch.amp.autocast('cuda'): feat, _ = smiles_encoder([input_ids, attn_mask]) feat = feat / feat.norm(dim=-1, keepdim=True) features_mmap[np.array(indices)] = feat.cpu().numpy().astype(precision) if __name__ == "__main__": parquet_path = "database/molecular_dataset/CSU-EI-DB_100M/id2molecule.parquet" config_path = "results/CSUEI_nist23_smil6_msl6_bs1024_DDP1_lr4.0e-04_smiesa1_msesa1_rm1_se1_2026-02-13_193603/config_2026-02-13_193603.json" distributed_extract_features( parquet_path, config_path, read_chunk_size=500000, per_gpu_batch=4096, precision='float16' )