"No module named 'mlflow.utils.async_logging'"

[t-wolfeadam]

When running the "hyperparameter tuning and monitoring" example, the job fails due to missing the async_logging module from mlflow. This is confirmed by the fact the module was added in v2.11 (link to release), yet the conda file specifies v2.4.1 (link).

I believe this is due to the unspecified version for at least the azureml-mlflow module, as evidenced by the traceback.

Traceback (most recent call last):
  File "/mnt/azureml/cr/j/0d777414ce404c88bded3e037c6d4d40/exe/wd/tune.py", line 208, in <module>
    run(args.num_tune_samples, args.env_name, args.scheduler)
  File "/mnt/azureml/cr/j/0d777414ce404c88bded3e037c6d4d40/exe/wd/tune.py", line 90, in run
    current_run = mlflow.start_run()
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/mlflow/tracking/fluent.py", line 291, in start_run
    client = MlflowClient()
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/mlflow/tracking/client.py", line 80, in __init__
    self._tracking_client = TrackingServiceClient(final_tracking_uri)
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/mlflow/tracking/_tracking_service/client.py", line 53, in __init__
    self.store
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/mlflow/tracking/_tracking_service/client.py", line 57, in store
    return utils._get_store(self.tracking_uri)
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/mlflow/tracking/_tracking_service/utils.py", line 217, in _get_store
    return _tracking_store_registry.get_store(store_uri, artifact_uri)
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/mlflow/tracking/_tracking_service/registry.py", line 39, in get_store
    return self._get_store_with_resolved_uri(resolved_store_uri, artifact_uri)
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/mlflow/tracking/_tracking_service/registry.py", line 49, in _get_store_with_resolved_uri
    return builder(store_uri=resolved_store_uri, artifact_uri=artifact_uri)
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/azureml/mlflow/entry_point_loaders.py", line 33, in azureml_store_builder
    from azureml.mlflow._store.tracking.store import AzureMLRestStore
  File "/azureml-envs/azureml_1d2d74dc70c5588f68061b70e19a7950/lib/python3.10/site-packages/azureml/mlflow/_store/tracking/store.py", line 23, in <module>
    from mlflow.utils.async_logging.run_operations import RunOperations
ModuleNotFoundError: No module named 'mlflow.utils.async_logging'

This is the conda spec I landed on as working, through some trial and error.

channels:
  - anaconda
  - conda-forge
dependencies:
  - python=3.10.11
  - pip=23.0.1
  - pip:
      - azureml-mlflow==1.55.0
      - azureml-defaults==1.55.0
      - mlflow==2.11.1
      - mlflow-skinny==2.11.1
      - ray-on-aml==0.2.4
      - ray[data]==2.9.3
      - ray[rllib]==2.9.3
      - torch>=2.0.1
      - tensorflow==2.15.0
      - tensorflow-probability==0.23.0
      - tensorboardX>=2.6.2
      - torch>=2.2.1
      - gymnasium==0.28.1
      - numpy==1.23.5