Add keras flavor for Keras 3 support #10830

chenmoneygithub · 2024-01-16T23:15:40Z

🛠 DevTools 🛠

Install mlflow from this PR

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

Checkout with GitHub CLI

gh pr checkout 10830

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Add keras flavor for Keras 3 support, this supports all backends Keras 3 supports, i.e., Tensorflow, PyTorch and JAX.

This code is adapted from tensorflow flavor, but mostly rewritten. Redundant logic is removed.

How is this PR tested?

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

Does this PR require documentation update?

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.

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

Components

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 · 2024-01-16T23:16:04Z

Documentation preview for 4f876fd will be available here when this CircleCI job completes successfully.

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/7733187485.

docs/source/deep-learning/keras/quickstart/quickstart_keras_core.ipynb

harupy · 2024-01-19T00:28:18Z

setup.py

@@ -96,7 +96,7 @@ def run(self):
        print("\n".join(dependencies))


-MINIMUM_SUPPORTED_PYTHON_VERSION = "3.8"
+MINIMUM_SUPPORTED_PYTHON_VERSION = "3.9"


Suggested change

MINIMUM_SUPPORTED_PYTHON_VERSION = "3.9"

MINIMUM_SUPPORTED_PYTHON_VERSION = "3.8"

harupy · 2024-01-22T06:16:36Z

mlflow/ml_package_versions.py

+        },
+        "models": {
+            "minimum": "3.0.2",
+            "maximum": "3.10.0", # Dummy version number to remove cap.


We can use the latest version. We don't know if the current implementation really works with 3.10.0.

harupy · 2024-01-22T06:17:02Z

mlflow/keras_core/__init__.py

@@ -1,3 +0,0 @@
-from mlflow.keras_core.callback import MLflowCallback
-
-__all__ = ["MLflowCallback"]


should we raise a deprecation warning?

it was marked as experimental, and I don't think anyone is using it, so I guess we can just delete it.

harupy · 2024-01-22T06:24:17Z

mlflow/keras/autolog.py

+                log_every_n_steps=log_every_n_steps,
+            )
+            callbacks = _add_mlflow_to_keras_callbacks(callbacks, mlflow_callback)
+            kwargs["callbacks"] = callbacks


Suggested change

kwargs["callbacks"] = callbacks

kwargs["callbacks"] = [*callbacks, mlflow_callback]

_add_mlflow_to_keras_callbacks can just check if callbacks contains an MLflowCallback object.

makes sense to me, I am renaming it to _check_existing_mlflow_callback and append the mlflow_callback beforehand.

harupy · 2024-01-22T06:24:32Z

mlflow/keras/autolog.py

+                    _logger.warning(f"Failed to log dataset information to MLflow. Reason: {e}")
+
+            # Add `MLflowCallback` to the callback list.
+            callbacks = args[5] if len(args) >= 6 else kwargs.get("callbacks", [])


Any chance that callbacks is None?

hmm there is no enforcement, but I doubt anyone will want to do that.

Suggested change

callbacks = args[5] if len(args) >= 6 else kwargs.get("callbacks", [])

callbacks = args[5] if len(args) >= 6 else kwargs.get("callbacks") or []

in case callbacks is None

harupy · 2024-01-22T06:24:42Z

mlflow/keras/callback.py

@@ -45,15 +45,13 @@ class MLflowCallback(keras.callbacks.Callback):
                label,
                batch_size=4,
                epochs=2,
-                callbacks=[mlflow.keras_core.MLflowCallback(run)],
+                callbacks=[mlflow.keras_core.MLflowCallback()],


Suggested change

callbacks=[mlflow.keras_core.MLflowCallback()],

callbacks=[mlflow.keras.MLflowCallback()],

good catch!

harupy · 2024-01-22T06:39:41Z

mlflow/keras/autolog.py

@@ -0,0 +1,261 @@
+# MLflow autologging support for Keras 3.


This comment is useless. Let's remove it.

Oh I made a mistake here, it should be a docstring instead of a comment. Basically according to Google python style guide, each module should have a top-level docstring.

harupy · 2024-01-22T06:41:04Z

mlflow/ml_package_versions.py

can you also mlflow/ml-package-versions.yml?

harupy · 2024-01-22T06:42:04Z

mlflow/keras/callback.py

-        self.metrics_logger.record_metrics(logs, epoch)
+        log_metrics(logs, step=epoch, synchronous=False)


curious why we need this change

good question - a few months ago we added async logging support to our fluent API, so we on longer need to use the specific logger as before.

harupy · 2024-01-22T06:43:11Z

tests/keras/test_autolog.py

+    AUTOLOGGING_INTEGRATIONS.pop("keras", None)
+
+
+def _create_keras_model():


Is it possible to test a pytorch model?

This Keras model can be a PyTorch, Tensorflow or JAX model, which is controlled by environment variable KERAS_BACKEND. It sounds a bit weird at first look, I suggest giving it a try, which i find pretty cool!

harupy · 2024-01-22T07:26:55Z

mlflow/keras/save.py

