client runner keeps restarting and closing the connection until my model is saved in the mlflow instance about nvflare HOT 9 CLOSED

   @flare.train
    def train(input_model=None):
        """
        Train UNet model with training images (X_train) and masks (y_train)
        and user-defined parameters.
        Use test dataset only to visualise performance before final evaluation.
        X data can be 3 channel or of different dimension (UNet will automatically be adapted).

        :param X_train: train images
        :param y_train_onehot: corresponding onehot encoded train masks
        :param X_test: test images
        :param y_test_onehot: corresponding onehot encoded test masks
        :param model_path: (Path) path to which model will be saved
        :return: model saved to provided path
        """
        start = time.time()
        # Check that provided path can be used for saving model
        assert (
            type(model_path) == PosixPath
            and model_path.suffix == ".h5py"
        ), f"Provided model path '{model_path}' is not a Path ending in '.h5py'!"

        # (5) loads model from NVFlare
        for k, v in input_model.params.items():
            model.get_layer(k).set_weights(v)
        cfg["epochs"] = 1
        _logger.info("Training model...")
        history = model.fit(
            X_train,
            y_train_onehot,
            batch_size=cfg["batch_size"],
            verbose=2,
            epochs=cfg["epochs"],
            validation_data=(X_test, y_test_onehot),
            callbacks=[CustomEpochLogger()],
        )

        # print('history is ', history.history)
        metrics = {
            metric: values[-1]
            for metric, values in history.history.items()
        }

        mlflow_writer.log_metrics(
            metrics=metrics, step=input_model.current_round
        )
        duration = time.time() - start
        _logger.info(
            f"Elapsed time during model training:\t{round(duration / 60, 2)} min"
        )

        # Save trained Model
        # model.save_weights(model_path)
        model.save(model_dir)
        _logger.info(f"Model saved to '{model_path}'.")
        save_model_info(model=model, model_dir=model_dir)

        # save json config to model directory
        # cp_conf(model_dir)
        _logger.info(
            f"Saved configuration of training run to {model_dir}"
        )

        # (3) send back the model to nvflare server
        output_model = flare.FLModel(
            params={
                layer.name: layer.get_weights()
                for layer in model.layers
            },
            metrics=history.history,
        )

        return output_model
while flare.is_running():
        input_model = flare.receive()
        input_model.current_round
        print(f"current_round={input_model.current_round}")

        # (optional) print system info
        system_info = flare.system_info()
        print(f"NVFlare system info: {system_info}")

        train(input_model=input_model)

while flare.is_running():
       input_model = flare.receive()
       current_round = input_model.current_round
        total_rounds = input_model.total_rounds
         print(f"Currently in round {current_round} out of a total of {total_rounds} rounds")