Hello everyone, does anyone know if the sandbox deployment allows a connection between mlflow and flyte? When I try to run a task to test it, the task fails. Here is the code that I'm running, as well as the error I'm getting: Code:

import mlflow.keras
import tensorflow as tf
from flytekit import task, workflow
from flytekitplugins.mlflow import mlflow_autolog

@task(disable_deck=False)
@mlflow_autolog(framework=mlflow.keras)
def train_model(epochs: int):
    # Refer to <https://www.tensorflow.org/tutorials/keras/classification>
    fashion_mnist = tf.keras.datasets.fashion_mnist
    (train_images, train_labels), (_, _) = fashion_mnist.load_data()
    train_images = train_images / 255.0

    model = tf.keras.Sequential(
        [
            tf.keras.layers.Flatten(input_shape=(28, 28)),
            tf.keras.layers.Dense(128, activation="relu"),
            tf.keras.layers.Dense(10),
        ]
    )

    model.compile(
        optimizer="adam",
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=["accuracy"],
    )
    model.fit(train_images, train_labels, epochs=epochs)

@workflow
def ml_pipeline(epochs: int):
    train_model(epochs=epochs)


if __name__ == "__main__":
    print(f"Running {__file__} main...")
    ml_pipeline(epochs=5)

Error:

[1/1] currentAttempt done. Last Error: USER::                                   │
│ ❱ 295 │   │   │   │   return func(args, *kwargs)                           │
│                                                                              │
│ /usr/local/lib/python3.10/site-packages/flytekit/core/python_auto_container. │
│ py:235 in load_task                                                          │
│                                                                              │
│ ❱ 235 │   │   task_module = importlib.import_module(name=task_module)  # typ │
│                                                                              │
│ /usr/local/lib/python3.10/importlib/init.py:126 in import_module         │
│                                                                              │
│ ❱ 126 │   return _bootstrap._gcd_import(name[level:], package, level)        │
│ in _gcd_import:1050                                                          │
│ in _find_and_load:1027                                                       │
│ in _find_and_load_unlocked:1006                                              │
│ in _load_unlocked:688                                                        │
│ in exec_module:883                                                           │
│ in _call_with_frames_removed:241                                             │
│                                                                              │
│ /root/test2.py:1 in <module>                                                 │
│                                                                              │
│ ❱  1 import mlflow.keras                                                     │
╰──────────────────────────────────────────────────────────────────────────────╯
ModuleNotFoundError: No module named 'mlflow'

whats the folder structure of your project?

cc @Ketan (kumare3)

I just created a PR for a new type of mlflow decorator and have the integration working

Thank you, if I have any follow up questions on this topic I'll ask them here