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

PyTorchJobTunerによるジョブチューニング例

このドキュメントは、PyTorchJobTunerを用いてKubernetesクラスタ上のPyTorchJobを実際にチューニングする例を説明します。 ここでは、intelligence/components/zenith-tune/examples/integration/kubernetesに含まれるMNISTの学習スクリプトに対して、データローダのワーカー数とOpenMPのスレッド数をチューニングすることで学習時間を削減する方法を示します。

前提条件​

ジョブチューニング例を実行する前に、以下を準備してください：

  1. PyTorchJobを投入・取得可能なKubernetesクラスタ環境
  2. kubectlがインストールされていること
  3. ZenithTuneがインストールされたPython環境

アクセス可能なKubernetesクラスタを所有していない場合は、minikubeを用いてローカル環境にクラスタを構築することでジョブチューニング例を実行できます。

# Minikubeクラスタの起動とコンテキスト設定
minikube start
kubectl config use-context minikube

# Kubeflow Training Operatorのインストール
kubectl apply --server-side -k "github.com/kubeflow/training-operator.git/manifests/overlays/standalone?ref=v1.8.1"

ジョブチューニングで使用するファイル​

FAIBのintelligence/components/zenith-tune/examples/integration/kubernetesに含まれる、ジョブチューニングで使用するファイルについて説明します。

job_training.yaml​

このYAMLファイルは、シンプルなMNISTの学習スクリプトを実行するPyTorchJobを定義しています。 なお、学習スクリプトはConfigMapで定義されています。

学習スクリプトの主な特徴：

  • DataLoaderワーカーを設定するための--num-workers引数を受け付ける
  • PyTorchスレッドを設定するためのOMP_NUM_THREADS環境変数を読み取る
  • 学習時間が出力される

tune_job_training.py​

このPythonスクリプトは、PyTorchJobTunerを用いてjob_training.yamlで定義されたMNIST学習のnum-workersとOMP_NUM_THREADSをチューニングします。

まず、PyTorchJobTunerはチューニング対象のジョブ名job_nameとその名前空間get_namespaceを受け取り、ジョブを特定します。 submit_namespaceはジョブを投入する際の名前空間でget_namespaceとは独立に設定することができます。 db_pathは既存のジョブチューニング結果を再開・分析するために指定するデータベースのパスです。

    tuner = PyTorchJobTuner(
        job_name=args.job_name,
        get_namespace=args.namespace,
        submit_namespace=args.namespace,
        db_path=args.db_path,
    )

次にユーザは、既存のPyTorchJobをチューニングジョブへと変換する関数job_converterと目的値を抽出する関数value_extractorを定義し、optimize関数に渡すことでジョブチューニングを実行します。 job_converterとvalue_extractorの実装例は以降で説明します。

    tuner.optimize(
        job_converter=job_converter,
        value_extractor=value_extractor,
        n_trials=args.n_trials,
    )

job_converterの実装例​

このスクリプトのjob_converterでは、チューニングパラメータomp_num_threadsとnum_workersを定義し、omp_num_threadsは環境変数に、num_workersはコマンドの引数に設定します。 環境変数の設定にはPyTorchJobのset_env関数を使用しています。

def job_converter(trial: Trial, job: PyTorchJob) -> PyTorchJob:
    """
    Update job definition with different OMP_NUM_THREADS and num_workers settings.

    Args:
        trial: Optuna trial object for suggesting parameters
        job: PyTorchJob object to update

    Returns:
        Updated PyTorchJob object
    """
    # Suggest number of threads and workers
    num_threads = trial.suggest_int("omp_num_threads", 1, 8)
    num_workers = trial.suggest_int("num_workers", 0, 4)

    # Set environment variable using convenient API
    job.set_env("OMP_NUM_THREADS", str(num_threads))

    ...

num_workersはコマンドの引数に追加するため、CommandBuilderクラスを用いてコマンドラインを変更します。 そのために、PyTorchJob.get_commandでオリジナルのコマンドを取得しますが、PyTorchJobのコマンドはArray形式で表現されているため、変更対象のコマンド文字列のみ抽出します。 今回の例では、コマンドが['sh', '-c', 'command']のように配列の2番目に格納されているため、2番目の文字列のみをCommandBuilderに与えています。

そして、CommandBuilderのappend関数を用いて--num-workersオプションを追加します。 今回は、既存のコマンドに--num-workersオプションが含まれていないことがわかっているためappend関数を使用していますが、すでに存在するオプションの値を更新したい場合はupdate関数を使用します。

    ...

    # Update command to include num_workers argument using CommandBuilder
    current_command = job.get_command()
    assert (
        current_command
        and len(current_command) >= 3
        and current_command[0] == "sh"
        and current_command[1] == "-c"
    ), f"Expected ['sh', '-c', 'command'] format, got: {current_command}"

    # Modify only the actual command part (index 2)
    actual_command = current_command[2]
    builder = CommandBuilder(actual_command)
    builder.append(f"--num-workers {num_workers}")
    ...

