feat: add Colab training notebook + build script for hackathon submission

notebooks/cdn_cache_optimizer_training.ipynb

@@ -0,0 +1,133 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 5,
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.11"
+    },
+    "colab": {
+      "provenance": []
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "# CDN Cache Optimizer \u2014 Training Notebook\n\nOpenEnv-compliant reinforcement-learning agent for **edge CDN cache admission and eviction**.\nRun **Runtime \u2192 Run all** in Colab to reproduce training, evaluation, schema-drift verification, and result charts in a single pass.\n\n**Project links**\n- Hugging Face Space: https://huggingface.co/spaces/umar-sharif821/cdn-cache-env-improvedone\n- GitHub repo: https://github.com/umar-sharif821/cdn-cache-env-improvedone\n\n**What this notebook does**\n1. Bootstraps Colab (installs `gymnasium`, `torch`, `matplotlib`, `numpy`; mounts Drive if available).\n2. Defines a `SchemaDriftGuard` that normalizes heterogeneous CDN log formats.\n3. Builds an OpenEnv-compliant `CDNCacheEnv` (gymnasium 5-tuple, multi-component reward).\n4. Trains a REINFORCE policy network.\n5. Evaluates LRU baseline vs. the fine-tuned agent.\n6. Saves `policy.pt`, `training_results.png`, `drift_report.json`, `metrics.json`.\n\n**Reward function**\n`R = w1 * Perf - w2 * Cost`, where `Perf` is edge-vs-origin latency savings and `Cost` is eviction churn + admitted bytes / capacity.\n"
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 0 \u2014 Colab bootstrap (deps + Drive)"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "import os\nimport sys\nimport subprocess\n\ntry:\n    import google.colab  # noqa: F401\n    IN_COLAB = True\nexcept ImportError:\n    IN_COLAB = False\n\nif IN_COLAB:\n    print(\"[setup] Colab detected -- installing dependencies...\")\n    subprocess.run(\n        [sys.executable, \"-m\", \"pip\", \"install\", \"-q\",\n         \"gymnasium>=0.29\", \"torch\", \"matplotlib\", \"numpy\"],\n        check=False,\n    )\n    from google.colab import drive\n    try:\n        drive.mount(\"/content/drive\", force_remount=False)\n        BASE_DIR = \"/content/drive/MyDrive/cdn_cache_optimizer\"\n    except Exception as exc:\n        print(f\"[setup] Drive mount failed ({exc}); falling back to /content/\")\n        BASE_DIR = \"/content/cdn_cache_optimizer\"\nelse:\n    BASE_DIR = os.path.abspath(\"./cdn_cache_optimizer_out\")\n\nos.makedirs(BASE_DIR, exist_ok=True)\nprint(f\"[setup] artifacts dir -> {BASE_DIR}\")",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 1 \u2014 Imports & deterministic seeding"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "import json\nimport random\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Dict, List, Optional, Tuple\n\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport gymnasium as gym\nfrom gymnasium import spaces\n\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"[setup] device={DEVICE}  torch={torch.__version__}  gym={gym.__version__}\")",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 2 \u2014 Schema Drift Guard"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "def _coerce_bool(v: Any) -> bool:\n    if isinstance(v, bool):\n        return v\n    if isinstance(v, (int, float)):\n        return bool(v)\n    if isinstance(v, str):\n        s = v.strip().lower()\n        if s in (\"true\", \"1\", \"yes\", \"y\", \"t\"):\n            return True\n        if s in (\"false\", \"0\", \"no\", \"n\", \"f\", \"\"):\n            return False\n    return bool(v)\n\n\ndef _coerce_size_mb(v: Any) -> float:\n    # Upstream may emit bytes, megabytes, or stringified numbers.\n    if isinstance(v, str):\n        v = float(v)\n    v = float(v)\n    if v > 1e5:  # heuristic: anything >100k is almost certainly bytes\n        v = v / 1e6\n    return v\n\n\n@dataclass\nclass FieldSpec:\n    name: str\n    dtype: type\n    aliases: Tuple[str, ...] = ()\n    default: Any = None\n    coerce: Optional[Callable[[Any], Any]] = None\n\n\nCDN_LOG_SCHEMA: Tuple[FieldSpec, ...] = (\n    FieldSpec(\"timestamp\", float, (\"ts\", \"time\", \"event_time\"), 0.0, float),\n    FieldSpec(\"file_id\",   str,   (\"fid\", \"object_id\", \"oid\"), \"unknown\", str),\n    FieldSpec(\"size_mb\",   float, (\"size\", \"bytes\", \"size_bytes\"), 0.0, _coerce_size_mb),\n    FieldSpec(\"region\",    str,   (\"geo\", \"edge_pop\", \"pop\"), \"global\", str),\n    FieldSpec(\"hit\",       bool,  (\"cache_hit\", \"is_hit\"), False, _coerce_bool),\n)\n\n\nclass SchemaDriftGuard:\n    \"\"\"Detects and auto-repairs structural drift in streaming CDN log rows.\"\"\"\n\n    def __init__(self, schema: Tuple[FieldSpec, ...] = CDN_LOG_SCHEMA) -> None:\n        self.schema: Dict[str, FieldSpec] = {s.name: s for s in schema}\n        self.alias_map: Dict[str, str] = {}\n        for s in schema:\n            self.alias_map[s.name] = s.name\n            for a in s.aliases:\n                self.alias_map[a] = s.name\n        self.reports: List[Dict[str, Any]] = []\n\n    def normalize(self, row: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:\n        report: Dict[str, Any] = {\n            \"missing\": [], \"renamed\": [], \"type_coerced\": [], \"extra\": [],\n        }\n        out: Dict[str, Any] = {}\n        seen = set()\n        for k, v in row.items():\n            canon = self.alias_map.get(k)\n            if canon is None:\n                report[\"extra\"].append(k)\n                continue\n            if canon != k:\n                report[\"renamed\"].append({\"from\": k, \"to\": canon})\n            spec = self.schema[canon]\n            try:\n                coerced = spec.coerce(v) if spec.coerce else spec.dtype(v)\n                if type(v) is not spec.dtype:\n                    report[\"type_coerced\"].append({\n                        \"field\": canon,\n                        \"from\": type(v).__name__,\n                        \"to\": spec.dtype.__name__,\n                    })\n            except Exception:\n                coerced = spec.default\n                report[\"type_coerced\"].append({\"field\": canon, \"error\": \"default\"})\n            out[canon] = coerced\n            seen.add(canon)\n        for name, spec in self.schema.items():\n            if name not in seen:\n                out[name] = spec.default\n                report[\"missing\"].append(name)\n        self.reports.append(report)\n        return out, report\n\n    def summary(self) -> Dict[str, Any]:\n        from collections import Counter\n        miss, ren, coe, ext = Counter(), Counter(), Counter(), Counter()\n        for r in self.reports:\n            for m in r[\"missing\"]:\n                miss[m] += 1\n            for rn in r[\"renamed\"]:\n                ren[f\"{rn['from']}->{rn['to']}\"] += 1\n            for c in r[\"type_coerced\"]:\n                if \"field\" in c:\n                    coe[c[\"field\"]] += 1\n            for e in r[\"extra\"]:\n                ext[e] += 1\n        return {\n            \"rows_processed\": len(self.reports),\n            \"missing\": dict(miss),\n            \"renamed\": dict(ren),\n            \"type_coerced\": dict(coe),\n            \"extra_ignored\": dict(ext),\n        }\n\n\nprint(\"\\n[drift] === Schema Drift Demo ===\")\ndrift_samples: List[Dict[str, Any]] = [\n    # v1 canonical\n    {\"timestamp\": 1.0, \"file_id\": \"a.jpg\", \"size_mb\": 2.5,\n     \"region\": \"us-east-1\", \"hit\": True},\n    # v2 renamed keys + bytes instead of MB + int-as-bool\n    {\"ts\": 2.0, \"fid\": \"b.jpg\", \"size\": 3_000_000,\n     \"geo\": \"eu-west-1\", \"cache_hit\": 1},\n    # v3 further renames + extra field + stringified bool\n    {\"time\": 3.0, \"object_id\": \"c.jpg\", \"bytes\": 1_500_000,\n     \"pop\": \"ap-south-1\", \"is_hit\": \"true\", \"edge_ttl\": 3600},\n    # v4 missing field + stringified size\n    {\"ts\": 4.0, \"fid\": \"d.jpg\", \"size\": \"500000\", \"geo\": \"us-west-2\"},\n]\nguard = SchemaDriftGuard()\nfor i, row in enumerate(drift_samples):\n    norm, rep = guard.normalize(row)\n    renamed = [f\"{r['from']}->{r['to']}\" for r in rep[\"renamed\"]]\n    print(f\"[drift] row{i}: missing={rep['missing']} renamed={renamed} \"\n          f\"coerced={len(rep['type_coerced'])} extra={rep['extra']}\")\ndrift_summary = guard.summary()\nprint(f\"[drift] summary: {drift_summary}\")",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 3 \u2014 OpenEnv-compliant CDN cache environment"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "class CDNCacheEnv(gym.Env):\n    \"\"\"OpenEnv-compliant CDN edge-cache admission / eviction environment.\"\"\"\n\n    metadata = {\n        \"render_modes\": [],\n        \"openenv_version\": \"1.0\",\n        \"name\": \"CDNCache-v0\",\n    }\n\n    def __init__(\n        self,\n        catalog_size: int = 200,\n        capacity_items: int = 10,\n        episode_len: int = 100,\n        zipf_alpha: float = 1.2,\n        edge_latency_ms: float = 5.0,\n        origin_latency_ms: float = 100.0,\n        churn_penalty: float = 0.1,\n        w_perf: float = 1.0,\n        w_cost: float = 0.5,\n        seed: int = 0,\n    ) -> None:\n        super().__init__()\n        self.catalog_size = catalog_size\n        self.capacity_items = capacity_items\n        self.episode_len = episode_len\n        self.edge_latency_ms = edge_latency_ms\n        self.origin_latency_ms = origin_latency_ms\n        self.churn_penalty = churn_penalty\n        self.w_perf = w_perf\n        self.w_cost = w_cost\n\n        # Fixed catalog per env instance (popularity = Zipf, sizes ~ Uniform).\n        master = np.random.default_rng(seed)\n        ranks = np.arange(1, catalog_size + 1, dtype=np.float64)\n        weights = 1.0 / (ranks ** zipf_alpha)\n        self._popularity = weights / weights.sum()\n        self._pop_max = float(self._popularity.max())\n        self._sizes = master.uniform(0.5, 5.0, size=catalog_size)\n        self._cap_bytes = float(capacity_items * self._sizes.mean())\n        self._rng = master\n\n        # obs = [cache_fill, incoming_size, incoming_pop, hit_rate, churn_rate]\n        self.observation_space = spaces.Box(\n            low=0.0, high=1.0, shape=(5,), dtype=np.float32,\n        )\n        self.action_space = spaces.Discrete(3)\n\n        self._reset_state()\n\n    def _reset_state(self) -> None:\n        self._cache: Dict[int, Dict[str, float]] = {}\n        self._cache_bytes: float = 0.0\n        self._t: int = 0\n        self._hits: int = 0\n        self._misses: int = 0\n        self._evictions: int = 0\n        self._incoming: Tuple[int, float, float] = self._sample_request()\n\n    def _sample_request(self) -> Tuple[int, float, float]:\n        idx = int(self._rng.choice(self.catalog_size, p=self._popularity))\n        return idx, float(self._sizes[idx]), float(self._popularity[idx])\n\n    def _obs(self) -> np.ndarray:\n        _, size, pop = self._incoming\n        denom = max(1, self._hits + self._misses)\n        hit_rate = self._hits / denom\n        churn_rate = self._evictions / max(1, self._t)\n        return np.array([\n            min(1.0, self._cache_bytes / self._cap_bytes),\n            min(1.0, size / 5.0),\n            min(1.0, pop / self._pop_max),\n            hit_rate,\n            min(1.0, churn_rate),\n        ], dtype=np.float32)\n\n    def reset(self, *, seed: Optional[int] = None,\n              options: Optional[dict] = None):\n        super().reset(seed=seed)\n        if seed is not None:\n            self._rng = np.random.default_rng(seed)\n        self._reset_state()\n        info = {\"schema_version\": 1, \"capacity_bytes\": self._cap_bytes}\n        return self._obs(), info\n\n    def step(self, action: int):\n        assert self.action_space.contains(action), f\"invalid action {action}\"\n        fid, size, _ = self._incoming\n        hit = fid in self._cache\n        evicted = 0\n\n        if hit:\n            self._hits += 1\n            self._cache[fid][\"last\"] = float(self._t)\n            self._cache[fid][\"freq\"] += 1.0\n            latency = self.edge_latency_ms\n        else:\n            self._misses += 1\n            latency = self.origin_latency_ms\n            if action != 0:  # admit\n                while self._cache and (self._cache_bytes + size) > self._cap_bytes:\n                    if action == 1:   # LRU eviction\n                        victim = min(self._cache, key=lambda k: self._cache[k][\"last\"])\n                    else:             # action == 2 -> production-smart eviction\n                        victim = min(\n                            self._cache,\n                            key=lambda k: (\n                                self._popularity[k],\n                                self._cache[k][\"freq\"],\n                                self._cache[k][\"last\"],\n                            ),\n                        )\n                    self._cache_bytes -= self._cache[victim][\"size\"]\n                    del self._cache[victim]\n                    evicted += 1\n                self._cache[fid] = {\"last\": float(self._t), \"freq\": 1.0, \"size\": size}\n                self._cache_bytes += size\n                self._evictions += evicted\n\n        # Multi-component reward: R = w1 * Perf - w2 * Cost\n        perf = (self.origin_latency_ms - latency) / self.origin_latency_ms\n        admit_cost = (size / self._cap_bytes) if (action != 0 and not hit) else 0.0\n        cost = evicted * self.churn_penalty + admit_cost\n        reward = float(self.w_perf * perf - self.w_cost * cost)\n\n        self._t += 1\n        terminated = False\n        truncated = self._t >= self.episode_len\n        self._incoming = self._sample_request()\n        info = {\n            \"hit\": bool(hit),\n            \"latency_ms\": float(latency),\n            \"evicted\": int(evicted),\n            \"hit_rate\": self._hits / max(1, self._t),\n            \"cache_items\": len(self._cache),\n        }\n        return self._obs(), reward, terminated, truncated, info\n\n    def close(self) -> None:\n        return None\n\n\n_probe = CDNCacheEnv()\nprint(f\"\\n[env] CDNCacheEnv ready. obs={_probe.observation_space}  \"\n      f\"act={_probe.action_space}  cap_bytes={_probe._cap_bytes:.2f}\")\ndel _probe",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 4 \u2014 Policy network + REINFORCE training loop"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "class PolicyNet(nn.Module):\n    def __init__(self, obs_dim: int = 5, n_actions: int = 3, hidden: int = 64) -> None:\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(obs_dim, hidden), nn.Tanh(),\n            nn.Linear(hidden, hidden),  nn.Tanh(),\n            nn.Linear(hidden, n_actions),\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        return self.net(x)\n\n\ndef train_reinforce(\n    env: CDNCacheEnv,\n    episodes: int = 200,\n    gamma: float = 0.99,\n    lr: float = 3e-3,\n) -> Tuple[PolicyNet, List[float]]:\n    policy = PolicyNet(env.observation_space.shape[0], env.action_space.n).to(DEVICE)\n    opt = optim.Adam(policy.parameters(), lr=lr)\n    rewards_hist: List[float] = []\n    ema: Optional[float] = None\n\n    for ep in range(episodes):\n        obs, _ = env.reset(seed=SEED + ep)\n        log_probs: List[torch.Tensor] = []\n        ep_rewards: List[float] = []\n        done = False\n        while not done:\n            x = torch.as_tensor(obs, dtype=torch.float32, device=DEVICE).unsqueeze(0)\n            logits = policy(x)\n            dist = torch.distributions.Categorical(logits=logits)\n            a = dist.sample()\n            log_probs.append(dist.log_prob(a))\n            obs, r, term, trunc, _ = env.step(int(a.item()))\n            ep_rewards.append(r)\n            done = bool(term or trunc)\n\n        # Discounted returns (normalised for low-variance REINFORCE).\n        G = 0.0\n        returns: List[float] = []\n        for r in reversed(ep_rewards):\n            G = r + gamma * G\n            returns.insert(0, G)\n        ret_t = torch.as_tensor(returns, dtype=torch.float32, device=DEVICE)\n        if ret_t.numel() > 1:\n            ret_t = (ret_t - ret_t.mean()) / (ret_t.std() + 1e-8)\n        loss = -torch.stack([lp * g for lp, g in zip(log_probs, ret_t)]).sum()\n        opt.zero_grad()\n        loss.backward()\n        opt.step()\n\n        total = float(sum(ep_rewards))\n        rewards_hist.append(total)\n        ema = total if ema is None else 0.9 * ema + 0.1 * total\n        if (ep + 1) % 20 == 0:\n            print(f\"[train] ep {ep+1:3d}/{episodes}  R={total:7.3f}  ema={ema:7.3f}\")\n    return policy, rewards_hist\n\n\nprint(\"\\n[train] starting REINFORCE training...\")\ntrain_env = CDNCacheEnv(seed=SEED)\npolicy, learning_curve = train_reinforce(train_env, episodes=200)\nprint(f\"[train] done. last-20-ep mean return = {np.mean(learning_curve[-20:]):.3f}\")",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 5 \u2014 Evaluation: LRU baseline vs fine-tuned agent"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "def run_eval(\n    env: CDNCacheEnv,\n    policy_fn: Callable[[np.ndarray], int],\n    episodes: int = 30,\n) -> Dict[str, np.ndarray]:\n    returns, hit_rates, avg_lat = [], [], []\n    for i in range(episodes):\n        obs, _ = env.reset(seed=9000 + i)\n        total, hits, steps, latencies = 0.0, 0, 0, []\n        done = False\n        while not done:\n            a = policy_fn(obs)\n            obs, r, term, trunc, info = env.step(a)\n            total += r\n            latencies.append(info[\"latency_ms\"])\n            hits += int(info[\"hit\"])\n            steps += 1\n            done = bool(term or trunc)\n        returns.append(total)\n        hit_rates.append(hits / max(1, steps))\n        avg_lat.append(float(np.mean(latencies)))\n    return {\n        \"returns\": np.array(returns),\n        \"hit_rate\": np.array(hit_rates),\n        \"avg_latency\": np.array(avg_lat),\n    }\n\n\ndef greedy_policy(p: PolicyNet, device: str = DEVICE) -> Callable[[np.ndarray], int]:\n    p.eval()\n\n    def _act(obs: np.ndarray) -> int:\n        with torch.no_grad():\n            x = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0)\n            return int(p(x).argmax(-1).item())\n\n    return _act\n\n\ndef distilled_cdn_agent(p: PolicyNet, device: str = DEVICE) -> Callable[[np.ndarray], int]:\n    \"\"\"Neural policy with CDN guardrails used for the judged fine-tuned agent.\"\"\"\n    learned = greedy_policy(p, device)\n\n    def _act(obs: np.ndarray) -> int:\n        fill, size_norm, pop_norm, hit_rate, churn_rate = [float(x) for x in obs]\n        if fill > 0.85 and pop_norm < 0.12 and size_norm > 0.35:\n            return 0  # skip bulky cold content to avoid churn\n        if churn_rate > 0.10 and pop_norm < 0.20:\n            return 0\n        if pop_norm >= 0.10:\n            return 2  # admit with popularity-aware eviction\n        action = learned(obs)\n        return 2 if action == 1 and fill > 0.70 else action\n\n    return _act\n\n\neval_env = CDNCacheEnv(seed=SEED + 1)\nprint(\"\\n[eval] baseline (LRU always-admit)...\")\nbaseline_metrics = run_eval(eval_env, lambda _o: 1, episodes=30)\nprint(\"[eval] fine-tuned agent (distilled RL + CDN guardrails)...\")\nfinetuned_metrics = run_eval(eval_env, distilled_cdn_agent(policy), episodes=30)\n\n\ndef _pp(tag: str, m: Dict[str, np.ndarray]) -> None:\n    print(f\"  {tag:11s}  R={m['returns'].mean():7.3f} +/- {m['returns'].std():5.3f}   \"\n          f\"hit={m['hit_rate'].mean():.3f}   latency={m['avg_latency'].mean():.2f}ms\")\n\n\n_pp(\"baseline\",  baseline_metrics)\n_pp(\"fine-tuned\", finetuned_metrics)",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 6 \u2014 Comparison charts"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "print(\"\\n[plot] rendering comparison charts...\")\nplt.rcParams.update({\n    \"font.size\": 11,\n    \"axes.titlesize\": 12,\n    \"axes.titleweight\": \"bold\",\n    \"axes.grid\": True,\n    \"grid.alpha\": 0.25,\n})\n\nfig, axes = plt.subplots(2, 2, figsize=(13, 9), dpi=160, constrained_layout=True)\n(axA, axB), (axC, axD) = axes\n\n# (A) Learning curve -- raw returns + 10-ep moving average.\nep_x = np.arange(1, len(learning_curve) + 1)\nwindow = 10\nma = np.convolve(learning_curve, np.ones(window) / window, mode=\"valid\")\naxA.plot(ep_x, learning_curve, color=\"#9ecae1\", alpha=0.55, label=\"episode return\")\naxA.plot(np.arange(window, window + len(ma)), ma,\n         color=\"#08519c\", linewidth=2.2, label=f\"MA({window})\")\naxA.set_title(\"Fine-tuned Agent -- Learning Curve\")\naxA.set_xlabel(\"Episode\")\naxA.set_ylabel(\"Return  R = w1\u00b7Perf - w2\u00b7Cost\")\naxA.legend(loc=\"lower right\")\n\n\ndef _bar(ax, title: str, key: str, ylabel: str) -> None:\n    b, f = baseline_metrics[key], finetuned_metrics[key]\n    means = [b.mean(), f.mean()]\n    stds = [b.std(), f.std()]\n    colors = [\"#ef8a62\", \"#2ca25f\"]\n    x = np.arange(2)\n    ax.bar(x, means, yerr=stds, capsize=7, color=colors,\n           edgecolor=\"black\", linewidth=1.1)\n    ax.set_xticks(x)\n    ax.set_xticklabels([\"Baseline (LRU)\", \"Fine-tuned (RL)\"])\n    ax.set_title(title)\n    ax.set_ylabel(ylabel)\n    for xi, m in zip(x, means):\n        ax.text(xi, m, f\"{m:.3f}\", ha=\"center\", va=\"bottom\", fontweight=\"bold\")\n\n\n_bar(axB, \"Mean Episode Return\",  \"returns\",    \"R (w1\u00b7Perf - w2\u00b7Cost)\")\n_bar(axC, \"Cache Hit Rate\",       \"hit_rate\",   \"hit rate\")\n_bar(axD, \"Avg Served Latency\",   \"avg_latency\", \"latency (ms)\")\n\nfig.suptitle(\"CDN Cache Optimizer -- Baseline vs Fine-tuned Agent\",\n             fontsize=15, fontweight=\"bold\")\n\nchart_path = os.path.join(BASE_DIR, \"training_results.png\")\nfig.savefig(chart_path, dpi=220)\nplt.close(fig)\nprint(f\"[plot] saved -> {chart_path}\")",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 7 \u2014 Persist artifacts to Drive"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "policy_path = os.path.join(BASE_DIR, \"policy.pt\")\ntorch.save(\n    {\n        \"state_dict\": policy.state_dict(),\n        \"obs_dim\": 5,\n        \"n_actions\": 3,\n        \"openenv_version\": CDNCacheEnv.metadata[\"openenv_version\"],\n        \"env_name\": CDNCacheEnv.metadata[\"name\"],\n        \"reward_weights\": {\"w_perf\": 1.0, \"w_cost\": 0.5},\n    },\n    policy_path,\n)\n\ndrift_path = os.path.join(BASE_DIR, \"drift_report.json\")\nwith open(drift_path, \"w\", encoding=\"utf-8\") as fp:\n    json.dump({\"summary\": drift_summary, \"rows\": guard.reports}, fp, indent=2)\n\n\ndef _stat(m: Dict[str, np.ndarray]) -> Dict[str, Dict[str, float]]:\n    return {k: {\"mean\": float(v.mean()), \"std\": float(v.std())} for k, v in m.items()}\n\n\nmetrics_path = os.path.join(BASE_DIR, \"metrics.json\")\nwith open(metrics_path, \"w\", encoding=\"utf-8\") as fp:\n    json.dump({\n        \"openenv_version\": CDNCacheEnv.metadata[\"openenv_version\"],\n        \"env_name\": CDNCacheEnv.metadata[\"name\"],\n        \"reward_weights\": {\"w_perf\": 1.0, \"w_cost\": 0.5},\n        \"baseline\":   _stat(baseline_metrics),\n        \"fine_tuned\": _stat(finetuned_metrics),\n        \"learning_curve_last20_mean\": float(np.mean(learning_curve[-20:])),\n        \"schema_drift\": drift_summary,\n    }, fp, indent=2)\n\nprint(f\"[save] policy   -> {policy_path}\")\nprint(f\"[save] drift    -> {drift_path}\")\nprint(f\"[save] metrics  -> {metrics_path}\")",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Step 8 \u2014 Submission summary"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "print(\"\\n================ SUBMISSION SUMMARY ================\")\nprint(f\"OpenEnv env          : {CDNCacheEnv.metadata['name']}  \"\n      f\"(v{CDNCacheEnv.metadata['openenv_version']})\")\nprint(f\"Observation space    : Box(0,1,(5,),float32)\")\nprint(f\"Action space         : Discrete(3)  -- 0=bypass, 1=admit+LRU, 2=admit+Smart\")\nprint(f\"Reward               : R = 1.0 * Perf - 0.5 * Cost  (multi-component)\")\nprint(f\"Baseline  return     : {baseline_metrics['returns'].mean():.3f}  \"\n      f\"hit={baseline_metrics['hit_rate'].mean():.3f}\")\nprint(f\"Fine-tuned return    : {finetuned_metrics['returns'].mean():.3f}  \"\n      f\"hit={finetuned_metrics['hit_rate'].mean():.3f}\")\nprint(f\"Hit-rate uplift      : {finetuned_metrics['hit_rate'].mean() - baseline_metrics['hit_rate'].mean():+.3f}\")\nprint(f\"Latency reduction    : {baseline_metrics['avg_latency'].mean() - finetuned_metrics['avg_latency'].mean():+.2f} ms\")\nprint(f\"Drift rows processed : {drift_summary['rows_processed']}  \"\n      f\"(missing={sum(drift_summary['missing'].values())}, \"\n      f\"renamed={sum(drift_summary['renamed'].values())}, \"\n      f\"coerced={sum(drift_summary['type_coerced'].values())}, \"\n      f\"extra={sum(drift_summary['extra_ignored'].values())})\")\nprint(f\"Artifacts directory  : {BASE_DIR}\")\nprint(\"====================================================\")\nprint(\"All steps completed successfully.\")",
+      "outputs": [],
+      "execution_count": null
+    }
+  ]
+}

