Deep Dive Into Flyte

Ketan Umare (Founder of Flyte, Co-founder & CEO Union.ai)

Haytham Abuelfutuh (Co-creator of Flyte, Co-founder & CTO Union.ai)

flyte.org

Agenda

• Introductions
• Challenges for ML Orchestration
• Flyte
• Tutorial I - Writing a Flyte Workflow
• Tutorial II - Extending Flyte
• Q/A

Ketan Umare

Introductions

• Father, Husband, Software Engineer
• 15 years software engineer, worked on
• Built many products ground up
• Across many industries - Finance, HFT, logistics, mapping, Cloud infrastructure and machine learning
• Started the Flyte project @Lyft 5 years ago
• Founder of flyte.org and Union.ai (also an accidental CEO)
• Passionate about large scale distributed systems and applied ML

​

Haytham Abuelfutuh

Introductions

Haytham is a father, husband, Co-Founder & CTO @ Union.Ai and a co-founder and a maintainer of the Flyte Open Source Project.

He has gained experience in building distributed systems and cloud native solutions through his tenure at Microsoft, Google and Lyft.

Where do we come from?

Kubernetes-native

Workflow Automation Platform

for Business-critical

Machine Learning and

Data Processes

at Scale

What is Flyte?

Kubernetes-native

Workflow Automation Platform

for Business-critical

Machine Learning and

Data Processes

at Scale

What is Flyte?

What is Flyte?

Define: Workflow

The sequence of industrial, administrative, or other processes through which a piece of work passes from initiation to completion.

Not quite a DAG

Directed Acyclic Graphs imply no loops/repeats. But complex processes may have repetitions. Runtime is still a DAG.

Workflows in Data & ML?

​

​

What is Flyte?

Flyte’s role in the Data / ML World!

Data Warehouse / Data Lake

ETL (Spark, Hadoop, Presto)

DataPrep (ETL?)

Serve Features

Ingest Data

Train Models

Batch Inference

Serve Models

Stream Data

Offline Feature Store

Model & Artifact Store

Model Monitoring (Drift etc)

Feature Monitoring

Store transformed Data

ML Transforms

Feature Service

Models & other artifacts (embeddings etc)

Record predictions

Record predictions

Get Ground truth data and store results to visualize

stream features

What is Flyte?

Real World Examples of Workflows

• Extract, transform and clean features from raw data sources
• Load features, train one model, validate the model and deploy it for production serving
• Train a model and perform batch predictions
• Use ground truth data to compare with prediction to determine model performance
• Use events to trigger a process — train / run predictions, etc.
• Perform traditional ETL / ELT processes
• Perform asynchronous data quality checks
• Perform feature drift checks
• …

​

These are only on ML and Data ecosystems. Workflows exist everywhere!

History Before the Dawn!

What is Flyte?

Open Source Journey

What is Flyte?

💡

It’s Really Day-1!

​

What is Flyte?

Integrations

Best Workflow and Pipeline Orchestration Tools, neptune.ai blog - April 2021

https://neptune.ai/blog/best-workflow-and-pipeline-orchestration-tools

​

25 Hot New Data Tools and What They DON’T Do

https://towardsdatascience.com/25-hot-new-data-tools-and-what-they-dont-do-31bf23bd8e56

Collaborators & Contributors

• Lyft
• Spotify
• Striveworks
• USU
• Freenome
• Level 5
• Intel
• GoJek
• Union.ai
• Blackshark.ai
• Wolt
• RunX
• AppliedAI
• LatchBio
• others

Challenges of ML Orchestration

What Flyte solves.

Challenge 1

Develop Incrementally & Constantly Iterate @ Scale

Scale the Job

e.g. one region -> all regions, more GPUs

​

Start with one Job, run it locally

e.g. spark job, a training job, a query etc

​

Create a pipeline, test it locally

e.g. Fetch data -> train model -> calculate metrics

​

Execute the pipeline on demand, at scale

e.g. Run a pipeline with parameters

​

​

Run the pipeline on a schedule or in-response to an event

e.g. Run every hour

​

​

​

​