その後、PyTorchJob.set_commandで新たなコマンドを設定し、更新後のjobを返すことでjob_converterの実装は完了です。

    ...

    # Replace the command part while keeping sh -c wrapper
    new_command = current_command.copy()
    new_command[2] = builder.get_command()
    job.set_command(new_command)

    print(
        f"Trial {trial.number}: OMP_NUM_THREADS={num_threads}, num_workers={num_workers}"
    )

    return job

value_extractorの実装例​

value_extractorはCommandOutputTunerの用法と同様、第1引数にログファイルのパスが与えられるため、ログから目的値を抽出する処理を実装します。 今回は、Elapsed time:以降に学習時間が出力されるため、この値を抽出して返します。 学習失敗やジョブ中断などによって値の抽出に失敗した場合はNoneを返すことで、そのTrialを棄却します。

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

    # Look for the line with elapsed time
    match = re.search(r"Elapsed time: ([0-9.]+) seconds", logs)
    if match:
        elapsed_time = float(match.group(1))
        return elapsed_time
    else:
        print(f"Could not find elapsed time in {log_path}")
        return None

ジョブチューニングの実行​

ジョブチューニングを実行するための既存ジョブとしてjob_training.yamlで定義されたPyTorchJobをKubernetesクラスタに投入します。

kubectl create -f job_training.yaml
  • Error from server (AlreadyExists): error when creating "job_training.yaml": configmaps "training-script" already existsのようなエラーが生じた場合、すでにjob_training.yamlは作成されているため無視して構いません。
  • 初回のcreateにはDocker imageのpullに数十分程度の時間を要する場合があります。

デプロイされたPyTorchJobの名前と学習時間を確認します。

$ kubectl get pytorchjobs
NAME                STATUS      AGE
job-training-76kxt  Succeeded   1m
$ kubectl logs job-training-76kxt-worker-0
...
Elapsed time: 78.9675 seconds

そして、上記で確認したジョブを対象にチューニングスクリプトを実行します。 試行回数は5回とします。

python tune_job_training.py --job-name job-training-76kxt --n-trials 5

チューニングが終了すると以下のようなログが出力されます（zenith-tuneのログ出力を一部省略）。 チューニングの結果{'omp_num_threads': 1, 'num_workers': 0}の組み合わせが、最も高速であることがわかりました。

[I 2025-08-25 14:27:10,064] Trial 0 finished with value: 15.9339 and parameters: {'omp_num_threads': 1, 'num_workers': 0}. Best is trial 0 with value: 15.9339.
[I 2025-08-25 14:28:14,611] Trial 1 finished with value: 27.2895 and parameters: {'omp_num_threads': 2, 'num_workers': 4}. Best is trial 0 with value: 15.9339.
[I 2025-08-25 14:29:37,069] Trial 2 finished with value: 58.804 and parameters: {'omp_num_threads': 4, 'num_workers': 3}. Best is trial 0 with value: 15.9339.
[I 2025-08-25 14:30:38,995] Trial 3 finished with value: 47.4847 and parameters: {'omp_num_threads': 3, 'num_workers': 4}. Best is trial 0 with value: 15.9339.
[I 2025-08-25 14:32:03,641] Trial 4 finished with value: 30.1066 and parameters: {'omp_num_threads': 3, 'num_workers': 0}. Best is trial 0 with value: 15.9339.
2025-08-25 14:32:03,653 - zenith-tune - INFO - Best trial: trial_id=1, value=15.9339, params={'omp_num_threads': 1, 'num_workers': 0}

チューニング結果の適用​

チューニングによって得られたパラメータをYAMLファイルに適用することで、今後実行される学習を高速化します。

job_training.yamlにパラメータを反映します。

    command:
    - sh
    - -c
    - |
-        python /scripts/train_mnist.py
+        OMP_NUM_THREADS=1 python /scripts/train_mnist.py --num-workers 0

パラメータを反映したPyTorchJobを実行し、実際に高速化されていることを確認します。

# configmaps "training-script" already existsは無視する
$ kubectl create -f job_training.yaml
pytorchjob.kubeflow.org/job-training-nbmtz created
Error from server (AlreadyExists): error when creating "job_training_new.yaml": configmaps "training-script" already exists

# 実行完了後
$ kubectl logs job-training-nbmtz-worker-0
...
Elapsed time: 16.0635 seconds

チューニング中の最適値15.9339 secondsと同程度の時間で学習が完了しており、ジョブチューニングによって元の学習時間78.9675 secondsと比較して、およそ4.9倍の高速化を達成しました。

このように、ZenithTuneが提供するPyTorchJobTunerを使用することで、手動での試行錯誤なしに最適なハイパーパラメータを効率的に発見し、学習時間を削減することができます。