Add missing type hints in checkpointing.py - #4153
Open
RudrenduPaul wants to merge 1 commit into
Open
Conversation
Adds parameter and return type hints to load_accelerator_state, save_accelerator_state, save_custom_state, and load_custom_state, mirroring the types already established for the equivalent save_accelerator_state parameters in the same file.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR does
Adds missing type hints to four functions in
src/accelerate/checkpointing.py:save_accelerator_state: adds the missing-> Pathreturn type (thefunction already returns
output_dir, aPath, but the annotation wasabsent even though every parameter was already typed).
load_accelerator_state: was entirely untyped. Addsinput_dir: str,models: list,optimizers: list,schedulers: list,dataloaders: list,process_index: int,scaler: Optional[GradScaler] = None,load_kwargs: Optional[dict] = None, and a-> dictreturntype. Each of these mirrors the type already established for the
equivalent parameter in the sibling
save_accelerator_statefunctionin the same file (e.g.
process_index: intandscaler: Optional[GradScaler] = Noneare copied verbatim), and the-> dictreturn type matches the function's own docstring (
Returns: dict) andthe actual returned value (
override_attributes = dict()).save_custom_state/load_custom_state: addspath: str(matchingthe
output_dir: str/input_dir: strconvention used elsewhere inthe file) and
-> None(both functions have no return statement).Scope note: intentionally left untyped
objinsave_custom_state/load_custom_state: this is aduck-typed parameter (only requires
.state_dict()/.load_state_dict()), and there's no established precedent for itstype in this file, so it was left unannotated rather than guess.
map_locationinload_accelerator_state: the docstring saysstr,but the actual runtime value is reassigned to a
PartialState().deviceobject when
"on_device"is passed, so a single accurate type isn'tobvious from reading the function. Left unannotated rather than
introduce an inaccurate hint.
load_model_func_kwargs(**kwargs): left untyped, consistent withhow
**kwargsis handled elsewhere in this typing-cleanup series.No behaviour change
Type-hint-only change. No logic was modified.
Testing
ruff check src/accelerate/checkpointing.py— 3 pre-existing findings(an unnecessary
dict()call, aenumerate()index-lookup nit, and ablind
except Exception), none on changed lines, none introduced bythis change (verified against upstream
mainbefore/after).ruff format --check src/accelerate/checkpointing.py— passes, filealready formatted.
python3 -m py_compileon the modified file — passes.Related PRs
Follows the same pattern as my other open type-hint PRs in this repo:
#4123, #4124, #4125.