04_ML_Ops.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ML Ops\n",
    "In this section of the demo, we will utilize Snowpark's Python client-side Dataframe API as well as the Snowpark server-side runtime to create an **ML ops pipeline**.  We will take the functions created by the Data Scientist and ML Engineer and create a set of functions that can be easily automated with the company's orchestration tools. \n",
    "\n",
    "The ML Engineer must create a pipeline to **automate deployment** of models and batch predictions where the business users can consume them easily from an ML-powered front-end application like Streamlit.  The predictions must be accompanied by an explanation of which features were most impactful for the prediction.  \n",
    "\n",
    "Most importantly no data should leave Snowflake.  The entire end-to-end workflow should be pushed-down to run where the data sits."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Input: Data in `trips` table.  Feature engineering, train, predict functions from data scientist.  \n",
    "Output: Automatable pipeline of feature engineering, train, predict."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Load  credentials and connect to Snowflake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dags.snowpark_connection import snowpark_connect\n",
    "session, state_dict = snowpark_connect('./include/state.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Setup Training and Inference Pipeline\n",
    "\n",
    "We will generate a unique identifier which we will use to provide lineage across all components of the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from snowflake.snowpark import functions as F\n",
    "import uuid\n",
    "run_date='2020_01_01'\n",
    "model_id = str(uuid.uuid1()).replace('-', '_')\n",
    "\n",
    "state_dict.update({'model_id': model_id})\n",
    "state_dict.update({'run_date': run_date})\n",
    "state_dict.update({'weather_database_name': 'WEATHER_NYC'})\n",
    "state_dict.update({'trips_table_name': 'TRIPS',\n",
    "                   'holiday_table_name': 'HOLIDAYS',\n",
    "                   'weather_table_name': state_dict['weather_database_name']+'.ONPOINT_ID.HISTORY_DAY',\n",
    "                   'weather_view_name': 'WEATHER_NYC_VW',\n",
    "                   'feature_table_name' : 'FEATURES_'+model_id,\n",
    "                   'pred_table_name': 'PREDS_'+model_id,\n",
    "                   'eval_table_name': 'EVALS_'+model_id,\n",
    "                   'forecast_table_name': 'FORECAST_'+model_id,\n",
    "                   'forecast_steps': 30,\n",
    "                   'train_udf_name': 'station_train_predict_udtf',\n",
    "                   'train_func_name': 'StationTrainPredictFunc',\n",
    "                   'eval_udf_name': 'eval_model_output_udf',\n",
    "                   'eval_func_name': 'eval_model_func',\n",
    "                   'model_stage_name': 'MODEL_STAGE',\n",
    "                  })\n",
    "\n",
    "import json\n",
    "with open('./include/state.json', 'w') as sdf:\n",
    "    json.dump(state_dict, sdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will deploy the model training and inference as a permanent [Python Snowpark User-Defined Table Function (UDTF)](https://docs.snowflake.com/en/developer-guide/udf/python/udf-python-tabular-functions.html). This will make the function available to not only our automated training/inference pipeline but also to any users needing the function for manually generated predictions.  \n",
    "  \n",
    "As a permanent function we will need a staging area."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql('CREATE STAGE IF NOT EXISTS ' + state_dict['model_stage_name']).collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For production we need to be able to reproduce results.  The `trips` table will change as new data is loaded each month so we need a point-in-time snapshot.  Snowflake [Zero-Copy Cloning](https://docs.snowflake.com/en/sql-reference/sql/create-clone.html) allows us to do this with copy-on-write features so we don't have multiple copies of the same data.  We will create a unique ID to identify each training/inference run as well as the features and predictions generated.  We can use [object tagging](https://docs.snowflake.com/en/user-guide/object-tagging.html) to tag each object with the `model_id` as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clone_table_name = 'TRIPS_CLONE_'+state_dict[\"run_date\"]\n",
    "state_dict.update({\"clone_table_name\":clone_table_name})\n",
    "\n",
    "_ = session.sql('CREATE OR REPLACE TABLE '+clone_table_name+\" CLONE \"+state_dict[\"trips_table_name\"]).collect()\n",
    "_ = session.sql('CREATE TAG IF NOT EXISTS model_id_tag').collect()\n",
    "_ = session.sql(\"ALTER TABLE \"+clone_table_name+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will start by importing the functions created by the ML Engineer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dags.mlops_pipeline import materialize_holiday_table\n",
    "from dags.mlops_pipeline import subscribe_to_weather_data\n",
    "from dags.mlops_pipeline import create_weather_view\n",
    "from dags.mlops_pipeline import deploy_pred_train_udf\n",
    "from dags.mlops_pipeline import deploy_eval_udf\n",
    "from dags.mlops_pipeline import create_forecast_table\n",
    "from dags.mlops_pipeline import create_feature_table\n",
    "from dags.mlops_pipeline import train_predict\n",
    "from dags.mlops_pipeline import evaluate_station_model\n",
    "from dags.mlops_pipeline import flatten_tables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pipeline will be orchestrated by our companies orchestration framework but we will test the steps here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "holiday_table_name = materialize_holiday_table(session=session, \n",
    "                                               holiday_table_name=state_dict['holiday_table_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weather_database_name = subscribe_to_weather_data(session=session, \n",
    "                                                  weather_database_name=state_dict['weather_database_name'], \n",
    "                                                  weather_listing_id=state_dict['weather_listing_id'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weather_view_name = create_weather_view(session, \n",
    "                                        weather_table_name=state_dict['weather_table_name'],\n",
    "                                        weather_view_name=state_dict['weather_view_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_udf_name = deploy_pred_train_udf(session=session, \n",
    "                                       udf_name=state_dict['train_udf_name'],\n",
    "                                       function_name=state_dict['train_func_name'],\n",
    "                                       model_stage_name=state_dict['model_stage_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_udf_name = deploy_eval_udf(session=session, \n",
    "                                udf_name=state_dict['eval_udf_name'],\n",
    "                                function_name=state_dict['eval_func_name'],\n",
    "                                model_stage_name=state_dict['model_stage_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_table_name = create_feature_table(session, \n",
    "                                          trips_table_name=state_dict['clone_table_name'], \n",
    "                                          holiday_table_name=state_dict['holiday_table_name'], \n",
    "                                          weather_view_name=state_dict['weather_view_name'],\n",
    "                                          feature_table_name=state_dict['feature_table_name'])\n",
    "\n",
    "_ = session.sql(\"ALTER TABLE \"+feature_table_name+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "forecast_table_name = create_forecast_table(session, \n",
    "                                            trips_table_name=state_dict['trips_table_name'],\n",
    "                                            holiday_table_name=state_dict['holiday_table_name'], \n",
    "                                            weather_view_name=state_dict['weather_view_name'], \n",
    "                                            forecast_table_name=state_dict['forecast_table_name'],\n",
    "                                            steps=state_dict['forecast_steps'])\n",
    "\n",
    "_ = session.sql(\"ALTER TABLE \"+forecast_table_name+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.use_warehouse(state_dict['compute_parameters']['train_warehouse'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_table_name = train_predict(session, \n",
    "                                station_train_pred_udf_name=state_dict['train_udf_name'], \n",
    "                                feature_table_name=state_dict['feature_table_name'], \n",
    "                                forecast_table_name=state_dict['forecast_table_name'],\n",
    "                                pred_table_name=state_dict['pred_table_name'])\n",
    "\n",
    "_ = session.sql(\"ALTER TABLE \"+pred_table_name+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#session.sql('ALTER WAREHOUSE IF EXISTS '+state_dict['compute_parameters']['train_warehouse']+' SUSPEND').collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.use_warehouse(state_dict['compute_parameters']['default_warehouse'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_table_name = evaluate_station_model(session, \n",
    "                                         run_date=state_dict['run_date'], \n",
    "                                         eval_model_udf_name=state_dict['eval_udf_name'], \n",
    "                                         pred_table_name=state_dict['pred_table_name'], \n",
    "                                         eval_table_name=state_dict['eval_table_name'])\n",
    "\n",
    "_ = session.sql(\"ALTER TABLE \"+eval_table_name+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flat_pred_table, flat_forecast_table, flat_eval_table = flatten_tables(session, \n",
    "                                                                       pred_table_name, \n",
    "                                                                       forecast_table_name, \n",
    "                                                                       eval_table_name)\n",
    "state_dict['flat_pred_table'] = flat_pred_table\n",
    "state_dict['flat_forecast_table'] = flat_forecast_table\n",
    "state_dict['flat_eval_table'] = flat_eval_table\n",
    "\n",
    "_ = session.sql(\"ALTER TABLE \"+flat_pred_table+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()\n",
    "_ = session.sql(\"ALTER TABLE \"+flat_forecast_table+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()\n",
    "_ = session.sql(\"ALTER TABLE \"+flat_eval_table+\" SET TAG model_id_tag = '\"+state_dict[\"model_id\"]+\"'\").collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Consolidate Ingest, Training, Inference and Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile dags/mlops_tasks.py\n",
    "\n",
    "def snowpark_database_setup(state_dict:dict)-> dict: \n",
    "    import snowflake.snowpark.functions as F\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.elt import reset_database\n",
    "\n",
    "    session, _ = snowpark_connect('./include/state.json')\n",
    "    reset_database(session=session, state_dict=state_dict, prestaged=True)\n",
    "\n",
    "    _ = session.sql('CREATE STAGE '+state_dict['model_stage_name']).collect()\n",
    "    _ = session.sql('CREATE TAG model_id_tag').collect()\n",
    "\n",
    "    session.close()\n",
    "\n",
    "    return state_dict\n",
    "\n",
    "def incremental_elt_task(state_dict: dict, files_to_download:list)-> dict:\n",
    "    from dags.ingest import incremental_elt\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    print('Ingesting '+str(files_to_download))\n",
    "    download_base_url=state_dict['connection_parameters']['download_base_url']\n",
    "\n",
    "    _ = session.use_warehouse(state_dict['compute_parameters']['load_warehouse'])\n",
    "\n",
    "    _ = incremental_elt(session=session, \n",
    "                        state_dict=state_dict, \n",
    "                        files_to_ingest=files_to_download,\n",
    "                        download_base_url=download_base_url,\n",
    "                        use_prestaged=True)\n",
    "\n",
    "    #_ = session.sql('ALTER WAREHOUSE IF EXISTS '+state_dict['compute_parameters']['load_warehouse']+\\\n",
    "    #                ' SUSPEND').collect()\n",
    "\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def initial_bulk_load_task(state_dict:dict)-> dict:\n",
    "    from dags.ingest import bulk_elt\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    _ = session.use_warehouse(state_dict['compute_parameters']['load_warehouse'])\n",
    "\n",
    "    print('Running initial bulk ingest from '+state_dict['connection_parameters']['download_base_url'])\n",
    "    \n",
    "    _ = bulk_elt(session=session, \n",
    "                 state_dict=state_dict, \n",
    "                 download_base_url=state_dict['connection_parameters']['download_base_url'],\n",
    "                 use_prestaged=True)\n",
    "\n",
    "    #_ = session.sql('ALTER WAREHOUSE IF EXISTS '+state_dict['compute_parameters']['load_warehouse']+\\\n",
    "    #                ' SUSPEND').collect()\n",
    "\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def materialize_holiday_task(state_dict: dict)-> dict:\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import materialize_holiday_table\n",
    "\n",
    "    print('Materializing holiday table.')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    _ = materialize_holiday_table(session=session, \n",
    "                                  holiday_table_name=state_dict['holiday_table_name'])\n",
    "\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def subscribe_to_weather_data_task(state_dict: dict)-> dict:\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import subscribe_to_weather_data\n",
    "\n",
    "    print('Subscribing to weather data')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    _ = subscribe_to_weather_data(session=session, \n",
    "                                  weather_database_name=state_dict['weather_database_name'], \n",
    "                                  weather_listing_id=state_dict['weather_listing_id'])\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def create_weather_view_task(state_dict: dict)-> dict:\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import create_weather_view\n",
    "\n",
    "    print('Creating weather view')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    _ = create_weather_view(session=session,\n",
    "                            weather_table_name=state_dict['weather_table_name'],\n",
    "                            weather_view_name=state_dict['weather_view_name'])\n",
    "    session.close()\n",
    "    return state_dict\n",
    "    \n",
    "def deploy_model_udf_task(state_dict:dict)-> dict:\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import deploy_pred_train_udf\n",
    "\n",
    "    print('Deploying station model')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    _ = session.sql('CREATE STAGE IF NOT EXISTS ' + state_dict['model_stage_name']).collect()\n",
    "\n",
    "    _ = deploy_pred_train_udf(session=session, \n",
    "                              udf_name=state_dict['train_udf_name'],\n",
    "                              function_name=state_dict['train_func_name'],\n",
    "                              model_stage_name=state_dict['model_stage_name'])\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def deploy_eval_udf_task(state_dict:dict)-> dict:\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import deploy_eval_udf\n",
    "\n",
    "    print('Deploying station model')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    _ = session.sql('CREATE STAGE IF NOT EXISTS ' + state_dict['model_stage_name']).collect()\n",
    "\n",
    "    _ = deploy_eval_udf(session=session, \n",
    "                        udf_name=state_dict['eval_udf_name'],\n",
    "                        function_name=state_dict['eval_func_name'],\n",
    "                        model_stage_name=state_dict['model_stage_name'])\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def generate_feature_table_task(state_dict:dict, \n",
    "                                holiday_state_dict:dict, \n",
    "                                weather_state_dict:dict)-> dict:\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import create_feature_table\n",
    "\n",
    "    print('Generating features for all stations.')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    session.use_warehouse(state_dict['compute_parameters']['fe_warehouse'])\n",
    "\n",
    "    _ = session.sql(\"CREATE OR REPLACE TABLE \"+state_dict['clone_table_name']+\\\n",
    "                    \" CLONE \"+state_dict['trips_table_name']).collect()\n",
    "    _ = session.sql(\"ALTER TABLE \"+state_dict['clone_table_name']+\\\n",
    "                    \" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "\n",
    "    _ = create_feature_table(session, \n",
    "                             trips_table_name=state_dict['clone_table_name'], \n",
    "                             holiday_table_name=state_dict['holiday_table_name'], \n",
    "                             weather_view_name=state_dict['weather_view_name'],\n",
    "                             feature_table_name=state_dict['feature_table_name'])\n",
    "\n",
    "    _ = session.sql(\"ALTER TABLE \"+state_dict['feature_table_name']+\\\n",
    "                    \" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def generate_forecast_table_task(state_dict:dict, \n",
    "                                 holiday_state_dict:dict, \n",
    "                                 weather_state_dict:dict)-> dict: \n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import create_forecast_table\n",
    "\n",
    "    print('Generating forecast features.')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    _ = create_forecast_table(session, \n",
    "                              trips_table_name=state_dict['trips_table_name'],\n",
    "                              holiday_table_name=state_dict['holiday_table_name'], \n",
    "                              weather_view_name=state_dict['weather_view_name'], \n",
    "                              forecast_table_name=state_dict['forecast_table_name'],\n",
    "                              steps=state_dict['forecast_steps'])\n",
    "\n",
    "    _ = session.sql(\"ALTER TABLE \"+state_dict['forecast_table_name']+\\\n",
    "                    \" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def bulk_train_predict_task(state_dict:dict, \n",
    "                            feature_state_dict:dict, \n",
    "                            forecast_state_dict:dict)-> dict: \n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import train_predict\n",
    "\n",
    "    state_dict = feature_state_dict\n",
    "\n",
    "    print('Running bulk training and forecast.')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    session.use_warehouse(state_dict['compute_parameters']['train_warehouse'])\n",
    "\n",
    "    pred_table_name = train_predict(session, \n",
    "                                    station_train_pred_udf_name=state_dict['train_udf_name'], \n",
    "                                    feature_table_name=state_dict['feature_table_name'], \n",
    "                                    forecast_table_name=state_dict['forecast_table_name'],\n",
    "                                    pred_table_name=state_dict['pred_table_name'])\n",
    "\n",
    "    _ = session.sql(\"ALTER TABLE \"+state_dict['pred_table_name']+\\\n",
    "                    \" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "    #_ = session.sql('ALTER WAREHOUSE IF EXISTS '+state_dict['compute_parameters']['train_warehouse']+\\\n",
    "    #                ' SUSPEND').collect()\n",
    "\n",
    "    session.close()\n",
    "    return state_dict\n",
    "\n",
    "def eval_station_models_task(state_dict:dict, \n",
    "                             pred_state_dict:dict,\n",
    "                             run_date:str)-> dict:\n",
    "\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import evaluate_station_model\n",
    "\n",
    "    print('Running eval UDF for model output')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    eval_table_name = evaluate_station_model(session, \n",
    "                                             run_date=run_date, \n",
    "                                             eval_model_udf_name=state_dict['eval_udf_name'], \n",
    "                                             pred_table_name=state_dict['pred_table_name'], \n",
    "                                             eval_table_name=state_dict['eval_table_name'])\n",
    "\n",
    "    _ = session.sql(\"ALTER TABLE \"+state_dict['eval_table_name']+\\\n",
    "                    \" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "    session.close()\n",
    "    return state_dict                                               \n",
    "\n",
    "def flatten_tables_task(pred_state_dict:dict, state_dict:dict)-> dict:\n",
    "    from dags.snowpark_connection import snowpark_connect\n",
    "    from dags.mlops_pipeline import flatten_tables\n",
    "\n",
    "    print('Flattening tables for end-user consumption.')\n",
    "    session, _ = snowpark_connect()\n",
    "\n",
    "    flat_pred_table, flat_forecast_table, flat_eval_table = flatten_tables(session,\n",
    "                                                                           pred_table_name=state_dict['pred_table_name'], \n",
    "                                                                           forecast_table_name=state_dict['forecast_table_name'], \n",
    "                                                                           eval_table_name=state_dict['eval_table_name'])\n",
    "    state_dict['flat_pred_table'] = flat_pred_table\n",
    "    state_dict['flat_forecast_table'] = flat_forecast_table\n",
    "    state_dict['flat_eval_table'] = flat_eval_table\n",
    "\n",
    "    _ = session.sql(\"ALTER TABLE \"+flat_pred_table+\" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "    _ = session.sql(\"ALTER TABLE \"+flat_forecast_table+\" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "    _ = session.sql(\"ALTER TABLE \"+flat_eval_table+\" SET TAG model_id_tag = '\"+state_dict['model_id']+\"'\").collect()\n",
    "\n",
    "    return state_dict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile dags/mlops_setup_pipeline.py\n",
    "from dags.mlops_tasks import snowpark_database_setup\n",
    "from dags.mlops_tasks import initial_bulk_load_task\n",
    "from dags.mlops_tasks import materialize_holiday_task\n",
    "from dags.mlops_tasks import subscribe_to_weather_data_task\n",
    "from dags.mlops_tasks import create_weather_view_task\n",
    "from dags.mlops_tasks import deploy_model_udf_task\n",
    "from dags.mlops_tasks import deploy_eval_udf_task\n",
    "from dags.mlops_tasks import generate_feature_table_task\n",
    "from dags.mlops_tasks import generate_forecast_table_task\n",
    "from dags.mlops_tasks import bulk_train_predict_task\n",
    "from dags.mlops_tasks import eval_station_models_task \n",
    "from dags.mlops_tasks import flatten_tables_task\n",
    "\n",
    "def citibikeml_setup_taskflow(run_date:str):\n",
    "    \"\"\"\n",
    "    End to end Snowflake ML Demo\n",
    "    \"\"\"\n",
    "    import uuid\n",
    "    import json\n",
    "\n",
    "    with open('./include/state.json') as sdf:\n",
    "        state_dict = json.load(sdf)\n",
    "    \n",
    "    model_id = str(uuid.uuid1()).replace('-', '_')\n",
    "\n",
    "    state_dict.update({'model_id': model_id})\n",
    "    state_dict.update({'run_date': run_date})\n",
    "    state_dict.update({'weather_database_name': 'WEATHER_NYC'})\n",
    "    state_dict.update({'load_table_name': 'RAW_',\n",
    "                       'trips_table_name': 'TRIPS',\n",
    "                       'load_stage_name': 'LOAD_STAGE',\n",
    "                       'model_stage_name': 'MODEL_STAGE',\n",
    "                       'weather_table_name': state_dict['weather_database_name']+'.ONPOINT_ID.HISTORY_DAY',\n",
    "                       'weather_view_name': 'WEATHER_NYC_VW',\n",
    "                       'holiday_table_name': 'HOLIDAYS',\n",
    "                       'clone_table_name': 'CLONE_'+model_id,\n",
    "                       'feature_table_name' : 'FEATURE_'+model_id,\n",
    "                       'pred_table_name': 'PRED_'+model_id,\n",
    "                       'eval_table_name': 'EVAL_'+model_id,\n",
    "                       'forecast_table_name': 'FORECAST_'+model_id,\n",
    "                       'forecast_steps': 30,\n",
    "                       'train_udf_name': 'station_train_predict_udtf',\n",
    "                       'train_func_name': 'StationTrainPredictFunc',\n",
    "                       'eval_udf_name': 'eval_model_output_udf',\n",
    "                       'eval_func_name': 'eval_model_func'\n",
    "                      })\n",
    "    \n",
    "    #Task order - one-time setup\n",
    "    setup_state_dict = snowpark_database_setup(state_dict)\n",
    "    load_state_dict = initial_bulk_load_task(setup_state_dict)\n",
    "    holiday_state_dict = materialize_holiday_task(setup_state_dict)\n",
    "    subscribe_state_dict = subscribe_to_weather_data_task(setup_state_dict)\n",
    "    weather_state_dict = create_weather_view_task(subscribe_state_dict)\n",
    "    model_udf_state_dict = deploy_model_udf_task(setup_state_dict)\n",
    "    eval_udf_state_dict = deploy_eval_udf_task(setup_state_dict)\n",
    "    feature_state_dict = generate_feature_table_task(load_state_dict, holiday_state_dict, weather_state_dict) \n",
    "    foecast_state_dict = generate_forecast_table_task(load_state_dict, holiday_state_dict, weather_state_dict)\n",
    "    pred_state_dict = bulk_train_predict_task(model_udf_state_dict, feature_state_dict, foecast_state_dict)\n",
    "    eval_state_dict = eval_station_models_task(eval_udf_state_dict, pred_state_dict, run_date)  \n",
    "    state_dict = flatten_tables_task(pred_state_dict, eval_state_dict)\n",
    "\n",
    "    return state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile dags/mlops_monthly_pipeline.py\n",
    "\n",
    "from dags.mlops_tasks import incremental_elt_task\n",
    "from dags.mlops_tasks import generate_feature_table_task\n",
    "from dags.mlops_tasks import generate_forecast_table_task\n",
    "from dags.mlops_tasks import bulk_train_predict_task\n",
    "from dags.mlops_tasks import eval_station_models_task \n",
    "from dags.mlops_tasks import flatten_tables_task\n",
    "\n",
    "def citibikeml_monthly_taskflow(files_to_download:list, run_date:str):\n",
    "    \"\"\"\n",
    "    End to end Snowflake ML Demo\n",
    "    \"\"\"\n",
    "    import uuid\n",
    "    import json\n",
    "\n",
    "    with open('./include/state.json') as sdf:\n",
    "        state_dict = json.load(sdf)\n",
    "    \n",
    "    model_id = str(uuid.uuid1()).replace('-', '_')\n",
    "\n",
    "    state_dict.update({'model_id': model_id})\n",
    "    state_dict.update({'run_date': run_date})\n",
    "    state_dict.update({'weather_database_name': 'WEATHER_NYC'})\n",
    "    state_dict.update({'load_table_name': 'RAW_',\n",
    "                       'trips_table_name': 'TRIPS',\n",
    "                       'load_stage_name': 'LOAD_STAGE',\n",
    "                       'model_stage_name': 'MODEL_STAGE',\n",
    "                       'weather_table_name': state_dict['weather_database_name']+'.ONPOINT_ID.HISTORY_DAY',\n",
    "                       'weather_view_name': 'WEATHER_NYC_VW',\n",
    "                       'holiday_table_name': 'HOLIDAYS',\n",
    "                       'clone_table_name': 'CLONE_'+model_id,\n",
    "                       'feature_table_name' : 'FEATURE_'+model_id,\n",
    "                       'pred_table_name': 'PRED_'+model_id,\n",
    "                       'eval_table_name': 'EVAL_'+model_id,\n",
    "                       'forecast_table_name': 'FORECAST_'+model_id,\n",
    "                       'forecast_steps': 30,\n",
    "                       'train_udf_name': 'station_train_predict_udtf',\n",
    "                       'train_func_name': 'StationTrainPredictFunc',\n",
    "                       'eval_udf_name': 'eval_model_output_udf',\n",
    "                       'eval_func_name': 'eval_model_func'\n",
    "                      })\n",
    "    \n",
    "    #Task order - monthlyl incremental\n",
    "    incr_state_dict = incremental_elt_task(state_dict, files_to_download)\n",
    "    feature_state_dict = generate_feature_table_task(incr_state_dict, incr_state_dict, incr_state_dict) \n",
    "    forecast_state_dict = generate_forecast_table_task(incr_state_dict, incr_state_dict, incr_state_dict)\n",
    "    pred_state_dict = bulk_train_predict_task(feature_state_dict, feature_state_dict, forecast_state_dict)\n",
    "    eval_state_dict = eval_station_models_task(pred_state_dict, pred_state_dict, run_date)\n",
    "    state_dict = flatten_tables_task(pred_state_dict, eval_state_dict)\n",
    "\n",
    "    return state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dags.snowpark_connection import snowpark_connect\n",
    "session, state_dict = snowpark_connect('./include/state.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from dags.mlops_setup_pipeline import citibikeml_setup_taskflow\n",
    "\n",
    "state_dict = citibikeml_setup_taskflow(run_date='2020_01_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from dags.mlops_monthly_pipeline import citibikeml_monthly_taskflow\n",
    "\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202001-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_02_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202002-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_03_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202003-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_04_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202004-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_05_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202005-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_06_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202006-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_07_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202007-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_08_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "state_dict = citibikeml_monthly_taskflow(files_to_download = ['202008-citibike-tripdata.csv.zip'], \n",
    "                                         run_date='2020_09_01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "authors": [
   {
    "name": "cforbe"
   }
  ],
  "kernelspec": {
   "display_name": "snowpark_0110:Python",
   "language": "python",
   "name": "conda-env-snowpark_0110-py"
  },
  "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.8.13"
  },
  "msauthor": "trbye"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}