Prefect deploy hangs indefinitely when flow imports AzureML MLClient and MlflowClient

[AMontgomerie]

Bug summary

After upgrading from prefect 2 to 3, running prefect --no-prompt deploy --prefect-file prefect.yaml --all gets stuck indefinitely. With log level set to debug, the only console output is:

❯ prefect --no-prompt deploy --prefect-file prefect.yaml --all
07:52:51.394 | DEBUG   | prefect.profiles - Using profile 'local'
07:52:51.488 | DEBUG   | prefect.client - Connecting to API at http://127.0.0.1:4200/api/
07:52:51.522 | DEBUG   | prefect.client - Connecting to API at http://127.0.0.1:4200/api/
07:52:51.557 | DEBUG   | prefect.client - Connecting to API at http://127.0.0.1:4200/api/

After lots of trial and error, I managed to track down the lines of code that are causing it to get stuck. If my flow imports both azure.ai.ml.MLClient and mlflow.tracking.MlflowClient then this issue occurs. If either class is not imported then deployment goes ahead normally.

For example:

# hello.py
from prefect import flow
from prefect.logging import get_run_logger

from azure.ai.ml import MLClient  # noqa
from mlflow.tracking import MlflowClient  # noqa


@flow
def hello():
    get_run_logger().info("hello")


if __name__ == "__main__":
    hello()

name: test-project
prefect-version: 3.0.5

deployments:
  - name: development
    tags: null
    schedule: null
    entrypoint: hello.py:hello
    work_pool:
      name: my-work-pool
      job_variables:
        image: myimage:latest

Version info (prefect version output)

❯ prefect version
07:59:09.542 | DEBUG   | prefect.profiles - Using profile 'local'
Version:             3.0.5
API version:         0.8.4
Python version:      3.11.9
Git commit:          1d1b92a7
Built:               Tue, Oct 8, 2024 5:55 PM
OS/Arch:             linux/x86_64
Profile:             local
Server type:         server
Pydantic version:    2.7.1
Integrations:
  prefect-kubernetes: 0.5.0

Additional context

Other dependencies:

• azure-ai-ml-1.20.0
• mlflow-skinny-2.16.2
• azureml-mlflow-1.57.0

I have a local self-hosted Prefect server running inside Docker Desktop Kubernetes, and also a Prefect Cloud account with a work pool that connects to an AKS cluster, and the same issue occurs when attempting to deploy to either.

[zzstoatzz]

hi @AMontgomerie - hmm interesting. I have reproduced but I'm not sure what's causing this to hang

I will say that you should be able to unblock yourself by deferring your imports, something like

# hello.py

from prefect import flow
from prefect.logging import get_run_logger


@flow
def hello():
    from azure.ai.ml import MLClient  # noqa
    from mlflow.tracking import MlflowClient  # noqa

    get_run_logger().info("hello")


if __name__ == "__main__":
    hello()

will try to understand more about what's going on

[AMontgomerie]

Thanks for the workaround! It was a little awkward to implement because these imports are inside a utility package that is imported throughout a lot of different flows, but I am able to deploy the flows now at least.

[lennartvandeguchte]

I'm running into the same error when importing mlflow in any of my flow scripts. @zzstoatzz did you already found a clue what is causing this behaviour? I don't want to use the workaround you proposed because mlflow is used all over the place in my codebase.

[jeanluciano]

I think this might be related to #16115. I was not able to reproduce on 3.1.13. Could you try reproducing it again?

[lennartvandeguchte]

@jeanluciano Yes, I can confirm that my issue has been resolved with version 3.1.13

[jeanluciano]

Awesome! will be closing this issue.