Automated_Retraining_Workflow_Example_with_MLflow.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyMp5PQTYtD2ioQjwEeY8F+6",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/OneFineStarstuff/OneFineStardust/blob/main/Automated_Retraining_Workflow_Example_with_MLflow.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install mlflow"
      ],
      "metadata": {
        "id": "ZEy6Ki9RvkT5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import mlflow\n",
        "import mlflow.pytorch\n",
        "import torch\n",
        "\n",
        "# Define your model\n",
        "class MyModel(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MyModel, self).__init__()\n",
        "        self.fc = nn.Linear(10, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.fc(x)\n",
        "\n",
        "# Instantiate the model\n",
        "model = MyModel()\n",
        "\n",
        "# Example input for the model\n",
        "input_example = torch.randn(1, 10).numpy()  # Convert to numpy.ndarray\n",
        "\n",
        "# Placeholder for model performance comparison\n",
        "new_model_performance = 0.85  # Example performance metric (e.g., accuracy)\n",
        "baseline_performance = 0.80  # Example baseline performance metric (e.g., accuracy)\n",
        "\n",
        "# Define the conda environment\n",
        "conda_env = {\n",
        "    'name': 'mlflow-env',\n",
        "    'channels': ['defaults', 'conda-forge'],\n",
        "    'dependencies': [\n",
        "        'python=3.8',\n",
        "        'pytorch=2.5.1',\n",
        "        'pip',\n",
        "        {\n",
        "            'pip': [\n",
        "                f'torch=={torch.__version__}',\n",
        "                'mlflow'\n",
        "            ]\n",
        "        }\n",
        "    ]\n",
        "}\n",
        "\n",
        "# Log model and parameters\n",
        "with mlflow.start_run():\n",
        "    mlflow.log_params({\"learning_rate\": 0.01, \"epochs\": 10})\n",
        "    mlflow.pytorch.log_model(model, \"model\", conda_env=conda_env, input_example=input_example)\n",
        "\n",
        "    # Log new model performance\n",
        "    mlflow.log_metric(\"new_model_performance\", new_model_performance)\n",
        "\n",
        "    # Evaluate performance and deploy if improved\n",
        "    if new_model_performance > baseline_performance:\n",
        "        print(\"New model performance improved. Ready to deploy.\")\n",
        "        # Remove existing directory if it exists\n",
        "        import shutil\n",
        "        shutil.rmtree(\"/content/production_model_path\", ignore_errors=True)\n",
        "        # Save the model to a production path\n",
        "        mlflow.pytorch.save_model(model, \"/content/production_model_path\")"
      ],
      "metadata": {
        "id": "cEL8BQ6XCP66"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}