Skip to content

Save WandB run information to model folder#3017

Merged
AlexEMG merged 10 commits intomainfrom
maxim/save_wandb_data
Aug 26, 2025
Merged

Save WandB run information to model folder#3017
AlexEMG merged 10 commits intomainfrom
maxim/save_wandb_data

Conversation

@maximpavliv
Copy link
Copy Markdown
Contributor

This Pull Request adds a run identifiers saving mechanism to the WandbLogger:

When WandbLogger is configured as a logger - the WandB run information (entity, project, run_id) are saved to a file named wandb_info.yaml in the model folder.
This information can be used later to recover the WandB run to which the training process has been logged.

@MMathisLab
Copy link
Copy Markdown
Member

hi @maximpavliv shall we add this to the docs?

@MMathisLab MMathisLab added the enhancement New feature or request label Jun 25, 2025
@MMathisLab MMathisLab requested a review from Copilot June 26, 2025 13:51
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

Adds a mechanism to persist WandB run identifiers to disk and tightens argument filtering when instantiating registered objects. Key changes:

  • WandbLogger now requires a train_folder, writes run info to wandb_info.yaml, and prints a confirmation.
  • Registry’s build_from_cfg filters extra constructor kwargs based on the function signature.
  • Training entrypoint passes train_folder from the data loader into the logger config.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File Description
deeplabcut/pose_estimation_pytorch/runners/logger.py Imported yaml; added _save_wandb_info() and updated CSVLogger to accept str paths
deeplabcut/pose_estimation_pytorch/registry.py Enhanced build_from_cfg to drop unexpected kwargs before calling constructors
deeplabcut/pose_estimation_pytorch/apis/training.py Modified LOGGER.build invocation to inject train_folder from loader.model_folder
Comments suppressed due to low confidence (2)

deeplabcut/pose_estimation_pytorch/runners/logger.py:399

  • The constructor only accepts str for train_folder but immediately converts to Path; consider annotating as str | Path to reflect the accepted types and improve clarity.
    def __init__(self, train_folder: str, log_filename: str) -> None:

deeplabcut/pose_estimation_pytorch/apis/training.py:102

  • The variable loader isn’t defined in this scope, which will raise a NameError. You should reference the actual training loader variable (e.g., train_loader.model_folder) or pass the folder path directly into train.
            {**logger_config, "model": model, "train_folder": loader.model_folder}

with open(output_path, "w") as f:
yaml.dump(wandb_info, f)

print(f"WandB run info saved to {output_path}")
Copy link

Copilot AI Jun 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] Using print in library code can clutter stdout; consider using the project’s logger (e.g., logging.info) so output can be controlled by the user’s log level.

Suggested change
print(f"WandB run info saved to {output_path}")
logging.info(f"WandB run info saved to {output_path}")

Copilot uses AI. Check for mistakes.
@maximpavliv
Copy link
Copy Markdown
Contributor Author

@MMathisLab I integrated Copilot's suggestion, and added a description of the feature to the docs.

@AlexEMG AlexEMG merged commit d9c856b into main Aug 26, 2025
7 of 8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants