simple_versioned_workflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bd97604c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kfp import dsl, components\n",
    "from disdat_kfp.caching_wrapper import Caching\n",
    "import kfp\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56e1d381",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_IMAGE = 'python:3.8.11-slim' # feel free to use other python images \n",
    "YOUR_S3_BUCKET = YOUR_BUCKET # for instance s3://hello-bucket. Make sure your EKS cluster has access to this bucket!\n",
    "PIPELINE_NAME = 'simple_versioned_workflow'\n",
    "SERVICE_ACCOUNT = YOUR_ROLE # for instance 'mlp-pipelines'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6448e9da",
   "metadata": {},
   "source": [
    "### User ContainerOp 1\n",
    "This component produces a file and return some numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fdd2122d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def file_producer(char_num: int, data: components.OutputTextFile(str)):\n",
    "    data.write('*' * char_num)\n",
    "# create a container op from pure python function \n",
    "# KFP compiler will generate the YAML for you. \n",
    "producer_op = components.create_component_from_func(file_producer,\n",
    "                                                    base_image=BASE_IMAGE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a267e20f",
   "metadata": {},
   "source": [
    "### User ContainerOp 2\n",
    "This component consumes the file produced by container 1 and count the number of chars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a8987d24",
   "metadata": {},
   "outputs": [],
   "source": [
    "def file_consumer(data_handle: components.InputTextFile(str)) -> int:\n",
    "    data = data_handle.read()\n",
    "    return len(data)\n",
    "consumer_op = components.create_component_from_func(file_consumer,\n",
    "                                                    base_image=BASE_IMAGE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0534a411",
   "metadata": {},
   "source": [
    "### User Container 3\n",
    "This is a random container to show you how enable_caching work with force dependencies "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2e595166",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hello(msg: str='123') -> str:\n",
    "    print('hello')\n",
    "    return msg\n",
    "\n",
    "hello_op = components.create_component_from_func(hello,\n",
    "                                                    base_image=BASE_IMAGE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffe972bd",
   "metadata": {},
   "source": [
    "### built workflow with disdat-kfp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "328013e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "@dsl.pipeline(\n",
    "    name=PIPELINE_NAME,\n",
    "    description=\"file passing pipeline\"\n",
    ")\n",
    "def pipeline():\n",
    "    caching = Caching(disdat_context=PIPELINE_NAME,\n",
    "                      disdat_repo_s3_url='s3://' + YOUR_S3_BUCKET, \n",
    "                      force_rerun_pipeline=False, \n",
    "                      use_verbose=True)\n",
    "    \n",
    "    \n",
    "    hello_msg = hello_op() \n",
    "    \n",
    "    # enable caching works like a decorator \n",
    "    cached_producer = caching.enable_caching(producer_op, \n",
    "                                             char_num=1000, \n",
    "                                             # if False to enable cacaching\n",
    "                                             _disdat_force_rerun=False, \n",
    "                                             # your own bundle name\n",
    "                                             _disdat_bundle='file_producer_bundle', \n",
    "                                             # you can force dependency like this \n",
    "                                             _after=[hello_msg])\n",
    "    \n",
    "    # pass outputs like you normally would, cached_producer will have the same \n",
    "    # output signature as file_producer\n",
    "    consumer_op(cached_producer.outputs['data'])\n",
    "    \n",
    "    # you can have two cascading cached components as well\n",
    "    cached_consumer = caching.enable_caching(consumer_op, \n",
    "                                             data_handle=cached_producer.outputs['data'], \n",
    "                                             _disdat_force_rerun=False, \n",
    "                                             _disdat_bundle='file_consumer_bundle')\n",
    "    # note that we have a cached consumer and a normal consumer. \n",
    "    # This is to show you that you can treat cached_container_op as a normal component"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8dd2c35",
   "metadata": {},
   "source": [
    "### Compile the workflow and extract the YAML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "aa8d3c0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "kfp.compiler.Compiler().compile(pipeline, 'pipeline.yaml')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93ba0d94",
   "metadata": {},
   "source": [
    "### Run the workflow \n",
    "I use Argo to execute KFP, you are free to choose other platforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6698ac06",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.system('argo submit --serviceaccount {} pipeline.yaml'.format(SERVICE_ACCOUNT))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13c82a68",
   "metadata": {},
   "source": [
    "### Results\n",
    "Below are images captured by Argo GUI, which is the default KFP platform at Intuit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab89c722",
   "metadata": {},
   "source": [
    "First execution:\n",
    "![run_1](docs/run_1.png)\n",
    "\n",
    "\n",
    "Second execution:\n",
    "![run_2](docs/run_2.png)\n",
    "\n",
    "If you rerun the pipeline with exactly the same parameters, file_producer and consumer will be skipped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d176b674",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c0abed",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}