ハングジョブ検出

RayCluster マニフェストにホスト IP 環境変数を追加する

RayCluster YAML workerGroupSpecsの headGroupSpecと の両方に次の環境変数を追加します。これにより、トレーニングコンテナ内で実行されている Python ライブラリがホストのジョブモニタリングサービスと通信できるようになります。

spec:
  headGroupSpec:
    template:
      spec:
        containers:
        - name: ray-head
          env:
          - name: HYPERPOD_JMA_HOST
            valueFrom:
              fieldRef:
                fieldPath: status.hostIP
  workerGroupSpecs:
  - template:
      spec:
        containers:
        - name: ray-worker
          env:
          - name: HYPERPOD_JMA_HOST
            valueFrom:
              fieldRef:
                fieldPath: status.hostIP

トレーニングコンテナイメージにツールキットライブラリをインストールする

PyPI ウェブサイトからトレーニングコンテナイメージに toolkit-for-ray-on-sagemaker-ai パッケージを追加します。ライブラリは SageMaker Distribution イメージにプリインストールされています。

pip install toolkit-for-ray-on-sagemaker-ai

トレーニング関数にモニタリングを追加する

有意義な作業を行う前に、トレーニング関数SageMakerLogMonitoring.start()内で を呼び出します。これにより、トレーニングの開始からモニタリングがアクティブになり、モデルのロード中または最初のフォワードパス中に発生するハングを検出できます。

from toolkit_for_ray_on_sagemaker_ai.log_monitoring import (
    SageMakerLogMonitoring,
    LogMonitorConfig,
)
from ray.train import RunConfig, FailureConfig, ScalingConfig
from ray.train.torch import TorchTrainer

def train_func():
    # Start monitoring BEFORE any meaningful work
    SageMakerLogMonitoring(config=LogMonitorConfig(
        enabled=True,
        rules=[
            {
                "name": "training_progress",
                "type": "log_pattern",
                "enabled": True,
                "log_pattern": "(Epoch|Step|Iteration) \\d+",
                "timeout_minutes": 10,
                "start_timeout_minutes": 30,
                "stop_pattern": "Training complete",
                "fault_on_match": False,
            },
            {
                "name": "oom_detection",
                "type": "log_pattern",
                "enabled": True,
                "log_pattern": "CUDA out of memory|OutOfMemoryError|OOM",
                "fault_on_match": True,
            },
        ],
        action="cancel",
    )).start()

    # ... your training loop
    for epoch in range(num_epochs):
        print(f"Epoch {epoch}")  # This output is what the rule monitors
        train_one_epoch(model, dataloader)

    print("Training complete")  # Matches stop_pattern, deactivates the rule

# Configure FailureConfig so Ray Train automatically restarts all workers
# when the cancel action terminates a hung worker.
trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
    run_config=RunConfig(
        failure_config=FailureConfig(max_failures=3),
    ),
)

cancel アクションがハングワーカーを終了すると、Ray ドキュメントの Ray Train の FailureConfig はワーカーの障害を検出し、最後のチェックポイントからすべてのワーカーを再起動します。ジョブが完全に失敗するまでに必要な自動復旧の試行回数max_failuresに設定します。を使用しない場合FailureConfig、1 つのワーカー終了でジョブ全体が失敗します。

Ray Train は、再起動時に最後に保存されたチェックポイントから復元するため、最新のチェックポイント以降の作業を超えてトレーニングの進行状況が失われることはありません。トレーニングコードが定期的にチェックポイントを保存する場合 (各エポック境界など）、ジョブは最後に保存された状態から自動的に再開されます。