メインコンテンツまでスキップ
バージョン: v2509

Kubernetes PyTorchJobのチューニング

ZenithTuneはKubernetes環境におけるPyTorchJobのチューニング機能（PyTorchJobTuner）を提供しています。 この機能は、既存のジョブを指定し、他のTunerと同じ要領でチューニングパラメータと目的関数を設定することで、クラスタ上でチューニングジョブを実行します。

PyTorchJobTunerの使い方​

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

ユーザは、ジョブの変換に必要なjob_converter関数と目的値を抽出するvalue_extractor関数を定義することで、ローカル環境での最適化とクラスタ上でのジョブ実行を反復的に繰り返します。

初期化​

チューニング対象である既存のジョブ名と、チューニングに使用する名前空間を指定して、PyTorchJobTunerのインスタンスを作成します。

from zenith_tune.integration.kubernetes import PyTorchJobTuner

# チューナーの初期化
tuner = PyTorchJobTuner(
    job_name="my-training-job",   # チューニング対象のPyTorchJob名
    namespace="default",          # ジョブを実行する名前空間（省略時は現在のコンテキストの名前空間）
)

job_converterの定義​

job_converter関数は、既存のジョブからチューニングジョブへと変換するユーザ定義関数です。 第1引数にはOptunaのTrial、第2引数にはPyTorchJobTunerインスタンス作成時に指定した既存ジョブが渡されます。

チューニングジョブへの変換は第2引数のPyTorchJobを介して、(a)コマンド文字列にチューニング対象のオプションを追加したり、(b)ジョブ定義にチューニング対象の環境変数を設定します。 PyTorchJobはset_env, set_commandなどの便利なメンバ関数に加えて、ジョブマニフェストをDict形式でアクセスできるため、ユーザは既存ジョブを任意のジョブへと変換することができます。 PyTorchJobの詳細はAPIリファレンスを参照してください。

また、ZenithTuneはコマンド文字列を操作するためにCommandBuilderクラスを提供しています。 CommandBuilderはコマンド文字列をパースし、既存のオプションの上書きや削除することができます。 詳細はCommandBuilderを参照してください。

from optuna.trial import Trial
from zenith_tune.command import CommandBuilder
from zenith_tune.integration.kubernetes import PyTorchJob

def job_converter(trial: Trial, job: PyTorchJob) -> PyTorchJob:
    # チューニングパラメータの定義
    learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-1, log=True)
    batch_size = trial.suggest_categorical("batch_size", [16, 32, 64, 128])
    omp_num_threads = trial.suggest_categorical("omp_num_threads", [1, 2, 4, 8])

    # 環境変数の設定
    job.set_env("OMP_NUM_THREADS", str(omp_num_threads))

    # コマンドの更新
    current_command = job.get_command()
    builder = CommandBuilder(current_command[2]) # Example: ["sh", "-c", "python train.py"]
    builder.update(f"--learning-rate {learning_rate}")
    builder.update(f"--batch-size {batch_size}")

    # 更新されたコマンドの設定
    new_command = current_command.copy()
    new_command[2] = builder.get_command()
    job.set_command(new_command)

    return job

value_extractorの定義​

value_extractor関数は、CommandOutputTunerと同様、ジョブの実行ログから目的の値（例：実行時間、精度、損失など）を抽出します。 第1引数には、PyTorchJobTunerが各試行完了後に自動的にダウンロードしたログのローカルファイルパスが与えられます。

import re
from typing import Optional

def value_extractor(log_path: str) -> Optional[float]:
    with open(log_path, "r") as f:
        logs = f.read()

    # 例1: 実行時間を抽出（最小化したい場合）
    match = re.search(r"Elapsed time: ([0-9.]+) seconds", logs)
    if match:
        elapsed_time = float(match.group(1))
        return elapsed_time

    # 例2: 検証精度を抽出（最大化したい場合）
    # match = re.search(r"Validation accuracy: ([0-9.]+)", logs)
    # if match:
    #     accuracy = float(match.group(1))
    #     return accuracy
    return None

チューニングの実行​

すべての準備が整ったら、optimizeメソッドを呼び出してチューニングを開始します。

# チューニングの実行
tuner.optimize(
    job_converter=job_converter,        # ジョブ変換関数
    value_extractor=value_extractor,    # 目的値抽出関数
    n_trials=20,                        # 試行回数
    default_params={                    # 初回試行のデフォルト値（オプション）
        "learning_rate": 0.001,
        "batch_size": 32,
        "omp_num_threads" : 2,
    }
)

チューニング中の詳細な動作は以下の通りです。

  1. 各試行ごとに新しいPyTorchJobが作成されます
  2. job_converterでパラメータが適用されます
  3. ジョブがKubernetesクラスタに送信されます
  4. ジョブの完了を待機します（タイムアウト設定あり）
  5. ジョブのログをKubernetesからローカルにダウンロードし、outputs/<study_name>/trial_<n>.txtとして保存します
  6. 保存されたログファイルに対してvalue_extractorを実行し、目的値を抽出します
  7. ジョブは自動的にクリーンアップされます
  8. Optunaが次の試行のパラメータを提案します

チューニングの中断と再開​

PyTorchJobTunerはチューニングセッションの中断と再開をサポートしています。 この機能は長時間のチューニングや、AWS環境などで認証時間が限られている場合に役立ちます。

チューニングセッションを永続化するためには、dp_pathパラメータを指定します。

tuner = PyTorchJobTuner(
    job_name="my-training-job",
    namespace="default",
    db_path="./study.db"  # 同じデータベースファイルを指定
)

再開時に前回のセッションで未完了のジョブが存在する場合、PyTorchJobTunerは自動的にジョブを検出し実行完了まで待機します。 未完了のジョブが存在しない場合は、新しい試行に基づくジョブを投入します。 この仕組みにより、予期しない中断があっても安全にジョブチューニングを再開できます。