Step FunctionsからEKSへPyTorchJobを投入できない？ConfigMapで華麗に解決する方法

Last updated at 2025-09-23Posted at 2025-09-23

はじめに

AWS Step Functionsを使用してEKSクラスター上のKubeflow PyTorchJobを動的に作成・実行する仕組みについて解説します。

はじめに
構成概要
- 全体アーキテクチャ図
事前準備
- ConfigMapのEKSクラスターへの登録
実装の仕組み
- ConfigMapによるPyTorchJobテンプレート定義
- Step FunctionsからのPyTorchJob投入
仕組みのポイント
注意点とベストプラクティス
Step FunctionsからPyTorchJob投入までの詳細フロー
トラブルシューティング
まとめ
関連ファイル
参考リンク
他の実装手法との比較
補足：なぜStep FunctionsはPyTorchJobを直接サポートしないのか？

背景と課題

AWS Step FunctionsにはEKS統合として標準のKubernetes Jobを投入する機能（arn:aws:states:::eks:runJob.sync）は提供されていますが、PyTorchJobのようなCustom Resource Definitionは直接サポートされていません。

AWS公式ドキュメントによると：

カスタムリソースやその他の Kubernetes リソースについては、eks:call サービス統合を使用してください。

引用元: AWS Step Functions での Amazon EKS の使用

Step FunctionsのEKS統合では、以下の制限があります：

eks:runJob: batch/v1 Jobのみサポート（.syncによる完了待機とログ取得が可能）
eks:call: カスタムリソースやその他のKubernetesリソース用の汎用API（ポーリングループが必要）

そのため、以下のような課題があります：

Step FunctionsからPyTorchJobを直接作成できない
分散深層学習ワークフローの自動化が困難
ML実験の動的なパラメータ制御が複雑

PyTorchJobについて

PyTorchJobは、Kubernetes上でPyTorch分散学習を実行するためのCustom Resource Definitionです。主要な実装として以下があります：

Kubeflow Training Operator: 現在の標準実装（kubeflow.org/v1）
TorchElastic Controller: PyTorch公式のElastic対応実装
その他: vSphere ML Extension、第三者フォークなど

本記事では、最も広く使用されているKubeflow Training OperatorのPyTorchJobを対象とします。

注意: 2025年現在、KubeflowはTrainer v2への移行を進めており、従来のPyTorchJob CRDから新しい統合API（TrainJob）への移行が計画されています。

解決アプローチ

本記事では、この制約を克服するため、Kubernetes Jobを介してPyTorchJobを間接的に作成する仕組みを実装します。ConfigMapをテンプレートとして使用し、環境変数による動的パラメータ注入でPyTorchJobを柔軟に作成できます。

構成概要

この仕組みでは、以下のコンポーネントを使用します：

AWS Step Functions: PyTorchJob投入のオーケストレーション
Amazon EKS: Kubernetesクラスター
Kubeflow PyTorchJob: 分散深層学習ジョブの実行
ConfigMap: PyTorchJobのテンプレート定義

全体アーキテクチャ図

以下の図は、Step FunctionsからPyTorchJobを投入する全体的な流れを示しています：

図に示されているように、処理は大きく4つの段階に分かれています：

Step Functions（上段）: 実行開始からPyTorchJob作成までのオーケストレーション
ConfigMap（左中段）: PyTorchJobテンプレートの定義
EKS Cluster（右中段）: Kubernetes Job → PyTorchJob → GPU Resourcesの流れ
詳細処理（下段）: Kubernetes Job内部でのStep-by-Stepな処理

事前準備

ConfigMapのEKSクラスターへの登録

PyTorchJobのテンプレートを使用するには、事前にConfigMapをEKSクラスターに登録する必要があります。

1. kubectlの設定

まず、EKSクラスターに接続するためのkubeconfigを設定します：

# AWS CLIを使用してEKSクラスターの認証情報を取得
aws eks update-kubeconfig --name eks-cluster-training --region us-west-2

# 接続確認
kubectl get nodes

2. ConfigMapの作成と登録

ConfigMapファイル（pytorch-job-template-configmap.yaml）を作成し、クラスターに登録：

