Test/torch utils unit tests #206
Conversation
- Add return type hint to get_tokenizer() - Add type hints and checkpoint validation to generate.py main() - Add parameter type hints to suppress_output() in torch/utils.py Improves IDE support and catches potential bugs early.
- Test suppress_output for different ranks - Test init_distributed with single/multi-GPU - Test CUDA availability and device validation - Test NCCL warmup and error handling - Test cleanup on exceptions
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
| mock_torch.cuda.is_available.return_value = False | ||
|
|
||
| with pytest.raises(RuntimeError, match="CUDA is not available"): | ||
| init_distributed() |
There was a problem hiding this comment.
The new tests around init_distributed expect the helper to raise descriptive RuntimeErrors when CUDA is unavailable and to perform cleanup on various failures (e.g. rank out of range, NCCL warm‑up errors). However, gpt_oss/torch/utils.py:init_distributed does not call torch.cuda.is_available, check torch.cuda.device_count, or wrap its warm‑up code in try/except, so it will happily return even when CUDA is mocked as unavailable and will never invoke dist.destroy_process_group. As written, test_init_distributed_cuda_not_available_raises (and the similar tests that follow) will fail because the function under test has no such error handling. Either extend init_distributed to implement these safety checks and cleanup logic or relax the assertions in these tests.
Useful? React with 👍 / 👎.
unit tests for torch distributed initialization utilities.
Tests Added: