コーディング × インフラ — asyncio GKE バッチワーカー改善

概要

問題

以下のコードの問題点を洗い出し、改善せよ。 このワーカーはGCP上のGKE AutopilotのPodで実行され、BigQueryからデータを取得して処理後にCloud Storageへ書き込む。

制約・前提条件

• Python 3.12 + asyncio を使用
• GKE Autopilot 上の Pod で実行
• Workload Identity Federation（サービスアカウントキーファイルなし）
• BigQuery・Cloud Storage のクライアントは google-cloud-python v3.x
• batch_id は外部システムから渡ってくるユーザー入力値

悪いコード (Before)

📋 コピー# worker.py（問題あり）
import asyncio
import json
from google.cloud import bigquery, storage
from datetime import datetime

BQ_CLIENT = bigquery.Client()  # グローバルに作成
GCS_CLIENT = storage.Client()

async def process_batch(batch_id):
    # BigQueryからデータ取得
    query = f"SELECT * FROM `my_project.my_dataset.orders` WHERE batch_id = '{batch_id}'"
    result = BQ_CLIENT.query(query).to_dataframe()

    # 処理
    processed = []
    for _, row in result.iterrows():
        processed.append({
            "order_id": row["order_id"],
            "total": row["amount"] * 1.1,  # 10%加算
            "processed_at": str(datetime.now())
        })

    # GCSに書き込み
    bucket = GCS_CLIENT.bucket("my-bucket")
    blob = bucket.blob(f"output/{batch_id}.json")
    blob.upload_from_string(json.dumps(processed))

    print(f"Done: {batch_id}")

async def main():
    batch_ids = ["batch_001", "batch_002", "batch_003"]
    tasks = [process_batch(bid) for bid in batch_ids]
    await asyncio.gather(*tasks)

if __name__ == "__main__":
    asyncio.run(main())

ヒント（段階的開示）

loop = asyncio.get_running_loop()
result = await loop.run_in_executor(
    None, lambda: client.query(query).to_dataframe()
)

job_config = bigquery.QueryJobConfig(
    query_parameters=[
        bigquery.ScalarQueryParameter("batch_id", "STRING", batch_id)
    ]
)
query = "SELECT * FROM `...` WHERE batch_id = @batch_id"

問題点分析

アーキテクチャ図 — GKE Autopilot + Workload Identity

模範解答

📋 コピー"""BigQuery orders バッチ処理ワーカー。

GKE Autopilot 上で動作し、Workload Identity Federation で
GCP リソースに認証する非同期バッチワーカー。
"""

from __future__ import annotations

import asyncio
import json
import logging
from datetime import UTC, datetime
from typing import Any

from google.cloud import bigquery, storage

# ─── 定数 ──────────────────────────────────────────────────────────────────
_PROJECT_ID = "my_project"
_DATASET = "my_dataset"
_TABLE = "orders"
_GCS_BUCKET = "my-bucket"
_GCS_OUTPUT_PREFIX = "output"
_SURCHARGE_RATE = 1.1  # 10% 加算レート（仕様書 §3.2 参照）

_QUERY = f"""
SELECT order_id, amount
FROM `{_PROJECT_ID}.{_DATASET}.{_TABLE}`
WHERE batch_id = @batch_id
"""

logger = logging.getLogger(__name__)


# ─── クライアントファクトリ（DI可能にする） ───────────────────────────────
def _make_bq_client() -> bigquery.Client:
    """BigQuery クライアントを生成する（Workload Identity 使用）。"""
    return bigquery.Client(project=_PROJECT_ID)


def _make_gcs_client() -> storage.Client:
    """Cloud Storage クライアントを生成する。"""
    return storage.Client(project=_PROJECT_ID)


# ─── 処理ロジック ────────────────────────────────────────────────────────
async def _fetch_orders(
    bq_client: bigquery.Client,
    batch_id: str,
) -> list[dict[str, Any]]:
    """BigQuery から指定バッチのオーダーを非同期取得する。

    Args:
        bq_client: BigQuery クライアント。
        batch_id: 処理対象バッチの識別子。

    Returns:
        オーダー行の辞書リスト。

    Raises:
        ValueError: batch_id が空の場合。
        google.api_core.exceptions.GoogleAPIError: BigQuery 呼び出し失敗時。
    """
    if not batch_id:
        raise ValueError("batch_id must not be empty")

    job_config = bigquery.QueryJobConfig(
        query_parameters=[
            bigquery.ScalarQueryParameter("batch_id", "STRING", batch_id)
        ]
    )

    loop = asyncio.get_running_loop()
    # BigQuery クライアントは同期 I/O → executor でブロッキングを回避
    df = await loop.run_in_executor(
        None,
        lambda: bq_client.query(_QUERY, job_config=job_config).to_dataframe(),
    )
    return df.to_dict(orient="records")


def _transform(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """オーダーに surcharge を付加して変換する。"""
    now_utc = datetime.now(UTC).isoformat()
    return [
        {
            "order_id": row["order_id"],
            "total": round(row["amount"] * _SURCHARGE_RATE, 2),
            "processed_at": now_utc,
        }
        for row in rows
    ]


async def _upload_result(
    gcs_client: storage.Client,
    batch_id: str,
    data: list[dict[str, Any]],
) -> None:
    """処理結果を GCS に非同期アップロードする。"""
    blob_name = f"{_GCS_OUTPUT_PREFIX}/{batch_id}.json"
    loop = asyncio.get_running_loop()
    await loop.run_in_executor(
        None,
        lambda: gcs_client.bucket(_GCS_BUCKET)
        .blob(blob_name)
        .upload_from_string(
            json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        ),
    )
    logger.info("Uploaded gs://%s/%s (%d rows)", _GCS_BUCKET, blob_name, len(data))


async def process_batch(
    batch_id: str,
    *,
    bq_client: bigquery.Client | None = None,
    gcs_client: storage.Client | None = None,
) -> None:
    """1バッチを取得・変換・アップロードする（DI対応）。"""
    bq = bq_client or _make_bq_client()
    gcs = gcs_client or _make_gcs_client()

    logger.info("Processing batch: %s", batch_id)
    try:
        rows = await _fetch_orders(bq, batch_id)
        if not rows:
            logger.warning("No orders found for batch: %s", batch_id)
            return
        transformed = _transform(rows)
        await _upload_result(gcs, batch_id, transformed)
    except Exception:
        logger.exception("Failed to process batch: %s", batch_id)
        raise


async def main(batch_ids: list[str]) -> None:
    """複数バッチを並行処理する。"""
    await asyncio.gather(*(process_batch(bid) for bid in batch_ids))


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    asyncio.run(main(["batch_001", "batch_002", "batch_003"]))

ポイント解説

BQ_CLIENT = bigquery.Client()  # モジュール副作用
GCS_CLIENT = storage.Client()

async def process_batch(batch_id):
    # グローバルを直接使う → モック不可
    result = BQ_CLIENT.query(query)...

async def process_batch(
    batch_id: str,
    *,
    bq_client=None,  # テストでモックを渡せる
    gcs_client=None,
) -> None:
    bq = bq_client or _make_bq_client()

実務への応用

PodのService AccountにIAMバインドを付与するだけでクレデンシャルを自動取得できる。Terraform での設定例:

📋 コピーresource "google_service_account_iam_member" "workload_identity" {
  service_account_id = google_service_account.worker.name
  role               = "roles/iam.workloadIdentityUser"
  member             = "serviceAccount:${var.project_id}.svc.id.goog[${var.namespace}/${var.ksa_name}]"
}

Pod の serviceAccountName を KSA に設定するだけで google.cloud クライアントが自動的に WIF で認証する（キーファイル不要）。

次のステップ

• 参考: Google Cloud Python クライアントライブラリの非同期サポートロードマップ
• 参考: Python asyncio.Semaphore を使った並行数制限

今日のまとめ

自己評価

自分の回答

気づき・メモ