[Train] Use user-defined train function name in training thread name and error stack traces

[avigyabb]

Description

When a user's training function raises an exception, the training thread (and the resulting stack trace) is named after Ray's internal wrapper train_fn_with_final_checkpoint_flush instead of the user-defined function.

The root cause is in RayTrainWorker.run_train_fn: the user's training function is wrapped by an inner train_fn_with_final_checkpoint_flush , and this wrapper is what gets handed to ThreadRunner. ThreadRunner names the training thread via get_callable_name(target), so it reads the wrapper's __name__ and produces
TrainingThread(train_fn_with_final_checkpoint_flush). We apply functools.wraps(train_fn) to the inner wrapper so it inherits the user function's __name__. This is similar to how construct_train_func preserves the user's name on train_fn via functools.wraps.

Related issues

Fixes TRAIN-744 on Jira

Testing

Added a unit test test_training_thread_name_matches_user_fn in python/ray/train/v2/tests/test_worker.py. It runs RayTrainWorker.run_train_fn with a user function named my_custom_train_fn and asserts the launched training thread is named TrainingThread(my_custom_train_fn) (rather than TrainingThread(train_fn_with_final_checkpoint_flush)). Without the functools.wraps fix this assertion fails, so the test guards the fix.

Reproduction

Repro script and before/after logs: https://gist.github.com/avigyabb/626151d1b05526f38b221198f6858c65

The script uses Ray's real ThreadRunner and get_callable_name and rebuilds the exact wrapper from worker.py, toggling @functools.wraps to show both states in a single run (no cluster required). It also has an --e2e mode that drives the real DataParallelTrainer.

Results (thread_name_before_after.log):

	Thread / stack-trace name
Before fix	TrainingThread(train_fn_with_final_checkpoint_flush)
After fix	TrainingThread(train_fn_per_worker)

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request imports wraps from functools and applies the @wraps(train_fn) decorator to the train_fn_with_final_checkpoint_flush helper function in worker.py to preserve the metadata of the original training function. There are no review comments, so I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[liulehui]

can this change be tested with a unit test?

could you also share a repro script so that:

1. repro the original issue with the script
2. show that this fix actually fix the issue,
3. link the repro script and corresponding logs in the PR description.