Skip to content

Add missing type hints in checkpointing.py - #4153

Open
RudrenduPaul wants to merge 1 commit into
huggingface:mainfrom
RudrenduPaul:add-checkpointing-type-hints
Open

Add missing type hints in checkpointing.py#4153
RudrenduPaul wants to merge 1 commit into
huggingface:mainfrom
RudrenduPaul:add-checkpointing-type-hints

Conversation

@RudrenduPaul

Copy link
Copy Markdown

What this PR does

Adds missing type hints to four functions in src/accelerate/checkpointing.py:

  • save_accelerator_state: adds the missing -> Path return type (the
    function already returns output_dir, a Path, but the annotation was
    absent even though every parameter was already typed).
  • load_accelerator_state: was entirely untyped. Adds input_dir: str,
    models: list, optimizers: list, schedulers: list,
    dataloaders: list, process_index: int, scaler: Optional[GradScaler] = None, load_kwargs: Optional[dict] = None, and a -> dict return
    type. Each of these mirrors the type already established for the
    equivalent parameter in the sibling save_accelerator_state function
    in the same file (e.g. process_index: int and scaler: Optional[GradScaler] = None are copied verbatim), and the -> dict
    return type matches the function's own docstring (Returns: dict) and
    the actual returned value (override_attributes = dict()).
  • save_custom_state / load_custom_state: adds path: str (matching
    the output_dir: str / input_dir: str convention used elsewhere in
    the file) and -> None (both functions have no return statement).

Scope note: intentionally left untyped

  • obj in save_custom_state / load_custom_state: this is a
    duck-typed parameter (only requires .state_dict() /
    .load_state_dict()), and there's no established precedent for its
    type in this file, so it was left unannotated rather than guess.
  • map_location in load_accelerator_state: the docstring says str,
    but the actual runtime value is reassigned to a PartialState().device
    object when "on_device" is passed, so a single accurate type isn't
    obvious from reading the function. Left unannotated rather than
    introduce an inaccurate hint.
  • load_model_func_kwargs (**kwargs): left untyped, consistent with
    how **kwargs is 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, a enumerate() index-lookup nit, and a
    blind except Exception), none on changed lines, none introduced by
    this change (verified against upstream main before/after).
  • ruff format --check src/accelerate/checkpointing.py — passes, file
    already formatted.
  • python3 -m py_compile on the modified file — passes.

Related PRs

Follows the same pattern as my other open type-hint PRs in this repo:
#4123, #4124, #4125.

Note: Claude Code was used to assist in drafting this change. All
changes were reviewed by the submitter.

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