Creating a workflow with different container images

In this example, we will create a workflow that will have different task configs, one with GPU and another with a non-GPU one. We will write a simple workflow to get the latest price of Bitcoin, Ethereum, and Litecoin, the main aim of this example is to understand the different types of configs. we will create two configs, one for GPU and one for CPU like this.

cpu_config = PythonTaskConfig(
    image=TaskPythonBuild(
        python_version="3.9",
        pip_packages=["truefoundry[workflow]==0.3.0rc13"],
    ),
    resources=Resources(cpu_request=0.45),
    service_account="<service-account>",  # replace with your service account
)

gpu_config = PythonTaskConfig(
    image=TaskPythonBuild(
        python_version="3.9",
        pip_packages=["truefoundry[workflow]==0.3.0rc10", "pynvml==11.5.0"],
        cuda_version="11.5-cudnn8",
    ),
    env={
        "NVIDIA_DRIVER_CAPABILITIES": "compute,utility",
        "NVIDIA_VISIBLE_DEVICES": "all",
    },
    resources=Resources(cpu_request=0.45, devices=[NvidiaGPU(name="T4", count=1)]),
    service_account="<service-account>",  # replace with your service account
)
...

Prerequsite

Before you proceed with the guide, make sure you have the following:

Truefoundry CLI: Set up and configure the TrueFoundry CLI tool on your local machine by following the Setup for CLI guide.
Workspace: To deploy your workflow, you'll need a workspace. If you don't have one, you can create it using this guide: Creating a Workspace or seek assistance from your cluster administrator.

Creating the workflow

Create a workflow.py where we will write the code for our workflow and then create a deploy.py file where we will store the truefoundry configuration for deploying workflow.

Your directory structure will then appear as follows:

.  
├── workflow.py  
└── deploy.py

workflow.py

from truefoundry.workflow import (
    task,
    workflow,
    PythonTaskConfig,
    TaskPythonBuild,
)
from truefoundry.deploy import Resources, NvidiaGPU
import requests
import json

cpu_config = PythonTaskConfig(
    image=TaskPythonBuild(
        python_version="3.9",
        pip_packages=["truefoundry[workflow]==0.3.0rc13"],
    ),
    resources=Resources(cpu_request=0.45),
    service_account="<service-account>",  # replace with your service account
)

gpu_config = PythonTaskConfig(
    image=TaskPythonBuild(
        python_version="3.9",
        pip_packages=["truefoundry[workflow]==0.3.0rc10", "pynvml==11.5.0"],
        cuda_version="11.5-cudnn8",
    ),
    env={
        "NVIDIA_DRIVER_CAPABILITIES": "compute,utility",
        "NVIDIA_VISIBLE_DEVICES": "all",
    },
    resources=Resources(cpu_request=0.45, devices=[NvidiaGPU(name="T4", count=1)]),
    service_account="<service-account>",  # replace with your service account
)


# Python Task: Fetch real-time prices for multiple cryptocurrencies from the CoinGecko API
@task(task_config=cpu_config)
def fetch_crypto_data() -> str:
    response = requests.get(
        "https://api.coingecko.com/api/v3/simple/price?ids=bitcoin,ethereum,litecoin&vs_currencies=usd"
    )
    return response.text


# Python Task: Process the data to extract prices
@task(task_config=gpu_config)
def extract_prices(data: str) -> dict:
    from pynvml import nvmlDeviceGetCount, nvmlInit

    nvmlInit()
    assert nvmlDeviceGetCount() > 0

    json_data = json.loads(data)
    prices = {
        "bitcoin": json_data["bitcoin"]["usd"],
        "ethereum": json_data["ethereum"]["usd"],
        "litecoin": json_data["litecoin"]["usd"],
    }
    return prices


# Workflow: Combine all tasks
@workflow
def crypto_workflow() -> dict:
    data = fetch_crypto_data()
    prices = extract_prices(data=data)
    return prices

As you can see for the fetch_crypto_data function we have defined the resource which does not use GPU whereas task extract_prices uses GPU and hence we have defined the env variable and we have also used the pynvml package to check whether GPU is present or not by asserting the condition that assert nvmlDeviceGetCount() > 0.

deploy.py

Now in deploy.py file write the below code to configure the deployment.

from truefoundry.deploy import Workflow, LocalSource


workflow = Workflow(
    name="get-crypto-data-example",  # your workflow application name
    workflow_file_path="workflow.py",
)
workflow.deploy(workspace_fqn="<workspace-fqn>")

Now run the below command in the terminal to deploy your workflow, replace <workfspace-fqn> with the workspace fqn which you can find on the UI.

python deploy.py

Run the above command from the same directory containing the deploy.py file.