[BUG] Task cannot return `keras.Model` - "Invalid filepath extension for saving."

[HansBambel]

Describe the bug

Following the example from https://www.union.ai/docs/flyte/user-guide/data-input-output/tensorflow/ I tried returning a keras.Model, but receive the following error due to a naming problem:

ValueError: Failed to convert outputs of task 'src.workflows.tutorial_example.train_model' at position 0.
    Failed to convert type <class 'keras.src.models.sequential.Sequential'> to type <class 'keras.src.models.model.Model'>.
    Error Message: Invalid filepath extension for saving. Please add either a `.keras` extension for the native Keras format (recommended) or a `.h5` extension. Use `model.export(filepath)` if you want to export a SavedModel for use with TFLite/TFServing/etc. Received: filepath=/tmp/flytek4ddfb4d/local_flytekit/6f804fd14418b087683c300d6a2547fe..

It seems that the path that flyte gives to keras' save function does not end with an allowed extension such as keras (which it can't infer anyway).

Is the workaround/fix to save the model and give back a folder/file instead? Then the docs should reflect that.

Expected behavior

To save the keras model as task output to feed it to another task.

Additional context to reproduce

1. Follow tutorial from https://www.union.ai/docs/flyte/user-guide/data-input-output/tensorflow/ (only requires the train_model-function)
2. Run the workflow

Screenshots

No response

Are you sure this issue hasn't been raised already?

• [x] Yes

Have you read the Code of Conduct?

• [x] Yes

[runllm[bot]]

It seems likely that the issue you're encountering is related to the file path extension used when saving a Keras model in Flyte. The error message suggests that the path provided by Flyte does not have a valid extension for saving Keras models, such as .keras or .h5. This is a known requirement for Keras model saving functions.

My understanding is that a potential workaround is to manually save the Keras model to a file with a valid extension and then return the file path or directory as a FlyteFile or FlyteDirectory. This approach ensures that the model is saved with the correct format and can be used as an output in Flyte tasks. Would you like more details on how to implement this workaround?

Sources:

• https://github.com/flyteorg/flyte/issues/6393
• https://www.union.ai/docs/flyte/user-guide/data-input-output/tensorflow/ Note: I wasn’t able to find highly relevant data sources, but above are a few potentially relevant links.

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @runllm to ask follow-up questions. See runllm.com for more!

[HansBambel]

As a workaround I am now using FlyteFile:

    keras_output_file = output_dir / "saved_model.keras"
    model.save(keras_output_file)
    return FlyteFile(keras_output_file)

And then in another task that wants to use this model:

    model = keras.models.load_model(model_file)

[davidmirror-ops]

@HansBambel I think this indeed is an issue with the example and using FlyteFile should be the alternative cc @ppiegaze

[github-actions[bot]]

"Hello 👋, this issue has been inactive for over 90 days. To help maintain a clean and focused backlog, we'll be marking this issue as stale and will close the issue if we detect no activity in the next 7 days. Thank you for your contribution and understanding! 🙏"

[HansBambel]

Bump

[github-actions[bot]]

"Hello 👋, this issue has been inactive for over 90 days. To help maintain a clean and focused backlog, we'll be marking this issue as stale and will close the issue if we detect no activity in the next 7 days. Thank you for your contribution and understanding! 🙏"

[HansBambel]

bump

[kumare3]

Uhh we missed this one

[dansola]

Proposed fix: https://github.com/flyteorg/flytekit/pull/3346