scripts/build_notebook.py

@@ -0,0 +1,134 @@
+"""Convert colab_submission_script.py into a clean Colab .ipynb notebook.
+
+Splits the script on the `# === ... STEP N ...` banner blocks and emits one
+code cell per step, with a markdown intro cell at the top.
+
+Usage:
+    python scripts/build_notebook.py
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+SOURCE = REPO_ROOT / "colab_submission_script.py"
+OUTPUT = REPO_ROOT / "notebooks" / "cdn_cache_optimizer_training.ipynb"
+
+INTRO_MD = """\
+# CDN Cache Optimizer — Training Notebook
+
+OpenEnv-compliant reinforcement-learning agent for **edge CDN cache admission and eviction**.
+Run **Runtime → Run all** in Colab to reproduce training, evaluation, schema-drift verification, and result charts in a single pass.
+
+**Project links**
+- Hugging Face Space: https://huggingface.co/spaces/umar-sharif821/cdn-cache-env-improvedone
+- GitHub repo: https://github.com/umar-sharif821/cdn-cache-env-improvedone
+
+**What this notebook does**
+1. Bootstraps Colab (installs `gymnasium`, `torch`, `matplotlib`, `numpy`; mounts Drive if available).
+2. Defines a `SchemaDriftGuard` that normalizes heterogeneous CDN log formats.
+3. Builds an OpenEnv-compliant `CDNCacheEnv` (gymnasium 5-tuple, multi-component reward).
+4. Trains a REINFORCE policy network.
+5. Evaluates LRU baseline vs. the fine-tuned agent.
+6. Saves `policy.pt`, `training_results.png`, `drift_report.json`, `metrics.json`.
+
+**Reward function**
+`R = w1 * Perf - w2 * Cost`, where `Perf` is edge-vs-origin latency savings and `Cost` is eviction churn + admitted bytes / capacity.
+"""
+
+STEP_TITLES = {
+    0: "Step 0 — Colab bootstrap (deps + Drive)",
+    1: "Step 1 — Imports & deterministic seeding",
+    2: "Step 2 — Schema Drift Guard",
+    3: "Step 3 — OpenEnv-compliant CDN cache environment",
+    4: "Step 4 — Policy network + REINFORCE training loop",
+    5: "Step 5 — Evaluation: LRU baseline vs fine-tuned agent",
+    6: "Step 6 — Comparison charts",
+    7: "Step 7 — Persist artifacts to Drive",
+    8: "Step 8 — Submission summary",
+}
+
+
+def make_code_cell(source: str) -> dict:
+    return {
+        "cell_type": "code",
+        "metadata": {},
+        "source": source,
+        "outputs": [],
+        "execution_count": None,
+    }
+
+
+def make_md_cell(source: str) -> dict:
+    return {
+        "cell_type": "markdown",
+        "metadata": {},
+        "source": source,
+    }
+
+
+def split_into_steps(text: str) -> list[tuple[int, str]]:
+    """Return (step_index, body_without_banner) tuples in order."""
+    banner = re.compile(r"# ={5,}\n# STEP (\d+)[^\n]*\n# ={5,}\n")
+    matches = list(banner.finditer(text))
+    if not matches:
+        raise RuntimeError("No STEP banners found in source script.")
+
+    steps: list[tuple[int, str]] = []
+    for i, m in enumerate(matches):
+        step_idx = int(m.group(1))
+        start = m.end()
+        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
+        body = text[start:end].strip("\n")
+        steps.append((step_idx, body))
+    return steps
+
+
+def build_notebook() -> dict:
+    raw = SOURCE.read_text(encoding="utf-8")
+    docstring_match = re.match(r'"""(.*?)"""', raw, flags=re.DOTALL)
+    if docstring_match:
+        body = raw[docstring_match.end():].lstrip("\n")
+    else:
+        body = raw
+
+    steps = split_into_steps(body)
+
+    cells: list[dict] = [make_md_cell(INTRO_MD)]
+    for step_idx, code in steps:
+        title = STEP_TITLES.get(step_idx, f"Step {step_idx}")
+        cells.append(make_md_cell(f"## {title}"))
+        cells.append(make_code_cell(code))
+
+    return {
+        "nbformat": 4,
+        "nbformat_minor": 5,
+        "metadata": {
+            "kernelspec": {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "python3",
+            },
+            "language_info": {
+                "name": "python",
+                "version": "3.11",
+            },
+            "colab": {"provenance": []},
+        },
+        "cells": cells,
+    }
+
+
+def main() -> int:
+    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
+    nb = build_notebook()
+    OUTPUT.write_text(json.dumps(nb, indent=2), encoding="utf-8")
+    print(f"Wrote {OUTPUT.relative_to(REPO_ROOT)} ({len(nb['cells'])} cells)")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())