# ConfigMapの作成
kubectl apply -f ConfigMaps/pytorch-job-template-configmap.yaml

# 作成されたConfigMapの確認
kubectl get configmap pytorch-job-template-configmap -n ml-training-pipeline

# ConfigMapの内容確認
kubectl describe configmap pytorch-job-template-configmap -n ml-training-pipeline

3. 必要なNamespaceの作成

PyTorchJobを実行するNamespaceが存在しない場合は作成：

# Namespaceの作成
kubectl create namespace ml-training-pipeline

# ServiceAccountの作成（必要な場合）
kubectl create serviceaccount ml-pipeline-sa -n ml-training-pipeline

実装の仕組み

1. ConfigMapによるPyTorchJobテンプレート定義

図の左中段に示されているように、PyTorchJobのテンプレートをConfigMapとして定義します。このConfigMapは環境変数による動的パラメータ注入機能を提供し、同じテンプレートから異なる設定のPyTorchJobを作成できます。

apiVersion: v1
kind: ConfigMap
metadata:
  name: pytorch-job-template-configmap
  namespace: ml-training-pipeline
data:
  pytorchjob.yaml: |
    apiVersion: kubeflow.org/v1
    kind: PyTorchJob
    metadata:
      name: ${PYTORCH_JOB_NAME}
      namespace: ml-training-pipeline
    spec:
      pytorchReplicaSpecs:
        Worker:
          replicas: 1
          restartPolicy: Never
          template:
            spec:
              nodeSelector:
                sagemaker.amazonaws.com/instance-group-name: gpu-nodegroup
              containers:
              - name: pytorch
                image: 123456789012.dkr.ecr.us-west-2.amazonaws.com/ml-training:v1.0-release
                imagePullPolicy: Always
                command:
                  - /bin/bash
                args:
                  - -c
                  - |
                    echo "=== ML Training Pipeline Start ==="
                    # 学習スクリプトの実行
                    python train.py --epochs ${NUM_EPOCHS} --batch-size ${BATCH_SIZE}
                resources:
                  limits:
                    nvidia.com/gpu: "8"
                env:
                  - name: NUM_EPOCHS
                    value: "${NUM_EPOCHS}"
                  - name: BATCH_SIZE
                    value: "${BATCH_SIZE}"
                  - name: EXPERIMENT_NAME
                    value: "${EXPERIMENT_NAME}"

2. Step FunctionsからのPyTorchJob投入

Step Functionsでは、以下のステップでPyTorchJobを動的に作成・投入します。

2.1 EKSクラスター情報の取得

EKSクラスターの認証情報を取得します：

{
  "Get EKS Cluster Info": {
    "Type": "Task",
    "Resource": "arn:aws:states:::aws-sdk:eks:describeCluster",
    "Parameters": {
      "Name": "eks-cluster-training"
    },
    "Credentials": {
      "RoleArn": "arn:aws:iam::123456789012:role/eks-cluster-role"
    },
    "ResultPath": "$.clusterInfo",
    "Next": "Create PyTorchJob"
  }
}

2.2 ConfigMapテンプレートを使用したPyTorchJob作成

図の右中段に示されているように、ConfigMapを参照してPyTorchJobを動的に作成します。この処理では、Kubernetes Job（青色）を使ってPyTorchJob（赤色）を間接的に作成し、最終的にGPUリソースで学習を実行します：

