PyTorchJobの自動チューニングスケジューラー

ZenithTuneは、Kubernetes環境においてPyTorchJobを自動的に発見し、ハイパーパラメータチューニングを実行するスケジューラー機能（PyTorchJobTuningScheduler）を提供しています。 この機能は、YAMLアノテーションによる設定のみで、コードを書くことなくクラスタ上でのジョブチューニングを自動化します。

PyTorchJobTuningSchedulerの使い方

PyTorchJobTuningSchedulerは以下のワークフローで動作します。

このスケジューラーは、Kubernetesクラスタ上で新しく作成されるPyTorchJobを自動的に監視し、特定の条件（アノテーション）を満たすジョブに対して自動的にハイパーパラメータチューニングを開始します。

アノテーションベースチューニング

PyTorchJobTuningSchedulerの最大の特徴は、PyTorchJobのYAMLアノテーションに最適化設定を記述するだけで、自動的にハイパーパラメータチューニングが開始されることです。

1. スケジューラーの起動

PyTorchJobTuningSchedulerを以下のように起動します。

from aibooster.intelligence.zenith_tune.integration.kubernetes import (
    JobFilter, PyTorchJobTuningScheduler, TuningConfig, TuningRule
)

# チューニングルールの定義
tuning_rules = [
    TuningRule(
        job_filter=JobFilter(),
        tuning_config=TuningConfig(),
    )
]

scheduler = PyTorchJobTuningScheduler(
    tuning_rules=tuning_rules,
    max_concurrent_tuning_per_namespace=1,  # 名前空間ごとに1つまで
)
scheduler.run()

2. チューニング対象ジョブの準備

スケジューラーが起動したら、チューニング対象となるPyTorchJobにzenith-tune/optimization-configアノテーションを追加します：

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  generateName: my-training-
  annotations:
    zenith-tune/optimization-config: |
      variables:
        - name: "learning_rate"
          type: "float"
          range: [0.001, 0.1]
          log: true
          target_env: "LEARNING_RATE"
        - name: "batch_size"
          type: "int"
          range: [16, 128]
          step: 16
          target_env: "BATCH_SIZE"
      objective:
        name: "loss"
        regex: "Final Loss: ([0-9]+\\.?[0-9]*)"
        direction: "minimize"
      n_trials: 10
spec:
  pytorchReplicaSpecs:
    Worker:
      replicas: 1
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch/pytorch:latest
              command:
                - python
                - train.py
                - --learning-rate
                - ${LEARNING_RATE:-0.01}
                - --batch-size
                - ${BATCH_SIZE:-32}

重要: target_envで指定した環境変数名（例：LEARNING_RATE、BATCH_SIZE）は、チューニング実行時に各トライアルで最適化された値が自動的に設定されます。上記の例では、${LEARNING_RATE:-0.01}のように環境変数を参照し、値が設定されていない場合のデフォルト値（0.01）も指定しています。

3. 自動チューニングの開始

スケジューラーが起動している状態で、zenith-tune/optimization-configアノテーションを持つPyTorchJobをクラスタに投入すると、自動的にチューニングが開始されます。

スケジューラーは新しく作成されたジョブを検出し、アノテーションに基づいて自動的にハイパーパラメータチューニングを実行します。

4. チューニングの停止方法

実行中のチューニングを停止したい場合は、オリジナルのPyTorchJobを削除することで停止できます：

kubectl delete pytorchjob my-training-xxxxx

オリジナルジョブが削除されると、そのジョブに対するチューニング処理は自動的に中断され、新しいトライアルジョブの生成も停止されます。既に実行中のトライアルジョブは完了まで実行されますが、新しいトライアルは開始されません。

アノテーション設定の詳細仕様

zenith-tune/optimization-configアノテーションのYAML形式について、厳密な仕様を説明します。

全体構造

metadata:
  annotations:
    zenith-tune/optimization-config: |
      variables:
        - name: <string>
          type: <string>
          target_env: <string>
          # type-specific fields
      objective:
        name: <string>
        regex: <string>
        direction: <string>
      n_trials: <integer>  # Optional

variables（チューニングパラメータ）

各変数で使用可能なフィールド:

フィールド	型	必須	説明
name	string	Yes	Optunaトライアルで使用する変数の一意識別子
type	string	Yes	変数型: float、int、または categorical
target_env	string	No	Workerコンテナに設定する環境変数名（未指定時はnameの大文字版を使用）

