Skip to content

Feat/mlflow nested run support#10

Open
illorens wants to merge 2 commits into
mainfrom
feat/mlflow-nested-run-support
Open

Feat/mlflow nested run support#10
illorens wants to merge 2 commits into
mainfrom
feat/mlflow-nested-run-support

Conversation

@illorens

@illorens illorens commented Jun 4, 2026

Copy link
Copy Markdown
Collaborator

Add nested run support to MLflow logger


Description

This PR extends MLflowLogger with support for nested MLflow runs, allowing child runs to be logged under a shared parent run. It also adds the ability to attach tags to both parent and child runs.

The logger now operates in two modes:

  • Standard mode (parent_run_name=None): existing behaviour is preserved — reuses an active run if one exists, otherwise creates a new run named run_name.
  • Nested mode (parent_run_name=<name>): creates or reuses a parent run named parent_run_name, then starts a nested child run named run_name. All metrics and parameters are logged to the child run.

New constructor parameters:

  • parent_run_name — enables nested mode when set
  • run_tags — tags applied to the child (or standard) run
  • parent_tags — tags applied to the parent run in nested mode

No breaking changes in MLflowLogger. All existing call sites that omit the new parameters continue to work identically.

⚠️ Breaking change: The Trainer constructor parameters train_loss_fct and val_loss_fct have been renamed to train_loss_fn and val_loss_fn respectively. Any call sites passing these as keyword arguments will need to be updated.


Checklist

  • Code is well documented and follows project conventions (linted, formatted)
  • CHANGELOG updated (if applicable)
  • License header is added to all new files
  • For a new contributors, their name has been added to the AUTHORS file
  • No sensitive data or secrets are exposed
  • Unit tests added covering the new code

@illorens illorens requested review from frazane and hlasco June 4, 2026 16:17
@illorens illorens self-assigned this Jun 4, 2026
@illorens illorens added the enhancement New feature or request label Jun 4, 2026
Comment thread obsweatherscale/logger.py
self,
experiment_name: str | None = None,
run_name: str | None = None,
parent_run_name: str | None = None,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

There's no guarantee than run names are unique, I think you should run_id instead of run name

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.

2 participants