model-registry-201/Logging_Models_by_Reference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/wandb/edu/blob/main/model-registry-201/Logging_Models_by_Reference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "<!--- @wandbcode{modelreg201-reference} -->"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Logging models by Reference\n",
    "<!--- @wandbcode{wandb201-reference} -->\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qqq wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "wandb.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import string\n",
    "from google.cloud import storage\n",
    "\n",
    "def generate_random_file(filename, size_in_kb=1):\n",
    "    \"\"\"Generate a random file with the given size in KB.\"\"\"\n",
    "    chars = ''.join([random.choice(string.ascii_letters + string.digits) for _ in range(1024)])  # 1 KB\n",
    "    with open(filename, 'w') as f:\n",
    "        for _ in range(size_in_kb):\n",
    "            f.write(chars)\n",
    "\n",
    "random_file_name = \"my_model.ckpt\"\n",
    "generate_random_file(random_file_name, 20)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Log References to Existing Object Stores\n",
    "\n",
    "- If your team manages a dedicated bucket (e.g. s3 or GCS) where all checkpoints are stored, you can have W&B log references to this bucket\n",
    "- W&B does not copy the checkpoints, we just manage a reference\n",
    "- Get all the benefits of an enterprise model registry on top of your existing buckets!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!export GOOGLE_APPLICATION_CREDENTIALS=\"/path/to/my_service_account.json\"\n",
    "!export GOOGLE_CLOUD_PROJECT=\"my_gcp_project\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "\n",
    "wandb.init(project=\"model-registry-walkthrough\", entity='smle-machine', job_type=\"log_gcs_reference\")\n",
    "\n",
    "storage_client = storage.Client()\n",
    "\n",
    "bucket = storage_client.bucket(\"my-checkpoint-demo-bucket\")\n",
    "blob = bucket.blob(\"my_model.ckpt\")\n",
    "blob.upload_from_filename(\"my_model.ckpt\")\n",
    "\n",
    "ref_art = wandb.Artifact(name='my_reference_checkpoint', type='model')\n",
    "ref_art.add_reference(\"gs://my-checkpoint-demo-bucket/my_model.ckpt\")\n",
    "wandb.log_artifact(ref_art)\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To consume the checkpoints you must have the right GCS/AWS/Azure credentials set\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}