docs/redis_model_store.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cbba56a9",
   "metadata": {},
   "source": [
    "![Redis](https://redis.io/wp-content/uploads/2024/04/Logotype.svg?auto=webp&quality=85,75&width=120)\n",
    "\n",
    "# ML Model Serialization in Redis\n",
    "\n",
    "The `ModelStore` class below implements the following logic:\n",
    "- Builds a model metadata index for model version management\n",
    "- Handles model chunking, serialization, and deserialization to/from Redis using Pickle\n",
    "\n",
    "Below we test with various Python ML-native data types and models.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0c3ead6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install scikit-learn torch tensorflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "307c25ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import redis\n",
    "\n",
    "from redis_model_store import ModelStore\n",
    "\n",
    "# Replace values below with your own if using Redis Cloud instance\n",
    "REDIS_HOST = os.getenv(\"REDIS_HOST\", \"localhost\") # ex: \"redis-18374.c253.us-central1-1.gce.cloud.redislabs.com\"\n",
    "REDIS_PORT = os.getenv(\"REDIS_PORT\", \"6379\")      # ex: 18374\n",
    "REDIS_PASSWORD = os.getenv(\"REDIS_PASSWORD\", \"\")  # ex: \"1TNxTEdYRDgIDKM2gDfasupCADXXXX\"\n",
    "\n",
    "# If SSL is enabled on the endpoint, use rediss:// as the URL prefix\n",
    "REDIS_URL = f\"redis://:{REDIS_PASSWORD}@{REDIS_HOST}:{REDIS_PORT}\"\n",
    "\n",
    "# Initialize Redis client\n",
    "redis_client = redis.Redis.from_url(REDIS_URL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dfe10ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the ModelStore\n",
    "model_store = ModelStore(redis_client, shard_size=1012*100) # ~100Kb sized keys in Redis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14a8c59f",
   "metadata": {},
   "source": [
    "## Test with simple Scikit-Learn model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "74cc62af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model accuracy: 1.00\n"
     ]
    }
   ],
   "source": [
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.datasets import load_iris\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "# Load a simple dataset and train a RandomForest model\n",
    "iris = load_iris()\n",
    "X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.2, random_state=42)\n",
    "\n",
    "# Train a RandomForestClassifier\n",
    "model = RandomForestClassifier()\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "# Evaluate the model\n",
    "y_pred = model.predict(X_test)\n",
    "accuracy = accuracy_score(y_test, y_pred)\n",
    "print(f\"Model accuracy: {accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "51a8f83f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:13.871 - model_store.store - INFO - Saving 'random_forest' model\n",
      "2025-01-22 10:56:13.873 - model_store.store - INFO - Starting model serialization and storage\n",
      "2025-01-22 10:56:13.880 - model_store.store - INFO - Stored model in 2 shards (0.0069s)\n",
      "2025-01-22 10:56:13.883 - model_store.store - INFO - Save operation completed in 0.0121s\n"
     ]
    }
   ],
   "source": [
    "# Save the trained model to Redis\n",
    "model_name = \"random_forest\"\n",
    "version = model_store.save_model(model, model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d0ce2b03",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:13.887 - model_store.store - INFO - Loading 'random_forest' model\n",
      "2025-01-22 10:56:13.889 - model_store.store - INFO - Starting model reconstruction from shards\n",
      "2025-01-22 10:56:13.899 - model_store.store - INFO - Loaded model from 2 shards (0.0105s)\n",
      "2025-01-22 10:56:13.899 - model_store.store - INFO - Load operation completed in 0.0130s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded model accuracy: 1.00\n"
     ]
    }
   ],
   "source": [
    "# Load the model from Redis\n",
    "loaded_model = model_store.load_model(model_name)\n",
    "\n",
    "# Verify that the loaded model works\n",
    "y_pred_loaded = loaded_model.predict(X_test)\n",
    "loaded_accuracy = accuracy_score(y_test, y_pred_loaded)\n",
    "print(f\"Loaded model accuracy: {loaded_accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c743a48d",
   "metadata": {},
   "source": [
    "## Test with 1Gb numpy array\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dd894c4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "desired_size_bytes = 1024*1024*1024\n",
    "num_elements = desired_size_bytes // 8\n",
    "large_array = np.random.rand(num_elements).astype(np.float64)\n",
    "\n",
    "model_name = \"numpy_array\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bcc6abd7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:14.435 - model_store.store - INFO - Saving 'numpy_array' model\n",
      "2025-01-22 10:56:14.438 - model_store.store - INFO - Starting model serialization and storage\n",
      "2025-01-22 10:56:21.996 - model_store.store - INFO - Stored model in 10611 shards (7.5575s)\n",
      "2025-01-22 10:56:22.004 - model_store.store - INFO - Save operation completed in 7.5689s\n"
     ]
    }
   ],
   "source": [
    "version = model_store.save_model(large_array, model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "12f57740",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:22.008 - model_store.store - INFO - Loading 'numpy_array' model\n",
      "2025-01-22 10:56:22.013 - model_store.store - INFO - Starting model reconstruction from shards\n",
      "2025-01-22 10:56:25.661 - model_store.store - INFO - Loaded model from 10611 shards (3.6480s)\n",
      "2025-01-22 10:56:25.702 - model_store.store - INFO - Load operation completed in 3.6941s\n"
     ]
    }
   ],
   "source": [
    "# Load the model from Redis\n",
    "loaded_model = model_store.load_model(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e29ce236",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.True_"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check if all elements match in the array\n",
    "sum(loaded_model == large_array) == num_elements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c452718",
   "metadata": {},
   "source": [
    "## Test with pytorch model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ccc06981",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# Define a simple model\n",
    "class SimpleModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleModel, self).__init__()\n",
    "        self.fc = nn.Linear(1, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(x)\n",
    "\n",
    "# Create model, define loss and optimizer\n",
    "model = SimpleModel()\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "# Dummy data\n",
    "x = torch.tensor([[1.0], [2.0], [3.0]], requires_grad=True)\n",
    "y = torch.tensor([[2.0], [4.0], [6.0]])\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(100):\n",
    "    optimizer.zero_grad()\n",
    "    output = model(x)\n",
    "    loss = criterion(output, y)\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f5e2dc3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction for input 4.0: 7.4401021003723145\n"
     ]
    }
   ],
   "source": [
    "# Make a simple prediction\n",
    "x_test = torch.tensor([[4.0]])\n",
    "prediction = model(x_test).item()\n",
    "print(f\"Prediction for input 4.0: {prediction}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "86dc9f9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:32.978 - model_store.store - INFO - Saving 'pytorch' model\n",
      "2025-01-22 10:56:32.982 - model_store.store - INFO - Starting model serialization and storage\n",
      "2025-01-22 10:56:32.983 - model_store.store - INFO - Stored model in 1 shards (0.0012s)\n",
      "2025-01-22 10:56:32.984 - model_store.store - INFO - Save operation completed in 0.0056s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'1.0'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save the trained model to Redis\n",
    "model_name = \"pytorch\"\n",
    "version = \"1.0\"\n",
    "model_store.save_model(model, model_name, version=version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "670f353f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:32.987 - model_store.store - INFO - Loading 'pytorch' model\n",
      "2025-01-22 10:56:32.988 - model_store.store - INFO - Starting model reconstruction from shards\n",
      "2025-01-22 10:56:32.990 - model_store.store - INFO - Loaded model from 1 shards (0.0012s)\n",
      "2025-01-22 10:56:32.990 - model_store.store - INFO - Load operation completed in 0.0027s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction for input 4.0 with loaded model: 7.4401021003723145\n"
     ]
    }
   ],
   "source": [
    "# Load the model from Redis\n",
    "loaded_model = model_store.load_model(model_name)\n",
    "\n",
    "prediction = loaded_model(x_test).item()\n",
    "print(f\"Prediction for input 4.0 with loaded model: {prediction}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6644d75",
   "metadata": {},
   "source": [
    "## Test with tensorflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e8bd2de2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<keras.src.callbacks.history.History at 0x16bd3a5d0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "# Define a simple model\n",
    "inputs = tf.keras.Input(shape=(1,))\n",
    "outputs = tf.keras.layers.Dense(1)(inputs)\n",
    "model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "# Compile the model\n",
    "model.compile(optimizer='sgd', loss='mse')\n",
    "\n",
    "# Dummy data\n",
    "x = tf.constant([[1.0], [2.0], [3.0]])\n",
    "y = tf.constant([[2.0], [4.0], [6.0]])\n",
    "\n",
    "# Train the model\n",
    "model.fit(x, y, epochs=100, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c26f35ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction for input 4.0: 7.760094165802002\n"
     ]
    }
   ],
   "source": [
    "# Make a simple prediction\n",
    "x_test = tf.constant([[4.0]])\n",
    "prediction = model(x_test).numpy()[0, 0]\n",
    "print(f\"Prediction for input 4.0: {prediction}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5aeaf0ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:36.920 - model_store.store - INFO - Saving 'tensorflow' model\n",
      "2025-01-22 10:56:36.922 - model_store.store - INFO - Starting model serialization and storage\n",
      "2025-01-22 10:56:36.932 - model_store.store - INFO - Stored model in 1 shards (0.0103s)\n",
      "2025-01-22 10:56:36.934 - model_store.store - INFO - Save operation completed in 0.0138s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'1.0'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save the trained model to Redis\n",
    "model_name = \"tensorflow\"\n",
    "version = \"1.0\"\n",
    "model_store.save_model(model, model_name, version=version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "117255d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-22 10:56:36.941 - model_store.store - INFO - Loading 'tensorflow' model\n",
      "2025-01-22 10:56:36.943 - model_store.store - INFO - Starting model reconstruction from shards\n",
      "2025-01-22 10:56:36.954 - model_store.store - INFO - Loaded model from 1 shards (0.0106s)\n",
      "2025-01-22 10:56:36.954 - model_store.store - INFO - Load operation completed in 0.0131s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction for input 4.0 with loaded model: 7.760094165802002\n"
     ]
    }
   ],
   "source": [
    "# Load the model from Redis\n",
    "loaded_model = model_store.load_model(model_name)\n",
    "\n",
    "prediction = loaded_model(x_test).numpy()[0, 0]\n",
    "print(f\"Prediction for input 4.0 with loaded model: {prediction}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db4db02e",
   "metadata": {},
   "source": [
    "# Model versioning "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "84136352",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['numpy_array', 'pytorch', 'random_forest', 'tensorflow']"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List all available models in the store\n",
    "models = model_store.list_models()\n",
    "models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "cadbec7e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ModelVersion(name='pytorch', description='', version='1.0', created_at=1737561392.98, shard_keys=['shard:pytorch:1.0:0'])]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List model versions for a model\n",
    "versions = model_store.get_all_versions(models[1])\n",
    "versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0962bc96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Delete a model version\n",
    "model_version = versions[0]\n",
    "model_store.delete_version(model_version.name, model_version.version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "28af2739",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10617"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Clear all versions for all models\n",
    "model_store.clear()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}