diff --git a/tests/functional_tests/shell_test_utils/_run_training.sh b/tests/functional_tests/shell_test_utils/_run_training.sh index cd22e9b1104..d3f166c88eb 100644 --- a/tests/functional_tests/shell_test_utils/_run_training.sh +++ b/tests/functional_tests/shell_test_utils/_run_training.sh @@ -187,10 +187,14 @@ DISTRIBUTED_ARGS=( --redirects "3" ) +FT_LAUNCHER_ARGS=( + --max-restarts=3 +) + # Start training if [[ "$IS_NEMO_TEST" == "true" ]]; then if [[ "$LAUNCHER" == "ft_launcher" ]]; then - ft_launcher ${DISTRIBUTED_ARGS[@]} \ + ft_launcher ${DISTRIBUTED_ARGS[@]} ${FT_LAUNCHER_ARGS[@]} \ --no-python /opt/venv/bin/$TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? else uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} \ @@ -198,7 +202,7 @@ if [[ "$IS_NEMO_TEST" == "true" ]]; then fi else if [[ "$LAUNCHER" == "ft_launcher" ]]; then - ft_launcher ${DISTRIBUTED_ARGS[@]} \ + ft_launcher ${DISTRIBUTED_ARGS[@]} ${FT_LAUNCHER_ARGS[@]} \ $TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? else uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} \