+        `save_model()` and `log_model()` produce a pip environment that, at minimum, contains these
+        requirements.
+    """
+    return [_get_pinned_requirement("keras")]


Is there a way to detect the backend framework (TF, torch, Jax)?

Very good point, let me add the backend framework information.

@chenmoneygithub can you add it?

I added it to flavor option above

From Haru: let's use this utility function to retrieve the backend package + version.

I tested and verified without any extra code, _get_pinned_requirement already retrieves the backend requirement, idk why tho.

Signed-off-by: chenmoneygithub <[email protected]>

chenmoneygithub · 2024-01-30T22:10:01Z

close and reopen to disable the CI cache.

Signed-off-by: chenmoneygithub <[email protected]>

harupy

LGTM

Signed-off-by: chenmoneygithub <[email protected]> Signed-off-by: ernestwong-db <[email protected]>

chenmoneygithub marked this pull request as draft January 16, 2024 23:16

github-actions bot added area/tracking Tracking service, tracking client APIs, autologging rn/feature Mention under Features in Changelogs. labels Jan 16, 2024

BenWilson2 reviewed Jan 17, 2024

View reviewed changes

docs/source/deep-learning/keras/quickstart/quickstart_keras_core.ipynb Outdated Show resolved Hide resolved

chenmoneygithub marked this pull request as ready for review January 17, 2024 21:29

chenmoneygithub mentioned this pull request Jan 18, 2024

Does tensorflow < 2.16 depend on keras or tf-keras? keras-team/keras#19069

Closed

harupy reviewed Jan 19, 2024

View reviewed changes

harupy reviewed Jan 22, 2024

View reviewed changes

chenmoneygithub force-pushed the keras-3-support branch from 69ddd1c to 66c2c9b Compare January 23, 2024 22:11

chenmoneygithub added 11 commits January 30, 2024 12:50

Add keras flavor

f4c78db

Signed-off-by: chenmoneygithub <[email protected]>

Fix style

41f61e0

Signed-off-by: chenmoneygithub <[email protected]>

small refactor

c23af45

Signed-off-by: chenmoneygithub <[email protected]>

support model exporting

ed09cbd

Signed-off-by: chenmoneygithub <[email protected]>

bump up minimum python version in setup.py

80db710

Signed-off-by: chenmoneygithub <[email protected]>

change keras requirements

045a660

Signed-off-by: chenmoneygithub <[email protected]>

fix lint

32f0a2e

Signed-off-by: chenmoneygithub <[email protected]>

fix comments

6c1aa33

Signed-off-by: chenmoneygithub <[email protected]>

Some tests fixation

ccc87d9

Signed-off-by: chenmoneygithub <[email protected]>

some more fixations

34e5454

Signed-off-by: chenmoneygithub <[email protected]>

fixes

7b4e0f8

Signed-off-by: chenmoneygithub <[email protected]>

chenmoneygithub force-pushed the keras-3-support branch from c3a8ec6 to 7b4e0f8 Compare January 30, 2024 21:28

change cross version testing config

1011625

Signed-off-by: chenmoneygithub <[email protected]>

chenmoneygithub closed this Jan 30, 2024

chenmoneygithub reopened this Jan 30, 2024

chenmoneygithub added 4 commits January 30, 2024 14:31

small

d42be76

Signed-off-by: chenmoneygithub <[email protected]>

move around export clause

a44104a

Signed-off-by: chenmoneygithub <[email protected]>

let's go

de920b8

Signed-off-by: chenmoneygithub <[email protected]>

change dataset logging logic

4f876fd

Signed-off-by: chenmoneygithub <[email protected]>

harupy approved these changes Feb 1, 2024

View reviewed changes

chenmoneygithub merged commit 093782b into mlflow:master Feb 1, 2024
39 checks passed

This was referenced Feb 5, 2024

Fix [Makefile:78: rsthtml] Error 1 #11002

Closed

Fix [Makefile:78: rsthtml] Error 1 (related to keras) #11003

Closed

ernestwong-db pushed a commit to ernestwong-db/mlflow that referenced this pull request Feb 6, 2024

Add keras flavor for Keras 3 support (mlflow#10830)

ecd2e36

Signed-off-by: chenmoneygithub <[email protected]> Signed-off-by: ernestwong-db <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add keras flavor for Keras 3 support #10830

Add keras flavor for Keras 3 support #10830

chenmoneygithub commented Jan 16, 2024 •

edited

Loading

github-actions bot commented Jan 16, 2024 •

edited

Loading

harupy Jan 19, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024 •

edited

Loading

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 23, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 22, 2024

chenmoneygithub Jan 22, 2024

harupy Jan 24, 2024

chenmoneygithub Jan 24, 2024

chenmoneygithub Jan 26, 2024

chenmoneygithub Jan 26, 2024

chenmoneygithub commented Jan 30, 2024

harupy left a comment

	MINIMUM_SUPPORTED_PYTHON_VERSION = "3.9"
	MINIMUM_SUPPORTED_PYTHON_VERSION = "3.8"

		@@ -1,3 +0,0 @@
		from mlflow.keras_core.callback import MLflowCallback

		__all__ = ["MLflowCallback"]

	kwargs["callbacks"] = callbacks
	kwargs["callbacks"] = [*callbacks, mlflow_callback]

	callbacks = args[5] if len(args) >= 6 else kwargs.get("callbacks", [])
	callbacks = args[5] if len(args) >= 6 else kwargs.get("callbacks") or []

	callbacks=[mlflow.keras_core.MLflowCallback()],
	callbacks=[mlflow.keras.MLflowCallback()],

		self.metrics_logger.record_metrics(logs, epoch)
		log_metrics(logs, step=epoch, synchronous=False)

		AUTOLOGGING_INTEGRATIONS.pop("keras", None)


		def _create_keras_model():

Add keras flavor for Keras 3 support #10830

Add keras flavor for Keras 3 support #10830

Conversation

chenmoneygithub commented Jan 16, 2024 • edited Loading

Install mlflow from this PR

Checkout with GitHub CLI

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

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

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

github-actions bot commented Jan 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chenmoneygithub Jan 22, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chenmoneygithub commented Jan 30, 2024

harupy left a comment

Choose a reason for hiding this comment

chenmoneygithub commented Jan 16, 2024 •

edited

Loading

github-actions bot commented Jan 16, 2024 •

edited

Loading

chenmoneygithub Jan 22, 2024 •

edited

Loading