Retrieve results for jobs / pipelines

1

3

5

2

4

6

Challenge 2

Tame Infrastructure & Self-serve

• Centrally managed Infrastructure
• Access resources — CPU/GPU/Mem, etc
• Framework/Library independence
• Multi-tenancy unaware
• Automatic scaling
• Knobs to Control costs

Challenge 3

Parameterize Executions & Dynamism

• In ML, experiments require:
• Similar algorithms with different parameters / inputs
• Altering behaviour based on inputs — usually scale
• Time is nothing special — usually time is used to test different hypothesis
• Even with dynamism, knowing the pipeline ahead of time is desirable

Challenge 4

Memoization, Recoverability & Reliability

• Ability to alter certain parts of the pipeline, without paying the penalty of the entire pipeline
• Recover from system failures, without having to pay the penalty of recomputations
• And finally in production, the executions should be timely and reliable

Input: x=m1

Input: x=m1

Challenge 5

Collaboration & Organizational Scaling

• Domain experts can work on their parts independently
• Other users in the organization can reuse
• Users are free to use the language of their choice in a unified pipeline
• Communication happens over a strong interface boundary

critical/complex algorithm

PipelineA

PipelineB

dataA

dataB

Team A

Team B

PipelineC

dataC

Composite Pipeline

• Composite pipeline is composed of TeamA, PipelineB + other tasks.
• PipelineC re-uses the shared critical task.

Challenge 6

Extend Simply

Flyte

Vendor A

Inhouse

Vendor B

Consistent API

Organizations want flexibility

Control costs (Migrate vendors, bring capabilities inhouse)

Users velocity and existing code should just work!

​

3

Users want flexibility

Add simple python extensions (Airflow operators)

Maybe only for their teams

​

​

1

Platform wants to keep adding new capabilities

Distributed training support, Spark, Streaming etc

Continue adding and controlling roll-out of features

​

2

Flytekit makes it easy to add new user customizations

Flyte also allows you to run just your own containers

Flyte backend plugins are independently deployed, maintained and are in the hosted service

Flyte control plane makes it possible to switch plugin associations and OSS makes it possible to migrate

Peek into Flyte

Building Blocks: Tasks

• Smallest Unit of work in Flyte
• Declarative, Versioned, language and framework independent.
• Maps to a Backend execution plugin (extendable)

E.g., Containers, SQL queries, Pods, WebAPI calls

• Strong Interface (Typed Inputs and outputs)

Workflows

• Declarative & composable
• Models data-flow through tasks
• Strongly typed interfaces
• Declarative, versioned and dynamically recursive
• DSL’s in python, Java, Scala.
• Python supports imperative model and bring your own DSL

inputs

Outputs

User Journey

Write your code

• Write code in python, with type annotations
• Use DataFrames, PySpark, Python datatypes, etc
• Annotate with @task / @workflow
• Optional: Enable caching
• Execute locally

@task(cache=True, cache_version=”1.0”)

def pay_multiplier(df: pandas.DataFrame, scalar: int) -> pandas.DataFrame:

df["col"] = 2 * df["col"]

return df

​

@task

def total_spend(df: pyspark.DataFrame) -> int:

return df.agg(F.sum("col")).collect()[0][0]

​

@workflow

def calculate_spend(emp_df: pandas.DataFrame) -> int:

return total_spend(df=pay_multiplier(df=emp_df, scalar=2))

​

# Execute

pay_multiplier(df=pandas.DataFrame())

​

calculate_spend(emp_df=pandas.DataFrame())

​

Get ready to scale!

• Specify resources for tasks
• Specify Spark cluster configuration
• Add resiliency — retries
• Optional: Configure one or more Schedule / Notification
• Of course, still execute it locally

@task(limits=Resources(cpu="2", mem="150Mi"))

def pay_multiplier(df: pandas.DataFrame, scalar: int) -> pandas.DataFrame:

df["col"] = 2 * df["col"]

return df

​

@task(task_config=Spark(

spark_conf={"spark.driver.memory": "1000M"}

), retries=2)

