data_loader.py

import pandas as pd
import numpy as np
import config

class AlibabaV2026DataLoader:
    def __init__(self, request_trace_path=config.TRAIN_DATA_PATH, model_meta_path=config.MODEL_META_PATH,
                 shard_id=0, num_shards=1):
        self.request_trace_path = request_trace_path
        self.model_meta_path = model_meta_path
        self.chunk_size = config.CHUNK_SIZE
        self.shard_id = shard_id
        self.num_shards = max(1, num_shards)
        self.model_metadata = self._load_model_metadata()

    def _load_model_metadata(self):
        try:
            df_meta = pd.read_csv(self.model_meta_path)
            df_meta['normalized_size'] = df_meta['size_gb'] / config.MAX_MODEL_SIZE_GB
            return df_meta.set_index('model_id')['normalized_size'].to_dict()
        except FileNotFoundError:
            return {}

    def get_model_size(self, model_id):
        return self.model_metadata.get(model_id, 0.05) 

    def stream_requests(self):
        chunk_iterator = pd.read_csv(
            self.request_trace_path, 
            chunksize=self.chunk_size,
            usecols=['checkpoint_model_version_id', 'prompt_length', 'num_inference_steps']
        )

        for chunk_index, chunk in enumerate(chunk_iterator):
            if chunk_index < config.CHUNK_START:
                continue
            if config.CHUNK_END is not None and chunk_index >= config.CHUNK_END:
                break
            if self.num_shards > 1 and (chunk_index % self.num_shards) != self.shard_id:
                continue
            chunk = chunk.dropna()
            
            chunk['norm_prompt'] = pd.to_numeric(chunk['prompt_length'], errors='coerce') / 256.0
            chunk['norm_completion'] = pd.to_numeric(chunk['num_inference_steps'], errors='coerce') / 100.0
            
            chunk = chunk.dropna(subset=['norm_prompt', 'norm_completion'])
            
            chunk['model_idx'] = chunk['checkpoint_model_version_id'].astype('category').cat.codes

            yield chunk

    def get_queue_demand(self, current_chunk, current_index, num_models,
                         window_size=config.QUEUE_WINDOW, tau=5.0):

        demand = np.zeros(num_models, dtype=np.float32)
        end = min(current_index + window_size, len(current_chunk))
        upcoming = current_chunk['model_idx'].values[current_index:end]
        for pos, m in enumerate(upcoming):
            m = int(m)
            if 0 <= m < num_models:
                demand[m] += np.exp(-pos / tau)
        peak = demand.max()
        if peak > 0:
            demand /= peak
        return demand

    def get_queue_window(self, current_chunk, current_index, window_size=config.QUEUE_WINDOW):
        remaining = len(current_chunk) - current_index
        if remaining >= window_size:
            window = current_chunk.iloc[current_index : current_index + window_size]
        else:
            window = current_chunk.iloc[current_index:]
            
        queue_features = window[['model_idx', 'norm_prompt', 'norm_completion']].astype(np.float32).values
        
        if len(queue_features) < window_size:
            padding = np.zeros((window_size - len(queue_features), 3), dtype=np.float32)
            queue_features = np.vstack((queue_features, padding))
            
        return queue_features.flatten()