drift/reporting/wandb.py

import pandas as pd
from config.types import RawConfig
from typing import Optional
from utils.helpers import weighted_average
from training.types import Stats

def launch_wandb(project_name:str, default_config: RawConfig, sweep:bool=False) -> Optional[object]:
    from wandb_setup import get_wandb
    wandb = get_wandb()   
    if wandb is None:
        raise Exception("Wandb can not be initalized, the environment variable WANDB_API_KEY is missing (can also use .env file)")  

    elif sweep:
        wandb.init(project=project_name, config = vars(default_config))             
        return wandb
    else:
        wandb.init(project=project_name, config = vars(default_config), reinit=True)
        return wandb
    
    
def override_config_with_wandb_values(wandb: Optional[object], raw_config: RawConfig) -> RawConfig:
    if wandb is None: return raw_config

    wandb_config: dict = wandb.config
    
    config_dict = vars(raw_config)
    for k in config_dict:
        config_dict[k] = wandb_config[k]

    return RawConfig(**config_dict)

def send_report_to_wandb(stats: Stats, wandb:Optional[object]):
    if wandb is None: return

    run = wandb.run
    run.save()

    for key, value in stats.items():
        run.log({ key: value })

    run.finish()