{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "47301b3d",
   "metadata": {},
   "source": [
    "# Model Deployment API Demo\n",
    "\n",
    "This notebook demonstrates how to use the Datamint API to deploy machine learning models as Docker images.\n",
    "\n",
    "> **IMPORTANT:**\n",
    "> Registered Model it not the same as a deployed model!\n",
    ">\n",
    "> A Registered Model is a model that has been uploaded to the platform, that might or might not be deployed.\n",
    "> Is a requirement to have a Registered Model before deploying it.\n",
    "> In this demo we will show how to deploy a registered model.\n",
    "\n",
    "## Features Covered:\n",
    "- Starting deployment jobs\n",
    "- Monitoring deployment status\n",
    "- Cancelling jobs\n",
    "- Managing deployed images\n",
    "- Checking image existence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb59d9f3",
   "metadata": {},
   "source": [
    "## Setup and Initialization\n",
    "\n",
    "First, import the necessary libraries and initialize the API client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "639e0b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from datamint import Api\n",
    "\n",
    "# Initialize the API client\n",
    "# The API key can be set via environment variable DATAMINT_API_KEY\n",
    "# or passed directly to the Api constructor\n",
    "api = Api()\n",
    "\n",
    "print(\"API client initialized successfully!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ebce00a",
   "metadata": {},
   "source": [
    "## List current registered models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f6a3ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all versions of all registered models\n",
    "from datetime import datetime\n",
    "\n",
    "from mlflow import MlflowClient\n",
    "\n",
    "client = MlflowClient()\n",
    "all_registered_models = client.search_registered_models()\n",
    "for rm in all_registered_models:\n",
    "    print(f\"Registered Model: {rm.name}\")\n",
    "    model_versions = client.search_model_versions(f\"name='{rm.name}'\")\n",
    "    # get all all attributes of each model version\n",
    "    for mv in model_versions:\n",
    "        created_datetime = datetime.fromtimestamp(mv.creation_timestamp / 1000.0) # converting from milliseconds\n",
    "        created_datetime = created_datetime.strftime(\"%Y-%m-%d %H:%M:%S\") # formatting\n",
    "        print(f\"  Version: {mv.version} | Created at: {created_datetime}\")\n",
    "        # INFO: do try to access ``mv.aliases``, since mlflow does not populate this field with the ``search_model_versions``.\n",
    "        # If you need aliases, use ``get_model_version(mv.name, mv.version)`` method instead."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca5f9fad",
   "metadata": {},
   "source": [
    "## Loading a Registered Model Locally\n",
    "\n",
    "Before deploying, you can load a model version directly in Python via `ModelVersion.load_model()` (a shortcut over `datamint.mlflow.flavors.load_model`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd55407f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a model version locally, without deploying it\n",
    "model = api.models.get_by_name(all_registered_models[1].name)\n",
    "latest = model.get_latest_version()\n",
    "\n",
    "loaded_model = latest.load_model()\n",
    "print(loaded_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9da152da",
   "metadata": {},
   "source": [
    "## Starting a Deployment Job\n",
    "\n",
    "Deploy a model by specifying the model name and optionally a version or alias."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e122f4df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start a deployment job\n",
    "# You can specify either model_version (int) or model_alias (str), but not both\n",
    "job = api.deploy.start(\n",
    "    model_name=all_registered_models[1].name,\n",
    "    model_alias=\"latest\",\n",
    "    # or use model_version=1\n",
    "    with_gpu=False,\n",
    ")\n",
    "\n",
    "print(\"Deployment job started!\")\n",
    "print(f\"Job ID: {job.id}\")\n",
    "print(f\"Status: {job.status}\")\n",
    "print(f\"Model: {job.model_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6481deaa",
   "metadata": {},
   "source": [
    "## Monitoring Job Status\n",
    "\n",
    "Check the status of a deployment job and monitor its progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bc975d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get job status by ID\n",
    "job_status = api.deploy.get_by_id(job.id)\n",
    "\n",
    "print(f\"Job Status: {job_status.status}\")\n",
    "print(f\"Progress: {job_status.progress_percentage}%\")\n",
    "print(f\"Current Step: {job_status.current_step}\")\n",
    "print(f\"Image Name: {job_status.image_name}\")\n",
    "print(f\"Image Tag: {job_status.image_tag}\")\n",
    "\n",
    "if job_status.error_message:\n",
    "    print(f\"Error: {job_status.error_message}\")\n",
    "\n",
    "if job_status.recent_logs:\n",
    "    print(\"\\nRecent Logs:\")\n",
    "    for log in job_status.recent_logs[-5:]:  # Show last 5 logs\n",
    "        print(f\"  {log}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b647d70e",
   "metadata": {},
   "source": [
    "## Polling for Job Completion\n",
    "\n",
    "Monitor a job until it completes (or fails)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25608260",
   "metadata": {},
   "outputs": [],
   "source": [
    "def wait_for_job_completion(job_id, max_wait_seconds=600, poll_interval=2):\n",
    "    \"\"\"\n",
    "    Wait for a deployment job to complete.\n",
    "    \n",
    "    Args:\n",
    "        job_id: The job ID to monitor\n",
    "        max_wait_seconds: Maximum time to wait (default 10 minutes)\n",
    "        poll_interval: How often to check status (default 2 seconds)\n",
    "    \"\"\"\n",
    "    start_time = time.time()\n",
    "    last_msg = \"\"\n",
    "    while time.time() - start_time < max_wait_seconds:\n",
    "        job_status = api.deploy.get_by_id(job_id)\n",
    "        \n",
    "        msg = f\"Status: {job_status.status} | Progress: {job_status.progress_percentage}% | Step: {job_status.current_step}\"\n",
    "        if msg != last_msg:\n",
    "            print(msg)\n",
    "        last_msg = msg\n",
    "        \n",
    "        if job_status.status in ['completed', 'failed', 'cancelled']:\n",
    "            return job_status\n",
    "        \n",
    "        time.sleep(poll_interval)\n",
    "    \n",
    "    print(\"Timeout waiting for job completion\")\n",
    "    return None\n",
    "\n",
    "# Wait for the job to complete\n",
    "final_status = wait_for_job_completion(job.id)\n",
    "\n",
    "if final_status:\n",
    "    print(f\"\\nFinal Status: {final_status.status}\")\n",
    "    if final_status.status == 'completed':\n",
    "        print(f\"Image built successfully: {final_status.image_name}:{final_status.image_tag}\")\n",
    "    elif final_status.status == 'failed':\n",
    "        print(f\"Error: {final_status.error_message}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b86955df",
   "metadata": {},
   "source": [
    "## Cancelling a Job\n",
    "\n",
    "Cancel a running deployment job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5a8a44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cancel a job\n",
    "cancelled = api.deploy.cancel(job.id)\n",
    "\n",
    "if cancelled:\n",
    "    print(f\"Job {job.id} cancelled successfully\")\n",
    "else:\n",
    "    print(f\"Job {job.id} could not be cancelled (may have already completed)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9db375cc",
   "metadata": {},
   "source": [
    "## Listing Active Jobs\n",
    "\n",
    "Get information about active deployment jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b697c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List active jobs\n",
    "active_jobs = api.deploy.list_active_jobs()\n",
    "\n",
    "print(f\"Active jobs count: {active_jobs['active_jobs_count']}\")\n",
    "active_jobs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "220d072a",
   "metadata": {},
   "source": [
    "## Managing Deployed Images\n",
    "\n",
    "List and manage Docker images that have been deployed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6022e586",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all deployed images\n",
    "all_images = api.deploy.list_images()\n",
    "\n",
    "print(f\"Total deployed images: {len(all_images)}\")\n",
    "for img in all_images[:5]:  # Show first 5\n",
    "    print(f\"  - {img['full_name']} ({img['size_mb']:.2f} MB)\")\n",
    "    print(f\"    Created: {img['created']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b0d1100",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List images for a specific model\n",
    "model_images = api.deploy.list_images(model_name=\"FracAtlas_adapted\")\n",
    "\n",
    "print(f\"Images for 'FracAtlas_adapted': {len(model_images)}\")\n",
    "for img in model_images:\n",
    "    print(f\"  - {img['name']}:{img['tag']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecaa3783",
   "metadata": {},
   "source": [
    "## Checking if an Image Exists\n",
    "\n",
    "Verify if a specific model image exists."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6eba6c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if an image exists\n",
    "exists = api.deploy.image_exists(\n",
    "    model_name=\"FracAtlas_adapted\",\n",
    "    tag=\"champion\"\n",
    ")\n",
    "\n",
    "if exists:\n",
    "    print(\"Image exists!\")\n",
    "else:\n",
    "    print(\"Image not found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9875937",
   "metadata": {},
   "source": [
    "## Removing Deployed Images\n",
    "\n",
    "Remove Docker images that are no longer needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cf2c108",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove a specific image tag\n",
    "result = api.deploy.remove_image(\n",
    "    model_name=\"my_model\",\n",
    "    tag=\"v1.0\"\n",
    ")\n",
    "\n",
    "print(f\"Removal {'successful' if result['success'] else 'failed'}\")\n",
    "print(f\"Message: {result['message']}\")\n",
    "if result.get('removed_tags'):\n",
    "    print(f\"Removed tags: {result['removed_tags']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8f2513",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove all images for a model\n",
    "result = api.deploy.remove_image(\n",
    "    model_name=\"my_model\"\n",
    "    # No tag specified = remove all tags\n",
    ")\n",
    "\n",
    "print(f\"Removal {'successful' if result['success'] else 'failed'}\")\n",
    "print(f\"Message: {result['message']}\")\n",
    "if result.get('removed_tags'):\n",
    "    print(f\"Removed tags: {', '.join(result['removed_tags'])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89701dd9",
   "metadata": {},
   "source": [
    "## Complete Workflow Example\n",
    "\n",
    "Here's a complete workflow from deployment to cleanup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da42289b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def deploy_and_monitor_model(model_name, model_alias=\"champion\"):\n",
    "    \"\"\"\n",
    "    Complete workflow: deploy, monitor, and verify a model.\n",
    "    \"\"\"\n",
    "    print(f\"1. Starting deployment for {model_name}@{model_alias}\")\n",
    "    job = api.deploy.start(\n",
    "        model_name=model_name,\n",
    "        model_alias=model_alias,\n",
    "        with_gpu=False\n",
    "    )\n",
    "    print(f\"   Job ID: {job.id}\")\n",
    "    \n",
    "    print(\"\\n2. Monitoring deployment progress...\")\n",
    "    final_status = wait_for_job_completion(job.id, poll_interval=5)\n",
    "    \n",
    "    if final_status and final_status.status == 'completed':\n",
    "        print(\"\\n3. Deployment completed successfully!\")\n",
    "        print(f\"   Image: {final_status.image_name}:{final_status.image_tag}\")\n",
    "        \n",
    "        # Verify image exists\n",
    "        exists = api.deploy.image_exists(\n",
    "            model_name=model_name,\n",
    "            tag=final_status.image_tag\n",
    "        )\n",
    "        print(f\"   Image verification: {'✓' if exists else '✗'}\")\n",
    "        \n",
    "        return final_status\n",
    "    else:\n",
    "        print(\"\\n3. Deployment failed or timed out\")\n",
    "        if final_status and final_status.error_message:\n",
    "            print(f\"   Error: {final_status.error_message}\")\n",
    "        return None\n",
    "\n",
    "# Run the complete workflow\n",
    "result = deploy_and_monitor_model(\"my_production_model\", \"champion\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3be1d9b",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook covered:\n",
    "- ✓ Starting deployment jobs with various configurations\n",
    "- ✓ Monitoring job progress and status\n",
    "- ✓ Cancelling running jobs\n",
    "- ✓ Listing and managing deployed images\n",
    "- ✓ Checking image existence\n",
    "- ✓ Removing images\n",
    "- ✓ Complete workflow examples"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}