Hi All We are using the Metaflow extension framework. We are experiencing some issues with the Metaflow UI visualising the DAG on certain flows. It seems to be related to the

metaflow_extension

library not being available in the UI. Has anyone else experienced this issue? Any guidance on solutions would be much appreciated. TIA! Here is the stacktrace:

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/root/services/ui_backend_service/data/cache/client/cache_server.py", line 307, in <module>
    cli(auto_envvar_prefix='MFCACHE')
  File "/opt/latest/lib/python3.11/site-packages/click/core.py", line 1128, in __call__
    return self.main(*args, **kwargs)
  File "/opt/latest/lib/python3.11/site-packages/click/core.py", line 1053, in main
    rv = self.invoke(ctx)
  File "/opt/latest/lib/python3.11/site-packages/click/core.py", line 1395, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/latest/lib/python3.11/site-packages/click/core.py", line 754, in invoke
    return __callback(*args, **kwargs)
  File "/root/services/ui_backend_service/data/cache/client/cache_server.py", line 301, in cli
    Scheduler(store, max_actions).loop()
  File "/root/services/ui_backend_service/data/cache/client/cache_server.py", line 196, in __init__
    self.pool = multiprocessing.Pool(
  File "/usr/local/lib/python3.11/multiprocessing/context.py", line 119, in Pool
    return Pool(processes, initializer, initargs, maxtasksperchild,
  File "/usr/local/lib/python3.11/multiprocessing/pool.py", line 215, in __init__
    self._repopulate_pool()
  File "/usr/local/lib/python3.11/multiprocessing/pool.py", line 306, in _repopulate_pool
    return self._repopulate_pool_static(self._ctx, self.Process,
  File "/usr/local/lib/python3.11/multiprocessing/pool.py", line 329, in _repopulate_pool_static
    w.start()
  File "/usr/local/lib/python3.11/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/usr/local/lib/python3.11/multiprocessing/context.py", line 281, in _Popen
    return Popen(process_obj)
  File "/usr/local/lib/python3.11/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/usr/local/lib/python3.11/multiprocessing/popen_fork.py", line 71, in _launch
    code = process_obj._bootstrap(parent_sentinel=child_r)
  File "/usr/local/lib/python3.11/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/local/lib/python3.11/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/local/lib/python3.11/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
  File "/root/services/ui_backend_service/data/cache/client/cache_worker.py", line 29, in execute_action
    execute(tempdir, action_cls, request)
  File "/root/services/ui_backend_service/data/cache/client/cache_worker.py", line 51, in execute
    res = action_cls.execute(
  File "/root/services/ui_backend_service/data/cache/generate_dag_action.py", line 89, in execute
    with streamed_errors(stream_output):
  File "/usr/local/lib/python3.11/contextlib.py", line 155, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/root/services/ui_backend_service/data/cache/utils.py", line 130, in streamed_errors
    get_traceback_str()
  File "/root/services/ui_backend_service/data/cache/utils.py", line 124, in streamed_errors
    yield
  File "/root/services/ui_backend_service/data/cache/generate_dag_action.py", line 93, in execute
    dag = DataArtifact("{}/_graph_info".format(param_step.task.pathspec)).data
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/latest/lib/python3.11/site-packages/metaflow/client/core.py", line 906, in data
    obj = filecache.get_artifact(ds_type, location[6:], meta, *components)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/latest/lib/python3.11/site-packages/metaflow/client/filecache.py", line 207, in get_artifact
    _, obj = next(
             ^^^^^
  File "/opt/latest/lib/python3.11/site-packages/metaflow/datastore/task_datastore.py", line 370, in load_artifacts
    yield name, pickle.loads(blob)
                ^^^^^^^^^^^^^^^^^^

ModuleNotFoundError: No module named 'metaflow_extensions'