{
  "Create PyTorchJob": {
    "Type": "Task",
    "Resource": "arn:aws:states:::eks:runJob.sync",
    "Parameters": {
      "ClusterName": "eks-cluster-training",
      "CertificateAuthority.$": "$.clusterInfo.Cluster.CertificateAuthority.Data",
      "Endpoint.$": "$.clusterInfo.Cluster.Endpoint",
      "Namespace": "ml-training-pipeline",
      "Job": {
        "apiVersion": "batch/v1",
        "kind": "Job",
        "metadata": {
          "generateName": "pytorch-job-creator-",
          "namespace": "ml-training-pipeline"
        },
        "spec": {
          "template": {
            "spec": {
              "containers": [
                {
                  "name": "job-creator",
                  "image": "bitnami/kubectl:latest",
                  "command": ["/bin/bash", "-c"],
                  "args": [
                    "# ConfigMapからテンプレートをコピー
" +
                    "cp /app/pytorchjob.yaml /tmp/pytorchjob.yaml
" +
                    "# 置換対象の環境変数をホワイトリスト形式で指定
" +
                    "# これにより、意図しない環境変数の置換を防ぐ
" +
                    "export VARS_TO_SUBSTITUTE='$PYTORCH_JOB_NAME,$NUM_EPOCHS,$BATCH_SIZE,$EXPERIMENT_NAME'
" +
                    "# envsubstで指定した環境変数のみを置換
" +
                    "envsubst \"${VARS_TO_SUBSTITUTE}\" < /tmp/pytorchjob.yaml > /tmp/pytorchjob_final.yaml
" +
                    "# デバッグ用：置換結果の確認（本番環境ではコメントアウト推奨）
" +
                    "echo '=== 置換された環境変数 ==='
" +
                    "grep -E 'name:|NUM_EPOCHS|BATCH_SIZE|EXPERIMENT_NAME' /tmp/pytorchjob_final.yaml
" +
                    "# PyTorchJobを作成
" +
                    "kubectl create -f /tmp/pytorchjob_final.yaml"
                  ],
                  "env": [
                    {
                      "name": "PYTORCH_JOB_NAME",
                      "value": "training-job-$(date +%s)"
                    },
                    {
                      "name": "NUM_EPOCHS",
                      "value": "10"
                    },
                    {
                      "name": "BATCH_SIZE", 
                      "value": "32"
                    },
                    {
                      "name": "EXPERIMENT_NAME",
                      "value": "experiment-001"
                    }
                  ],
                  "volumeMounts": [
                    {
                      "name": "pytorch-job-template",
                      "mountPath": "/app",
                      "readOnly": true
                    }
                  ]
                }
              ],
              "volumes": [
                {
                  "name": "pytorch-job-template",
                  "configMap": {
                    "name": "pytorch-job-template-configmap"
                  }
                }
              ],
              "restartPolicy": "Never"
            }
          }
        }
      }
    },
    "End": true
  }
}

仕組みのポイント

1. 環境変数による動的パラメータ注入

Step Functionsの実行時パラメータを環境変数としてPyTorchJobに渡すことで、同じテンプレートから異なる設定のジョブを作成できます。

主な動的パラメータ：

NUM_EPOCHS: 学習エポック数
BATCH_SIZE: バッチサイズ
EXPERIMENT_NAME: 実験名

2. ConfigMapによるテンプレート管理

PyTorchJobの定義をConfigMapで管理することで：

テンプレートの再利用性向上
バージョン管理の簡素化
設定変更時の影響範囲の明確化

3. 間接的なPyTorchJob作成

図に示されているように、Step FunctionsからPyTorchJobを直接作成するのではなく、Kubernetes Job（青色）を介してPyTorchJob（赤色）を作成する方式を採用。これにより：

Step Functionsの制約を回避
kubectlコマンドによる柔軟な操作
エラーハンドリングの改善

注意点とベストプラクティス

1. エラーハンドリング

各ステップで適切なエラーハンドリングを実装：

{
  "Catch": [
    {
      "ErrorEquals": ["States.ALL"],
      "Comment": "Continue to success summary even if marker update fails",
      "Next": "Summarize Success"
    }
  ]
}

2. リトライ戦略

重要な処理にはリトライを設定：

{
  "Retry": [
    {
      "ErrorEquals": ["States.TaskFailed"],
      "IntervalSeconds": 2,
      "MaxAttempts": 3,
      "BackoffRate": 2
    }
  ]
}

3. タイムアウト設定

長時間実行されるジョブには適切な待機時間を設定：

{
  "Wait 600 Seconds": {
    "Type": "Wait",
    "Seconds": 600,
    "Next": "Check PyTorchJob Status"
  }
}

Step FunctionsからPyTorchJob投入までの詳細フロー

基本的な処理フロー

