Add type hints to top-level public API functions#8079
Conversation
Annotate the public API entry points in deepspeed/__init__.py (initialize, init_inference, default_inference_config, add_config_arguments, tp_model_init, set_optimizer_flags) with precise parameter and return types, addressing deepspeedai#8074. Scoped first increment for the top-level deepspeed.* public API surface; no runtime behavior changes. Signed-off-by: Arun Sharma <sharm485@umn.edu>
arunshar
force-pushed
the
typing/public-api-hints
branch
from
c1711bb to
d33fe7f
Compare
![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.
Reviewed commit: d33fe7ff58
ℹ️ 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| config: Optional[Union[str, Dict[str, Any]]] = None, | ||
| mesh_param: Any = None, | ||
| config_params: Optional[Union[str, Dict[str, Any]]] = None | ||
| ) -> Tuple[DeepSpeedEngine, Optional[Optimizer], Optional[DataLoader], Optional[_LRScheduler]]: |
There was a problem hiding this comment.
For calls that enable common DeepSpeed wrappers, this tuple annotation is narrower than what initialize() actually returns: ZeRO/fp16/bf16 paths replace engine.optimizer with DeepSpeedOptimizer wrappers rather than torch.optim.Optimizer, DeepSpeedEngine.deepspeed_io() returns a DeepSpeedDataLoader when training_data is supplied rather than a torch.utils.data.DataLoader, and JSON-configured schedulers can be DeepSpeed scheduler objects rather than _LRScheduler. Since this is the public annotation rendered for deepspeed.initialize, typed callers will see false errors or assume PyTorch-only members for valid DeepSpeed configurations; use DeepSpeed wrapper/protocol types (or a wider union) for the returned tuple.
Useful? React with 👍 / 👎.
…pper types The Codex reviewer correctly noted the return tuple over-claimed PyTorch-only types. deepspeed.initialize() can return DeepSpeed wrappers, not the torch base classes: ZeRO/fp16/bf16 paths set engine.optimizer to a DeepSpeedOptimizer (which subclasses object, not torch.optim.Optimizer); engine.training_dataloader is a DeepSpeedDataLoader (not a torch DataLoader); and JSON-configured schedulers are DeepSpeed objects (not _LRScheduler). Widen the return to Tuple[DeepSpeedEngine, Optional[Union[Optimizer, DeepSpeedOptimizer]], Optional[DeepSpeedDataLoader], Any] so typed callers do not see false errors for valid DeepSpeed configurations. Signed-off-by: Arun Sharma <sharm485@umn.edu>
|
Thanks for the review. Good catch — the return tuple over-claimed PyTorch-only types. |
1 participant
What
Adds precise type hints to the top-level public API functions in
deepspeed/__init__.py, addressing #8074:initialize— annotates the previously-untypedargs,mpu,collate_fn,config,mesh_param,config_paramsparameters and adds the return type.init_inference,default_inference_config,add_config_arguments,tp_model_init,set_optimizer_flags.Why
#8074 asks for type annotations on the public API. This is a scoped first increment covering the top-level
deepspeed.*entry points (the functions users call directly), which keeps the diff small and reviewable; follow-ups can extend to deeper modules.Notes
pre-commit run --files deepspeed/__init__.pyis green (yapf, flake8, check-torchdist, check-license, codespell).initializereturnsTuple[DeepSpeedEngine, Optional[Optimizer], Optional[DataLoader], Optional[_LRScheduler]]— bothDeepSpeedHybridEngineandPipelineEnginesubclassDeepSpeedEngine.args,mpu,mesh_param) are annotatedAny.Opened as a draft for early feedback on the scope and the chosen annotations.
Addresses #8074