型別フィールド

Float型の場合:

フィールド	型	必須	説明
range	[float, float]	Yes	最小値と最大値 [low, high]
log	boolean	No	対数スケールを使用 (default: false)

- name: "learning_rate"
  type: "float"
  range: [0.0001, 0.1]
  log: true
  target_env: "LEARNING_RATE"

Integer型の場合:

フィールド	型	必須	説明
range	[int, int]	Yes	最小値と最大値 [low, high]
step	int	No	ステップサイズ (default: 1)

- name: "batch_size"
  type: "int"
  range: [16, 128]
  step: 16
  target_env: "BATCH_SIZE"

Categorical型の場合:

フィールド	型	必須	説明
choices	list	Yes	選択可能な値のリスト

- name: "optimizer"
  type: "categorical"
  choices: ["adam", "sgd", "adamw"]
  target_env: "OPTIMIZER_TYPE"

objective（最適化目標）

objectiveセクションで使用するフィールド:

フィールド	型	必須	説明
name	string	Yes	メトリクスの可読名
regex	string	Yes	値を抽出するための正規表現パターン
direction	string	Yes	最適化方向: minimize または maximize
selector	string	No	複数マッチ時の値選択方法（デフォルト: last）

正規表現パターンの規則

• 数値をキャプチャする1つのキャプチャグループ () が必要
• キャプチャグループは浮動小数点数にマッチ: ([0-9]+\.?[0-9]*)
• 特殊文字は適切にエスケープする

selectorオプション

正規表現パターンが複数の値にマッチした場合の選択方法を指定します：

値	説明
first	最初にマッチした値を使用
last	最後にマッチした値を使用（デフォルト）
min	マッチした値の最小値を使用
max	マッチした値の最大値を使用
mean	マッチした値の平均値を使用
median	マッチした値の中央値を使用

Validation Lossの最終値を最小化する場合の設定例です：

objective:
  name: "validation_loss"
  regex: "Validation Loss: ([0-9]+\\.?[0-9]*)"
  direction: "minimize"
  selector: "last"

custom_patches（カスタムパッチ）

custom_patchesを使用して、トライアルジョブのマニフェストに任意のパッチを適用できます。これにより、priorityClassName, nodeSelector などを設定できます。

metadata:
  annotations:
    zenith-tune/optimization-config: |
      variables:
        - name: "learning_rate"
          type: "float"
          range: [0.001, 0.1]
          target_env: "LEARNING_RATE"
      objective:
        name: "loss"
        regex: "Final Loss: ([0-9]+\\.?[0-9]*)"
        direction: "minimize"
      custom_patches:
        spec:
          pytorchReplicaSpecs:
            Worker:
              template:
                spec:
                  priorityClassName: low-priority

高度な設定

TuningRuleによるジョブのルーティング

PyTorchJob Tuning Schedulerでは、ジョブの種類や用途に応じて異なるチューニング設定を適用できます。 例えば、学習ジョブには多めのトライアル数を設定し、テストジョブには少なめのトライアル数を設定したり、ジョブごとに最適化対象のハイパーパラメータや目的関数を個別に定義したりすることが可能です。 このようなジョブごとのルーティングはTuningRuleを使って実現します。

TuningRuleは、ジョブの条件を定義するJobFilterと、チューニング設定を定義するTuningConfigを組み合わせたルールです。 複数のルールを定義した場合、最初にマッチしたルールが使用されます。

tuning_rules = [
    # 学習ジョブ: 多めのトライアル数
    TuningRule(
        job_filter=JobFilter(
            name_pattern=r"training-.*",
            annotations={"zenith-tune/optimization-config": None}
        ),
        tuning_config=TuningConfig(
            n_trials=20,
        )
    ),
    # テストジョブ: 少なめのトライアル数
    TuningRule(
        job_filter=JobFilter(
            name_pattern=r"test-.*",
            annotations={"zenith-tune/optimization-config": None}
        ),
        tuning_config=TuningConfig(
            n_trials=5,
        )
    ),
    # その他: デフォルト設定
    TuningRule(
        job_filter=JobFilter(
            annotations={"zenith-tune/optimization-config": None}
        ),
        tuning_config=TuningConfig()
    ),
]