図の下段に示されている詳細処理を参照しながら、Step-by-Stepで以下の流れで実行されます：

Step Functions実行開始: 入力パラメータ（エポック数、バッチサイズなど）を受け取り
EKSクラスター情報取得: 対象クラスターの認証情報を取得
Kubernetes Job作成: ConfigMapを使用してPyTorchJob作成用のJobを投入
PyTorchJob作成: Jobコンテナ内でenvsubstによる環境変数展開とkubectl createを実行

Kubernetes Job内部での詳細ステップ

図の下段「Kubernetes Job 内部処理の詳細」で示されているように、7つのステップで処理が実行されます：

Step 1-2: ConfigMapからテンプレート読み込み・Job名生成
Step 3: 環境変数展開（envsubst実行）
Step 4-5: PyTorchJob作成・ラベル追加
Step 6-7: PyTorchJob起動・学習実行（8GPU並列処理）

環境変数の流れ

図の下段中央にある「環境変数」ノート（グレー）に示されているように、以下の流れでパラメータが渡されます：

Step Functions Input
    ↓
{
  "numEpochs": 10,
  "batchSize": 32,
  "experimentName": "test-001"
}
    ↓
Kubernetes Job環境変数
    ↓
{
  "NUM_EPOCHS": "10",
  "BATCH_SIZE": "32",
  "EXPERIMENT_NAME": "test-001"
}
    ↓
envsubst実行でConfigMap展開（Step 3）
    ↓
PyTorchJob YAML作成・適用（Step 4）

データとストレージの流れ

図の最下段「データとストレージの流れ」セクションに示されているように、以下のストレージコンポーネントが連携します：

S3 Bucket: 学習データの保存
FSx for Lustre: 高性能ファイルシステム（S3 → FSx → S3）
ConfigMap: PyTorchJobテンプレートの管理

トラブルシューティング

ConfigMapが見つからない場合

# ConfigMapの存在確認
kubectl get configmap -n ml-training-pipeline

# ConfigMapの再作成
kubectl delete configmap pytorch-job-template-configmap -n ml-training-pipeline
kubectl apply -f ConfigMaps/pytorch-job-template-configmap.yaml

PyTorchJobが作成されない場合

# Kubernetes Jobのログ確認
kubectl logs -n ml-training-pipeline job/ml-pipeline-xxx

# PyTorchJobの一覧確認
kubectl get pytorchjob -n ml-training-pipeline

まとめ

AWS Step FunctionsからKubeflow PyTorchJobを動的に投入する仕組みにより：

テンプレート管理: ConfigMapによる再利用可能なジョブ定義
動的パラメータ: 環境変数による柔軟な設定変更
間接作成: Kubernetes Jobを介したPyTorchJob作成

この仕組みは、機械学習ワークフローの自動化とスケーラビリティの向上に有効です。

参考リンク

他の実装手法との比較

Step FunctionsからPyTorchJobを投入する様々な手法を比較検討します：

手法	実装難易度	保守性	パフォーマンス	セキュリティ	エラーハンドリング	推奨度
ConfigMap + Kubernetes Job （本記事の手法）	⭐⭐⭐ 中	⭐⭐⭐⭐ 高	⭐⭐⭐⭐ 高	⭐⭐⭐⭐ 高	⭐⭐⭐⭐ 高	⭐⭐⭐⭐⭐ 最推奨
`eks:call`直接操作	⭐⭐ 低	⭐⭐ 低	⭐⭐⭐ 中	⭐⭐⭐ 中	⭐⭐ 低	⭐⭐ 非推奨
AWS Lambda + kubectl	⭐⭐⭐⭐ 高	⭐⭐⭐ 中	⭐⭐⭐ 中	⭐⭐⭐ 中	⭐⭐⭐⭐ 高	⭐⭐⭐ 条件付推奨
Argo Workflows経由	⭐⭐⭐⭐⭐ 最高	⭐⭐⭐ 中	⭐⭐⭐⭐ 高	⭐⭐⭐⭐ 高	⭐⭐⭐⭐ 高	⭐⭐⭐ 条件付推奨
Helm Chart + AWS CLI	⭐⭐⭐⭐ 高	⭐⭐⭐⭐ 高	⭐⭐⭐ 中	⭐⭐⭐ 中	⭐⭐⭐ 中	⭐⭐⭐ 条件付推奨
Kubernetes API直接	⭐⭐⭐⭐⭐ 最高	⭐⭐ 低	⭐⭐⭐⭐ 高	⭐⭐ 低	⭐⭐ 低	⭐ 非推奨
SQS + Operator Pattern	⭐⭐⭐⭐⭐ 最高	⭐⭐ 低	⭐⭐⭐⭐⭐ 最高	⭐⭐⭐ 中	⭐⭐⭐ 中	⭐⭐ 非推奨

