Hung job detection

Add the host IP environment variable to your RayCluster manifest

Add the following environment variable to both headGroupSpec and workerGroupSpecs in your RayCluster YAML. This allows the Python library running inside the training container to communicate with the job monitoring service on the host.

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

Install the toolkit library in your training container image

Add the toolkit-for-ray-on-sagemaker-ai package from the PyPI website to your training container image. The library is preinstalled in SageMaker Distribution images.

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

Add monitoring to your training function

Call SageMakerLogMonitoring.start() inside your training function before doing any meaningful work. This ensures monitoring is active from the start of training and can detect hangs that occur during model loading or the first forward pass.

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),
    ),
)

When the cancel action terminates a hung worker, Ray Train's FailureConfig in the Ray documentation detects the worker failure and restarts all workers from the last checkpoint. Set max_failures to the number of automatic recovery attempts you want before the job fails permanently. Without FailureConfig, a single worker termination fails the entire job.

Because Ray Train restores from the last saved checkpoint on restart, no training progress is lost beyond the work since the most recent checkpoint. If your training code saves checkpoints periodically (for example, at each epoch boundary), the job resumes from the last saved state automatically.