agent/backend/utils.py

import torch
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from functools import partial
import pandas as pd
from typing import Dict, Union, List
import itertools
import solara

from .data import ExplorationDataset
from .models import Perceptron, NetSingleHiddenLayer
from .loss import loss_mape

def predict_dict(model, ds, inputs: Dict[str, Union[List[int], List[float]]]):
    combinations = itertools.product(*inputs.values())
    input_cols = list(inputs.keys())

    input_for_df = {col: [] for col in input_cols}
    for combination in combinations:
        for input_col, value in zip(input_cols, combination):
            input_for_df[input_col].append(value)

    input_df = pd.DataFrame(input_for_df)
    input_df_transformed = ds.input_transform(input_df)
    #print(input_df)
    #print(input_df_transformed)
    ds_new = ExplorationDataset(input_df, input_cols=ds.input_cols, 
                output_cols=[], transform_dict=ds.transform_dict)
    dl_new = DataLoader(ds_new, batch_size=len(input_df), shuffle=False)
    for x, y in dl_new:
        output = model.forward(x)
    output_for_df = {col_name: output[:,col_index].detach().numpy() for col_index, col_name in enumerate(ds.output_cols)}
    output_df = pd.DataFrame(output_for_df)
    output_df_transformed = ds.output_inv_transform(output_df)
    #print(output_df)
    #print(output_df_transformed)
    return input_df, output_df_transformed

def estimate_metrics(model, ds, cur_state):
    input_ranges = {key: [value] for key, value in cur_state.items()}
    input_df, output_df = predict_dict(model, ds, input_ranges)
    est_metrics = {metric: output_df.loc[0,metric] for metric in output_df.columns}
    return est_metrics
                
def read_metrics(df, cur_state):
    cols = list(df.columns)

    dff = df[cols]
    for col, value in cur_state.items():
        #print(f'{col} = {value}', dff.columns)
        dff = dff.query(f'{col} == {value}')
    if len(dff) == 0:
        return None
    output_df = dff.sample(1)
    for feature in ['replica','cpu','expected_tps']:
        if feature in cols:
            cols.remove(feature)
    output_df = output_df[cols].reset_index(drop=True)
    metrics = {metric: output_df.loc[0,metric] for metric in output_df.columns}
    return metrics

def train(ds: ExplorationDataset, model_name, trn_ratio, 
          batch_size_trn, batch_size_val, optimizer_name, learning_rate,
          max_epoch, loss_name, seed):
    torch.manual_seed(seed)
    input_cols, output_cols = ds.input_cols, ds.output_cols
    df = ds.df
    if model_name == "Perceptron":
        model = Perceptron(in_features=len(input_cols), out_features=len(output_cols))
    elif model_name == "NetSingleHiddenLayer":
        # TODO: make hidden_size adjustable in ui
        model = NetSingleHiddenLayer(in_features=len(input_cols), out_features=len(output_cols), hidden_size=10)

    if loss_name == "mape":
        loss_fn = loss_mape

    trn_size = int(len(ds)*trn_ratio)
    val_size = len(ds) - trn_size
    generator = torch.Generator().manual_seed(seed)
    ds_trn, ds_val = torch.utils.data.random_split(ds, [trn_size, val_size], generator=generator)
    dl_trn = DataLoader(ds_trn, batch_size=batch_size_trn, shuffle=True)
    dl_val = DataLoader(ds_val, batch_size=batch_size_val, shuffle=True)

    if optimizer_name == "Adam":
        optimizer_fn = partial(torch.optim.Adam,lr=learning_rate)
    print('backend training ...')
    print('training in progress...', len(df))
    print('data columns', list(df.columns))
    print('input columns', input_cols)
    print('output columns', output_cols)
    print('training ratio', trn_ratio)
    print('batch size trainig', batch_size_trn)
    print('batch size validation', batch_size_val)
    print(f'Number of samples {len(ds)}')
    print(f'Number of samples in training {len(ds_trn)}')
    print(f'Number of samples in validation {len(ds_val)}')
    print(f'Learning rate: {learning_rate}')
    print(f'Optimizer {optimizer_name}')
    print(f'Max epoch: {max_epoch}')
    print(f'random seed',seed)

    x, y = ds[0]
    in_features = x.shape[0]
    out_features = y.shape[0]


    optimizer = optimizer_fn(model.parameters())

    #epochbar = tqdm(range(max_epoch))
    for ep in range(max_epoch):
        model.train()
        for x, y in dl_trn:
            optimizer.zero_grad()
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            loss.backward()
            optimizer.step()

        trn_loss = evaluate(model, dl_trn, loss_fn)
        val_loss = evaluate(model, dl_val, loss_fn)
        #epochbar.set_postfix(epoch=ep+1,loss=loss.item(),val_loss=val_loss)
        yield ep, trn_loss, val_loss, model
        
    return ep, trn_loss, val_loss, model

