生产级编排AI工作流套件：Flyte全面使用指南 — Core concepts Caching & Named outputs & ImageSpec-CSDN博客

本文链接：https://blog.csdn.net/weixin_47339916/article/details/148030323

生产级编排AI工作流套件：Flyte全面使用指南 — Core concepts Caching & Named outputs & ImageSpec

Flyte 是一个开源编排器，用于构建生产级数据和机器学习流水线。它以 Kubernetes 作为底层平台，注重可扩展性和可重复性。借助 Flyte，用户团队可以使用 Python SDK 构建流水线，并将其无缝部署在云端和本地环境中，从而实现分布式处理和高效的资源利用。

文中内容仅限技术学习与代码实践参考，市场存在不确定性，技术分析需谨慎验证，不构成任何投资建议。

Flyte

缓存

Flyte 允许缓存节点（任务、子工作流和子启动计划）的输出，以加速后续执行。

当相同代码需要多次使用相同输入执行时，缓存功能非常有用。

以下视频简要解释并演示了任务缓存功能：https://www.youtube.com/watch?v=WNkThCp-gqo

输入缓存

Flyte 的输入缓存功能允许任务自动缓存执行所需的输入数据。该功能在需要重新执行任务时（如故障重试或用户手动触发）尤其有用。通过缓存输入数据，Flyte 能优化工作流性能和资源使用，避免不必要的输入重新计算。

输出缓存

Flyte 的输出缓存功能允许用户缓存任务结果以避免冗余计算。该功能对于执行昂贵或耗时操作且结果不频繁变化的任务特别有价值。

在工作流有向无环图（DAG）中，所有节点均可独立启用缓存
节点包括：任务、子工作流（在工作流中直接调用的工作流）和子启动计划（在工作流中调用的启动计划）
顶层工作流或启动计划（通过 UI 或 CLI 调用）不支持缓存
默认所有任务、子工作流和子启动计划均禁用缓存，以避免副作用任务的意外缓存结果。需要显式启用所需节点的缓存

启用和配置缓存

通过设置 @fl.task 装饰器（任务）的 cache 参数或 with_overrides 方法（子工作流/子启动计划）为 Cache 对象来启用缓存。Cache 对象参数用于配置缓存行为：

import flytekit as fl

# 定义一个任务并启用缓存
@fl.task(cache=fl.Cache(version="1.0", serialize=True, ignored_inputs=["a"]))
def sum(a: int, b: int, c: int) -> int:
    return a + b + c

# 定义子工作流
@fl.workflow
def child_wf(a: int, b: int, c: int) -> list[int]:
    return [\
        sum(a=a, b=b, c=c)\
        for _ in range(5)\
    ]

# 创建子启动计划
child_lp = fl.LaunchPlan.get_or_create(child_wf)

# 父工作流使用子工作流
@fl.workflow
def parent_wf_with_subwf(input: int = 0):
    return [\
        # 在子工作流上启用缓存
        child_wf(a=input, b=3, c=4).with_overrides(cache=fl.Cache(version="1.0", serialize=True, ignored_inputs=["a"]))\
        for i in [1, 2, 3]\
    ]

# 父工作流使用子启动计划
@fl.workflow
def parent_wf_with_sublp(input: int = 0):
    return [\
        child_lp(a=input, b=1, c=2).with_overrides(cache=fl.Cache(version="1.0", serialize=True, ignored_inputs=["a"]))\
        for i in [1, 2, 3]\
    ]

此示例展示了多级缓存配置：

任务级：sum 任务的 @fl.task 装饰器
工作流级：child_wf 调用的 with_overrides 方法
启动计划级：child_lp 调用的 with_overrides 方法

Cache 对象

Cache 对象包含以下参数：

version (Optional[str]): 缓存键组成部分。当任务函数变更时，修改版本号可使 Flyte 忽略旧缓存结果
serialize (bool): 启用/禁用缓存序列化。启用后确保节点单实例执行
ignored_inputs (Union[Tuple[str, ...], str]): 计算缓存哈希时忽略的输入变量
policies (Optional[Union[List[CachePolicy], CachePolicy]]): 生成版本哈希的缓存策略列表
salt (str): 哈希生成使用的加密盐值

overwrite-cache 标志

执行工作流、启动计划或任务时，可使用 overwrite-cache 标志强制重新执行。

命令行覆盖缓存

