Encountered AssertionError regarding distutils when running fx plan initialize

[psfoley]

Discussed in #1016

^{Originally posted by hwpang July 31, 2024}
Hi, I'm a new user to OpenFL and am going through the Quick Start example at https://openfl.readthedocs.io/en/latest/get_started/quickstart.html. I was able to follow the example up until fx plan initialize, where I encountered the following assertion error:

/anaconda/envs/openfl-env/lib/python3.8/site-packages/_distutils_hack/__init__.py:11: UserWarning: Distutils was imported before Setuptools, but importing Setuptools also replaces the `distutils` module in `sys.modules`. This may lead to undesirable behaviors or errors. To avoid these issues, avoid using distutils directly, ensure that setuptools is installed in the traditional way (e.g. not an editable install), and/or make sure that setuptools is always imported before distutils.
  warnings.warn(
/anaconda/envs/openfl-env/lib/python3.8/site-packages/_distutils_hack/__init__.py:26: UserWarning: Setuptools is replacing distutils.
  warnings.warn("Setuptools is replacing distutils.")
EXCEPTION : /anaconda/envs/openfl-env/lib/python3.8/distutils/core.py
Traceback (most recent call last):
  File "/anaconda/envs/openfl-env/bin/fx", line 8, in <module>
    sys.exit(entry())
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/openfl/interface/cli.py", line 268, in entry
    error_handler(e)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/openfl/interface/cli.py", line 195, in error_handler
    raise error
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/openfl/interface/cli.py", line 266, in entry
    cli(max_content_width=120)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/click/decorators.py", line 33, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/openfl/interface/plan.py", line 129, in initialize
    data_loader = plan.get_data_loader(collaborator_cname)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/openfl/federated/plan/plan.py", line 390, in get_data_loader
    self.loader_ = Plan.build(**defaults)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/openfl/federated/plan/plan.py", line 194, in build
    module = import_module(module_path)
  File "/anaconda/envs/openfl-env/lib/python3.8/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1014, in _gcd_import
  File "<frozen importlib._bootstrap>", line 991, in _find_and_load
  File "<frozen importlib._bootstrap>", line 975, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 671, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 843, in exec_module
  File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
  File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/hpang7/code/Users/hpang/Projects/Federated_learning/openfl/my_workspace/src/ptmnist_inmemory.py", line 7, in <module>
    from .mnist_utils import load_mnist_shard
  File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/hpang7/code/Users/hpang/Projects/Federated_learning/openfl/my_workspace/src/mnist_utils.py", line 10, in <module>
    from torchvision import datasets
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torchvision/__init__.py", line 10, in <module>
    from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils  # usort:skip
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torchvision/models/__init__.py", line 2, in <module>
    from .convnext import *
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torchvision/models/convnext.py", line 8, in <module>
    from ..ops.misc import Conv2dNormActivation, Permute
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torchvision/ops/__init__.py", line 23, in <module>
    from .poolers import MultiScaleRoIAlign
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torchvision/ops/poolers.py", line 10, in <module>
    from .roi_align import roi_align
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torchvision/ops/roi_align.py", line 7, in <module>
    from torch._dynamo.utils import is_compile_supported
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torch/_dynamo/__init__.py", line 2, in <module>
    from . import convert_frame, eval_frame, resume_execution
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 48, in <module>
    from . import config, exc, trace_rules
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torch/_dynamo/exc.py", line 12, in <module>
    from .utils import counters
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 1063, in <module>
    if has_triton_package():
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/torch/utils/_triton.py", line 9, in has_triton_package
    import triton
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/__init__.py", line 8, in <module>
    from .runtime import (
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/runtime/__init__.py", line 1, in <module>
    from .autotuner import (Autotuner, Config, Heuristics, autotune, heuristics)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/runtime/autotuner.py", line 9, in <module>
    from ..testing import do_bench, do_bench_cudagraph
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/testing.py", line 7, in <module>
    from . import language as tl
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/language/__init__.py", line 4, in <module>
    from . import math
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/language/math.py", line 1, in <module>
    from . import core
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/language/core.py", line 10, in <module>
    from ..runtime.jit import jit
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/runtime/jit.py", line 12, in <module>
    from ..runtime.driver import driver
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/runtime/driver.py", line 1, in <module>
    from ..backends import backends
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/backends/__init__.py", line 50, in <module>
    backends = _discover_backends()
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/backends/__init__.py", line 44, in _discover_backends
    driver = _load_module(name, os.path.join(root, name, 'driver.py'))
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/backends/__init__.py", line 12, in _load_module
    spec.loader.exec_module(module)
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/backends/amd/driver.py", line 7, in <module>
    from triton.runtime.build import _build
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/triton/runtime/build.py", line 8, in <module>
    import setuptools
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/setuptools/__init__.py", line 8, in <module>
    import _distutils_hack.override  # noqa: F401
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/_distutils_hack/override.py", line 1, in <module>
    __import__('_distutils_hack').do_override()
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/_distutils_hack/__init__.py", line 70, in do_override
    ensure_local_distutils()
  File "/anaconda/envs/openfl-env/lib/python3.8/site-packages/_distutils_hack/__init__.py", line 57, in ensure_local_distutils
    assert '_distutils' in core.__file__, core.__file__
AssertionError: /anaconda/envs/openfl-env/lib/python3.8/distutils/core.py

Information that may be relevant:

Python version
3.8.19

OpenFL version
1.5

Operating system
Ubuntu 20.04

My environment packages from conda list:

# packages in environment at /anaconda/envs/openfl-env:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                        main  
_openmp_mutex             5.1                       1_gnu  
absl-py                   2.1.0                    pypi_0    pypi
anyio                     4.4.0                    pypi_0    pypi
argon2-cffi               23.1.0                   pypi_0    pypi
argon2-cffi-bindings      21.2.0                   pypi_0    pypi
arrow                     1.3.0                    pypi_0    pypi
asttokens                 2.4.1                    pypi_0    pypi
async-lru                 2.0.4                    pypi_0    pypi
attrs                     23.2.0                   pypi_0    pypi
babel                     2.15.0                   pypi_0    pypi
backcall                  0.2.0                    pypi_0    pypi
beautifulsoup4            4.12.3                   pypi_0    pypi
bleach                    6.1.0                    pypi_0    pypi
bzip2                     1.0.8                h5eee18b_6  
ca-certificates           2024.7.2             h06a4308_0  
cachetools                5.4.0                    pypi_0    pypi
certifi                   2024.7.4                 pypi_0    pypi
cffi                      1.16.0                   pypi_0    pypi
charset-normalizer        3.3.2                    pypi_0    pypi
click                     8.1.7                    pypi_0    pypi
cloudpickle               3.0.0                    pypi_0    pypi
comm                      0.2.2                    pypi_0    pypi
cryptography              43.0.0                   pypi_0    pypi
debugpy                   1.8.2                    pypi_0    pypi
decorator                 5.1.1                    pypi_0    pypi
defusedxml                0.7.1                    pypi_0    pypi
docker                    7.1.0                    pypi_0    pypi
dynaconf                  3.2.5                    pypi_0    pypi
exceptiongroup            1.2.2                    pypi_0    pypi
executing                 2.0.1                    pypi_0    pypi
fastjsonschema            2.20.0                   pypi_0    pypi
filelock                  3.15.4                   pypi_0    pypi
flatten-json              0.1.14                   pypi_0    pypi
fqdn                      1.5.1                    pypi_0    pypi
fsspec                    2024.6.1                 pypi_0    pypi
google-auth               2.32.0                   pypi_0    pypi
google-auth-oauthlib      1.0.0                    pypi_0    pypi
grpcio                    1.65.2                   pypi_0    pypi
h11                       0.14.0                   pypi_0    pypi
httpcore                  1.0.5                    pypi_0    pypi
httpx                     0.27.0                   pypi_0    pypi
idna                      3.7                      pypi_0    pypi
importlib-metadata        8.2.0                    pypi_0    pypi
importlib-resources       6.4.0                    pypi_0    pypi
ipykernel                 6.29.5                   pypi_0    pypi
ipython                   8.12.3                   pypi_0    pypi
isoduration               20.11.0                  pypi_0    pypi
jedi                      0.19.1                   pypi_0    pypi
jinja2                    3.1.4                    pypi_0    pypi
joblib                    1.4.2                    pypi_0    pypi
json5                     0.9.25                   pypi_0    pypi
jsonpointer               3.0.0                    pypi_0    pypi
jsonschema                4.23.0                   pypi_0    pypi
jsonschema-specifications 2023.12.1                pypi_0    pypi
jupyter-client            8.6.2                    pypi_0    pypi
jupyter-core              5.7.2                    pypi_0    pypi
jupyter-events            0.10.0                   pypi_0    pypi
jupyter-lsp               2.2.5                    pypi_0    pypi
jupyter-server            2.14.2                   pypi_0    pypi
jupyter-server-terminals  0.5.3                    pypi_0    pypi
jupyterlab                4.2.4                    pypi_0    pypi
jupyterlab-pygments       0.3.0                    pypi_0    pypi
jupyterlab-server         2.27.3                   pypi_0    pypi
ld_impl_linux-64          2.38                 h1181459_1  
libffi                    3.4.4                h6a678d5_1  
libgcc-ng                 11.2.0               h1234567_1  
libgomp                   11.2.0               h1234567_1  
libstdcxx-ng              11.2.0               h1234567_1  
libuuid                   1.41.5               h5eee18b_0  
markdown                  3.6                      pypi_0    pypi
markdown-it-py            3.0.0                    pypi_0    pypi
markupsafe                2.1.5                    pypi_0    pypi
matplotlib-inline         0.1.7                    pypi_0    pypi
mdurl                     0.1.2                    pypi_0    pypi
mistune                   3.0.2                    pypi_0    pypi
mpmath                    1.3.0                    pypi_0    pypi
nbclient                  0.10.0                   pypi_0    pypi
nbconvert                 7.16.4                   pypi_0    pypi
nbformat                  5.10.4                   pypi_0    pypi
ncurses                   6.4                  h6a678d5_0  
nest-asyncio              1.6.0                    pypi_0    pypi
networkx                  3.1                      pypi_0    pypi
notebook-shim             0.2.4                    pypi_0    pypi
numpy                     1.24.4                   pypi_0    pypi
nvidia-cublas-cu12        12.1.3.1                 pypi_0    pypi
nvidia-cuda-cupti-cu12    12.1.105                 pypi_0    pypi
nvidia-cuda-nvrtc-cu12    12.1.105                 pypi_0    pypi
nvidia-cuda-runtime-cu12  12.1.105                 pypi_0    pypi
nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
nvidia-cufft-cu12         11.0.2.54                pypi_0    pypi
nvidia-curand-cu12        10.3.2.106               pypi_0    pypi
nvidia-cusolver-cu12      11.4.5.107               pypi_0    pypi
nvidia-cusparse-cu12      12.1.0.106               pypi_0    pypi
nvidia-nccl-cu12          2.20.5                   pypi_0    pypi
nvidia-nvjitlink-cu12     12.5.82                  pypi_0    pypi
nvidia-nvtx-cu12          12.1.105                 pypi_0    pypi
oauthlib                  3.2.2                    pypi_0    pypi
openfl                    1.5                      pypi_0    pypi
openssl                   3.0.14               h5eee18b_0  
overrides                 7.7.0                    pypi_0    pypi
packaging                 24.1                     pypi_0    pypi
pandas                    2.0.3                    pypi_0    pypi
pandocfilters             1.5.1                    pypi_0    pypi
parso                     0.8.4                    pypi_0    pypi
pexpect                   4.9.0                    pypi_0    pypi
pickleshare               0.7.5                    pypi_0    pypi
pillow                    10.4.0                   pypi_0    pypi
pip                       24.0             py38h06a4308_0  
pkgutil-resolve-name      1.3.10                   pypi_0    pypi
platformdirs              4.2.2                    pypi_0    pypi
prometheus-client         0.20.0                   pypi_0    pypi
prompt-toolkit            3.0.47                   pypi_0    pypi
protobuf                  3.20.3                   pypi_0    pypi
psutil                    6.0.0                    pypi_0    pypi
ptyprocess                0.7.0                    pypi_0    pypi
pure-eval                 0.2.3                    pypi_0    pypi
pyasn1                    0.6.0                    pypi_0    pypi
pyasn1-modules            0.4.0                    pypi_0    pypi
pycparser                 2.22                     pypi_0    pypi
pygments                  2.18.0                   pypi_0    pypi
python                    3.8.19               h955ad1f_0  
python-dateutil           2.9.0.post0              pypi_0    pypi
python-json-logger        2.0.7                    pypi_0    pypi
pytz                      2024.1                   pypi_0    pypi
pyyaml                    6.0.1                    pypi_0    pypi
pyzmq                     26.0.3                   pypi_0    pypi
readline                  8.2                  h5eee18b_0  
referencing               0.35.1                   pypi_0    pypi
requests                  2.32.3                   pypi_0    pypi
requests-oauthlib         2.0.0                    pypi_0    pypi
rfc3339-validator         0.1.4                    pypi_0    pypi
rfc3986-validator         0.1.1                    pypi_0    pypi
rich                      13.7.1                   pypi_0    pypi
rpds-py                   0.19.1                   pypi_0    pypi
rsa                       4.9                      pypi_0    pypi
scikit-learn              1.3.2                    pypi_0    pypi
scipy                     1.10.1                   pypi_0    pypi
send2trash                1.8.3                    pypi_0    pypi
setuptools                69.5.1           py38h06a4308_0  
six                       1.16.0                   pypi_0    pypi
sniffio                   1.3.1                    pypi_0    pypi
soupsieve                 2.5                      pypi_0    pypi
sqlite                    3.45.3               h5eee18b_0  
stack-data                0.6.3                    pypi_0    pypi
sympy                     1.13.1                   pypi_0    pypi
tensorboard               2.14.0                   pypi_0    pypi
tensorboard-data-server   0.7.2                    pypi_0    pypi
tensorboardx              2.6                      pypi_0    pypi
terminado                 0.18.1                   pypi_0    pypi
threadpoolctl             3.5.0                    pypi_0    pypi
tinycss2                  1.3.0                    pypi_0    pypi
tk                        8.6.14               h39e8969_0  
tomli                     2.0.1                    pypi_0    pypi
torch                     2.4.0                    pypi_0    pypi
torchvision               0.19.0                   pypi_0    pypi
tornado                   6.4.1                    pypi_0    pypi
tqdm                      4.66.4                   pypi_0    pypi
traitlets                 5.14.3                   pypi_0    pypi
triton                    3.0.0                    pypi_0    pypi
types-python-dateutil     2.9.0.20240316           pypi_0    pypi
typing-extensions         4.12.2                   pypi_0    pypi
tzdata                    2024.1                   pypi_0    pypi
uri-template              1.3.0                    pypi_0    pypi
urllib3                   2.2.2                    pypi_0    pypi
wcwidth                   0.2.13                   pypi_0    pypi
webcolors                 24.6.0                   pypi_0    pypi
webencodings              0.5.1                    pypi_0    pypi
websocket-client          1.8.0                    pypi_0    pypi
werkzeug                  3.0.3                    pypi_0    pypi
wheel                     0.43.0           py38h06a4308_0  
xz                        5.4.6                h5eee18b_1  
zipp                      3.19.2                   pypi_0    pypi
zlib                      1.2.13               h5eee18b_1  

[kta-intel]

@hwpang, thanks for raising this issue!

It is a strange one. To my understanding, it seems to be caused when triton imports setuptools. More specifically, it is because setuptools is being imported after distutils (I'd have to follow the stack trace deeper to see where this is occuring, though). setuptools attempts to override distutils with a local copy. If distutils is imported beforehand, though, this can result in a conflict like the one we are seeing.

We'll do a deeper investigation to see to what extent we can resolve this on our end.

For now, there are two quick fixes you can do:

Option 1.
Down grade to torch==2.3.1 and torchvision==0.18.1 this will automatically install an earlier version of triton that doesn't seem to have this issue. The workspace should automatically do this, so if you uninstall torch and torchvision, you should be able to just rerun:
fx workspace create --template torch_cnn_mnist --prefix my_workspace

-or-

Option 2.
Trying setting export SETUPTOOLS_USE_DISTUTILS=stdlib if you need the latest torch/triton. This will disable using the local distutils. This worked on my end. Hopefully on yours too