train.py

import torch
import torch.optim as optim
import torch.nn.functional as F
import copy
import heapq
import random
import numpy as np
import os

import config
from env import GenAIServerlessEnv
from data_loader import AlibabaV2026DataLoader
from models import (
    GenAIPPOArchitecture,
    compute_lookahead,
    compute_pseudo_gradient,
    apply_pseudo_gradient,
    calculate_fedprox_penalty,
)

class EdgeClient:
    def __init__(self, client_id, num_models):
        self.client_id = client_id
        if config.NON_IID:
            self.loader = AlibabaV2026DataLoader(shard_id=client_id, num_shards=config.NUM_CLIENTS)
        else:
            self.loader = AlibabaV2026DataLoader()
        self.env = GenAIServerlessEnv(self.loader)
        self.model = GenAIPPOArchitecture(self.env.state_dim, num_models)
        self.optimizer = optim.Adam(self.model.parameters(), lr=config.LEARNING_RATE)
        self.prev_weights = None
        self.state = None

    def compute_gae(self, rewards, values, dones):
        advantages = []
        last_adv = 0
        for t in reversed(range(len(rewards))):
            next_val = 0.0 if t == len(rewards) - 1 else values[t + 1]
            delta = rewards[t] + config.GAMMA * next_val * (1 - dones[t]) - values[t]
            last_adv = delta + config.GAMMA * config.GAE_LAMBDA * (1 - dones[t]) * last_adv
            advantages.insert(0, last_adv)
        return torch.tensor(advantages, dtype=torch.float32)

    def local_train(self, global_weights):
        theta_k = copy.deepcopy(global_weights)
        if self.prev_weights is None:
            self.prev_weights = copy.deepcopy(theta_k)
        extrapolated = compute_lookahead(theta_k, self.prev_weights, config.LOOKAHEAD_ALPHA)
        self.model.load_state_dict(extrapolated)

        states, actions, rewards, values, dones, masks_list = [], [], [], [], [], []

        if self.state is None:
            self.state = self.env.reset()
        state = self.state
        for _ in range(config.ROLLOUT_STEPS):
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            macro_mask = self.env.get_macro_action_mask()

            with torch.no_grad():
                action_idx, _, value = self.model.get_action(state_tensor, macro_mask)

            next_state, reward, done, _ = self.env.step(self.env.resolve_macro_action(action_idx))

            states.append(state)
            actions.append(action_idx)
            rewards.append(reward)
            values.append(value.item())
            dones.append(done)
            masks_list.append(macro_mask)

            state = next_state
            if done:
                state = self.env.reset()

        self.state = state

        advantages = self.compute_gae(rewards, values, dones)
        returns = advantages + torch.tensor(values, dtype=torch.float32)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        actions_tensor = torch.tensor(actions, dtype=torch.long)
        masks_tensor = torch.stack(masks_list)
        state_batch = torch.FloatTensor(np.array(states))

        from torch.distributions import Categorical

        old_log_probs = None
        for epoch in range(config.LOCAL_EPOCHS):
            logits, current_values = self.model(state_batch)
            logits = logits.masked_fill(~masks_tensor, -1e9)
            dist = Categorical(logits=logits)
            new_log_probs = dist.log_prob(actions_tensor)

            if epoch == 0:
                old_log_probs = new_log_probs.detach()

            ratio = torch.exp(new_log_probs - old_log_probs)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1.0 - config.PPO_CLIP, 1.0 + config.PPO_CLIP) * advantages
            actor_loss = -torch.min(surr1, surr2).mean()

            critic_loss = torch.nn.functional.mse_loss(current_values.squeeze(-1), returns)
            entropy = dist.entropy().mean()

            total_loss = actor_loss + (0.5 * critic_loss) - (config.ENTROPY_COEF * entropy)

            if config.FEDPROX_MU > 0.0:
                total_loss = total_loss + calculate_fedprox_penalty(self.model, theta_k, config.FEDPROX_MU)

            self.optimizer.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)
            self.optimizer.step()

        trained_weights = copy.deepcopy(self.model.state_dict())
        pseudo_grad = compute_pseudo_gradient(extrapolated, trained_weights)
        self.prev_weights = theta_k
        return pseudo_grad, len(rewards)