pyflyte run --remote --overwrite-cache example.py wf

UI 界面覆盖缓存

在启动对话框勾选 Override：

UI 覆盖缓存

编程方式覆盖缓存

使用 FlyteRemote.execute 的 overwrite_cache 参数：

from flytekit.configuration import Config
from flytekit.remote import FlyteRemote

remote = FlyteRemote(
    config=Config.auto(),
    default_project="flytesnacks",
    default_domain="development"
)

wf = remote.fetch_workflow(name="workflows.example.wf")
execution = remote.execute(wf, inputs={"name": "Kermit"}, overwrite_cache=True)

缓存工作原理

缓存条目键值由以下要素组成：

项目：跨项目不共享缓存
域：测试/生产环境隔离
缓存版本：功能变更时更新版本
节点签名：包含名称、输入输出参数类型
输入值：相同输入保证确定性输出

显式缓存版本

代码变更时通过更新版本号使缓存失效：

@fl.task(cache=fl.Cache(version="1.1"))
def t(n: int) -> int:
    return n * n + 1

节点签名

修改节点签名（增删改输入参数或输出类型）会使缓存失效。

本地运行缓存

本地缓存机制与远程相同，但键值不包含项目和域。缓存存储在 ~/.flyte/local-cache/，可通过以下命令清除：

pyflyte local-cache clear

缓存序列化

确保同一时刻只执行单个可缓存任务实例：

@fl.task(cache=fl.Cache(version="1.1", serialize=True))
def t(n: int) -> int:
    return n * n

卸载对象缓存

通过注解自定义哈希方法实现复杂对象缓存：

def hash_pandas_dataframe(df: pandas.DataFrame) -> str:
    return str(pandas.util.hash_pandas_object(df))

@fl.task
def foo_1(a: int, b: str) -> Annotated[pandas.DataFrame, HashMethod(hash_pandas_dataframe)]:
    df = pandas.DataFrame(...)
    return df

@fl.task(cache=True)
def bar_1(df: pandas.DataFrame) -> int:
    ...

@fl.workflow
def wf_1(a: int, b: str):
    df = foo(a=a, b=b)
    v = bar(df=df)

完整示例：

def hash_pandas_dataframe(df: pandas.DataFrame) -> str:
    return str(pandas.util.hash_pandas_object(df))

@fl.task
def uncached_data_reading_task() -> Annotated[pandas.DataFrame, HashMethod(hash_pandas_dataframe)]:
    return pandas.DataFrame({"column_1": [1, 2, 3]})

@fl.task(cache=True)
def cached_data_processing_task(df: pandas.DataFrame) -> pandas.DataFrame:
    time.sleep(1)
    return df * 2

@fl.task
def compare_dataframes(df1: pandas.DataFrame, df2: pandas.DataFrame):
    assert df1.equals(df2)

@fl.workflow
def cached_dataframe_wf():
    raw_data = uncached_data_reading_task()
    
    t1_node = create_node(cached_data_processing_task, df=raw_data)
    t2_node = create_node(cached_data_processing_task, df=raw_data)
    t1_node >> t2_node
    
    compare_dataframes(df1=t1_node.o0, df2=t2_node.o0)

if __name__ == "__main__":
    df1 = cached_dataframe_wf()
    print(f"Running cached_dataframe_wf once : {df1}")

命名输出

默认情况下，Flyte采用标准化命名约定为任务或工作流的输出分配名称。每个输出会按顺序标记为o1、o2、o3等。

但您可以通过使用NamedTuple来自定义这些输出名称。

首先导入必要的依赖项：

# basics/named_outputs.py
from typing import NamedTuple
import flytekit as fl

这里我们定义一个NamedTuple并将其分配给名为slope的任务作为输出：

slope_value = NamedTuple("slope_value", [("slope", float)])

@fl.task
def slope(x: list[int], y: list[int]) -> slope_value:
    sum_xy = sum([x[i] * y[i] for i in range(len(x))])
    sum_x_squared = sum([x[i] ** 2 for i in range(len(x))])
    n = len(x)
    return (n * sum_xy - sum(x) * sum(y)) / (n * sum_x_squared - sum(x) ** 2)

类似地，我们定义另一个NamedTuple并将其分配给intercept任务的输出：

intercept_value = NamedTuple("intercept_value", [("intercept", float)])