def total_spend(df: pyspark.DataFrame) -> int:

return df.agg(F.sum("col")).collect()[0][0]

​

@workflow

def calculate_spend(emp_df: pandas.DataFrame) -> int:

return total_spend(df=pay_multiplier(df=emp_df, scalar=2))

​

LaunchPlan.get_or_create(name="...",

workflow=calculate_spend,

schedule=FixedRate(duration=timedelta(minutes=10)),

notifications=[

Email(

phases=[WorkflowExecutionPhase.FAILED],

recipients_email=[...])]),

)

Ship & Execute on a Cluster

• Ship the code to remote — versioned
• Interactively execute
• Retrieve execution outputs
• OR use the CLI
• OR the UI

# Package and ship

$ pyflyte --pkgs myapp.workflows package --image ...

​

# Execute interactively

execution = remote.execute(

Flyte_entity,

inputs={...}, wait=True).sync()

print(execution.outputs)

​

# Execute using CLI

$ flytectl create execution --project flytesnacks --domain development --execFile exec_spec.yaml

​

Workflow modalities

@Workflow - Deferred evaluation -deferred to Launch!

@dynamic - Deferred to the return statement

@dynamic(cache=True, cache_version="0.1", limits=Resources(mem="600Mi"))

def parallel_fit_predict(

multi_train: typing.List[pd.DataFrame],

multi_val: typing.List[pd.DataFrame],

multi_test: typing.List[pd.DataFrame],

) -> typing.List[typing.List[float]]:

preds = []

​

for loc, train, val, test in zip(LOCATIONS, multi_train, multi_val, multi_test):

model = fit(loc=loc, train=train, val=val)

preds.append(predict(test=test, model_ser=model))

​

return preds

@workflow

def calculate_spend(emp_df: pandas.DataFrame) -> int:

return total_spend(df=pay_multiplier(df=emp_df, scalar=2))

UX Overview - UI (rendered graphs)

UX Overview - UI (error traces)

Concepts

Built by Platform Engineers for Platform Teams!

Flyte Component Layer Cake

Part IIa: Hello World!

Tutorial: Writing a Flyte Workflow

Let’s go!

• Clone the Workshop repo!

�

• Install FlyteCTL

​

Or: curl -sL https://ctl.flyte.org/install | sudo bash -s -- -b /usr/local/bin

​

• In a virtual environment install python requirements

​

• Start flyte sandbox cluster

​

$ git clone https://github.com/flyteorg/odsc-2021

$ cd odsc-2021

$ brew install flyteorg/homebrew-tap/flytectl

OR

$ curl -sL https://ctl.flyte.org/install | sudo bash -s -- -b /usr/local/bin

$ pip install -r 1_hello_world/requirements.txt

$ flytectl sandbox start --source=./

Local Execution

Local execution is a very desirable property, but is usually useful only for testing. In cases when exotic hardware is required or when data is protected, running in a production environment is the only solution.

​

But Flyte supports local execution — Simple Python script!

Why Registration?

• Automated tracking
• Full reproducibility — code snapshot captured, connected with a version (usually git-SHA)
• Ability to have multiple active versions at the same time
• Backtrack from predictions to actual code
• The strongly typed system of Flyte, makes it possible to auto-generate a Launch form, verify inputs and ensure correct execution behavior

Architecture overview

Process of Registration

Fast Registration

Code artifact

Re-use container

Blob store

Domains + Registration — DevOps power!

1

2

3

• At Lyft, everytime a new PR is created:
• A docker container is built
• And workflows & tasks are registered with Flyte
• If users modify their dependencies, then PR is first required to build containers (securely)
• For just code changes, users can use fast-register, to snapshot the code from laptops or dev environments
• Everything is automatically tracked!
• Multiple users can create isolated PR’s and test independently - @Scale!

Domains + Registration — DevOps power!

1

2

3

• Once the code is ready and passes code review, merge to master
• The users simulate a deployment through a deployment pipeline.
• At each stage of the deployment, workflows and tasks are registered with specific domain in Flyte
• Each domain, may change the data-directory, associated roles or some meta attributes like labels, annotations
• At Lyft, only production domain allows schedules

