kubeflow pipeline conditional and parallel task

We can easily run ML pipeline using aiplatfor sdk using the following setup - notebook 

PROJECT_ID = "project-your-id"  # @param {type:"string"}

LOCATION = "us-central1"  # @param {type:"string"}

BUCKET_URI = f"gs://your-bucket-name-{PROJECT_ID}-unique"  # @param {type:"string"}

! gcloud storage buckets create --location={LOCATION} --project={PROJECT_ID} {BUCKET_URI}

SERVICE_ACCOUNT = "ml-pipeline-dev-sa"  # @param {type:"string"}

import sys

IS_COLAB = "google.colab" in sys.modules

if (

    SERVICE_ACCOUNT == ""

    or SERVICE_ACCOUNT is None

    or SERVICE_ACCOUNT == "ml-pipeline-dev-sa"

):

    # Get your service account from gcloud

    if not IS_COLAB:

        shell_output = !gcloud auth list 2>/dev/null

        SERVICE_ACCOUNT = shell_output[2].replace("*", "").strip()

    if IS_COLAB:

        shell_output = ! gcloud projects describe  $PROJECT_ID

        project_number = shell_output[-1].split(":")[1].strip().replace("'", "")

        SERVICE_ACCOUNT = f"{project_number}-compute@developer.gserviceaccount.com"

    print("Service Account:", SERVICE_ACCOUNT)

Adding IAM policy for our service account 

! gcloud storage buckets add-iam-policy-binding $BUCKET_URI --member="serviceAccount:{SERVICE_ACCOUNT}" --role="roles/storage.objectCreator"

! gcloud storage buckets

add-iam-policy-binding $BUCKET_URI --member=serviceAccount:{SERVICE_ACCOUNT}

--role=roles/storage.objectViewer

import json

from google.cloud import aiplatform

from kfp import compiler, dsl

from kfp.dsl import component

PIPELINE_ROOT = "{}/pipeline_root/control".format(BUCKET_URI)

aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)

My pipelineline components:-

@component

def args_generator_op() -> list:

    return [{"cats": "1", "dogs": "2"}, {"cats": "10", "dogs": "20"}]

@component

def print_op(msg: str):

    print(msg)

@component

def flip_coin_op() -> str:

    """Flip a coin and return heads or tails randomly."""

    import random

    result = "heads" if random.randint(0, 1) == 0 else "tails"

    return result

The construction of the pipeline

import json

from kfp import dsl

ITEMS = [

    {"snakes": "anaconda", "lizards": "anole",

     "bunnies": [{"cottontail": "bugs"}, {"cottontail": "thumper"}]},

    {"snakes": "cobra", "lizards": "gecko",

     "bunnies": [{"cottontail": "roger"}]},

    {"snakes": "boa", "lizards": "iguana",

     "bunnies": [{"cottontail": "fluffy"},

                 {"fuzzy_lop": "petunia", "cottontail": "peter"}]},

]

@dsl.pipeline(name="control", pipeline_root=PIPELINE_ROOT)

def pipeline(

    loop_items: list = ITEMS,                               # real list, for looping

    json_string: str = json.dumps(ITEMS, sort_keys=True),   # string, for printing

):

    flip1 = flip_coin_op()

    with dsl.Condition(flip1.output != "no-such-result", name="alwaystrue"):

        args_generator = args_generator_op()

        with dsl.ParallelFor(args_generator.output) as item:

            print_op(msg=json_string)

            with dsl.Condition(flip1.output == "heads", name="heads"):

                print_op(msg=item.cats)

            with dsl.Condition(flip1.output == "tails", name="tails"):

                print_op(msg=item.dogs)

    with dsl.ParallelFor(loop_items) as item:

        with dsl.Condition(item.snakes == "boa", name="snakes"):

            print_op(msg=item.snakes)

            print_op(msg=item.lizards)

            print_op(msg=item.bunnies)

    with dsl.ParallelFor(loop_items) as item:

        with dsl.ParallelFor(item.bunnies) as item_bunnies:

            print_op(msg=item_bunnies.cottontail)

Generating the yaml pipeline 

compiler.Compiler().compile(

    pipeline_func=pipeline, package_path="control_pipeline.yaml"

)

And we run the pipeline and that's it. 

DISPLAY_NAME = "control"

job = aiplatform.PipelineJob(

    display_name=DISPLAY_NAME,

    template_path="control_pipeline.yaml",

    pipeline_root=PIPELINE_ROOT,

)

job.run()

And this is what it looks like