Skip to content

Type-fix init_wandb key parameter and _save_wandb_id return#96

Open
lonexreb wants to merge 1 commit intoNVlabs:mainfrom
lonexreb:fix/wandb-utils-key-optional-and-return-types
Open

Type-fix init_wandb key parameter and _save_wandb_id return#96
lonexreb wants to merge 1 commit intoNVlabs:mainfrom
lonexreb:fix/wandb-utils-key-optional-and-return-types

Conversation

@lonexreb
Copy link
Copy Markdown
Contributor

@lonexreb lonexreb commented May 4, 2026

Problem

Two small typing nits in src/alpamayo_r1/common/wandb_utils.py:

1. init_wandb declares key: str but accepts None.

def init_wandb(
    key: str,
    ...
) -> None:
    """Initialize wandb."""
    if key is not None:
        # login wandb, otherwise it will load from the home directory
        wandb.login(key=key)

The body explicitly branches on if key is not None and the inline comment ("otherwise it will load from the home directory") describes the None case as the supported "use cached credentials" path. Type checkers flag any init_wandb(None, ...) call site even though that's documented behavior.

2. _save_wandb_id has no return annotation.

def _save_wandb_id(output_dir: str, wandb_id: str):
    """Save the wandb id to the output_dir."""

It returns None. The neighboring _check_wandb_id is fully annotated (-> str | None); _save_wandb_id is the odd one out.

Fix

Pure typing/docstring change. No runtime behavior change.

  • Update init_wandb parameter annotation to key: str | None.
  • Expand the init_wandb one-word docstring into an Args block — the function is called from external launchers and was previously self-documenting only via parameter names.
  • Add -> None to _save_wandb_id.

Two small typing nits in src/alpamayo_r1/common/wandb_utils.py:

- init_wandb declares `key: str` but the body branches on
  `if key is not None`, and the comment "otherwise it will load from
  the home directory" makes clear that None is the supported "use
  cached credentials" path. Type checkers flag any `init_wandb(None,
  ...)` call site, even though that's the documented behavior.
  Update the annotation to `str | None`.

  While here, expand the one-word docstring to actually describe each
  parameter (the function is called from external launchers and was
  previously self-documenting only via the parameter names).

- _save_wandb_id has no return annotation. The function returns None
  (writes to disk and stops). Add `-> None` for consistency with
  `_check_wandb_id` immediately below it (which already has its
  return type annotated).

Pure typing/docstring change. No runtime behavior change.

Signed-off-by: lonexreb <[email protected]>
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