Domains + Registration — DevOps power!

1

2

3

• For production deployments, logs and metrics are automatically tracked.
• Models are automatically promoted to serving infrastructure (multiple options depending on the team)
• Users can use interactive notebooks to retrieve intermediate or final outputs, analyze data and also automate various monitoring tasks

Ok, but then what is Sandbox?

Cheatsheet

• Install requirements:�mkvirtualenv odsc --python=python3�pip install -r 1_hello_world/requirements.txt
• Run Locally:�python 1_hello_world/step1.py
• Start sandbox:�flytectl sandbox start --source=./
• Build docker image:�flytectl sandbox exec -- docker build -t hello_world:1 -f 1_hello_world.Dockerfile .
• Serialize Tasks and workflows:�pyflyte --pkgs 1_hello_world package --image hello_world:1
• Register Tasks and workflows:�flytectl register files flyte-package.tgz -p flytesnacks -d development --archive --version v2

Part IIb: Evolve a Model Training Pipeline

Tutorial: House Price Prediction

Part III: Custom Plugins

Tutorial: Writing a Flyte Plugin

Why?

Extensibility & Flexibility

• It is impossible to build an orchestration platform that is 100% fit
• As we built it and used by multiple different organizations, we realized that users want to
• Extend
• Decorate
• Contribute
• This led to highly extensible design of Flyte,
• Flyte is extensible in every component - the programming SDK, backend, UI etc
• Flyte also wanted to offer the convenience of a service oriented architecture to data orchestration - no more dealing with couple client code and infrastructure.

Use Cases for Extending Flyte

Extensibility & Flexibility

Use case

Python flytekit plugin

User container plugin

Prebuilt container plugin

Flytekit Type Transformer

Backend Plugin

Meta DSL on flytekit

K8s Plugin

WebAPI plugin

Fancy plugin

Golang? Multi-language support? High performance plugin

Python? Custom extensions, try before invest in backend plugin

Library, with user defined

extensions

Provide prebuilt container — plug & play

Custom specialized domain specific types

Write new experiences for your users

New language SDK

Flyte Service API

Contributions :D

Customized interface, no-code solutions etc

Extensibility & Flexibility

FlyteKit Type Transformers

Allows you to create domain specific types and let Flyte understand them. They can be configured to be auto-loaded.

​

FlyteKit Only Task Plugins

Allows syntactic sugar to be provided like a library. Flyte executes a python container, so you can potentially do anything in Flytekit itself.

​

​

FlyteKit data persistence plugins

Allows you to persist data to various stores by automatically using URIs, e.g., s3://, gcs://, bq://...

​

$ pip install flytekitplugins-*

$ pip install flytekitplugins-data-*

Extensibility & Flexibility

Write your own DSL

​

​

Flytekit has tools to write your own DSL in python

Simplified domain specific language, e.g., one class to do Load Data -> Train model -> predict OR using a YAML

Flytekit-learn coming soon!

​

class ModelTrainer(abc.ABC):

"""

This class can be derived to create an implicit pipeline

-- EXAMPLE ONLY --

"""

​

@abstractmethod

def load(params: Parameters) -> FlyteSchema:

...

​

@abstractmethod

def train(data: FlyteSchema) -> FlyteFile:

...

​

@abstractmethod

def predict(model: FlyteFile, datum: FlyteSchema) -> float:

...

​

Extensibility & Flexibility

DSLs in other languages

​

​

Use core protobuf to write SDK in any new language

JAVA/Scala already available (incubating) contributed by Spotify

​

case class GreetTaskInput(name: String)

case class GreetTaskOutput(greeting: String)

​

class GreetTask

extends SdkRunnableTask(

SdkScalaType[GreetTaskInput],

SdkScalaType[GreetTaskOutput]

) {

​

override def run(input: GreetTaskInput): GreetTaskOutput = GreetTaskOutput(s"Welcome, ${input.name}!")

}

​

object GreetTask {

def apply(name: SdkBindingData): SdkTransform =

new GreetTask().withInput("name", name)

}

