Check unit test pass/failure and handle timeout error in halo test

[amd-sriram]

Motivation

Halo test takes more than 10 hours to run - https://github.com/ROCm/apex/actions/runs/23252552706/job/67612063034?pr=320

There is no graceful way to exit in case of timeout
run - https://github.com/ROCm/apex/actions/runs/23505553410/job/68420328568

Technical Details

To check for unit test failure, assert statement similar to nvidia apex is added -
https://github.com/NVIDIA/apex/blob/master/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py#L134.

torch.testing.assert_close(list_y, list_y2, msg=memory_format_str) 

The following changes are made to address timeout error:

1. Shorter timeout on init_process_group so failures surface faster, and timeout for the unit test
2. Try/except around test functions to catch DistBackendError
3. Proper cleanup with destroy_process_group in a finally block

Test Plan

Run the CI and check how much time the halo test takes with this change.

Test Result

Run - https://github.com/ROCm/apex/actions/runs/23668144565
The halo test takes ... hours.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.