01_standalone_examples/image-segmentation/Image Segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "43e3aa9c-9a20-4468-9846-4c6f5902b7c4",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "<img src=\"https://lakefs.io/wp-content/uploads/2022/09/lakeFS-Logo.svg\" alt=\"lakeFS logo\" width=200/>\n",
    "\n",
    "# ML Data Version Control and Reproducibility at Scale\n",
    "\n",
    "### In the ever-evolving landscape of machine learning (ML), data stands as the cornerstone upon which triumphant models are built. However, as ML projects expand and encompass larger and more complex datasets, the challenge of efficiently managing and controlling data at scale becomes more pronounced.\n",
    "\n",
    "### Breaking Down Conventional Approaches:\n",
    "##### The Copy/Paste Predicament: In the world of data science, it's commonplace for data scientists to extract subsets of data to their local environments for model training. This method allows for iterative experimentation, but it introduces challenges that hinder the seamless evolution of ML projects.\n",
    "\n",
    "##### Reproducibility Constraints: Traditional practices of copying and modifying data locally lack the version control and audit-ability crucial for reproducibility. Iterating on models with various data subsets becomes a daunting task.\n",
    "\n",
    "##### Inefficient Data Transfer: Regularly shuttling data between the central repository and local environments strains resources and time, especially when choosing different subsets of data for each training run.\n",
    "\n",
    "##### Limited Compute Power: Operating within a local environment hampers the ability to harness the full power of parallel computing, as well as the distributed prowess of systems like Apache Spark.\n",
    "\n",
    "\n",
    "In this sample, you'll learn how to use lakeFS for scalable data version control and reproducibility in your ML workflows. This notebook demonstrates how to create branches for different experiments, work with your data either distributed or locally, and efficiently manage large datasets on the cloud without needing to copy data. We will leverage the technology stack of AWS S3, Databricks Delta Lake, PyTorch, and MLflow to ensure seamless data processing, experiment tracking, and model training at scale using the medallion architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b1a8d403-ec02-4dcc-8fc3-38da45734113",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "source": [
    "## Target Architecture\n",
    "\n",
    "<img src=\"https://www.databricks.com/sites/default/files/inline-images/db-277-blog-img-3.png\" alt=\"target architecture\" width=800/>\n",
    "\n",
    "#### Source: Databricks Blogs:\n",
    "##### [Accelerating Your Deep Learning with PyTorch Lightning on Databricks](https://www.databricks.com/blog/2022/09/07/accelerating-your-deep-learning-pytorch-lightning-databricks.html)\n",
    "##### [Image Segmentation with Databricks](https://florent-brosse.medium.com/image-segmentation-with-databricks-6db19d23725d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### You can run this same notebook in local container or on the Databricks cluster. This picture explains the full procees:\n",
    "<img src=\"./files/images/ImageSegmentation/ImageSegmentation.png\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0f3e1e04-ef13-45c9-8d5a-d637080fc5af",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "72db9c37-4dab-4712-8d8d-a79b4006354a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Change lakeFS endpoint and credentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "eec8f521-9a2c-42f1-979d-6e948febd9aa",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "lakefsEndPoint = 'http://lakefs:8000' # e.g. 'https://username.aws_region_name.lakefscloud.io' \n",
    "lakefsAccessKey = 'AKIAIOSFOLKFSSAMPLES'\n",
    "lakefsSecretKey = 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "c377a867-3690-4c7e-b899-7f5d376c0bd8",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### You can change repo name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d1144e38-f35e-43eb-9435-a656f5560af5",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "repo_name = \"image-segmentation-repo\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "f0c1b9f9-af1a-4714-a5a0-087c67440674",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Storage Information\n",
    "\n",
    "Change the Storage Namespace to a location in the bucket you’ve configured. The storage namespace is a location in the underlying storage where data for lakeFS repository will be stored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "cda4323b-32f5-4b34-90bc-68f9def3f5d7",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "storageNamespace = 's3://example' # e.g. \"s3://bucket\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6153f52a-e41e-4d6b-82d1-fbd57e496648",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Are you running this demo in LOCAL container or in Databricks DISTRIBUTED cluster?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "76793842-1c0d-4ba0-87a6-a7fb82b69135",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "localOrDistributedComputing = \"LOCAL\" # LOCAL or DISTRIBUTED"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b68d5e15-9705-4368-a40c-029362117e53",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Number of images to use for each experiment (use small number for LOCAL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "76793842-1c0d-4ba0-87a6-a7fb82b69135",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "imagesPerExperiment = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "79fdb63c-ee70-4991-84dc-bfa1c1f12e08",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Download demo dataset from [Kaggle](https://www.kaggle.com/c/airbus-ship-detection) and upload to \"airbus-ship-detection\" folder in your S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d6913f40-6a57-4ac3-a1b9-51abbf89c4dd",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "bucketName = '<S3 Bucket Name>'\n",
    "awsRegion = '<AWS Region>'\n",
    "prefix = \"airbus-ship-detection/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "cc4b13e3-8752-46aa-a88d-c723e5185972",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Provide your AWS credentials to access demo dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "06bb47c8-ced5-4cd9-8b15-ce406e1a978f",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "aws_access_key_id = 'aaaaaaaaaaaaa'\n",
    "aws_secret_access_key = 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "d9b7d4ff-2a76-40a4-8629-8818f667bfb8",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "source": [
    "## Setup\n",
    "\n",
    "**(you shouldn't need to change anything in this section, just run it)**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "496b546d-467a-48f0-91dd-1bc4aa34d7d8",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### If running LOCAL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ddc2a99a-da97-441f-a734-663a6617f88f",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    %run ./ImageSegmentationSetup.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b4aeff19-92b4-4252-a871-7842364984bd",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### If running DISTRIBUTED on Databricks cluster otherwise skip this cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bedb7140-7a81-4f32-9064-c5f73236252e",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%run ./ImageSegmentationSetup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### If running LOCAL then create an empty Git repository and configure Git. Git will version control your code while lakeFS will version control your data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    !git init {repo_name}\n",
    "    !git config --global user.email \"you@example.com\"\n",
    "    !git config --global user.name \"Your Name\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "2ed0c446-1d2f-4b0a-a59b-4b4e7e33ae1a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# Main demo starts here 🚦 👇🏻"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b873d753-63a9-4146-bf81-22d4ffd3c23c",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Import training data to experiment branch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "3a1d6958-521d-4fc8-b422-ac64e5b7b34a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Create branch for each experiment\n",
    "#### If running LOCAL then create a Git branch also"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "aaf98d56-be45-4093-bff1-e6dea24c76a8",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "experimentBranchN = experimentBranch+\"-1\"\n",
    "\n",
    "try:\n",
    "    repo.branch(experimentBranchN).head\n",
    "    branchExperimentBranchN = repo.branch(experimentBranchN)\n",
    "    print(f\"{experimentBranchN} already exists\")\n",
    "except NotFoundException as f:\n",
    "    if localOrDistributedComputing == \"LOCAL\":\n",
    "        !cd {repo_name} && git checkout -b {experimentBranchN}\n",
    "    branchExperimentBranchN = repo.branch(experimentBranchN).create(source_reference=emptyBranch)\n",
    "    print(f\"{experimentBranchN} branch created\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "347877ed-a861-47f0-8a32-5cd46700d30f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Get the list of images from S3 in the training dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6e4e5a53-8974-4e7b-8560-52e68cdc4a51",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "file_list = list_images()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6205cdae-dd4d-49b1-affa-0cba69e1fbc6",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Randomly select subset of the training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1b59ed9a-491c-4ec6-ad49-d4bdb1b572c4",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "file_list_random = random.choices(file_list, k=imagesPerExperiment)\n",
    "print(len(file_list_random))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "66e5aeb9-d444-46ed-b45d-755e54642cd8",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Import subset of the training data to lakeFS repo\n",
    "#### This is [zero-copy operation](https://docs.lakefs.io/understand/performance-best-practices.html#use-zero-copy-import)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "dfec7383-400e-4a21-aae6-aa0683c3350d",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import_images(file_list_random)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "7f854144-6757-437e-b4ba-34629b42e535",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Work locally with smaller dataset or work with bigger dataset in Databricks cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "72ec25f3-3903-4684-b098-4063980d4000",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    repo_path = f\"{repo_name}/lakefs_local\"\n",
    "elif localOrDistributedComputing == \"DISTRIBUTED\":\n",
    "    repo_path = f\"lakefs://{repo_name}/{experimentBranchN}\"\n",
    "\n",
    "raw_data_path = f\"{repo_path}/{raw_data_folder}\"\n",
    "training_data_path = f\"{raw_data_path}{training_data_folder}\"\n",
    "bronze_data_path = f\"{repo_path}/{bronze_data_folder}\"\n",
    "silver_data_path = f\"{repo_path}/{silver_data_folder}\"\n",
    "gold_data_path = f\"{repo_path}/{gold_data_folder}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fd319b42-6da9-42d0-b681-8e970527b9f1",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Clone experiment branch with smaller dataset locally\n",
    "#### This will download images locally. You will notice \"image-segmentation-repo/lakefs_local\" folder in Jupyter File Browser on the left side panel. You can browse the files inside this folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "67daa1c8-db4f-4816-9969-3673bab22d28",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    lakeFSLocalCommand = f\"lakectl local clone lakefs://{repo.id}/{experimentBranchN}/ {repo_path}\"\n",
    "    response = ! $lakeFSLocalCommand\n",
    "    print_lakectl_response(response, 8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's review \".gitignore\" file and \".lakefs_ref.yaml\" file created by previous \"lakectl local clone\" command.\n",
    "#### You will notice in .gitignore file that Git will not commit any data files in \"lakefs_local\" folder but will commit \".lakefs_ref.yaml\" file which includes lakeFS commit information. This way code as well as commit information about data will be kept together in Git repo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    !cat {repo_name}/.gitignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    !cat {repo_path}/.lakefs_ref.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete images smaller than 75KB in size locally. Add few images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    !find {training_data_path} -type f -name \"*.jpg\" -size -75k -delete\n",
    "    !cp ./data/*.jpg {training_data_path}\n",
    "    diff_branch(repo.id, repo_path, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add changes to Git repo and perform initial commit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    !cd {repo_name} && git add -A && git status\n",
    "    !cd {repo_name} && git commit -m \"Initial commit\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Commit local changes to lakeFS repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    commitMessage = 'Deleted images smaller than 75KB in size and added few images'\n",
    "    commit(repo.id, repo_path, experimentBranchN, commitMessage)\n",
    "    lakefs_set_tag(repo.id, f\"{tagPrefix}-{experimentBranchN}-raw-images\", experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "145e8baf-6bb3-4c9b-807d-27120c6ee422",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Verify that you can read the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "895b0ea9-855f-4368-a6ac-b71a7e1ffbb3",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = spark.read.format(\"image\").load(training_data_path)\n",
    "df.select(\"image.origin\", \"image.width\", \"image.height\").show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "62d1893d-9d3a-4901-859d-137f55465ffa",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Build the data pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "95dea8a4-ea85-4e56-9e80-c943822b8f8d",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Ingest raw images as bronze data set and save as Delta table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fa00f559-bc3e-4205-bcd8-04ba1bdce974",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_bronze_images = bronze_images()\n",
    "df_bronze_images.write.format(\"delta\").mode(\"overwrite\").save(f\"{bronze_data_path}/{training_data_folder}\")\n",
    "diff_branch(repo.id, repo_path, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "84a134fb-c324-4b8b-8496-ae7568256f80",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Commit bronze dataset to the lakeFS repository and tag it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "67487bb2-cc9e-4a06-997c-076ddad5173d",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "commitMessage = 'Converted raw images to binary content and saved as Delta table'\n",
    "commit(repo.id, repo_path, experimentBranchN, commitMessage)\n",
    "lakefs_set_tag(repo.id, f\"{tagPrefix}-{experimentBranchN}-bronze-images\", experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "9a7d4dd7-5581-4a12-a14e-6a959c2910cc",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Enrich dataset and save as silver dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "aa14323a-3e8c-4a4f-945e-1ab673b215dd",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_silver_images = silver_images(df_bronze_images)\n",
    "df_silver_images.write.format(\"delta\").mode(\"overwrite\").save(f\"{silver_data_path}/{training_data_folder}\")\n",
    "diff_branch(repo.id, repo_path, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1eebe466-316b-4f47-a178-47a3ac55a03e",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Commit silver dataset to the lakeFS repository and tag it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "28866c55-18a6-438b-a22d-a4336df25964",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "commitMessage = 'Enriched dataset and saved as silver dataset'\n",
    "commit(repo.id, repo_path, experimentBranchN, commitMessage)\n",
    "lakefs_set_tag(repo.id, f\"{tagPrefix}-{experimentBranchN}-silver-images\", experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "851ccf08-2446-44dd-b448-bd7a992d5982",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Load the raw image mask as bronze dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "437c2c0c-70cd-4c05-907a-6035467bb4fd",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_bronze_mask = bronze_mask()\n",
    "df_bronze_mask.write.format(\"delta\").mode(\"overwrite\").save(f\"{bronze_data_path}/{mask_data_folder}\")\n",
    "diff_branch(repo.id, repo_path, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "3a07d7f1-2841-4793-9bfd-cc17e957e77f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Commit bronze mask dataset to the lakeFS repository and tag it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "9b43ca23-d714-45dc-9d1f-5a3337427d48",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "commitMessage = 'Loaded the raw image mask and saved as Delta table'\n",
    "commit(repo.id, repo_path, experimentBranchN, commitMessage)\n",
    "lakefs_set_tag(repo.id, f\"{tagPrefix}-{experimentBranchN}-bronze-mask\", experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1e6c26b1-8895-4782-81f1-5864c5098d9a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Transform masks into images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "98d60ab0-7fd1-4811-82ce-8fdbb31b453b",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_silver_mask = silver_mask(df_bronze_mask)\n",
    "df_silver_mask.write.format(\"delta\").mode(\"overwrite\").save(f\"{silver_data_path}/{mask_data_folder}\")\n",
    "diff_branch(repo.id, repo_path, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "68e44b5f-b7ec-4de9-9813-ce6758c2d01b",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Commit silver mask dataset to the lakeFS repository and tag it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2c25bcbe-cb1b-4e7c-a96a-f44ab1dfc8d4",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "commitMessage = 'Transformed masks into images'\n",
    "commit(repo.id, repo_path, experimentBranchN, commitMessage)\n",
    "lakefs_set_tag(repo.id, f\"{tagPrefix}-{experimentBranchN}-silver-mask\", experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6b86d14d-bf0d-44e5-9daa-9408cbac586c",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### To verify that pipeline ran successfully, join image and mask both as the gold layer and select top 10 images with maximum number of boats/ships"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a7ad71cf-1d2f-4106-a025-caeb4f7051fb",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_gold_images = gold_images(df_silver_images, df_silver_mask)\n",
    "display_gold_images(df_gold_images.orderBy(desc(\"boat_number\")).limit(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "047b304a-ecfe-4a2a-8134-6e108dee85b0",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Save gold dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7ab42aed-8871-4a87-8072-4a477d28bd96",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_gold_images.write.format(\"delta\").mode(\"overwrite\").save(f\"{gold_data_path}/{training_data_folder}\")\n",
    "diff_branch(repo.id, repo_path, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "8e4b936a-c919-4959-8482-ccea61105f2e",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Commit gold dataset to the lakeFS repository and tag it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "60ecfc47-ff9f-48b7-af4b-e36a36a2ffd6",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "commitMessage = 'Joined image and mask both as the gold layer'\n",
    "commit(repo.id, repo_path, experimentBranchN, commitMessage)\n",
    "goldDatasetTagID = f\"{tagPrefix}-{experimentBranchN}-gold-images\"\n",
    "lakefs_set_tag(repo.id, goldDatasetTagID, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b8276798-7707-44a4-afe1-d455811fe079",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Build the Image Segmentation model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "ce0cb5aa-91c9-4241-8220-2f176bcdf158",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Split data as train/test datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f469b84f-bcdd-4a6f-ab15-a730db89bba2",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "gold_images_df = spark.read.format(\"delta\").load(f\"{gold_data_path}/{training_data_folder}\")\n",
    "(images_train, images_test) = gold_images_df.randomSplit(weights = [0.8, 0.2], seed = 42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "5e602614-98df-4ec3-a72f-3efcff223a6f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Prepare the dataset in PyTorch format by using Petastorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e826d87f-5a38-46df-ae0a-f04f7a351b5f",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set the converter cache folder to petastorm_path\n",
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    petastorm_path = 'file:///home/jovyan/petastorm/cache'\n",
    "elif localOrDistributedComputing == \"DISTRIBUTED\":\n",
    "    dbutils.fs.rm(\"dbfs:/tmp/petastorm\",True)\n",
    "    petastorm_path = 'file:///dbfs/tmp/petastorm/cache'\n",
    "\n",
    "spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, petastorm_path)\n",
    "# convert the image for pytorch\n",
    "converter_train = make_spark_converter(images_train.coalesce(4)) # You can increase number of partitions from 4 if parquet file sizes generated by Petastorm are more than 50 MB\n",
    "converter_test = make_spark_converter(images_test.coalesce(4))\n",
    "print(f\"Images in training dataset: {len(converter_train)}, Images in test dataset: {len(converter_test)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "61ccfd3a-743a-40cd-825e-d45a333d2760",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Train the base Model\n",
    "\n",
    "### If running LOCAL then train the model once with \"FPN\" architecture, \"resnet34\" encoder and learning rate of \"0.0001\"\n",
    "\n",
    "#### Model will return Intersection over Union (IoU) metric which is a widely-used evaluation metric in object detection and image segmentation tasks\n",
    "#### IoU measures the overlap between predicted bounding boxes and ground truth boxes, with scores ranging from 0 to 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "05cfc88d-6ecd-49fa-a1f5-be1821d76e42",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    valid_per_image_iou = train_model(\"FPN\", \"resnet34\", 0.0001)\n",
    "    print(f\"Intersection over Union (IoU) metric value: {valid_per_image_iou}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "db48f5dd-cd58-4de4-ac5f-25e1b9d44e52",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### If using Databricks cluster then fine-tune hyperparameters with Hyperopt\n",
    "\n",
    "#### This will crash when running LOCAL due to out-of-memory issues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "9dfd45f9-b244-42b0-a98a-da30b795e67a",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"DISTRIBUTED\":\n",
    "    # define hyperparameter search space\n",
    "    search_space = {\n",
    "        'lr': hp.loguniform('lr', -10, -4),\n",
    "        'segarch': hp.choice('segarch', ['Unet', 'FPN', 'deeplabv3plus', 'unetplusplus']),\n",
    "        'encoder_name': hp.choice('encoder_name', ['resnet50', 'resnet101', 'resnet152', 'resnet34'])}\n",
    "\n",
    "\n",
    "    # define training function to return results as expected by hyperopt\n",
    "    def train_fn(params):\n",
    "        arch = params['segarch']\n",
    "        encoder_name = params['encoder_name']\n",
    "        lr = params['lr']\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        valid_per_image_iou = train_model(arch, encoder_name, lr, nested=True)\n",
    "        return {'loss': 1 - valid_per_image_iou, 'status': STATUS_OK}\n",
    "\n",
    "    if localOrDistributedComputing == \"LOCAL\":\n",
    "        parallelism = 2\n",
    "    elif localOrDistributedComputing == \"DISTRIBUTED\":    \n",
    "        parallelism = int(spark.sparkContext.getConf().get('spark.databricks.clusterUsageTags.clusterWorkers'))\n",
    "\n",
    "    trials = SparkTrials(parallelism=parallelism) if parallelism > 1 else Trials()\n",
    "\n",
    "    # perform distributed hyperparameter tuning. Real training would go with max_eval > 20 \n",
    "    #mlflow.autolog(log_models=False)\n",
    "    with mlflow.start_run() as run:\n",
    "        argmin = fmin(fn=train_fn, space=search_space, algo=tpe.suggest, max_evals=3, trials=trials)\n",
    "        params = space_eval(search_space, argmin)\n",
    "        for p in params:\n",
    "            mlflow.log_param(p, params[p])\n",
    "        mlflow.set_tag(\"lakefs_demos\", \"image_segmentation\")\n",
    "        run_id = run.info.run_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1e273b5a-f534-4a08-9b51-7270e3f7d73d",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Save the best model to the MLflow registry (as a new version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fbbdb268-4c0c-4e47-82ed-557303c616bf",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# get the best model from the registry\n",
    "best_model = \\\n",
    "mlflow.search_runs(filter_string='attributes.status = \"FINISHED\" and tags.lakefs_demos = \"image_segmentation\"',\n",
    "                   order_by=[\"metrics.valid_per_image_iou DESC\"], max_results=1).iloc[0]\n",
    "model_registered = mlflow.register_model(\"runs:/\" + best_model.run_id + \"/model\", \"lakefs_demos_image_segmentation\")\n",
    "print(model_registered)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "928d2cd5-82e2-4880-9e6a-d68dff4910c5",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Save the best model information in the lakeFS repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c24b4e76-0868-47fd-bfc0-633d7f509962",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    pd.set_option('display.max_colwidth', -1)\n",
    "    f = open(f\"{repo_path}/best_model.txt\", \"w\")\n",
    "    f.write(best_model.to_string())\n",
    "    f.close()\n",
    "elif localOrDistributedComputing == \"DISTRIBUTED\":    \n",
    "    contentToUpload = best_model.to_string().encode('utf-8')\n",
    "    repo.branch(experimentBranchN).object('best_model.txt').upload(data=contentToUpload, mode='wb', pre_sign=False)\n",
    "commitMetadata = commit_metadata_for_best_model(best_model, model_registered)\n",
    "diff_branch(repo.id, repo_path, experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "bcacd8e0-24ac-4b3d-80db-018e4c117072",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Commit the best model information to the lakeFS repository\n",
    "#### Commit log in the lakeFS repository also URL to go to best registered model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "57b9af28-46e7-4213-b28a-4a2ac9099bb0",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "commitMessage = 'Information on best model'\n",
    "commit_id = commit(repo.id, repo_path, experimentBranchN, commitMessage, commitMetadata)\n",
    "lakefs_set_tag(repo.id, f\"{tagPrefix}-{experimentBranchN}-best-model\", experimentBranchN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e6c0b394-7567-4c92-bd87-58536399de2a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Flag the best model version as production-ready"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a53b2666-d536-4268-8361-ff1fc8b00421",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "client = mlflow.tracking.MlflowClient()\n",
    "print(\"registering model version \" + str(model_registered.version) + \" as production model\")\n",
    "client.transition_model_version_stage(name=\"lakefs_demos_image_segmentation\", version=model_registered.version,\n",
    "                                      stage=\"Production\", archive_existing_versions=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Copy notebooks (code) to Git repo. The git add command adds changes in the working directory to the staging area.\n",
    "#### Git doesn't add data files to staging area while adds \".lakefs_ref.yaml\" file which includes lakeFS commit information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    !cp -t {repo_name} 'Image Segmentation.ipynb' 'ImageSegmentationSetup.ipynb'\n",
    "    !cd {repo_name} && git add -A && git status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## If you are running LOCAL and want to access MLflow UI then open [start-mlflow-ui](./start-mlflow-ui.ipynb) notebook, start MLflow server and go to [MLflow UI](http://127.0.0.1:5001/).\n",
    "\n",
    "## If you are using Databricks then go to [Models page](https://dbc-8ada78b6-3a6d.cloud.databricks.com/#mlflow/models)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run following cell to generate the hyperlink to go to the commit page in lakeFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "md(f\"<br/>👉🏻 **Go to [the commit page in lakeFS]({lakefsUIEndPoint}/repositories/{repo_name}/commits/{commit_id}) \\\n",
    "to see the commit made to the repository along with information for the best model.<br>Click on 'Open Registered Model UI' button on the commit page to \\\n",
    "open the best model in MLflow UI.<br>Click on 'Source Run' link in MLflow UI to get run details including model pickle file(python_model.pkl).**\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "03ff8dba-f5d6-4aa9-8c93-7a74833abf3d",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## More Questions?\n",
    "\n",
    "[<img src=\"https://lakefs.io/wp-content/uploads/2023/06/Join-slack.svg\" alt=\"lakeFS logo\" width=700/>](https://lakefs.io/slack)"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "Image Segmentation",
   "widgets": {}
  },
  "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}