​

Extensibility & Flexibility

Backend Plugins

​

​

True power of Flyte!

Powerful, stateful plugins — starting multiple containers for a cluster, calling external APIs, performing complex auth flows, etc

unified API across languages

Maintain easily: patch-fix without deploying code fixes to users

Migrate seamlessly

​

​

@task(

task_config=MPIJob(

num_workers=2,

num_launcher_replicas=1,

slots=1,

),

retries=3, cache=True, cache_version="0.1",

requests=Resources(cpu='1', mem="300Mi"),

limits=Resources(cpu='2'),

)

def horovod_train_task(batch_size: int, buffer_size: int, dataset_size: int) -> FlyteDirectory:

hvd.init()

...

Extensibility & Flexibility

XGBoost plugin

​

​

Let’s write a plugin that will allow users to train an XGBoost model

Using a CSV, libSVM, or any DataFrame object

​

Why?

Easy to demo and support multiple data formats, without re-writing code. Get it type-safe!

​

​

# Define Task

xgboost_trainer = XGBoostTrainerTask(

name="xgboost_trainer",

config=XGBoostParameters(

hyper_parameters=HyperParameters(

max_depth=2, eta=1, objective="binary:logistic", verbosity=2

),

),

dataset_type=FlyteFile,

validate=True,

)

​

​

# Invoke it

xgboost_trainer(train=train, test=test, validation=validation,params=params)

​

Extensibility & Flexibility

How?

​

​

XGBoost trainer - executes a predefined function

Hence, we will make it of type PythonInstanceTask

Now implement the logic in execute method

XGBoostParameters is the config for this task - it follows a flytekit convention

​

class XGBoostTrainerTask(PythonInstanceTask[XGBoostParameters]):

​

def __init__(...):

super(XGBoostTrainerTask, self).__init__(

name,

task_type=self._TASK_TYPE,

task_config=config,

interface=Interface(inputs=inputs, outputs=outputs),

**kwargs,

)

​

def execute(self, **kwargs) -> Any:

...

​

Extensibility & Flexibility

Steps

​

​

• In the same repo, change directory to

​

​

• Now the code

​

• Edit the file task.py (hint - solution in task_sln.py) - follow `Step 1/2/3`
• Edit the file __init__.py
• Run the example

$ cd $repo-root/3_extend/flytekit-xgboost/flytekitplugins/xgboost

$ cd ./flytekitplugins/xgboost

$ cd $repo-root/3_extend

$ cd flytekit-xgboost

$ pip install -e .

$ cd ../../../xgboost_example

$ python example.py

Extensibility & Flexibility

Bonus round

Run it on the sandbox environment

Spark - Flyte - k8s interaction diagram

Use Spark Ephemeral clusters.

Environment is locked in the container.

Isolation is per execution

Backend plugin - Spark on Flyte + Horovod

Kubernetes cluster

Spark Operator

K8s

SparkDriver

SparkExecutors

SparkExecutors

SparkExecutors

Load Data from SQL stores -> Horovod

Your data may exist in SQL Data warehouses / Data Lakes

You want to continue your existing workflow - query from Snowflake and train a model using horovod.

Flyte comes to the rescue

Backend plugin - SQL + MPI

Query

Preprocess

Validation / Predictions

Flyte Workflow

Extensibility & Flexibility

PHEW!

Q & A

Roadmap

Part 1

Beyond

• Plugins
• Alternative implementations of the specification
• Performance and cross cloud
• ...

​

Q4’ 2021

• Streamline UX in Python (repo creation, etc.)
• UI Improvements
• Improvements in documentation
• Performance & scale improvements
• Improved schema support
• Improved deployment guides: GCP, OnPrem, Azure
• Bazel build support
• Better debugging support

​

Early 2022

• Low latency Serverless workflows
• Java, Scala SDK improvements
• Out of core backend plugins
• Other language SDKs (R?)
• Advanced debugging support
• Reactive workflows
• UI plugins: quality, ML monitoring, etc
• More plugins: hosted services
• flytekit-learn

​

Questions

Thank you!