# pep8: disable=E501 import functools import json import os import pickle import sys from pathlib import Path from unittest.mock import patch import numpy as np import pytest import tensorflow as tf import yaml from packaging.version import Version from tensorflow.keras import layers import mlflow from mlflow import MlflowClient from mlflow.exceptions import MlflowException from mlflow.models import Model from mlflow.models.utils import _read_example from mlflow.tensorflow import load_checkpoint from mlflow.tensorflow.autologging import _TensorBoard from mlflow.tensorflow.callback import MlflowCallback from mlflow.tracking.fluent import _shut_down_async_logging from mlflow.types.utils import _infer_schema from mlflow.utils.autologging_utils import ( AUTOLOGGING_INTEGRATIONS, autologging_is_disabled, ) from mlflow.utils.file_utils import local_file_uri_to_path from mlflow.utils.process import _exec_cmd np.random.seed(1337) @pytest.fixture(autouse=True) def clear_session(): yield _shut_down_async_logging() tf.keras.backend.clear_session() @pytest.fixture def random_train_data(): return np.random.random((150, 4)) @pytest.fixture def random_one_hot_labels(): n = 150 n_class = 3 classes = np.random.randint(0, n_class, n) labels = np.zeros((n, n_class)) labels[np.arange(n), classes] = 1 return labels @pytest.fixture def random_train_dict_mapping(random_train_data): def _generate_features(pos): return [v[pos] for v in random_train_data] return { "a": np.array(_generate_features(0)), "b": np.array(_generate_features(1)), "c": np.array(_generate_features(2)), "d": np.array(_generate_features(3)), } def _create_model_for_dict_mapping(): inputs = { "a": tf.keras.Input(shape=(1,), name="a"), "b": tf.keras.Input(shape=(1,), name="b"), "c": tf.keras.Input(shape=(1,), name="c"), "d": tf.keras.Input(shape=(1,), name="d"), } concatenated = layers.Concatenate()(list(inputs.values())) x = layers.Dense(16, activation="relu", input_shape=(4,))(concatenated) outputs = layers.Dense(3, activation="softmax")(x) model = tf.keras.Model(inputs=inputs, outputs=outputs) model.compile( optimizer=tf.keras.optimizers.Adam(), loss="categorical_crossentropy", metrics=["accuracy"] ) return model @pytest.fixture def fashion_mnist_tf_dataset(): train, _ = tf.keras.datasets.fashion_mnist.load_data() images, labels = train images = images / 255.0 labels = labels.astype(np.int32) fmnist_train_ds = tf.data.Dataset.from_tensor_slices((images, labels)) return fmnist_train_ds.shuffle(5000).batch(32) @pytest.fixture def fashion_mnist_tf_dataset_eval(): _, eval_dataset = tf.keras.datasets.fashion_mnist.load_data() images, labels = eval_dataset images = images / 255.0 labels = labels.astype(np.int32) fmnist_train_ds = tf.data.Dataset.from_tensor_slices((images, labels)) return fmnist_train_ds.shuffle(5000).batch(32) def _create_fashion_mnist_model(): model = tf.keras.Sequential([ tf.keras.Input((28, 28)), tf.keras.layers.Flatten(), tf.keras.layers.Dense(10), ]) model.compile( optimizer=tf.keras.optimizers.Adam(), loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["accuracy"], ) return model @pytest.fixture def keras_data_gen_sequence(random_train_data, random_one_hot_labels): class DataGenerator(tf.keras.utils.Sequence): def __len__(self): return 128 def __getitem__(self, index): x = random_train_data y = random_one_hot_labels return x, y return DataGenerator() @pytest.fixture(autouse=True) def clear_fluent_autologging_import_hooks(): """ Clears import hooks for MLflow fluent autologging (`mlflow.autolog()`) between tests to ensure that interactions between fluent autologging and TensorFlow / tf.keras can be tested successfully """ mlflow.utils.import_hooks._post_import_hooks.pop("tensorflow", None) mlflow.utils.import_hooks._post_import_hooks.pop("keras", None) @pytest.fixture(autouse=True) def clear_autologging_config(): """ Clears TensorFlow autologging config, simulating a fresh state where autologging has not been previously enabled with any particular configuration """ AUTOLOGGING_INTEGRATIONS.pop(mlflow.tensorflow.FLAVOR_NAME, None) def create_tf_keras_model(): model = tf.keras.Sequential() model.add(tf.keras.Input(shape=(4,), dtype="float64")) model.add(layers.Dense(16, activation="relu")) model.add(layers.Dense(3, activation="softmax")) model.compile( optimizer=tf.keras.optimizers.Adam(), loss="categorical_crossentropy", metrics=["accuracy"] ) return model def test_tf_keras_autolog_ends_auto_created_run(random_train_data, random_one_hot_labels): mlflow.tensorflow.autolog() data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() model.fit(data, labels, epochs=10) assert mlflow.active_run() is None def test_extra_tags_tensorflow_autolog(random_train_data, random_one_hot_labels): mlflow.tensorflow.autolog(extra_tags={"test_tag": "tf_autolog"}) data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() model.fit(data, labels, epochs=10) run = mlflow.last_active_run() assert run.data.tags["test_tag"] == "tf_autolog" assert run.data.tags[mlflow.utils.mlflow_tags.MLFLOW_AUTOLOGGING] == "tensorflow" @pytest.mark.parametrize("log_models", [True, False]) def test_tf_keras_autolog_log_models_configuration( random_train_data, random_one_hot_labels, log_models ): mlflow.tensorflow.autolog(log_models=log_models) data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() model.fit(data, labels, epochs=10) assert (mlflow.last_logged_model() is not None) == log_models @pytest.mark.parametrize("log_models", [True, False]) @pytest.mark.parametrize("log_datasets", [True, False]) def test_tf_keras_autolog_log_datasets_configuration_with_numpy( random_train_data, random_one_hot_labels, log_datasets, log_models ): mlflow.tensorflow.autolog(log_datasets=log_datasets, log_models=log_models) data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() model.fit(data, labels, epochs=10) client = MlflowClient() run_inputs = client.get_run(mlflow.last_active_run().info.run_id).inputs dataset_inputs = run_inputs.dataset_inputs if log_datasets: assert len(dataset_inputs) == 1 feature_schema = _infer_schema(data) target_schema = _infer_schema(labels) assert dataset_inputs[0].dataset.schema == json.dumps({ "mlflow_tensorspec": { "features": feature_schema.to_json(), "targets": target_schema.to_json(), } }) else: assert len(dataset_inputs) == 0 logged_model_inputs = run_inputs.model_inputs logged_model = mlflow.last_logged_model() if log_models: if log_datasets: assert len(logged_model_inputs) == 1 assert logged_model_inputs[0].model_id == logged_model.model_id else: assert logged_model is not None assert logged_model.source_run_id == mlflow.last_active_run().info.run_id else: assert len(logged_model_inputs) == 0 assert logged_model is None @pytest.mark.parametrize("log_datasets", [True, False]) def test_tf_keras_autolog_log_datasets_configuration_with_tensor( random_train_data, random_one_hot_labels, log_datasets ): mlflow.tensorflow.autolog(log_datasets=log_datasets) data_as_tensor = tf.convert_to_tensor(random_train_data) labels_as_tensor = tf.convert_to_tensor(random_one_hot_labels) model = create_tf_keras_model() model.fit(data_as_tensor, labels_as_tensor, epochs=10) client = MlflowClient() dataset_inputs = client.get_run(mlflow.last_active_run().info.run_id).inputs.dataset_inputs if log_datasets: assert len(dataset_inputs) == 1 feature_schema = _infer_schema(data_as_tensor.numpy()) target_schema = _infer_schema(labels_as_tensor.numpy()) assert dataset_inputs[0].dataset.schema == json.dumps({ "mlflow_tensorspec": { "features": feature_schema.to_json(), "targets": target_schema.to_json(), } }) else: assert len(dataset_inputs) == 0 @pytest.mark.parametrize("log_datasets", [True, False]) def test_tf_keras_autolog_log_datasets_configuration_with_tf_dataset( fashion_mnist_tf_dataset, log_datasets ): mlflow.tensorflow.autolog(log_datasets=log_datasets) fashion_mnist_model = _create_fashion_mnist_model() fashion_mnist_model.fit(fashion_mnist_tf_dataset) client = MlflowClient() dataset_inputs = client.get_run(mlflow.last_active_run().info.run_id).inputs.dataset_inputs if log_datasets: assert len(dataset_inputs) == 1 numpy_data = next(fashion_mnist_tf_dataset.as_numpy_iterator()) assert dataset_inputs[0].dataset.schema == json.dumps({ "mlflow_tensorspec": { "features": _infer_schema({ str(i): data_element for i, data_element in enumerate(numpy_data) }).to_json(), "targets": None, } }) else: assert len(dataset_inputs) == 0 def test_tf_keras_autolog_log_datasets_with_validation_data( fashion_mnist_tf_dataset, fashion_mnist_tf_dataset_eval ): mlflow.tensorflow.autolog(log_datasets=True) fashion_mnist_model = _create_fashion_mnist_model() fashion_mnist_model.fit(fashion_mnist_tf_dataset, validation_data=fashion_mnist_tf_dataset_eval) client = MlflowClient() dataset_inputs = client.get_run(mlflow.last_active_run().info.run_id).inputs.dataset_inputs assert len(dataset_inputs) == 2 assert dataset_inputs[0].tags[0].value == "train" assert dataset_inputs[1].tags[0].value == "eval" def test_tf_keras_autolog_log_datasets_with_validation_data_as_numpy_tuple( fashion_mnist_tf_dataset, fashion_mnist_tf_dataset_eval ): mlflow.tensorflow.autolog(log_datasets=True) fashion_mnist_model = _create_fashion_mnist_model() X_eval, y_eval = next(fashion_mnist_tf_dataset_eval.as_numpy_iterator()) fashion_mnist_model.fit(fashion_mnist_tf_dataset, validation_data=(X_eval, y_eval)) client = MlflowClient() dataset_inputs = client.get_run(mlflow.last_active_run().info.run_id).inputs.dataset_inputs assert len(dataset_inputs) == 2 assert dataset_inputs[0].tags[0].value == "train" assert dataset_inputs[1].tags[0].value == "eval" def test_tf_keras_autolog_log_datasets_with_validation_data_as_tf_tuple( fashion_mnist_tf_dataset, fashion_mnist_tf_dataset_eval ): mlflow.tensorflow.autolog(log_datasets=True) fashion_mnist_model = _create_fashion_mnist_model() # convert tensorflow dataset into tensors X_eval, y_eval = next(fashion_mnist_tf_dataset_eval.as_numpy_iterator()) X_eval_tensor = tf.convert_to_tensor(X_eval) y_eval_tensor = tf.convert_to_tensor(y_eval) fashion_mnist_model.fit( fashion_mnist_tf_dataset, validation_data=(X_eval_tensor, y_eval_tensor) ) client = MlflowClient() dataset_inputs = client.get_run(mlflow.last_active_run().info.run_id).inputs.dataset_inputs assert len(dataset_inputs) == 2 assert dataset_inputs[0].tags[0].value == "train" assert dataset_inputs[1].tags[0].value == "eval" def test_tf_keras_autolog_persists_manually_created_run(random_train_data, random_one_hot_labels): mlflow.tensorflow.autolog() with mlflow.start_run() as run: data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() model.fit(data, labels, epochs=10) assert mlflow.active_run() assert mlflow.active_run().info.run_id == run.info.run_id @pytest.fixture def tf_keras_random_data_run(random_train_data, random_one_hot_labels, initial_epoch): mlflow.tensorflow.autolog() data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() history = model.fit( data, labels, epochs=initial_epoch + 10, steps_per_epoch=1, initial_epoch=initial_epoch ) client = MlflowClient() return client.get_run(client.search_runs(["0"])[0].info.run_id), history @pytest.mark.parametrize("initial_epoch", [0, 10]) def test_tf_keras_autolog_logs_expected_data(tf_keras_random_data_run): run, history = tf_keras_random_data_run data = run.data assert "accuracy" in data.metrics assert "loss" in data.metrics # Testing explicitly passed parameters are logged correctly assert "epochs" in data.params assert data.params["epochs"] == str(history.epoch[-1] + 1) assert "steps_per_epoch" in data.params assert data.params["steps_per_epoch"] == "1" # Testing default parameters are logged correctly assert "initial_epoch" in data.params assert data.params["initial_epoch"] == str(history.epoch[0]) # Testing unwanted parameters are not logged assert "callbacks" not in data.params assert "validation_data" not in data.params # Testing optimizer parameters are logged assert "opt_name" in data.params assert data.params["opt_name"].lower() == "adam" assert "opt_learning_rate" in data.params assert "opt_beta_1" in data.params assert "opt_beta_2" in data.params assert "opt_epsilon" in data.params assert "opt_amsgrad" in data.params assert data.params["opt_amsgrad"] == "False" client = MlflowClient() all_epoch_acc = client.get_metric_history(run.info.run_id, "accuracy") num_of_epochs = len(history.history["loss"]) assert len(all_epoch_acc) == num_of_epochs == 10 artifacts = client.list_artifacts(run.info.run_id) artifacts = (x.path for x in artifacts) assert "model_summary.txt" in artifacts def __example_tf_dataset(batch_size): a = tf.data.Dataset.range(1) b = tf.data.Dataset.range(1) ds = tf.data.Dataset.zip((a, b)) return ds.batch(batch_size) class __ExampleSequence(tf.keras.utils.Sequence): def __init__(self, batch_size, with_sample_weights=False): self.batch_size = batch_size self.with_sample_weights = with_sample_weights def __len__(self): return 10 def __getitem__(self, idx): x = np.array([idx] * self.batch_size) y = np.array([-idx] * self.batch_size) if self.with_sample_weights: w = np.array([1] * self.batch_size) return x, y, w return x, y def __generator(data, target, batch_size): data_batches = np.split(data, data.shape[0] // batch_size) target_batches = np.split(target, target.shape[0] // batch_size) yield from zip(data_batches, target_batches) class __GeneratorClass: def __init__(self, data, target, batch_size): self.data = data self.target = target self.batch_size = batch_size self.ptr = 0 def __next__(self): if self.ptr >= len(self.data): raise StopIteration idx = self.ptr % len(self.data) self.ptr += 1 return self.data[idx : idx + self.batch_size], self.target[idx : idx + self.batch_size] def __iter__(self): return self @pytest.mark.parametrize( "generate_data", [ __example_tf_dataset, __ExampleSequence, functools.partial(__ExampleSequence, with_sample_weights=True), functools.partial(__generator, np.array([[1]] * 10), np.array([[1]] * 10)), pytest.param( functools.partial(__GeneratorClass, np.array([[1]] * 10), np.array([[1]] * 10)), marks=pytest.mark.skipif( Version(tf.__version__).release >= (2, 15) and "TF_USE_LEGACY_KERAS" not in os.environ, reason="does not support", ), ), ], ) @pytest.mark.parametrize("batch_size", [5, 10]) def test_tf_keras_autolog_implicit_batch_size_works(generate_data, batch_size): mlflow.autolog() model = tf.keras.Sequential() model.add(tf.keras.layers.Dense(1, input_shape=(1,))) model.compile(loss="mse") # 'x' passed as arg model.fit(generate_data(batch_size), verbose=0) assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size) # 'x' passed as kwarg model.fit(x=generate_data(batch_size), verbose=0) assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size) def __tf_dataset_multi_input(batch_size): a = tf.data.Dataset.range(1) b = tf.data.Dataset.range(1) c = tf.data.Dataset.range(1) ds = tf.data.Dataset.zip(((a, b), c)) return ds.batch(batch_size) class __SequenceMultiInput(tf.keras.utils.Sequence): def __init__(self, batch_size): self.batch_size = batch_size def __len__(self): return 10 def __getitem__(self, idx): return (np.random.rand(self.batch_size), np.random.rand(self.batch_size)), np.random.rand( self.batch_size ) def __generator_multi_input(data, target, batch_size): data_batches = np.split(data, data.shape[1] // batch_size, axis=1) target_batches = np.split(target, target.shape[0] // batch_size) for inputs, output in zip(data_batches, target_batches): yield tuple(inputs), output class __GeneratorClassMultiInput: def __init__(self, data, target, batch_size): self.data = data self.target = target self.batch_size = batch_size self.ptr = 0 def __next__(self): if self.ptr >= len(self.data): raise StopIteration idx = self.ptr % len(self.data) self.ptr += 1 return ( self.data[idx : idx + self.batch_size, 0], self.data[idx : idx + self.batch_size, 1], ), self.target[idx : idx + self.batch_size] def __iter__(self): return self @pytest.mark.parametrize( "generate_data", [ __tf_dataset_multi_input, __SequenceMultiInput, functools.partial(__generator_multi_input, np.random.rand(2, 10), np.random.rand(10)), functools.partial(__GeneratorClassMultiInput, np.random.rand(10, 2), np.random.rand(10, 1)), ], ) @pytest.mark.parametrize("batch_size", [5, 10]) def test_tf_keras_autolog_implicit_batch_size_works_multi_input(generate_data, batch_size): mlflow.tensorflow.autolog() input1 = tf.keras.Input(shape=(1,)) input2 = tf.keras.Input(shape=(1,)) concat = tf.keras.layers.Concatenate()([input1, input2]) output = tf.keras.layers.Dense(1, activation="sigmoid")(concat) model = tf.keras.models.Model(inputs=[input1, input2], outputs=output) model.compile(loss="mse") # 'x' passed as arg model.fit(generate_data(batch_size), verbose=0) assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size) # 'x' passed as kwarg model.fit(x=generate_data(batch_size), verbose=0) assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size) @pytest.mark.skipif( Version(tf.__version__) < Version("2.1.4"), reason="Does not support passing of generator classes as `x` in `fit`", ) @pytest.mark.parametrize( "generator", [ __generator, pytest.param( __GeneratorClass, marks=pytest.mark.skipif( Version(tf.__version__).release >= (2, 15) and "TF_USE_LEGACY_KERAS" not in os.environ, reason="does not support", ), ), ], ) @pytest.mark.parametrize("batch_size", [2, 3, 6]) def test_tf_keras_autolog_implicit_batch_size_for_generator_dataset_without_side_effects( generator, batch_size, ): from tensorflow.keras.layers import Dense from tensorflow.keras.models import Sequential data = np.array([[1, 2, 3], [3, 2, 1], [2, 2, 2], [10, 20, 30], [30, 20, 10], [20, 20, 20]]) target = np.array([[1], [3], [2], [11], [13], [12]]) model = Sequential() model.add( Dense( 5, input_dim=3, activation="relu", kernel_initializer="zeros", bias_initializer="zeros" ) ) model.add(Dense(1, kernel_initializer="zeros", bias_initializer="zeros")) model.compile(loss="mae", optimizer="adam", metrics=["mse"]) mlflow.autolog() actual_mse = model.fit(generator(data, target, batch_size), verbose=0).history["mse"][-1] mlflow.autolog(disable=True) expected_mse = model.fit(generator(data, target, batch_size), verbose=0).history["mse"][-1] np.testing.assert_allclose(actual_mse, expected_mse, atol=1) assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size) def test_tf_keras_autolog_succeeds_for_tf_datasets_lacking_batch_size_info(): X_train = np.random.rand(100, 100) y_train = np.random.randint(0, 10, 100) train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)) train_ds = train_ds.batch(50) train_ds = train_ds.cache().prefetch(buffer_size=5) assert not hasattr(train_ds, "_batch_size") model = tf.keras.Sequential() model.add(tf.keras.Input((100,))) model.add(tf.keras.layers.Dense(256, activation="relu")) model.add(tf.keras.layers.Dropout(rate=0.4)) model.add(tf.keras.layers.Dense(10, activation="sigmoid")) model.compile( loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False), optimizer="Adam", metrics=["accuracy"], ) mlflow.tensorflow.autolog() model.fit(train_ds, epochs=100) assert mlflow.last_active_run().data.params["batch_size"] == "None" def test_tf_keras_autolog_records_metrics_for_last_epoch(random_train_data, random_one_hot_labels): num_training_epochs = 17 mlflow.tensorflow.autolog(log_every_epoch=True) model = create_tf_keras_model() with mlflow.start_run() as run: model.fit( random_train_data, random_one_hot_labels, epochs=num_training_epochs, initial_epoch=0, ) client = MlflowClient() run_metrics = client.get_run(run.info.run_id).data.metrics assert "accuracy" in run_metrics all_epoch_acc = client.get_metric_history(run.info.run_id, "accuracy") assert len(all_epoch_acc) == num_training_epochs def test_tf_keras_autolog_logs_metrics_for_single_epoch_training( random_train_data, random_one_hot_labels ): """ tf.Keras exhibits inconsistent epoch indexing behavior in comparison with other TF2 APIs (e.g., tf.Estimator). tf.Keras uses zero-indexing for epochs, while other APIs use one-indexing. Accordingly, this test verifies that metrics are produced in the boundary case where a model is trained for a single epoch, ensuring that we don't miss the zero index in the tf.Keras case. """ mlflow.tensorflow.autolog() model = create_tf_keras_model() with mlflow.start_run() as run: model.fit(random_train_data, random_one_hot_labels, epochs=1) client = MlflowClient() run_metrics = client.get_run(run.info.run_id).data.metrics assert "accuracy" in run_metrics assert "loss" in run_metrics def test_tf_keras_autolog_names_positional_parameters_correctly( random_train_data, random_one_hot_labels ): mlflow.tensorflow.autolog() data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() with mlflow.start_run(): # Pass `batch_size` as a positional argument for testing purposes model.fit(data, labels, 8, epochs=10, steps_per_epoch=1) run_id = mlflow.active_run().info.run_id client = MlflowClient() run_info = client.get_run(run_id) assert run_info.data.params.get("batch_size") == "8" @pytest.mark.parametrize("initial_epoch", [0, 10]) def test_tf_keras_autolog_model_can_load_from_artifact(tf_keras_random_data_run, random_train_data): run, _ = tf_keras_random_data_run client = MlflowClient() artifacts = client.list_artifacts(run.info.run_id) artifacts = (x.path for x in artifacts) assert "tensorboard_logs" in artifacts model = mlflow.tensorflow.load_model("runs:/" + run.info.run_id + "/model") model.predict(random_train_data) def get_tf_keras_random_data_run_with_callback( random_train_data, random_one_hot_labels, callback, restore_weights, patience, initial_epoch, log_models, ): mlflow.tensorflow.autolog(log_models=log_models) data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() if callback == "early": # min_delta is set as such to guarantee early stopping callback = tf.keras.callbacks.EarlyStopping( monitor="loss", patience=patience, min_delta=99999999, restore_best_weights=restore_weights, verbose=1, ) else: class CustomCallback(tf.keras.callbacks.Callback): def on_train_end(self, logs=None): pass callback = CustomCallback() history = model.fit( data, labels, epochs=initial_epoch + 10, callbacks=[callback], initial_epoch=initial_epoch ) client = MlflowClient() return client.get_run(client.search_runs(["0"])[0].info.run_id), history, callback @pytest.fixture def tf_keras_random_data_run_with_callback( random_train_data, random_one_hot_labels, callback, restore_weights, patience, initial_epoch, log_models, ): return get_tf_keras_random_data_run_with_callback( random_train_data, random_one_hot_labels, callback, restore_weights, patience, initial_epoch, log_models=log_models, ) @pytest.mark.parametrize("log_models", [True, False]) @pytest.mark.parametrize("restore_weights", [True]) @pytest.mark.parametrize("callback", ["early"]) @pytest.mark.parametrize("patience", [0, 1, 5]) @pytest.mark.parametrize("initial_epoch", [0, 10]) def test_tf_keras_autolog_early_stop_logs( tf_keras_random_data_run_with_callback, initial_epoch, log_models ): run, history, callback = tf_keras_random_data_run_with_callback metrics = run.data.metrics params = run.data.params assert "patience" in params assert params["patience"] == str(callback.patience) assert "monitor" in params assert params["monitor"] == "loss" assert "verbose" not in params assert "mode" not in params assert "stopped_epoch" in metrics assert "restored_epoch" in metrics restored_epoch = int(metrics["restored_epoch"]) # In this test, the best epoch is always the first epoch because the early stopping callback # never observes a loss improvement due to an extremely large `min_delta` value assert restored_epoch == initial_epoch assert "loss" in history.history client = MlflowClient() metric_history = client.get_metric_history(run.info.run_id, "loss") # Check that MLflow has logged the metrics of the "best" model, in addition to per-epoch metrics loss = history.history["loss"] assert len(metric_history) == len(loss) + 1 steps, values = map(list, zip(*[(m.step, m.value) for m in metric_history])) # Check that MLflow has logged the correct steps assert steps == [*history.epoch, callback.stopped_epoch + 1] # Check that MLflow has logged the correct metric values np.testing.assert_allclose(values, [*loss, callback.best]) artifacts = [f.path for f in client.list_artifacts(run.info.run_id)] assert "tensorboard_logs" in artifacts # Check metrics are logged to the LoggedModel if log_models: logged_model = mlflow.last_logged_model() assert logged_model is not None assert {metric.key: metric.value for metric in logged_model.metrics} == metrics @pytest.mark.parametrize("log_models", [False]) @pytest.mark.parametrize("restore_weights", [True]) @pytest.mark.parametrize("callback", ["early"]) @pytest.mark.parametrize("patience", [11]) @pytest.mark.parametrize("initial_epoch", [0, 10]) def test_tf_keras_autolog_early_stop_no_stop_does_not_log(tf_keras_random_data_run_with_callback): run, history, callback = tf_keras_random_data_run_with_callback metrics = run.data.metrics params = run.data.params assert "patience" in params assert params["patience"] == str(callback.patience) assert "monitor" in params assert params["monitor"] == "loss" assert "verbose" not in params assert "mode" not in params assert "stopped_epoch" not in metrics assert "restored_epoch" not in metrics assert "loss" in history.history num_of_epochs = len(history.history["loss"]) client = MlflowClient() metric_history = client.get_metric_history(run.info.run_id, "loss") # Check the test epoch numbers are correct assert num_of_epochs == 10 assert len(metric_history) == num_of_epochs @pytest.mark.parametrize("log_models", [False]) @pytest.mark.parametrize("restore_weights", [False]) @pytest.mark.parametrize("callback", ["early"]) @pytest.mark.parametrize("patience", [5]) @pytest.mark.parametrize("initial_epoch", [0, 10]) def test_tf_keras_autolog_early_stop_no_restore_doesnt_log(tf_keras_random_data_run_with_callback): run, history, callback = tf_keras_random_data_run_with_callback metrics = run.data.metrics params = run.data.params assert "patience" in params assert params["patience"] == str(callback.patience) assert "monitor" in params assert params["monitor"] == "loss" assert "verbose" not in params assert "mode" not in params assert "stopped_epoch" in metrics assert "restored_epoch" not in metrics assert "loss" in history.history num_of_epochs = len(history.history["loss"]) client = MlflowClient() metric_history = client.get_metric_history(run.info.run_id, "loss") # Check the test epoch numbers are correct assert num_of_epochs == callback.patience + 1 assert len(metric_history) == num_of_epochs @pytest.mark.parametrize("log_models", [False]) @pytest.mark.parametrize("restore_weights", [False]) @pytest.mark.parametrize("callback", ["not-early"]) @pytest.mark.parametrize("patience", [5]) @pytest.mark.parametrize("initial_epoch", [0, 10]) def test_tf_keras_autolog_non_early_stop_callback_no_log(tf_keras_random_data_run_with_callback): run, history = tf_keras_random_data_run_with_callback[:-1] metrics = run.data.metrics params = run.data.params assert "patience" not in params assert "monitor" not in params assert "verbose" not in params assert "mode" not in params assert "stopped_epoch" not in metrics assert "restored_epoch" not in metrics assert "loss" in history.history num_of_epochs = len(history.history["loss"]) client = MlflowClient() metric_history = client.get_metric_history(run.info.run_id, "loss") # Check the test epoch numbers are correct assert num_of_epochs == 10 assert len(metric_history) == num_of_epochs @pytest.mark.parametrize("positional", [True, False]) def test_tf_keras_autolog_does_not_mutate_original_callbacks_list( tmp_path, random_train_data, random_one_hot_labels, positional ): """ TensorFlow autologging passes new callbacks to the `fit()` / `fit_generator()` function. If preexisting user-defined callbacks already exist, these new callbacks are added to the user-specified ones. This test verifies that the new callbacks are added to the without permanently mutating the original list of callbacks. """ mlflow.tensorflow.autolog() tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=tmp_path) callbacks = [tensorboard_callback] model = create_tf_keras_model() data = random_train_data labels = random_one_hot_labels if positional: model.fit(data, labels, None, 10, 1, callbacks) else: model.fit(data, labels, epochs=10, callbacks=callbacks) assert len(callbacks) == 1 assert callbacks == [tensorboard_callback] def test_tf_keras_autolog_does_not_delete_logging_directory_for_tensorboard_callback( tmp_path, random_train_data, random_one_hot_labels ): tensorboard_callback_logging_dir_path = str(tmp_path.joinpath("tb_logs")) tensorboard_callback = tf.keras.callbacks.TensorBoard( tensorboard_callback_logging_dir_path, histogram_freq=0 ) mlflow.tensorflow.autolog() data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() model.fit(data, labels, epochs=10, callbacks=[tensorboard_callback]) assert os.path.exists(tensorboard_callback_logging_dir_path) def test_tf_keras_autolog_logs_to_and_deletes_temporary_directory_when_tensorboard_callback_absent( tmp_path, random_train_data, random_one_hot_labels ): from mlflow.tensorflow import _TensorBoardLogDir mlflow.tensorflow.autolog() mock_log_dir_inst = _TensorBoardLogDir( location=str(tmp_path.joinpath("tb_logging")), is_temp=True ) with patch("mlflow.tensorflow._TensorBoardLogDir", autospec=True) as mock_log_dir_class: mock_log_dir_class.return_value = mock_log_dir_inst data = random_train_data labels = random_one_hot_labels model = create_tf_keras_model() model.fit(data, labels, epochs=10) assert not os.path.exists(mock_log_dir_inst.location) def get_text_vec_model(train_samples): # Taken from: https://github.com/mlflow/mlflow/issues/3910 try: from tensorflow.keras.layers.experimental.preprocessing import TextVectorization except ModuleNotFoundError: from tensorflow.keras.layers import TextVectorization VOCAB_SIZE = 10 SEQUENCE_LENGTH = 16 EMBEDDING_DIM = 16 vectorizer_layer = TextVectorization( max_tokens=VOCAB_SIZE, output_mode="int", output_sequence_length=SEQUENCE_LENGTH, ) vectorizer_layer.adapt(train_samples) model = tf.keras.Sequential([ vectorizer_layer, tf.keras.layers.Embedding( VOCAB_SIZE, EMBEDDING_DIM, name="embedding", mask_zero=True, ), tf.keras.layers.GlobalAveragePooling1D(), tf.keras.layers.Dense(16, activation="relu"), tf.keras.layers.Dense(1, activation="tanh"), ]) model.compile(optimizer="adam", loss="mse", metrics=["mae"]) return model @pytest.mark.skipif( Version(tf.__version__) < Version("2.3.0"), reason=( "Deserializing a model with `TextVectorization` and `Embedding` " "fails in tensorflow < 2.3.0. See this issue: " "https://github.com/tensorflow/tensorflow/issues/38250." ), ) def test_autolog_text_vec_model(tmp_path): """ Verifies autolog successfully saves a model that can't be saved in the H5 format """ mlflow.tensorflow.autolog() train_samples = tf.convert_to_tensor(["this is an example", "another example"]) train_labels = np.array([0.4, 0.2]) model = get_text_vec_model(train_samples) with mlflow.start_run() as run: model.fit(train_samples, train_labels, epochs=1) loaded_model = mlflow.tensorflow.load_model("runs:/" + run.info.run_id + "/model") np.testing.assert_array_equal(loaded_model.predict(train_samples), model.predict(train_samples)) def test_tf_keras_model_autolog_registering_model(random_train_data, random_one_hot_labels): registered_model_name = "test_autolog_registered_model" mlflow.tensorflow.autolog(registered_model_name=registered_model_name) with mlflow.start_run(): model = create_tf_keras_model() model.fit(random_train_data, random_one_hot_labels, epochs=10) registered_model = MlflowClient().get_registered_model(registered_model_name) assert registered_model.name == registered_model_name def test_fluent_autolog_with_tf_keras_logs_expected_content( random_train_data, random_one_hot_labels ): """ Guards against previously-exhibited issues where using the fluent `mlflow.autolog()` API with `tf.keras` Models did not work due to conflicting patches set by both the `mlflow.tensorflow.autolog()` and the `mlflow.keras.autolog()` APIs. """ mlflow.autolog() model = create_tf_keras_model() with mlflow.start_run() as run: model.fit(random_train_data, random_one_hot_labels, epochs=10) client = MlflowClient() run_data = client.get_run(run.info.run_id).data assert "accuracy" in run_data.metrics assert "epochs" in run_data.params def test_callback_is_picklable(): cb = MlflowCallback() pickle.dumps(cb) tb = _TensorBoard() pickle.dumps(tb) @pytest.mark.skipif( Version(tf.__version__) < Version("2.1.0"), reason="This test requires tensorflow >= 2.1.0" ) def test_tf_keras_autolog_distributed_training(random_train_data, random_one_hot_labels): # Ref: https://www.tensorflow.org/tutorials/distribute/keras mlflow.tensorflow.autolog() with tf.distribute.MirroredStrategy().scope(): model = create_tf_keras_model() fit_params = {"epochs": 10, "batch_size": 10} with mlflow.start_run() as run: model.fit(random_train_data, random_one_hot_labels, **fit_params) client = MlflowClient() assert client.get_run(run.info.run_id).data.params.keys() >= fit_params.keys() def test_import_tensorflow_with_fluent_autolog_enables_tensorflow_autologging(): mlflow.autolog() import tensorflow # noqa: F401 assert not autologging_is_disabled(mlflow.tensorflow.FLAVOR_NAME) def _assert_autolog_infers_model_signature_correctly(input_sig_spec, output_sig_spec): logged_model = mlflow.last_logged_model() artifact_path = local_file_uri_to_path(logged_model.artifact_location) ml_model_path = os.path.join(artifact_path, "MLmodel") with open(ml_model_path) as f: data = yaml.safe_load(f) assert data is not None assert "signature" in data signature = data["signature"] assert signature is not None assert "inputs" in signature assert "outputs" in signature assert json.loads(signature["inputs"]) == input_sig_spec assert json.loads(signature["outputs"]) == output_sig_spec def _assert_keras_autolog_input_example_load_and_predict_with_nparray(random_train_data): logged_model = mlflow.last_logged_model() model_conf = Model.load(logged_model.model_uri) input_example = _read_example(model_conf, logged_model.model_uri) np.testing.assert_array_almost_equal(input_example, random_train_data[:5]) pyfunc_model = mlflow.pyfunc.load_model(logged_model.model_uri) pyfunc_model.predict(input_example) def test_keras_autolog_input_example_load_and_predict_with_nparray( random_train_data, random_one_hot_labels ): mlflow.tensorflow.autolog(log_input_examples=True, log_model_signatures=True) initial_model = create_tf_keras_model() with mlflow.start_run(): initial_model.fit(random_train_data, random_one_hot_labels) _assert_keras_autolog_input_example_load_and_predict_with_nparray(random_train_data) def test_keras_autolog_infers_model_signature_correctly_with_nparray( random_train_data, random_one_hot_labels ): mlflow.tensorflow.autolog(log_model_signatures=True) initial_model = create_tf_keras_model() with mlflow.start_run(): initial_model.fit(random_train_data, random_one_hot_labels) _assert_autolog_infers_model_signature_correctly( [{"type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1, 4]}}], [{"type": "tensor", "tensor-spec": {"dtype": "float32", "shape": [-1, 3]}}], ) @pytest.mark.skipif( Version(tf.__version__) < Version("2.1.0"), reason="tf.data.Dataset inputs are unsupported for input example logging in TensorFlow < 2.1.0", ) def test_keras_autolog_input_example_load_and_predict_with_tf_dataset(fashion_mnist_tf_dataset): mlflow.tensorflow.autolog(log_input_examples=True, log_model_signatures=True) fashion_mnist_model = _create_fashion_mnist_model() with mlflow.start_run(): fashion_mnist_model.fit(fashion_mnist_tf_dataset) logged_model = mlflow.last_logged_model() model_conf = Model.load(logged_model.model_uri) input_example = _read_example(model_conf, logged_model.model_uri) pyfunc_model = mlflow.pyfunc.load_model(logged_model.model_uri) pyfunc_model.predict(input_example) @pytest.mark.skipif( Version(tf.__version__) < Version("2.1.0"), reason="tf.data.Dataset inputs are unsupported for signature logging in TensorFlow < 2.1.0", ) def test_keras_autolog_infers_model_signature_correctly_with_tf_dataset(fashion_mnist_tf_dataset): mlflow.tensorflow.autolog(log_model_signatures=True) fashion_mnist_model = _create_fashion_mnist_model() with mlflow.start_run(): fashion_mnist_model.fit(fashion_mnist_tf_dataset) _assert_autolog_infers_model_signature_correctly( [{"type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1, 28, 28]}}], [{"type": "tensor", "tensor-spec": {"dtype": "float32", "shape": [-1, 10]}}], ) def test_keras_autolog_input_example_load_and_predict_with_dict( random_train_dict_mapping, random_one_hot_labels ): mlflow.tensorflow.autolog(log_input_examples=True, log_model_signatures=True) model = _create_model_for_dict_mapping() with mlflow.start_run(): model.fit(random_train_dict_mapping, random_one_hot_labels) logged_model = mlflow.last_logged_model() model_conf = Model.load(logged_model.model_uri) input_example = _read_example(model_conf, logged_model.model_uri) for k, v in random_train_dict_mapping.items(): np.testing.assert_array_almost_equal(input_example[k], np.take(v, range(0, 5))) pyfunc_model = mlflow.pyfunc.load_model(logged_model.model_uri) pyfunc_model.predict(input_example) def test_keras_autolog_infers_model_signature_correctly_with_dict( random_train_dict_mapping, random_one_hot_labels ): mlflow.tensorflow.autolog(log_model_signatures=True) model = _create_model_for_dict_mapping() with mlflow.start_run(): model.fit(random_train_dict_mapping, random_one_hot_labels) _assert_autolog_infers_model_signature_correctly( [ {"name": "a", "type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1]}}, {"name": "b", "type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1]}}, {"name": "c", "type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1]}}, {"name": "d", "type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1]}}, ], [{"type": "tensor", "tensor-spec": {"dtype": "float32", "shape": [-1, 3]}}], ) def test_keras_autolog_input_example_load_and_predict_with_keras_sequence(keras_data_gen_sequence): mlflow.tensorflow.autolog(log_input_examples=True, log_model_signatures=True) model = create_tf_keras_model() with mlflow.start_run(): model.fit(keras_data_gen_sequence) _assert_keras_autolog_input_example_load_and_predict_with_nparray( keras_data_gen_sequence[:][0][:5] ) def test_keras_autolog_infers_model_signature_correctly_with_keras_sequence( keras_data_gen_sequence, ): mlflow.tensorflow.autolog(log_model_signatures=True) initial_model = create_tf_keras_model() with mlflow.start_run(): initial_model.fit(keras_data_gen_sequence) _assert_autolog_infers_model_signature_correctly( [{"type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1, 4]}}], [{"type": "tensor", "tensor-spec": {"dtype": "float32", "shape": [-1, 3]}}], ) def test_keras_autolog_load_saved_hdf5_model(keras_data_gen_sequence): mlflow.tensorflow.autolog(keras_model_kwargs={"save_format": "h5"}) model = create_tf_keras_model() with mlflow.start_run(): model.fit(keras_data_gen_sequence) logged_model = mlflow.last_logged_model() artifact_path = local_file_uri_to_path(logged_model.artifact_location) assert Path(artifact_path, "data", "model.h5").exists() def test_keras_autolog_logs_model_signature_by_default(keras_data_gen_sequence): mlflow.autolog() initial_model = create_tf_keras_model() initial_model.fit(keras_data_gen_sequence) logged_model = mlflow.last_logged_model() artifact_path = local_file_uri_to_path(logged_model.artifact_location) mlmodel_path = os.path.join(artifact_path, "MLmodel") with open(mlmodel_path) as f: mlmodel_contents = yaml.safe_load(f) assert "signature" in mlmodel_contents.keys() signature = mlmodel_contents["signature"] assert signature is not None assert "inputs" in signature assert "outputs" in signature assert json.loads(signature["inputs"]) == [ {"type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1, 4]}} ] assert json.loads(signature["outputs"]) == [ {"type": "tensor", "tensor-spec": {"dtype": "float32", "shape": [-1, 3]}} ] def test_extract_tf_keras_input_example_unsupported_type_returns_None(): from mlflow.tensorflow.autologging import extract_tf_keras_input_example extracted_data = extract_tf_keras_input_example([1, 2, 4, 5]) assert extracted_data is None, ( "Keras input data extraction function should have " "returned None as input type is not supported." ) def test_extract_input_example_from_tf_input_fn_unsupported_type_returns_None(): from mlflow.tensorflow.autologging import extract_tf_keras_input_example extracted_data = extract_tf_keras_input_example(lambda: [1, 2, 4, 5]) assert extracted_data is None, ( "Tensorflow's input_fn training data extraction should have" " returned None as input type is not supported." ) @pytest.mark.skipif( Version(tf.__version__) < Version("2.6.0"), reason=("TensorFlow only has a hard dependency on Keras in version >= 2.6.0"), ) def test_import_keras_model_trigger_import_tensorflow(): # This test is for guarding importing keras model will trigger importing tensorflow # Because in Keras>=2.6, the keras autologging patching is installed by # `mlflow.tensorflow.autolog`, suppose user enable autolog by `mlflow.autolog()`, # and then import keras, if keras does not trigger importing tensorflow, # then the keras autologging patching cannot be installed. py_executable = sys.executable _exec_cmd([ py_executable, "-c", "from keras import Model; import sys; assert 'tensorflow' in sys.modules", ]) def test_autolog_throw_error_on_explicit_mlflow_callback(keras_data_gen_sequence): mlflow.tensorflow.autolog() model = create_tf_keras_model() with mlflow.start_run() as run: with pytest.raises(MlflowException, match="MLflow autologging must be turned off*"): model.fit(keras_data_gen_sequence, callbacks=[MlflowCallback(run)]) def test_autolog_correct_logging_frequency(random_train_data, random_one_hot_labels): logging_freq = 5 num_epochs = 2 batch_size = 10 mlflow.tensorflow.autolog(log_every_epoch=False, log_every_n_steps=logging_freq) initial_model = create_tf_keras_model() with mlflow.start_run() as run: initial_model.fit( random_train_data, random_one_hot_labels, batch_size=batch_size, epochs=num_epochs, ) client = MlflowClient() loss_history = client.get_metric_history(run.info.run_id, "loss") assert len(loss_history) == num_epochs * (len(random_train_data) // batch_size) // logging_freq def test_automatic_checkpoint_per_epoch_callback(random_train_data, random_one_hot_labels): mlflow.tensorflow.autolog( checkpoint=True, checkpoint_monitor=None, checkpoint_mode=None, checkpoint_save_best_only=False, checkpoint_save_weights_only=False, checkpoint_save_freq="epoch", ) model = create_tf_keras_model() with mlflow.start_run() as run: model.fit(random_train_data, random_one_hot_labels, epochs=1) run_id = run.info.run_id logged_metrics = mlflow.artifacts.load_dict( f"runs:/{run_id}/checkpoints/epoch_0/checkpoint_metrics.json" ) assert set(logged_metrics) == {"epoch", "loss", "accuracy", "global_step"} assert logged_metrics["epoch"] == 0 assert logged_metrics["global_step"] == 5 pred_result = model.predict(random_train_data) pred_result2 = load_checkpoint(run_id=run_id).predict(random_train_data) np.testing.assert_array_almost_equal(pred_result, pred_result2) pred_result3 = load_checkpoint(run_id=run_id, epoch=0).predict(random_train_data) np.testing.assert_array_almost_equal(pred_result, pred_result3) def test_automatic_checkpoint_per_epoch_save_weight_only_callback( random_train_data, random_one_hot_labels ): mlflow.tensorflow.autolog( checkpoint=True, checkpoint_monitor=None, checkpoint_mode=None, checkpoint_save_best_only=False, checkpoint_save_weights_only=True, checkpoint_save_freq="epoch", ) model = create_tf_keras_model() with mlflow.start_run() as run: model.fit(random_train_data, random_one_hot_labels, epochs=1) run_id = run.info.run_id logged_metrics = mlflow.artifacts.load_dict( f"runs:/{run_id}/checkpoints/epoch_0/checkpoint_metrics.json" ) assert set(logged_metrics) == {"epoch", "loss", "accuracy", "global_step"} assert logged_metrics["epoch"] == 0 assert logged_metrics["global_step"] == 5 model2 = create_tf_keras_model() pred_result = model.predict(random_train_data) pred_result2 = load_checkpoint(model=model2, run_id=run_id).predict(random_train_data) np.testing.assert_array_almost_equal(pred_result, pred_result2) def test_automatic_checkpoint_per_3_steps_callback(random_train_data, random_one_hot_labels): mlflow.tensorflow.autolog( checkpoint=True, checkpoint_monitor=None, checkpoint_mode=None, checkpoint_save_best_only=False, checkpoint_save_weights_only=False, checkpoint_save_freq=3, ) model = create_tf_keras_model() with mlflow.start_run() as run: model.fit(random_train_data, random_one_hot_labels, epochs=1) run_id = run.info.run_id logged_metrics = mlflow.artifacts.load_dict( f"runs:/{run_id}/checkpoints/global_step_3/checkpoint_metrics.json" ) assert set(logged_metrics) == {"epoch", "loss", "accuracy", "global_step"} assert logged_metrics["epoch"] == 0 assert logged_metrics["global_step"] == 3 assert isinstance(load_checkpoint(run_id=run_id), tf.keras.Sequential) assert isinstance(load_checkpoint(run_id=run_id, global_step=3), tf.keras.Sequential) def test_automatic_checkpoint_per_3_steps_save_best_only_callback( random_train_data, random_one_hot_labels ): mlflow.tensorflow.autolog( checkpoint=True, checkpoint_monitor="loss", checkpoint_mode="min", checkpoint_save_best_only=True, checkpoint_save_weights_only=False, checkpoint_save_freq=3, ) model = create_tf_keras_model() with mlflow.start_run() as run: model.fit( random_train_data, random_one_hot_labels, epochs=1, ) run_id = run.info.run_id logged_metrics = mlflow.artifacts.load_dict( f"runs:/{run_id}/checkpoints/latest_checkpoint_metrics.json" ) assert set(logged_metrics) == {"epoch", "loss", "accuracy", "global_step"} assert logged_metrics["epoch"] == 0 assert logged_metrics["global_step"] == 3 assert isinstance(load_checkpoint(run_id=run_id), tf.keras.Sequential)