def predict(model, df, input_cols, output_cols, trn_ratio, 
            batch_size_trn, batch_size_val, seed):
    torch.manual_seed(seed)
    ds = ExplorationDataset(df, input_cols=input_cols, output_cols=output_cols)
    trn_size = int(len(ds)*trn_ratio)
    val_size = len(ds) - trn_size
    generator = torch.Generator().manual_seed(seed)
    ds_trn, ds_val = torch.utils.data.random_split(ds, [trn_size, val_size], generator=generator)
    dl_trn = DataLoader(ds_trn, batch_size=batch_size_trn, shuffle=True)
    dl_val = DataLoader(ds_val, batch_size=batch_size_val, shuffle=True)

    trn_pred, trn_target = predict_dataloader(model, dl_trn)
    val_pred, val_target = predict_dataloader(model, dl_val)

    results = {}
    for col, col_name in enumerate(output_cols):
        trn_df = pd.DataFrame(torch.cat([trn_pred[:,[col]], trn_target[:,[col]]],dim=1))
        trn_df = trn_df.rename(columns={0:'prediction',1:'target'})
        val_df = pd.DataFrame(torch.cat([val_pred[:,[col]], val_target[:,[col]]],dim=1))
        val_df = val_df.rename(columns={0:'prediction',1:'target'})
        results[col_name] = {'training': trn_df, 'validation': val_df}
    return results


    
def predict_dataloader(model, dataloader):
    with torch.no_grad():
        predictions = torch.empty(0, model.out_features)
        targets = torch.empty(predictions.shape)
        for x, y in dataloader:
            y_pred = model.forward(x)
            predictions = torch.cat([predictions, y_pred], dim=0)
            targets = torch.cat([targets, y], dim=0)
        return predictions, targets
        

def evaluate(model, dataloader, loss_fn):
    with torch.no_grad():
        avg_loss = 0
        for x, y in dataloader:
            y_pred = model.forward(x)
            loss = loss_fn(y_pred, y)
            avg_loss += loss.item()
        avg_loss = avg_loss / len(dataloader) 
        return avg_loss
    

def update_policy(model, rewards, log_probabilities, gamma, learning_rate, optimizer):
    discounted_rewards = []

    for t in range(len(rewards)):
        gt = 0
        pw = 0
        for r in rewards[t:]:
            gt = gt + gamma ** pw * r
            pw = pw + 1
        discounted_rewards.append(gt)

    discounted_rewards = torch.tensor(discounted_rewards)
    # normalize discounted rewards
    discounted_rewards = (discounted_rewards - discounted_rewards.mean()) / (discounted_rewards.std(0) + 1e-9)

    policy_gradient = []
    for log_probability, gt in zip(log_probabilities, discounted_rewards):
        policy_gradient.append(-log_probability * gt)
        # policy_gradient.append(1.0 / log_probability * gt)

    model.optimizer.zero_grad()
    policy_gradient = torch.stack(policy_gradient).sum()
    # policy_gradient.backward()
    policy_gradient.backward(retain_graph=True)
    optimizer.step()


@solara.component
def Plot1D(x: List, y: List, title='title',xlabel='xlabel', ylabel='ylabel', force_render=0):
    options = {
        'title': {
            'text': title,
            'left': 'center'},
        'tooltip': {
            'trigger': 'axis',
            'axisPointer': {
                'type': 'cross'
            }
        },
        'xAxis': {
            'axisTick': {
                'alignWithLabel': True
            },
            'data': x,
            'name': xlabel,
            'nameLocation': 'middle',
            'nameTextStyle': {'verticalAlign': 'top','padding': [10, 0, 0, 0]}
        },
        'yAxis': [
            {
                'type': 'value',
                'name': ylabel,
                'position': 'left',
                'alignTicks': True,
                'axisLine': {
                    'show': True,
                    'lineStyle': {'color': 'green'}}
            },
        ],
        'series': [
            {
            'name': ylabel,
            'data': y,
            'type': 'line',
            'yAxisIndex': 1
            },
        ],
    }
    solara.FigureEcharts(option=options)