Getting Started with Flytekit Python SDK

Hello world with Flytekit

4 min readOct 9, 2023

Flytekit Python SDK is a robust library that allows developers to design Flyte workflows and tasks using Python. This guide will help you understand the fundamental concepts of Flytekit and guide you through setting up your environment.

Overview

Flytekit introduces two primary concepts:

Tasks: These are units of work that can be executed. While tasks can be crafted in any language, Flytekit predominantly supports Python-based tasks. Dive deeper into Tasks for a comprehensive understanding.
Workflows: These are sequences of tasks executed in a particular order. Written in Python, workflows maintain strong typing. Explore Workflows to understand their intricacies.

See Flytekit for a detailed overview of the Flytekit Python SDK.

Now, get started with Flytekit by setting up your environment.

Setting Up Your Environment

Before delving into Flytekit, ensure your development environment is adequately configured.

Prerequisites

Docker: Flyte utilizes Docker for local task execution. Ensure Docker is installed and operational.
Python: FlyteKit is compatible with Python version 3.8 and above. This guide has been verified with Python 3.10. To check your Python version:
On Windows:

python -V

On OSX:

python3 -V

You should see an output similar to:

Python 3.10.9

Flyte CLI: Install Flyte CLI using brew:

brew install flyteorg/homebrew-tap/flytectl

Flyte SDK: Install the SDK using pip:

pip install flytekit

Booting Up the Demo Flyte Cluster

Initialize a local Flyte cluster with:

flytectl demo start

Upon successful startup, you’ll receive a URL for the Flyte Console:

Flyte Console: http://localhost:30080/console

Replace $USERNAME with your user directory name.

To set up sandbox environment variables for flytectl, use:

export FLYTECTL_CONFIG=/Users/$USERNAME/.flyte/config-sandbox.yaml

Now that your environment is set up, let’s dive into Flytekit.

Crafting Your First Flytekit Workflow

Flyte allows you to write workflows in Python using Flytekit. Let’s create a simple workflow that returns “Hello, World!”. When you’re done, you’ll be able to run the workflow locally or remotely on the Flyte cluster.

Create a hello_world.py file:
In your IDE, create a file named hello_world.py in your working directory.

touch hello_world.py

Import Necessary Modules:
Import the task and workflow modules from the flytekit library.

from flytekit import task, workflow

Define a Task:

Write a task that returns “Hello, World!”.

@task
def say_hello() -> str:
    return "Hello, World!"

Tasks are defined using the @task decorator and act as primary building blocks in Flyte. They are versioned, strongly typed, declarative, independently executable, and suitable for unit testing. A task's body is executed at runtime.

Design a Workflow:

Write a workflow that calls the task.

@workflow
def hello_world_wf() -> str:
    res = say_hello()
    return res

Workflows orchestrate tasks. Unlike tasks, workflows is executed at registration time, during the workflow’s registration process. Registration involves uploading the packed (serialized) workflow to the Flyte backend, which is then executed by the Flyte engine.

Execute the Workflow:

In the main function, run the workflow.

if __name__ == "__main__":
    print(f"Running hello_world_wf() {hello_world_wf()}")

Choose one of the following methods to execute the workflow:

Run the following command to execute the workflow locally:

pyflyte run hello_world.py hello_world_wf

You should see the following output:

Hello, World!

2. Alternatively, run the following command to execute it remotely on the Flyte cluster:

pyflyte run --remote hello_world.py hello_world_wf

Go to the Flyte Console to see the workflow execution.
Open the Flyte Console in your browser.
Select View Inputs & Outputs to see the workflow inputs and outputs.
Select Outputs to see the workflow output.
You should see the following output:

{"o0":"Hello, World!"}

Testing

Testing is an integral part of the development process. Flytekit provides a testing framework to test tasks and workflows. When using Flytekit, you have the advantage of a specialized testing framework designed to test tasks and workflows efficiently.

Using task_mock

You can use task_mock to mock a task declaration and mock any task in FlyteKit as long as it has a Python interface associated with it.

import pytest
from flytekit.testing import task_mock
from hello_world import hello_world_wf, say_hello

def test_workflow_returns_hello_world():
    with task_mock(say_hello) as mock_task:
        # Define what the mock should return when the task is called
        mock_task.side_effect = lambda: "Hello, World!"
        result = hello_world_wf()
        assert result == "Hello, World!"

In this example, test the hello_world_wf workflow by mocking the say_hello task.

The task_mock context manager ensures that the specified task (say_hello in this case) is mocked during the test, leaving other tasks unaffected. The side_effect attribute allows you to customize the mocked task's return behavior.

Ensure that you’ve imported necessary modules and the tasks or workflows you’re testing.

Now run pytest to execute the test.

You should see something similar to the following:

================================================== test session starts ===================================================
rootdir: /Users/user/flytectl/
configfile: pyproject.toml
plugins: timeout-2.1.0
collected 1 itemtests/hello_world_test.py::test_workflow_returns_hello_world PASSED                                                [100%]=================================================== 1 passed in 0.48s ====================================================

You’ve successfully tested your workflow!

Next Steps

Congratulations! You’ve successfully run your first Flyte workflow.

Workflows and tasks are the building blocks of Flyte.