Modernize metrics module with type hints and generic types #1993

sugatmahanti · 2025-11-05T00:13:15Z

SUMMARY:
This is part of #1927

Modernize type annotations using | operator and built-in generics in the metrics module as part of codebase modernization effort.

I also fixed typo for the function name SystemLoggingWraper -> SystemLoggingWrapper as part of this PR.

TEST PLAN:

make style
make quality
make tests

Notes:
Happy to address any comments! Thank you!

github-actions · 2025-11-05T00:13:23Z

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

gemini-code-assist · 2025-11-05T00:13:26Z

Summary of Changes

Hello @sugatmahanti, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request modernizes the type annotations within the metrics module, aligning them with contemporary Python syntax by utilizing the | operator for unions and optionals, and built-in generic types. This effort contributes to a broader codebase modernization initiative, enhancing readability and maintainability. Additionally, a minor but important typo in a class name has been corrected.

Highlights

Type Annotation Modernization: Updated type annotations across the metrics module to use the | operator for Union and Optional types, and built-in generic types like dict and list.
Typo Correction: Fixed a typo in the class name SystemLoggingWraper, renaming it to SystemLoggingWrapper.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request does a great job of modernizing the type hints in the metrics module, which significantly improves code clarity and maintainability. The typo fix for SystemLoggingWrapper is also a welcome correction. I've identified a few areas where the type hints could be further improved for better type safety and consistency. Specifically, several log_hyperparams methods could use a more specific dict type, and some deprecated methods in LoggerManager have type hints that could lead to runtime errors. Additionally, there's an opportunity to simplify some redundant type hints in frequency_manager.py. My detailed comments and suggestions are below.

gemini-code-assist · 2025-11-05T00:15:29Z

src/llmcompressor/metrics/logger.py

-        level: Optional[Union[int, str]] = None,
+        step: int | None = None,
+        wall_time: float | None = None,
+        log_types: str | list[str] | None = ALL_TOKEN,


The type hint for log_types allows None, but the called function self.metric.log_scalar does not accept None for its log_types parameter. This could lead to a runtime TypeError if None is passed explicitly. The | None should be removed to match the signature of the wrapped function.

Suggested change

log_types: str | list[str] | None = ALL_TOKEN,

log_types: str | list[str] = ALL_TOKEN,

gemini-code-assist · 2025-11-05T00:15:29Z

src/llmcompressor/metrics/logger.py

+        values: dict[str, float],
+        step: int | None = None,
+        wall_time: float | None = None,
+        log_types: str | list[str] | None = ALL_TOKEN,


The type hint for log_types allows None, but the called function self.metric.log_scalars does not accept None for its log_types parameter. This could lead to a runtime TypeError if None is passed explicitly. The | None should be removed to match the signature of the wrapped function.

Suggested change

log_types: str | list[str] | None = ALL_TOKEN,

log_types: str | list[str] = ALL_TOKEN,

gemini-code-assist · 2025-11-05T00:15:30Z

src/llmcompressor/metrics/logger.py

-        log_types: Union[str, List[str]] = ALL_TOKEN,
-        level: Optional[Union[int, str]] = None,
+        params: dict,
+        log_types: str | list[str] | None = ALL_TOKEN,


The type hint for log_types allows None, but the called function self.metric.log_hyperparams does not accept None for its log_types parameter. This could lead to a runtime TypeError if None is passed explicitly. The | None should be removed to match the signature of the wrapped function.

Suggested change

log_types: str | list[str] | None = ALL_TOKEN,

log_types: str | list[str] = ALL_TOKEN,

gemini-code-assist · 2025-11-05T00:15:30Z

src/llmcompressor/metrics/logger.py

        self,
-        params: Dict,
-        level: Optional[Union[int, str]] = None,
+        params: dict,


For better type safety and consistency with the base class BaseLogger, it's recommended to use a more specific type for params. The BaseLogger.log_hyperparams method defines it as dict[str, float], and the _lambda_func this method calls also expects dict[str, float] | None for its values argument. Using dict is equivalent to dict[Any, Any] which is too broad.

Suggested change

params: dict,

params: dict[str, float],

gemini-code-assist · 2025-11-05T00:15:30Z

src/llmcompressor/metrics/logger.py

        return self._loggers