各手法の詳細比較

1. ConfigMap + Kubernetes Job（本記事の手法）

メリット:

Step Functionsのeks:runJob.syncでログ取得と完了待機が可能
ConfigMapによるテンプレート管理で再利用性が高い
環境変数による動的パラメータ注入
AWS公式サポートの範囲内

デメリット:

間接的な作成のため若干の複雑性
Kubernetes Jobの追加リソース消費

2. `eks:call`直接操作

メリット:

最もシンプルなアプローチ
中間層が不要

デメリット:

.syncが使用できないためポーリングループが必要
ジョブの完了監視が複雑
エラーハンドリングが困難

3. AWS Lambda + kubectl

メリット:

Pythonでの柔軟なロジック実装
詳細なエラーハンドリング
複雑な前処理・後処理が可能

デメリット:

Lambda関数の開発・保守が必要
実行時間制限（15分）
kubectlバイナリの管理

4. Argo Workflows経由

メリット:

Kubernetes-nativeなワークフロー
複雑なDAGの表現が可能
GitOpsとの親和性

デメリット:

Argo Workflowsの追加インストール・管理
学習コストが高い
Step Functionsとの二重管理

5. Helm Chart + AWS CLI

メリット:

Helm Chartによる標準化されたテンプレート管理
バージョン管理の容易さ
パッケージング機能

デメリット:

Helmの設定・管理が必要
AWS CLIの実行環境構築
複雑性の増加

6. Kubernetes API直接呼び出し

メリット:

最もダイレクトなアプローチ
AWSサービスに依存しない
細かい制御が可能

デメリット:

認証・認可の複雑な管理
エラーハンドリングの実装が困難
セキュリティリスク

7. SQS + Operator Pattern

メリット:

非同期処理によるスケーラビリティ
疎結合アーキテクチャ
高い拡張性

デメリット:

カスタムOperatorの開発・運用コスト
アーキテクチャの複雑化
デバッグの困難さ

選択指針

本記事の手法（ConfigMap + Kubernetes Job）を推奨する理由:

バランスの良さ: 実装難易度、保守性、パフォーマンス、セキュリティのバランスが最適
AWS公式サポート: eks:runJob.syncによる公式サポートの範囲内
運用性: ログ取得、完了待機、エラーハンドリングが簡潔
拡張性: ConfigMapテンプレートによる柔軟な拡張

他手法を検討すべきケース:

Lambda + kubectl: 複雑な前処理・後処理が必要な場合
Argo Workflows: 既にArgoを使用しており、複雑なワークフローが必要な場合
Helm Chart: 既存のHelmベースの運用体制がある場合

補足：なぜStep FunctionsはPyTorchJobを直接サポートしないのか？

AWS Step FunctionsがPyTorchJobのようなCustom Resource Definition（CRD）を直接サポートしていない理由について、技術的背景を考察します：

1. Kubernetesの標準仕様との整合性

AWS Step Functionsのeks:runJobは、Kubernetes標準のbatch/v1 Job APIに特化して設計されています。これは以下の理由によるものです：

標準化: batch/v1 Jobは全てのKubernetesクラスターで利用可能な標準リソース
API安定性: Kubernetes v1.0からサポートされている成熟したAPI仕様
互換性: ベンダー固有の実装に依存しない汎用性

2. ライフサイクル管理の複雑性

PyTorchJobは標準のJobと比較して、より複雑なライフサイクルを持ちます：