def _greedy_hit_rate(model, loader, num_models, steps=2000):
    env = GenAIServerlessEnv(loader, num_models=num_models)
    s = env.reset()
    hits = 0
    model.eval()
    for i in range(steps):
        macro_mask = env.get_macro_action_mask()
        st = torch.FloatTensor(s).unsqueeze(0)
        with torch.no_grad():
            logits, _ = model(st)
            a = int(logits.masked_fill(~macro_mask, -1e9).argmax())
        s, _, done, info = env.step(env.resolve_macro_action(a))
        hits += info['cache_hit']
        if done:
            break
    model.train()
    return hits / (i + 1)


def dispatch_to(client, global_model, global_step):
    return {
        "weights": copy.deepcopy(global_model.state_dict()),
        "version": global_step,
    }


def main():
    os.makedirs(config.CHECKPOINT_DIR, exist_ok=True)
    seed = getattr(config, "SEED", 0)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    print(f"Seed: {seed}")
    clients = [EdgeClient(i, config.NUM_MODELS) for i in range(config.NUM_CLIENTS)]
    state_dim = clients[0].env.state_dim
    global_model = GenAIPPOArchitecture(state_dim, config.NUM_MODELS)

    for c in clients:
        c.model.load_state_dict(global_model.state_dict())
        c.optimizer = optim.Adam(c.model.parameters(), lr=config.LEARNING_RATE)

    global_step = 0
    skips = 0

    event_queue = []
    pending = {}
    sim_time = 0.0

    print("Starting Asynchronous Federated Training (AFedPG)...")
    for client in clients:
        pending[client.client_id] = dispatch_to(client, global_model, global_step)
        delay = random.uniform(config.CLIENT_MIN_DELAY, config.CLIENT_MAX_DELAY)
        heapq.heappush(event_queue, (sim_time + delay, client.client_id))

    eval_loader = AlibabaV2026DataLoader()
    best_path = os.path.join(config.CHECKPOINT_DIR, "checkpoint_best.pth")
    best_hit = _greedy_hit_rate(global_model, eval_loader, config.NUM_MODELS, steps=3000)
    torch.save(global_model.state_dict(), best_path)
    print(f"Initial (BC) hit-rate {best_hit * 100:.1f}% -> saved {best_path}")

    while global_step < config.TOTAL_GLOBAL_STEPS:
        finish_time, cid = heapq.heappop(event_queue)
        sim_time = finish_time
        client = clients[cid]

        dispatched = pending[cid]
        dispatched_weights = dispatched["weights"]
        dispatched_version = dispatched["version"]

        pseudo_grad, _ = client.local_train(dispatched_weights)

        staleness = global_step - dispatched_version
        applied = apply_pseudo_gradient(
            global_model,
            pseudo_grad,
            server_lr=config.SERVER_LR,
            staleness=staleness,
            use_staleness=config.STALENESS_ENABLED,
        )
        global_step += 1
        skips += int(not applied)

        flag = "" if applied else "  [SKIPPED non-finite gradient]"
        print(
            f"[step {global_step:>4}/{config.TOTAL_GLOBAL_STEPS}] "
            f"client={cid} staleness={staleness} sim_time={sim_time:6.1f}{flag}"
        )

        pending[cid] = dispatch_to(client, global_model, global_step)
        delay = random.uniform(config.CLIENT_MIN_DELAY, config.CLIENT_MAX_DELAY)
        heapq.heappush(event_queue, (sim_time + delay, cid))

        if global_step % config.CHECKPOINT_EVERY == 0:
            save_path = os.path.join(
                config.CHECKPOINT_DIR, f"checkpoint_step_{global_step}.pth"
            )
            torch.save(global_model.state_dict(), save_path)
            hit = _greedy_hit_rate(global_model, eval_loader, config.NUM_MODELS, steps=3000)
            if hit > best_hit:
                best_hit = hit
                torch.save(global_model.state_dict(), best_path)
                print(f"Saved Checkpoint: {save_path} | hit {hit * 100:.1f}% -> NEW BEST")
            else:
                print(f"Saved Checkpoint: {save_path} | hit {hit * 100:.1f}% (best {best_hit * 100:.1f}%)")

    final_path = os.path.join(config.CHECKPOINT_DIR, "checkpoint_final.pth")
    torch.save(global_model.state_dict(), final_path)
    print(f"\nTraining complete. Final: {final_path} | "
          f"Best (hit {best_hit * 100:.1f}%): {best_path}")
    print(f"Non-finite gradient skips: {skips}/{config.TOTAL_GLOBAL_STEPS} "
          f"({100 * skips / max(1, config.TOTAL_GLOBAL_STEPS):.1f}%)")


if __name__ == "__main__":
    main()