@fl.task
def intercept(x: list[int], y: list[int], slope: float) -> intercept_value:
    mean_x = sum(x) / len(x)
    mean_y = sum(y) / len(y)
    intercept = mean_y - slope * mean_x
    return intercept

虽然可以直接在代码中创建NamedTuple，但显式声明通常是更好的做法。这有助于避免mypy等工具可能产生的lint错误。

def slope() -> NamedTuple("slope_value", slope=float):
    pass

您可以直接在工作流中解包NamedTuple输出。此外，您也可以让工作流返回一个NamedTuple作为输出。

请注意，我们通过解引用操作来提取单个任务执行输出。这是必要的，因为NamedTuple本质上是元组（tuple），需要进行解引用操作。

slope_and_intercept_values = NamedTuple("slope_and_intercept_values", [("slope", float), ("intercept", float)])

@fl.workflow
def simple_wf_with_named_outputs(x: list[int] = [-3, 0, 3], y: list[int] = [7, 4, -2]) -> slope_and_intercept_values:
    slope_value = slope(x=x, y=y)
    intercept_value = intercept(x=x, y=y, slope=slope_value.slope)
    return slope_and_intercept_values(slope=slope_value.slope, intercept=intercept_value.intercept)

ImageSpec

在本节中，您将了解Flyte如何利用Docker镜像在底层构建容器，并学习如何创建自定义镜像以包含任务或工作流所需的所有依赖项。

您将探索如何执行带有自定义命令的原始容器，在单个工作流中指定多个容器镜像，并深入了解ImageSpec的方方面面！

ImageSpec允许您无需Dockerfile即可为Flyte任务定制容器镜像。通过复用PyPI和APT缓存中已下载的包，ImageSpec可加速构建过程。

默认情况下，ImageSpec使用默认的Docker构建器构建镜像，但您也可以指定自定义构建器，例如使用flytekitplugins-envd通过envd构建ImageSpec。

对于每个flytekit.PythonFunctionTask任务或用@task装饰器装饰的任务，您都可以指定容器镜像绑定规则。默认情况下，flytekit将所有任务绑定到单个容器镜像，即默认Docker镜像。要修改此行为，可使用flytekit.task装饰器中的container_image参数，并传入ImageSpec定义。

在构建镜像前，flytekit会检查容器注册表以确认镜像是否已存在。若镜像不存在，flytekit将在注册工作流前构建镜像，并将任务模板中的镜像名称替换为新构建的镜像名称。

前提条件

确保本地机器已运行docker服务
在ImageSpec中使用注册表时，需执行docker login以推送镜像

安装Python或APT包

您可以在ImageSpec中指定Python包和APT包。这些指定包将被添加到默认镜像之上，该默认镜像可在flytekit的Dockerfile中找到。具体来说，flytekit会调用DefaultImages.default_image()函数，该函数根据Python版本和flytekit版本确定并返回默认镜像。例如，若使用Python 3.8和flytekit 1.6.0，默认镜像将为ghcr.io/flyteorg/flytekit:py3.8-1.6.0。

前提条件

将ghcr.io/flyteorg替换为您可发布的容器注册表。若要在演示集群的本地注册表上传镜像，需通过registry参数将注册表指定为localhost:30000。

from flytekit import ImageSpec

sklearn_image_spec = ImageSpec(
  packages=["scikit-learn", "tensorflow==2.5.0"],
  apt_packages=["curl", "wget"],
)

安装Conda包

通过定义ImageSpec从指定conda频道安装包：

image_spec = ImageSpec(
  conda_packages=["langchain"],
  conda_channels=["conda-forge"],  # 指定从中拉取包的conda频道列表
)

在镜像中使用不同Python版本

通过在ImageSpec中指定python_version来构建不同Python版本的镜像：

image_spec = ImageSpec(
  packages=["pandas"],
  python_version="3.9",
)

在特定imageSpec环境中导入模块

is_container()方法用于判断任务是否使用由ImageSpec构建的镜像。若任务确实使用该镜像，则返回true。此方法有助于减少模块加载时间，并避免在单个镜像中安装不必要的依赖。

以下示例中，task1和task2都会导入pandas模块，但Tensorflow仅在task2中导入：

from flytekit import ImageSpec, task
import pandas as pd

pandas_image_spec = ImageSpec(
  packages=["pandas"],
  registry="ghcr.io/flyteorg",
)

