feat(types): add type hints to public comm API functions#8118
Conversation
Adds type hints to the most commonly used public API functions in deepspeed.comm: - init_distributed() - is_initialized() - destroy_process_group() - get_world_size() - get_rank() - get_local_rank() - get_global_rank() This is a first pass - more functions can be typed incrementally. Relates to deepspeedai#8074 Signed-off-by: NovusEdge <novusedge0@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7039a88198
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| timeout: timedelta = default_pg_timeout, | ||
| init_method: Optional[str] = None, | ||
| dist_init_required: Optional[bool] = None, | ||
| config: Optional[Dict[str, Any]] = None, |
There was a problem hiding this comment.
When a caller follows this new public annotation and passes a plain dict for config, init_distributed() immediately calls configure(deepspeed_config=config), and configure() dereferences deepspeed_config.comms_config; dictionaries do not have that attribute, so initialization fails before reaching the distributed setup. The supported object here is the parsed DeepSpeed config carrying .comms_config (or the parameter should be normalized before use), not Dict[str, Any].
Useful? React with 👍 / 👎.
Summary
First pass at adding type hints to the public API in
deepspeed.comm:DFunctions typed:
init_distributed()is_initialized()destroy_process_group()get_world_size()get_rank()get_local_rank()get_global_rank()Why
Better IDE autocomplete, mypy support, and easier onboarding for new contributors. Types are basically free documentation that doesnt go stale.
Notes
This is incremental — more functions can be typed in follow-up PRs. Started with the most commonly used ones.
Relates to #8074