notebooks/3_lineage_tracking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "656db1e0",
   "metadata": {},
   "source": [
    "# Lineage Tracking\n",
    "This notebook uses the ml-lineage-helper repo to track the lineage of data, code and ML models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "def7d2da",
   "metadata": {},
   "source": [
    "**Note:** The lineage tracking functionality requires the use of [https://github.com/aws-samples/ml-lineage-helper](https://github.com/aws-samples/ml-lineage-helper) repo. This is done by pip installing it directly from it's github repo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88032165",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "995020c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "import logging\n",
    "import boto3\n",
    "import sys\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d2e6ba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import from a different path\n",
    "path = Path(os.path.abspath(os.getcwd()))\n",
    "package_dir = f'{str(path.parent)}/utils'\n",
    "print(package_dir)\n",
    "sys.path.insert(0, package_dir)\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ecc22d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# pip install directly from the github repo for ml-lineage-helper\n",
    "!pip3 install git+https://github.com/aws-samples/ml-lineage-helper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80201523",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ml_lineage_helper import MLLineageHelper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "483eb7fb",
   "metadata": {},
   "source": [
    "## Setup Logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a58714df",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger('__name__')\n",
    "logging.basicConfig(format=\"%(asctime)s,%(filename)s,%(funcName)s,%(lineno)s,%(levelname)s,p%(process)s,%(message)s\", level=logging.INFO)       \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "934ec277",
   "metadata": {},
   "source": [
    "## Setup Config Variables\n",
    "Read the metadata (feature group name, model endpoint name etc.) produced by the previous notebooks so that they can be provided as inputs to the lineage tracking module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8043704c",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = utils.read_param(\"endpoint_name\")\n",
    "customer_inputs_fg_name = utils.read_param(\"customer_inputs_fg_name\")\n",
    "destinations_fg_name = utils.read_param(\"destinations_fg_name\")\n",
    "customer_inputs_fg_query_string = utils.read_param(\"customer_inputs_fg_query_string\")\n",
    "query_string = utils.read_param(\"query_string\")\n",
    "training_job_name = utils.read_param(\"training_job_name\")\n",
    "logger.info(f\"endpoint_name={endpoint_name}, customer_inputs_fg_name={customer_inputs_fg_name},\\n\"\n",
    "            f\"customer_inputs_fg_query_string={customer_inputs_fg_query_string}, training_job_name={training_job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4416fd26",
   "metadata": {},
   "source": [
    "## Setup lineage tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2d4d979",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model name is same as endpoint name in this example\n",
    "ml_lineage = MLLineageHelper()\n",
    "lineage = ml_lineage.create_ml_lineage(training_job_name, \n",
    "                                       model_name=endpoint_name,\n",
    "                                       query=customer_inputs_fg_query_string,\n",
    "                                       feature_group_names=[customer_inputs_fg_name, destinations_fg_name], \n",
    "                                       sagemaker_processing_job_description=None\n",
    "                                      )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "717eb842",
   "metadata": {},
   "source": [
    "## Lineage Information\n",
    "Lineage information is provided both in tabular form and graphical form as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eb3a22b",
   "metadata": {},
   "outputs": [],
   "source": [
    "lineage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bddd31b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visual Representation of the ML Lineage\n",
    "ml_lineage.graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7637f95",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}