tensorflow_image_spec = ImageSpec(
  packages=["tensorflow", "pandas"],
  registry="ghcr.io/flyteorg",
)

# 当且仅当任务使用tensorflow_image_spec构建的镜像时返回true
if tensorflow_image_spec.is_container():
  import tensorflow as tf

@task(container_image=pandas_image_spec)
def task1() -> pd.DataFrame:
  return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [1, 22]})


@task(container_image=tensorflow_image_spec)
def task2() -> int:
  num_gpus = len(tf.config.list_physical_devices('GPU'))
  print("Num GPUs Available: ", num_gpus)
  return num_gpus

在镜像中安装CUDA

有以下几种方式安装CUDA：

使用Nvidia docker镜像

Nvidia docker镜像已预装CUDA。可在ImageSpec中指定基础镜像：

image_spec = ImageSpec(
  base_image="nvidia/cuda:12.6.1-cudnn-devel-ubuntu22.04",
  packages=["tensorflow", "pandas"],
  python_version="3.9",
)

从额外索引安装包

通过指定pip_extra_index_url在ImageSpec中安装CUDA：

image_spec = ImageSpec(
  name="pytorch-mnist",
  packages=["torch", "torchvision", "flytekitplugins-kfpytorch"],
  pip_extra_index_url=["https://download.pytorch.org/whl/cu118"],
)

构建不同架构的镜像

通过platform参数指定架构（如linux/arm64或darwin/arm64）构建镜像：

image_spec = ImageSpec(
  packages=["pandas"],
  platform="linux/arm64",
)

从GitHub安装flytekit

当更新flytekit时，您可能希望使用特定提交哈希进行测试：

new_flytekit = "git+https://github.com/flyteorg/flytekit@90a4455c2cc2b3e171dfff69f605f47d48ea1ff1"
new_spark_plugins = f"git+https://github.com/flyteorg/flytekit.git@90a4455c2cc2b3e171dfff69f605f47d48ea1ff1#subdirectory=plugins/flytekit-spark"

image_spec = ImageSpec(
  apt_packages=["git"],
  packages=[new_flytekit, new_spark_plugins],
  registry="ghcr.io/flyteorg",
)

自定义镜像标签

通过tag_format参数自定义镜像标签格式。以下示例将生成<spec_hash>-dev格式的标签：

image_spec = ImageSpec(
  name="my-image",
  packages=["pandas"],
  tag_format="{spec_hash}-dev",
)

复制额外文件或目录

通过copy参数指定要复制到容器/root目录的文件或目录，用户可访问所需文件。目录结构将匹配相对路径。Docker仅支持相对路径，不允许使用绝对路径或超出当前工作目录的路径（如包含"…/"的路径）：

from flytekit import task, workflow, ImageSpec

image_spec = ImageSpec(
    name="image_with_copy",
    copy=["files/input.txt"],
)

@task(container_image=image_spec)
def my_task() -> str:
    with open("/root/files/input.txt", "r") as f:
        return f.read()

在YAML文件中定义ImageSpec

可通过向pyflyte run或pyflyte register命令提供ImageSpec YAML文件来覆盖容器镜像。例如：

# imageSpec.yaml
python_version: 3.11
packages:
  - sklearn
env:
  Debug: "True"

使用pyflyte注册工作流：

pyflyte run --remote --image image.yaml image_spec.py wf

仅构建镜像不注册工作流

使用pyflyte build命令仅构建镜像：

pyflyte build --remote image_spec.py wf

强制推送镜像

如需强制重建镜像（即使ImageSpec未更改），可向pyflyte命令传入FLYTE_FORCE_PUSH_IMAGE_SPEC=True：

FLYTE_FORCE_PUSH_IMAGE_SPEC=True pyflyte run --remote image_spec.py wf

也可在Python代码中调用force_push()方法强制推送：

image = ImageSpec(packages=["pandas"]).force_push()

将源代码文件注入ImageSpec

在Flyte后端运行时，通常通过快速注册机制将源代码文件注入任务镜像。如果ImageSpec构造函数指定了source_root且copy参数设置为CopyFileDetection.NO_COPY以外的值，无论快速注册状态如何都会复制文件。若source_root和copy字段为空，则是否复制源代码文件取决于是否使用快速注册。详见运行代码的完整说明。

由于文件有时会被复制到构建的镜像中，ImageSpec发布的标签将根据快速注册是否启用以及复制文件的内容而变化。