Skip to content

feat(types): add type hints to public comm API functions#8118

Open
NovusEdge wants to merge 1 commit into
deepspeedai:masterfrom
NovusEdge:add-type-hints
Open

feat(types): add type hints to public comm API functions#8118
NovusEdge wants to merge 1 commit into
deepspeedai:masterfrom
NovusEdge:add-type-hints

Conversation

@NovusEdge

Copy link
Copy Markdown

Summary

First pass at adding type hints to the public API in deepspeed.comm :D

Functions 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

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>
@NovusEdge NovusEdge requested a review from GuanhuaWang as a code owner July 3, 2026 14:29

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread deepspeed/comm/comm.py
timeout: timedelta = default_pg_timeout,
init_method: Optional[str] = None,
dist_init_required: Optional[bool] = None,
config: Optional[Dict[str, Any]] = None,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Correct the config type hint

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 👍 / 👎.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant