Keras automatic checkpoint

[WeichenXu123]

🛠 DevTools 🛠

Install mlflow from this PR

pip install git+https://github.com/mlflow/mlflow.git@refs/pull/11197/merge

Checkout with GitHub CLI

gh pr checkout 11197

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Keras automatic checkpoint implementation.

This PR replaces the old PR: #10955

How is this PR tested?

• Existing unit/integration tests
• New unit/integration tests
• Manual tests

Does this PR require documentation update?

• No. You can skip the rest of this section.
• Yes. I've updated:
  • Examples
  • API references
  • Instructions

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

Several automatic checkpoint arguments are added into keras autologging.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/deployments: MLflow Deployments client APIs, server, and third-party Deployments integrations
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[github-actions]

Documentation preview for 6d9ad09 will be available when this CircleCI job
completes successfully.

• Top page
• Changed pages

More info

• Ignore this comment if this PR does not change the documentation.
• It takes a few minutes for the preview to be available.
• The preview is updated when a new commit is pushed to this PR.
• This comment was created by https://github.com/mlflow/mlflow/actions/runs/8046328422.

[chenmoneygithub]

Thanks Weichen, great work! Left some comments.

[chenmoneygithub]

can we use the fluent API instead of the client API?

[WeichenXu123]

I prefer to use client API, it can set tracking uri explicitly (we need to set tracking uri here, for distributed training case). Any reason to use fluent API ?

[chenmoneygithub]

I think fluent api is our standard user interface, so it is more robust, also we can shorten the arglist. For distributed training scenario, is mlflow.set_tracking_uri() enough?

But anyway, I don't think this is a blocking issue, just want to open the discussion to align the API design.

[WeichenXu123]

oh, we can remove client arguements from MlflowModelCheckpointCallback constructor.

[WeichenXu123]

I also removed run_id from MlflowModelCheckpointCallback constructor and use fluent apis instead

[chenmoneygithub]

I'm a little concerned about this line - we are setting a class attribute in on_train_batch_end, which is called multiple times. Do we need this line?

[WeichenXu123]

I can make it only set it at the first time.

The purpose of setting this attribute is for usage in save_checkpoint, otherwise we need to figure out a way to pass trainer object to save_checkpoint function. Note save_checkpoint is called in base class but its implementation is in sub-class.

[chenmoneygithub]

can we set it in the constructor or in on_train_start hook?

[WeichenXu123]

Moved it into

    @rank_zero_only
    def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        self.trainer = trainer

We can't put it in the constructor, because the callbacks argument is also defined in Trainer constructor, and we need to support usage like:

        mlflow_checkpoint_callback = MLflowModelCheckpointCallback()  # we can't pass trainer object here.

        trainer = Trainer(
            callbacks=[mlflow_checkpoint_callback]
        )

[chenmoneygithub]

Thanks Weichen! We are pretty close, left a few non-blocking comments.

[chenmoneygithub]

Although at this time we don't expect users to directly use this class, since this is a public class, can we add a code example on how to use it?

[WeichenXu123]

I prefer to make it a private class _MlflowModelCheckpointCallbackBase

[chenmoneygithub]

ohh we cannot make it a private class because we are importing it in a different file.

[chenmoneygithub]

Thanks Weichen! Approved with a comment

[harupy]

LGTM

[harupy]

This line looks a bit scary. .split(".")[-2] fails if the result of os.path.basename(downloaded_checkpoint_filepath) doesn't contain a dot.

[WeichenXu123]

Improved error handling for this case:

        artifact_name = os.path.basename(downloaded_checkpoint_filepath)
        artifact_name_splits = artifact_name.split(".")
        if len(artifact_name_splits) < 2:
            raise MlflowException(
                f"The model checkpoint artifact file name '{artifact_name}' is malformed."
            )

[WeichenXu123]

if a str doesn't contain dot, 'xx'.split(".") gets ['xx']

[harupy]

I'd use endswith or regex so we don't need os.path.basename or split.

[WeichenXu123]

Personally I prefer split than regex, if we incautiously write a bad regex it might cause performance issue when matching a long string, split is safer to me.

[harupy]

@WeichenXu123 #11250

[WeichenXu123]

Test failures in tensorflow / dev are not relevent.