[BUG]: ModuleNotFoundError: No module named 'pytest' and 'decorator' in hpcaitech/colossalai:0.3.0 docker image

[lshmouse]

🐛 Describe the bug

When getting started with hpcaitech/colossalai:0.3.0 docker image, the following errors are encountered.

docker run -ti --gpus all --rm --ipc=host hpcaitech/colossalai:0.3.0

# colossalai check -i
Traceback (most recent call last):
  File "/opt/conda/bin/colossalai", line 5, in <module>
    from colossalai.cli import cli
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/__init__.py", line 1, in <module>
    from .cli import cli
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/cli.py", line 3, in <module>
    from .benchmark import benchmark
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/benchmark/__init__.py", line 5, in <module>
    from .benchmark import run_benchmark
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/benchmark/benchmark.py", line 13, in <module>
    from colossalai.testing import free_port
  File "/opt/conda/lib/python3.9/site-packages/colossalai/testing/__init__.py", line 9, in <module>
    from .pytest_wrapper import run_on_environment_flag
  File "/opt/conda/lib/python3.9/site-packages/colossalai/testing/pytest_wrapper.py", line 7, in <module>
    import pytest
ModuleNotFoundError: No module named 'pytest'

# colossalai check -i
Traceback (most recent call last):
  File "/opt/conda/lib/python3.9/site-packages/fabric/connection.py", line 7, in <module>
    from invoke.vendor.decorator import decorator
ModuleNotFoundError: No module named 'invoke.vendor.decorator'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/conda/bin/colossalai", line 5, in <module>
    from colossalai.cli import cli
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/__init__.py", line 1, in <module>
    from .cli import cli
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/cli.py", line 5, in <module>
    from .launcher import run
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/launcher/__init__.py", line 5, in <module>
    from .run import launch_multi_processes
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/launcher/run.py", line 12, in <module>
    from .multinode_runner import MultiNodeRunner
  File "/opt/conda/lib/python3.9/site-packages/colossalai/cli/launcher/multinode_runner.py", line 5, in <module>
    import fabric
  File "/opt/conda/lib/python3.9/site-packages/fabric/__init__.py", line 3, in <module>
    from .connection import Config, Connection
  File "/opt/conda/lib/python3.9/site-packages/fabric/connection.py", line 9, in <module>
    from decorator import decorator
ModuleNotFoundError: No module named 'decorator'

After installing these modules manually, the error are fixed. It's better to include these python modules in docker image.

# pip3 install pytest decorator
Collecting pytest
  Downloading pytest-7.4.3-py3-none-any.whl (325 kB)
     |████████████████████████████████| 325 kB 562 kB/s 
Collecting decorator
  Downloading decorator-5.1.1-py3-none-any.whl (9.1 kB)
Collecting exceptiongroup>=1.0.0rc8
  Downloading exceptiongroup-1.1.3-py3-none-any.whl (14 kB)
Requirement already satisfied: packaging in /opt/conda/lib/python3.9/site-packages (from pytest) (23.0)
Collecting iniconfig
  Downloading iniconfig-2.0.0-py3-none-any.whl (5.9 kB)
Requirement already satisfied: pluggy<2.0,>=0.12 in /opt/conda/lib/python3.9/site-packages (from pytest) (1.0.0)
Collecting tomli>=1.0.0
  Downloading tomli-2.0.1-py3-none-any.whl (12 kB)
Installing collected packages: tomli, iniconfig, exceptiongroup, pytest, decorator
Successfully installed decorator-5.1.1 exceptiongroup-1.1.3 iniconfig-2.0.0 pytest-7.4.3 tomli-2.0.1
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
# colossalai check -i 
#### Installation Report ####

Environment

No response