分散処理: Master/Workerの複数Pod管理
弾性スケーリング: 動的なWorker数調整
障害処理: 個別Pod障害時の自動復旧
状態管理: 分散学習特有の状態遷移

Step Functionsの.syncによる同期待機メカニズムは、単一Jobの完了を前提として設計されており、こうした複雑な状態管理には対応していません。

3. サードパーティ依存の回避

PyTorchJobはKubeflow Training Operatorという外部プロジェクトによって提供されます：

依存関係: AWS管理外のCRDに依存するリスク
バージョン互換性: Operatorのアップデートによる破壊的変更
サポート範囲: AWS公式サポートの責任範囲外

AWSは自社でコントロールできる標準Kubernetes APIのみを公式サポートすることで、SLA保証とサポート品質を維持しています。

4. エコシステムの多様性

機械学習フレームワークの多様性も要因の一つです：

PyTorchJob: Kubeflow Training Operator
TFJob: TensorFlow分散学習
MPIJob: MPI並列処理
XGBoostJob: XGBoost分散学習

特定のCRDのみをサポートすることは、他のフレームワークとの公平性を欠き、エコシステム全体の健全性を損なう可能性があります。

5. 設計思想：汎用性 vs 専門性

AWS Step Functionsは汎用的なワークフロー管理サービスとして設計されており：

汎用統合: 様々なAWSサービスとの統合を重視
抽象化レベル: 低レベルなKubernetes APIの詳細を隠蔽
運用負荷: 複雑なCRD固有の運用知識を要求しない

一方、PyTorchJobは高度に専門化されたML固有のリソースであり、両者の設計思想にギャップがあります。

6. 将来的な展望

ただし、以下のような技術的進歩により、将来的にはサポートが拡張される可能性もあります：

Kubernetes Job API v2: CRDとの統一的な管理機能
AWS容認化: Kubeflow OperatorのAWS公式サポート
標準化: PyTorchJobの業界標準化

現実的な対応策

現在の制約下では、本記事で紹介したような「間接的な作成手法」が実用的な解決策となります。この手法により：

AWS公式サポートの範囲内での実装
PyTorchJobの柔軟な制御
運用負荷の最小化
を同時に実現できます。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up

Step FunctionsからEKSへPyTorchJobを投入できない？ConfigMapで華麗に解決する方法

はじめに

目次

背景と課題

PyTorchJobについて

解決アプローチ

構成概要

全体アーキテクチャ図

事前準備

ConfigMapのEKSクラスターへの登録

1. kubectlの設定

2. ConfigMapの作成と登録

3. 必要なNamespaceの作成

実装の仕組み

1. ConfigMapによるPyTorchJobテンプレート定義

2. Step FunctionsからのPyTorchJob投入

2.1 EKSクラスター情報の取得

2.2 ConfigMapテンプレートを使用したPyTorchJob作成

仕組みのポイント

1. 環境変数による動的パラメータ注入

2. ConfigMapによるテンプレート管理

3. 間接的なPyTorchJob作成

注意点とベストプラクティス

1. エラーハンドリング

2. リトライ戦略

3. タイムアウト設定

Step FunctionsからPyTorchJob投入までの詳細フロー

基本的な処理フロー

Kubernetes Job内部での詳細ステップ

環境変数の流れ

データとストレージの流れ

トラブルシューティング

ConfigMapが見つからない場合

PyTorchJobが作成されない場合

まとめ

関連ファイル

参考リンク

他の実装手法との比較

各手法の詳細比較

1. ConfigMap + Kubernetes Job（本記事の手法）

2. eks:call直接操作

3. AWS Lambda + kubectl

4. Argo Workflows経由

5. Helm Chart + AWS CLI

6. Kubernetes API直接呼び出し

7. SQS + Operator Pattern

選択指針

補足：なぜStep FunctionsはPyTorchJobを直接サポートしないのか？

1. Kubernetesの標準仕様との整合性

2. ライフサイクル管理の複雑性

3. サードパーティ依存の回避

4. エコシステムの多様性

5. 設計思想：汎用性 vs 専門性

6. 将来的な展望

現実的な対応策

2. `eks:call`直接操作