将 Ray on Vertex AI 与 BigQuery 搭配使用

在 Vertex AI 上运行 Ray 应用时，请将 BigQuery 用作云数据库。本部分介绍如何从 Vertex AI 上的 Ray 集群对 BigQuery 数据库进行读取和写入操作。本部分中的步骤假定您使用的是 Vertex AI SDK for Python。

如需从 BigQuery 数据集读取数据，请新建 BigQuery 数据集或使用现有数据集。

导入并初始化 Ray on Vertex AI 客户端

如果您已连接到 Vertex AI 上的 Ray 集群，请重启内核并运行以下代码。连接时必须采用 runtime_env 变量才能运行 BigQuery 命令。

import ray
from google.cloud import aiplatform

# The CLUSTER_RESOURCE_NAME is the one returned from vertex_ray.create_ray_cluster.
address = 'vertex_ray://{}'.format(CLUSTER_RESOURCE_NAME)

runtime_env = {
    "pip":
       ["google-cloud-aiplatform[ray]","ray==2.47.1"]
  }

ray.init(address=address, runtime_env=runtime_env)

从 BigQuery 中读取数据

从 BigQuery 数据集读取数据。必须通过一个 Ray 任务来执行读取操作。

aiplatform.init(project=PROJECT_ID, location=LOCATION)

@ray.remote
def run_remotely():
    import vertex_ray
    dataset = DATASET
    parallelism = PARALLELISM
    query = QUERY

    ds = vertex_ray.data.read_bigquery(
        dataset=dataset,
        parallelism=parallelism,
        query=query
    )
    ds.materialize()

其中：

  • PROJECT_ID： Google Cloud 项目 ID。您可以在 Google Cloud 控制台的欢迎页面中找到项目 ID。

  • LOCATION：Dataset 的存储位置。例如 us-central1。

  • DATASET：BigQuery 数据集。必须采用 dataset.table 格式。如果您提供查询，则设置为 None。

  • PARALLELISM：一个整数，用于指定并行创建的读取任务数。创建的读取流数可能少于您请求的数量。

  • QUERY：一个字符串，其中包含用于从 BigQuery 数据库读取数据的 SQL 查询。如果不需要查询，则设置为 None。

转换数据

使用 pyarrow 或 pandas 更新和删除 BigQuery 表中的行和列。如果您想使用 pandas 转换，建议将输入保留为 pyarrow 类型，在用户定义的函数 (UDF) 中再将其转换为 pandas 类型，以便您可以捕获 UDF 中的任何 pandas 转换类型错误。必须通过一个 Ray 任务来执行转换。

@ray.remote
def run_remotely():
    # BigQuery Read first
    import pandas as pd
    import pyarrow as pa

    def filter_batch(table: pa.Table) -> pa.Table:
        df = table.to_pandas(types_mapper={pa.int64(): pd.Int64Dtype()}.get)
        # PANDAS_TRANSFORMATIONS_HERE
        return pa.Table.from_pandas(df)

    ds = ds.map_batches(filter_batch, batch_format="pyarrow").random_shuffle()
    ds.materialize()

    # You can repartition before writing to determine the number of write blocks
    ds = ds.repartition(4)
    ds.materialize()

将数据写入 BigQuery

将数据插入 BigQuery 数据集。 必须通过一个 Ray 任务来执行写入。

@ray.remote
def run_remotely():
    # BigQuery Read and optional data transformation first
    dataset=DATASET
    vertex_ray.data.write_bigquery(
        ds,
        dataset=dataset
    )

其中：

  • DATASET：BigQuery 数据集。数据集必须采用 dataset.table 格式。

后续步骤