-    def log_hyperparams(self, params: Dict, level: Optional[Union[int, str]] = None):
+    def log_hyperparams(self, params: dict, level: int | str | None = None):


For consistency and better type safety, params should be typed as dict[str, float] instead of the general dict. This aligns with the BaseLogger.log_hyperparams signature and ensures that the wrapped loggers receive correctly typed data.

Suggested change

def log_hyperparams(self, params: dict, level: int | str | None = None):

def log_hyperparams(self, params: dict[str, float], level: int | str | None = None):

gemini-code-assist · 2025-11-05T00:15:30Z

src/llmcompressor/metrics/logger.py

-        params: Dict,
-        log_types: Union[str, List[str]] = ALL_TOKEN,
-        level: Optional[Union[int, str]] = None,
+        params: dict,


For better type safety and consistency, it's recommended to use a more specific type for params. The BaseLogger.log_hyperparams method defines it as dict[str, float]. Using dict is equivalent to dict[Any, Any] which is too broad and less safe.

Suggested change

params: dict,

params: dict[str, float],

gemini-code-assist · 2025-11-05T00:15:30Z

src/llmcompressor/metrics/utils/frequency_manager.py

+    current_log_step: LogStepType | None,
+    last_log_step: LogStepType | None,
+    log_frequency: LogStepType | None,
+    last_model_update_step: LogStepType | None = None,


The type hint LogStepType | None is redundant because LogStepType is already an alias for int | float | None. Using just LogStepType would be cleaner and less confusing.

Suggested change

current_log_step: LogStepType | None,

last_log_step: LogStepType | None,

log_frequency: LogStepType | None,

last_model_update_step: LogStepType | None = None,

current_log_step: LogStepType,

last_log_step: LogStepType,

log_frequency: LogStepType,

last_model_update_step: LogStepType = None,

yeah this comment looks appropriate

sugatmahanti · 2025-11-05T00:16:15Z

I will update Structural Pattern Matching for this module in a follow up commit

sugatmahanti · 2025-11-05T00:17:22Z

I will also address the comments by Gemini as well in a follow up commit.

brian-dellabetta

thanks for the contribution. The gemini comments look appropriate to me

brian-dellabetta · 2025-11-06T18:17:52Z

src/llmcompressor/metrics/utils/frequency_manager.py

+    current_log_step: LogStepType | None,
+    last_log_step: LogStepType | None,
+    log_frequency: LogStepType | None,
+    last_model_update_step: LogStepType | None = None,


yeah this comment looks appropriate

brian-dellabetta · 2025-11-06T18:18:10Z

src/llmcompressor/metrics/logger.py

        self,
-        params: Dict,
-        level: Optional[Union[int, str]] = None,
+        params: dict,


Modernize metrics module with type hints and generic types

bf791db

gemini-code-assist bot reviewed Nov 5, 2025

View reviewed changes

sugatmahanti mentioned this pull request Nov 5, 2025

[Python 3.10+] Modernize codebase with Python 3.10+ idioms and features #1927

Open

20 tasks

brian-dellabetta reviewed Nov 6, 2025

View reviewed changes

	log_types: str \| list[str] \| None = ALL_TOKEN,
	log_types: str \| list[str] = ALL_TOKEN,

	def log_hyperparams(self, params: dict, level: int \| str \| None = None):
	def log_hyperparams(self, params: dict[str, float], level: int \| str \| None = None):

Modernize metrics module with type hints and generic types #1993

Are you sure you want to change the base?

Modernize metrics module with type hints and generic types #1993

Conversation

sugatmahanti commented Nov 5, 2025

Uh oh!

github-actions bot commented Nov 5, 2025

Uh oh!

gemini-code-assist bot commented Nov 5, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

brian-dellabetta Nov 6, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

brian-dellabetta Nov 6, 2025

Choose a reason for hiding this comment

Uh oh!

sugatmahanti commented Nov 5, 2025

Uh oh!

sugatmahanti commented Nov 5, 2025

Uh oh!

brian-dellabetta left a comment

Choose a reason for hiding this comment

Uh oh!

brian-dellabetta Nov 6, 2025

Choose a reason for hiding this comment

Uh oh!

brian-dellabetta Nov 6, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants