build_workflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "import sys\n",
    "!{sys.executable} -m pip install --upgrade stepfunctions\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "import boto3\n",
    "from botocore.client import Config\n",
    "import logging\n",
    "import sagemaker\n",
    "import stepfunctions\n",
    "\n",
    "from sagemaker.amazon.amazon_estimator import get_image_uri\n",
    "from sagemaker.estimator import Estimator\n",
    "from sagemaker.sklearn.estimator import SKLearn\n",
    "from sagemaker.s3 import S3Uploader\n",
    "from sagemaker.transformer import Transformer\n",
    "\n",
    "from stepfunctions.inputs import ExecutionInput\n",
    "from stepfunctions import steps\n",
    "from stepfunctions.steps import Parallel\n",
    "from stepfunctions.steps.sagemaker import TrainingStep, ModelStep, TransformStep\n",
    "from stepfunctions.steps.compute import LambdaStep\n",
    "from stepfunctions.workflow import Workflow\n",
    "\n",
    "import os\n",
    "import zipfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define boto3 clients\n",
    "s3_client = boto3.client('s3', config=Config(signature_version='s3v4'))\n",
    "lambda_client = boto3.client('lambda')\n",
    "resource_client = boto3.client('resourcegroupstaggingapi')\n",
    "\n",
    "# define catch-all execution role\n",
    "hbomax_datascience_service_role = 'arn:aws:iam::613630599026:role/hbomax-datascience-service-role'\n",
    "\n",
    "# set logging\n",
    "stepfunctions.set_stream_logger(level=logging.INFO)\n",
    "\n",
    "# collect session info\n",
    "region = boto3.Session().region_name\n",
    "acount_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "\n",
    "# define s3 bucket\n",
    "resources_bucket = 'hbomax-datascience-development-dev'\n",
    "\n",
    "# sagemaker session\n",
    "sagemaker_session = sagemaker.Session(default_bucket=resources_bucket)\n",
    "\n",
    "# name the Stepfunctions pipeline\n",
    "pipeline_name = 'FTInferenceRoutine'\n",
    "\n",
    "# XGBoost image\n",
    "#xgboost_image = get_image_uri(region, 'xgboost', repo_version='latest')\n",
    "xgboost_image = '613630599026.dkr.ecr.us-east-1.amazonaws.com/temp-backgrounddata/hbomax-xgboost-shap:1.0-1-cpu-py3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the data sources\n",
    "    # raw data\n",
    "s3_train_raw = sagemaker.s3_input(s3_data=f's3://{resources_bucket}/free_trial_model/snowflake-hbomax-staging/train/raw', content_type='text/csv')\n",
    "s3_test_raw = sagemaker.s3_input(s3_data=f's3://{resources_bucket}/free_trial_model/snowflake-hbomax-staging/test/raw', content_type='text/csv')\n",
    "s3_new_raw = sagemaker.s3_input(s3_data=f's3://{resources_bucket}/lifecycle/free-trial-propensity-model/inference/raw', content_type='text/csv')\n",
    "\n",
    "    # transformed data\n",
    "s3_train_transformed = sagemaker.s3_input(s3_data=f's3://{resources_bucket}/free_trial_model/snowflake-hbomax-staging/train/transformed', content_type='text/csv')\n",
    "s3_test_transformed = sagemaker.s3_input(s3_data=f's3://{resources_bucket}/free_trial_model/snowflake-hbomax-staging/test/transformed', content_type='text/csv')\n",
    "s3_new_transformed = sagemaker.s3_input(s3_data=f's3://{resources_bucket}/free_trial_model/snowflake-hbomax-staging/new/transformed', content_type='text/csv')\n",
    "\n",
    "s3_out_data = f's3://{resources_bucket}/free_trial_model/snowflake-hbomax-staging/model_output'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def writeLambda(function_name, role, description):\n",
    "    zip_name = f'{function_name}.zip'\n",
    "    lambda_source_code = f'lambda_code/{function_name}.py'\n",
    "\n",
    "    zf = zipfile.ZipFile(zip_name, mode='w')\n",
    "    zf.write(lambda_source_code, arcname=lambda_source_code.split('/')[-1])\n",
    "    zf.close()\n",
    "\n",
    "    #s3_client.copy_object(zip_name, Bucket='datascience-hbo-users', zip_name) # ExtraArgs={\"ServerSideEncryption\": \"aws:kms\"},\n",
    "    \n",
    "    S3Uploader.upload(local_path=zip_name, \n",
    "                      desired_s3_uri=f's3://datascience-hbo-users/lambda_code', ## UPDATE!!\n",
    "                      #kms_key='alias/aws/s3',\n",
    "                      #kms_key='aws:kms',\n",
    "                      session=sagemaker_session)\n",
    "\n",
    "    lambda_client = boto3.client('lambda')\n",
    "\n",
    "    # delete the existing function if necessary\n",
    "    response = resource_client.get_resources(\n",
    "    TagFilters=[\n",
    "                    {\n",
    "                        'Key': 'function_name',\n",
    "                        'Values': [\n",
    "                            function_name\n",
    "                        ]\n",
    "                    },\n",
    "                ]\n",
    "            )\n",
    "    \n",
    "    if len(response['ResourceTagMappingList']) > 0:\n",
    "        lambda_client.delete_function(FunctionName=function_name)\n",
    "\n",
    "    # create the function\n",
    "    response = lambda_client.create_function(\n",
    "        FunctionName=function_name,\n",
    "        Runtime='python3.7',\n",
    "        Role=role,\n",
    "        Handler=f'{function_name}.lambda_handler',\n",
    "        Code={\n",
    "            'S3Bucket': 'datascience-hbo-users', ## UPDATE!!\n",
    "            'S3Key': 'lambda_code/{}'.format(zip_name)\n",
    "        },\n",
    "        Description=description,\n",
    "        Timeout=15,\n",
    "        MemorySize=128,\n",
    "        Tags={\n",
    "        'function_name': function_name\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # delete the zip archive\n",
    "    os.remove(zip_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for f in [{'name':'free_trial_transform_new_data', 'description':'Collect the most recent model and transform new data for inference.'},\n",
    "          {'name':'free_trial_query_sagemaker_job', 'description':'Query for status of the transform new data job.'},\n",
    "          {'name':'free_trial_xgboost_transform', 'description':'Query for status of the transform new data job.'},\n",
    "          {'name':'free_trial_collect_dayofweek', 'description':'Collect the current day of the week.'},\n",
    "          {'name':'free_trial_postprocess_shap', 'description':'Postprocess SHAP values.'},\n",
    "          {'name':'free_trial_execute_workflow', 'description':'Execute the StepFunctions state machine.'}\n",
    "         ]:\n",
    "    \n",
    "    writeLambda(function_name=f['name'], role=hbomax_datascience_service_role, description=f['description'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# associate event with S3 trigger location (console)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define runtime input.  SageMaker expects unique names for each job, model and endpoint. \n",
    "execution_input = ExecutionInput(schema={\n",
    "    'SKLearnFeaturizerJobName': str,\n",
    "    'TransformTrainJobName': str,\n",
    "    'TransformTestJobName': str,\n",
    "    'FeaturizerModelName': str,\n",
    "    'XGBModelName': str,\n",
    "    'TrainXGBoostJobName': str,\n",
    "    'PipelineModelName': str,\n",
    "    'TransformNewJobName': str,\n",
    "    'TransformXGBoostJobName': str,\n",
    "    'TimestampPrefix': str\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "day_of_week_step = LambdaStep(\n",
    "    'Collect Day of Week',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_collect_dayofweek'\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the SKLearn Preprocessing Estimator\n",
    "sklearn_featurizer = SKLearn(\n",
    "    source_dir='sklearn_featurizer',\n",
    "    entry_point='featurizer.py',\n",
    "    role=hbomax_datascience_service_role,\n",
    "    output_kms_key='alias/aws/s3',\n",
    "    train_instance_type=\"ml.c4.xlarge\",\n",
    "    hyperparameters = {'resources_bucket': resources_bucket},\n",
    "    sagemaker_session=sagemaker_session)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_featurizer_step = TrainingStep(\n",
    "    'Fit Featurizer', \n",
    "    estimator = sklearn_featurizer,\n",
    "    data={\n",
    "        'train': s3_train_raw.config['DataSource']['S3DataSource']['S3Uri'],\n",
    "    },\n",
    "    tags= {'model': 'free_trial_sklearn_featurizer'},\n",
    "    job_name=execution_input['SKLearnFeaturizerJobName']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_featurizer_model_step = ModelStep(\n",
    "    'Create Featurizer Model', \n",
    "    model = fit_featurizer_step.get_expected_model(),\n",
    "    tags= {'model': 'free_trial_sklearn_featurizer'},\n",
    "    model_name=execution_input['FeaturizerModelName']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transformer = Transformer(\n",
    "    model_name=execution_input['FeaturizerModelName'],\n",
    "    instance_count=1,\n",
    "    instance_type='ml.m4.2xlarge',\n",
    "    strategy='MultiRecord',\n",
    "    assemble_with='Line',\n",
    "    output_kms_key='alias/aws/s3',\n",
    "    accept='text/csv',\n",
    "    output_path=s3_train_transformed.config['DataSource']['S3DataSource']['S3Uri']\n",
    ")\n",
    "\n",
    "transform_train_step = TransformStep(\n",
    "    'Transform Training Data',\n",
    "    transformer=train_transformer,\n",
    "    job_name=execution_input['TransformTrainJobName'],\n",
    "    model_name=execution_input['FeaturizerModelName'],\n",
    "    data=s3_train_raw.config['DataSource']['S3DataSource']['S3Uri'],\n",
    "    content_type= 'text/csv',\n",
    "    split_type='Line',\n",
    "    wait_for_completion=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_transformer = Transformer(\n",
    "    model_name=execution_input['FeaturizerModelName'],\n",
    "    instance_count=1,\n",
    "    instance_type='ml.m4.2xlarge',\n",
    "    strategy='MultiRecord',\n",
    "    assemble_with='Line',\n",
    "    output_kms_key='alias/aws/s3',\n",
    "    accept='text/csv',\n",
    "    output_path=s3_test_transformed.config['DataSource']['S3DataSource']['S3Uri']\n",
    ")\n",
    "\n",
    "transform_test_step = TransformStep(\n",
    "    'Transform Test Data',\n",
    "    transformer=test_transformer,\n",
    "    job_name=execution_input['TransformTestJobName'],\n",
    "    model_name=execution_input['FeaturizerModelName'],\n",
    "    data=s3_test_raw.config['DataSource']['S3DataSource']['S3Uri'],\n",
    "    content_type= 'text/csv',\n",
    "    split_type='Line',\n",
    "    wait_for_completion=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parallel_transform_step = Parallel(\n",
    "    state_id=\"Branch Transformations\"\n",
    ")\n",
    "\n",
    "parallel_transform_step.add_branch(transform_test_step)\n",
    "parallel_transform_step.add_branch(transform_train_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the XGBoost Model Estimator\n",
    "xgboost_estimator = Estimator(image_name = xgboost_image,\n",
    "                          role = hbomax_datascience_service_role, \n",
    "                          train_instance_count = 1, \n",
    "                          train_instance_type='ml.m4.4xlarge',\n",
    "                          output_path = s3_out_data,\n",
    "                          output_kms_key = 'alias/aws/s3',\n",
    "                          metric_definitions = [{'Name': 'validation:auc', 'Regex': \"validation-auc:(.*?)$\"},\n",
    "                                                {'Name': 'train:auc', 'Regex': \"train-auc:(.*?)#\"}],\n",
    "                          hyperparameters = {\n",
    "                                             'eval_metric':'auc'\n",
    "                                            , 'alpha':1.218487609\n",
    "                                            , 'eta':0.225242353\n",
    "                                            , 'max_depth':10\n",
    "                                            , 'min_child_weight':2.284773815\n",
    "                                            , 'num_round':2\n",
    "                                            , 'objective':'binary:logistic'\n",
    "                                            , 'rate_drop':0.3\n",
    "                                            , 'tweedie_variance_power':1.4\n",
    "                                          },\n",
    "                          sagemaker_session=sagemaker_session)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_xgboost_step = TrainingStep(\n",
    "    'Train XGBoost', \n",
    "    estimator = xgboost_estimator,\n",
    "    data={\n",
    "        'train': s3_train_transformed,\n",
    "        'validation': s3_test_transformed\n",
    "    },\n",
    "    tags= {'model': 'free_trial_xgboost'},\n",
    "    job_name=execution_input['TrainXGBoostJobName']\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_xgboost_model_step = ModelStep(\n",
    "    'Create XGBoost Model', \n",
    "    model = train_xgboost_step.get_expected_model(),\n",
    "    model_name=execution_input['XGBModelName'],\n",
    "    tags= {'model': 'free_trial_xgboost'},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_train_xgboost_step = LambdaStep(\n",
    "    'Query Train New Data',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_query_sagemaker_job',\n",
    "        'Payload':{\n",
    "            'job_type': 'Train',\n",
    "            'job_name': execution_input['TrainXGBoostJobName']\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_new_data_step = LambdaStep(\n",
    "    'Transform New Data',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_transform_new_data',\n",
    "        'Payload':{\n",
    "            \"bucket\": resources_bucket,\n",
    "            'TransformJobName': execution_input['TransformNewJobName']\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add catch for xgboost training failure\n",
    "train_xgboost_step.add_catch(steps.states.Catch(\n",
    "    error_equals=['States.TaskFailed'],\n",
    "    next_step=transform_new_data_step\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_transform_new_data_step = LambdaStep(\n",
    "    'Query Transform New Data',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_query_sagemaker_job',\n",
    "        'Payload':{\n",
    "            'job_type': 'Transform',\n",
    "            'job_name': execution_input['TransformNewJobName']\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_new_data_complete_step = steps.states.Choice(\n",
    "    'Transform new data Complete?'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "succeed_state = steps.Succeed(\n",
    "    state_id=\"Success\"             \n",
    ")\n",
    "\n",
    "transform_new_data_fail_state = steps.Fail(\n",
    "    state_id=\"Transform New Data Fail\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_new_data_step = LambdaStep(\n",
    "    'Infer New Data',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_xgboost_transform',\n",
    "        'Payload':{\n",
    "            \"bucket\": resources_bucket,\n",
    "            'TransformJobName': execution_input['TransformXGBoostJobName']\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_new_data_wait_state = steps.Wait(\n",
    "    state_id=\"Wait for transform new data\",\n",
    "    seconds=60\n",
    ")\n",
    "\n",
    "transform_new_data_complete_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_transform_new_data_step.output()['Payload'][\"JobStatus\"], value='Completed'), \n",
    "    next_step=infer_new_data_step\n",
    ")\n",
    "\n",
    "transform_new_data_complete_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_transform_new_data_step.output()['Payload'][\"JobStatus\"], value='Failed'), \n",
    "    next_step=transform_new_data_fail_state\n",
    ")\n",
    "\n",
    "transform_new_data_complete_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_transform_new_data_step.output()['Payload'][\"JobStatus\"], value='InProgress'), \n",
    "    next_step=transform_new_data_wait_state\n",
    ")\n",
    "\n",
    "transform_new_data_wait_state.next(query_transform_new_data_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_infer_new_data_step = LambdaStep(\n",
    "    'Query Infer New Data',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_query_sagemaker_job',\n",
    "        'Payload':{\n",
    "            'job_type': 'Transform',\n",
    "            'job_name': execution_input['TransformXGBoostJobName']\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_new_data_complete_step = steps.states.Choice(\n",
    "    'Infer new data Complete?'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_new_data_fail_state = steps.Fail(\n",
    "    state_id=\"Infer New Data Fail\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "postprocess_step = LambdaStep(\n",
    "    'Postprocess SHAP',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_postprocess_shap',\n",
    "        'Payload':{\n",
    "            'bucket': resources_bucket\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_new_data_wait_state = steps.Wait(\n",
    "    state_id=\"Wait for infer new data\",\n",
    "    seconds=60\n",
    ")\n",
    "\n",
    "infer_new_data_complete_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_infer_new_data_step.output()['Payload'][\"JobStatus\"], value='Completed'), \n",
    "    next_step=postprocess_step\n",
    ")\n",
    "\n",
    "infer_new_data_complete_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_infer_new_data_step.output()['Payload'][\"JobStatus\"], value='Failed'), \n",
    "    next_step=infer_new_data_fail_state\n",
    ")\n",
    "\n",
    "infer_new_data_complete_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_infer_new_data_step.output()['Payload'][\"JobStatus\"], value='InProgress'), \n",
    "    next_step=infer_new_data_wait_state\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "infer_new_data_wait_state.next(query_infer_new_data_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "complete_retrain_step = steps.states.Choice(\n",
    "    'Retrain?'\n",
    ")\n",
    "\n",
    "complete_retrain_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=day_of_week_step.output()['Payload'][\"retrain\"], value='True'), \n",
    "    next_step=fit_featurizer_step\n",
    ")\n",
    "\n",
    "complete_retrain_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=day_of_week_step.output()['Payload'][\"retrain\"], value='False'), \n",
    "    next_step=transform_new_data_step\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_xgboost_auc_step = LambdaStep(\n",
    "    'Query New Model AUC',\n",
    "    parameters={  \n",
    "        \"FunctionName\": 'free_trial_query_sagemaker_job',\n",
    "        'Payload':{\n",
    "            'job_type': 'Train',\n",
    "            'job_name': execution_input['TrainXGBoostJobName']\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_validation_step = steps.states.Choice(\n",
    "    'New Model Meet Threshold?'\n",
    ")\n",
    "\n",
    "model_validation_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_xgboost_auc_step.output()['Payload']['JobStatus']['MeetsThreshold'], value='True'), \n",
    "    next_step=transform_new_data_step\n",
    ")\n",
    "\n",
    "model_validation_step.add_choice(\n",
    "    rule=steps.choice_rule.ChoiceRule.StringEquals(variable=query_xgboost_auc_step.output()['Payload']['JobStatus']['MeetsThreshold'], value='False'), \n",
    "    next_step=steps.Fail('Threshold Not Met')\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_definition = steps.Chain([day_of_week_step, complete_retrain_step])\n",
    "\n",
    "steps.Chain([fit_featurizer_step, create_featurizer_model_step, parallel_transform_step, train_xgboost_step, create_xgboost_model_step, query_xgboost_auc_step, model_validation_step])\n",
    "\n",
    "steps.Chain([transform_new_data_step, query_transform_new_data_step, transform_new_data_complete_step])\n",
    "\n",
    "steps.Chain([infer_new_data_step, query_infer_new_data_step, infer_new_data_complete_step])\n",
    "\n",
    "#workflow_definition = steps.Chain([postprocess_step]) \n",
    "\n",
    "workflow = Workflow(\n",
    "    name=pipeline_name,\n",
    "    definition=workflow_definition,\n",
    "    role=hbomax_datascience_service_role,\n",
    "    execution_input=execution_input\n",
    ")\n",
    "\n",
    "workflow.create()\n",
    "workflow.update(workflow_definition)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}