scheduler = PyTorchJobTuningScheduler(tuning_rules=tuning_rules)
scheduler.run()

JobFilter

JobFilterは、チューニング対象となるジョブの条件を定義します：

job_filter = JobFilter(
    namespace_pattern=r"^ml-.*",  # 名前空間の正規表現パターン
    name_pattern=r"training-.*",  # ジョブ名の正規表現パターン
    labels={"team": "ai-research"},  # ラベル（完全一致）
    annotations={"zenith-tune/optimization-config": None}  # アノテーション（Noneはキー存在チェック）
)

TuningConfig

TuningConfigは、チューニングの実行設定を定義します：

tuning_config = TuningConfig(
    n_trials=20,                             # トライアル数（Noneでアノテーションから取得）
    output_dir="outputs",                    # ログ出力ディレクトリ
    maximize=False,                          # 目的値を最大化するか
    wait_resources=True,                     # リソース待機を有効化
    submit_namespace="tuning-ns",            # トライアルジョブの投入先名前空間（Noneの場合はオリジナルを引き継ぐ）
    default_params={"learning_rate": 0.01},  # 初回トライアルのデフォルトパラメータ
    default_custom_patches=None,             # トライアルジョブに適用するカスタムパッチ
    sampler=None,                            # Optunaのサンプラー
)

default_custom_patchesを使うと、トライアルジョブのKubernetesマニフェストをカスタマイズできます（custom_patchesを参照）。例えば、Kueueを使用している環境でトライアルジョブを特定のLocalQueueに割り当てる場合：

tuning_config = TuningConfig(
    default_custom_patches={
        "metadata": {
            "labels": {
                "kueue.x-k8s.io/queue-name": "tuning-queue"
            }
        }
    },
)

※ジョブのアノテーションにcustom_patchesが設定されている場合、アノテーションの設定が優先されてマージされます。

並行実行数の調整

並行実行数を調整するための2つのパラメータがあります：

• max_concurrent_tuning: クラスタ全体で同時実行可能なチューニング数の上限（Noneで無制限）
• max_concurrent_tuning_per_namespace: 各名前空間で同時実行可能なチューニング数の上限（Noneで無制限）

scheduler = PyTorchJobTuningScheduler(
    tuning_rules=tuning_rules,
    max_concurrent_tuning=10,             # クラスタ全体で最大10
    max_concurrent_tuning_per_namespace=2,  # 名前空間ごとに最大2
)

チューニングメタデータ環境変数

トライアルジョブには以下の環境変数が自動的に設定されます：

環境変数名	説明
ZENITH_TUNE_ENABLED	チューニングが有効であることを示す（値: 1）
ZENITH_TUNE_TRIAL_ID	トライアル番号（0から開始）
ZENITH_TUNE_PARAMS_LABEL	パラメータのラベル文字列（例: learning_rate=0.01,batch_size=32）

これらの環境変数は、トライアルジョブ内からログ出力や監視目的で参照できます。

YAMLファイルによるチューニング設定の記述

スケジューラーの設定はYAMLファイルでも記述することができます。

from aibooster.intelligence.zenith_tune.integration.kubernetes import PyTorchJobTuningScheduler

scheduler = PyTorchJobTuningScheduler.from_yaml("config.yaml")
scheduler.run()

設定ファイル例（config.yaml）:

scheduler:
  max_concurrent_tuning_per_namespace: 1  # 名前空間ごとの並行チューニング上限

tuning_rules:
  - job_filter:
      annotations:
        zenith-tune/optimization-config: null
    tuning_config:
      default_custom_patches:
        spec:
          pytorchReplicaSpecs:
            Worker:
              template:
                spec:
                  priorityClassName: low-priority

注意事項とベストプラクティス

• スケジューラー起動前から存在するジョブは自動的に除外されます
• チューニングにより生成されたジョブ（zenith-tune/created-by: PyTorchJobTunerアノテーション付き）は再度チューニング対象にはなりません
• max_concurrent_tuningとmax_concurrent_tuning_per_namespaceはクラスタのリソース容量に応じて適切に設定してください
• アノテーションのYAML形式が正しくない場合、そのジョブはスキップされます

関連ドキュメント

• PyTorchJobTuner
• PyTorchJobTunerによるジョブチューニング例
• PyTorchJobTuningSchedulerによる自動チューニング例