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

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

PyTorchJobTuningSchedulerの使い方

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

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

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

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

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

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

from zenith_tune.integration.kubernetes import JobFilter, PyTorchJobTuningScheduler

# Only target jobs with zenith-tune/optimization-config annotation
job_filter = JobFilter(
    annotations={"zenith-tune/optimization-config": None}  # Key existence check
)

scheduler = PyTorchJobTuningScheduler(
    submit_namespace="default",
    job_filter=job_filter,
    max_concurrent_tuning=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をクラスタに投入すると、自動的にチューニングが開始されます。

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

高度な設定

JobFilterによるジョブの選択

JobFilterを使用することで、より細かい条件でチューニング対象を絞り込むことができます。

# 特定のnamespaceのジョブのみを対象
job_filter = JobFilter(
    namespace_pattern=r"^ml-.*",  # "ml-"で始まるnamespace
    annotations={"zenith-tune/optimization-config": None}
)

# 特定のラベルを持つジョブのみを対象
job_filter = JobFilter(
    labels={"team": "ai-research"},
    annotations={"zenith-tune/optimization-config": None}
)

# 複数条件の組み合わせ
job_filter = JobFilter(
    namespace_pattern=r"^production-.*",
    labels={"environment": "production"},
    annotations={"zenith-tune/optimization-config": None}
)

並行実行数の調整

max_concurrent_tuningパラメータでクラスタリソースに応じた並行実行数を調整できます：

scheduler = PyTorchJobTuningScheduler(
    submit_namespace="default",
    tuning_config=tuning_config,
    job_filter=job_filter,
    max_concurrent_tuning=5,  # 最大5つのチューニングジョブを同時実行
)

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

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

正規表現パターンの規則

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

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

チューニングの停止方法

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

kubectl delete pytorchjob my-training-xxxxx

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

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

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

関連ドキュメント

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