# Copyright 2021 MosaicML. All Rights Reserved.

import os
import sys
import textwrap

import setuptools
from setuptools import setup


def package_files(directory: str):
    # from https://stackoverflow.com/a/36693250
    paths = []
    for (path, _, filenames) in os.walk(directory):
        for filename in filenames:
            paths.append(os.path.join('..', path, filename))
    return paths


with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

install_requires = [
    "pyyaml>=5.4.1",
    "tqdm>=4.62.3",
    "torchmetrics>=0.6.0",
    "torch_optimizer==0.1.0",
    "torchvision>=0.9.0",
    "torch>=1.9",
    "yahp>=0.0.14",
    "numpy==1.21.5",
]
extra_deps = {}

extra_deps['base'] = []

extra_deps['dev'] = [
    "custom_inherit==2.3.2",
    'junitparser>=2.1.1',
    'coverage[toml]>=6.1.1',
    'fasteners>=0.16.3',  # run_directory_uploader tests require fasteners
    'pytest>=6.2.0',
    'yapf>=0.13.0',
    'isort>=5.9.3',
    'ipython>=7.29.0',
    'ipykernel>=6.5.0',
    'jupyter>=1.0.0',
    'yamllint>=1.26.2',
    'pytest-timeout>=1.4.2',
    'recommonmark>=0.7.1',
    'sphinx>=4.2.0',
    'sphinx_copybutton>=0.4.0',
    'sphinx_markdown_tables>=0.0.15',
    'sphinx-argparse>=0.3.1',
    'sphinxcontrib.katex>=0.8.6',
    'sphinxext.opengraph>=0.4.2',
    'sphinxemoji>=0.2.0',
    'sphinx_rtd_theme>=1.0.0',
    'testbook>=0.4.2',
    'myst-parser>=0.15.2',
]
extra_deps['logging'] = ['wandb>=0.12.2', 'apache-libcloud>=3.4.1']

extra_deps['nlp'] = [
    'transformers>=4.11.3',
    'datasets>=1.14.0',
]

extra_deps['unet'] = [
    'monai>=0.7.0',
    'scikit-learn>=1.0.1',
]

extra_deps['deepspeed'] = [
    'deepspeed>=0.5.5',
]

extra_deps['all'] = set(dep for deps in extra_deps.values() for dep in deps)

setup(
    name="mosaicml",
    version="0.3.1",
    author="MosaicML",
    author_email="team@mosaicml.com",
    description="composing methods for ML training efficiency",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/mosaicml/composer",
    include_package_data=True,
    package_data={
        "composer": ['py.typed'],
        "": package_files('composer/yamls'),
    },
    packages=setuptools.find_packages(exclude=["tests*"]),
    classifiers=[
        "Programming Language :: Python :: 3",
    ],
    install_requires=install_requires,
    entry_points={
        'console_scripts': ['composer = composer.cli.launcher:main',],
    },
    extras_require=extra_deps,
    dependency_links=['https://developer.download.nvidia.com/compute/redist'],
    python_requires='>=3.7',
    ext_package="composer",
)

# only visible if user installs with verbose -v flag
# Printing to stdout as not to interfere with setup.py CLI flags (e.g. --version)
print("*" * 20, file=sys.stderr)
print(textwrap.dedent("""NOTE: For best performance, we recommend installing Pillow-SIMD
    for accelerated image processing operations. To install:
    \t pip uninstall pillow && pip install pillow-simd"""),
      file=sys.stderr)
print("*" * 20, file=sys.stderr)

setup.py
# Copyright 2021 MosaicML. All Rights Reserved.

import os
import site
import sys
import textwrap

import setuptools
from setuptools import setup

# From https://github.com/pypa/pip/issues/7953#issuecomment-645133255
site.ENABLE_USER_SITE = "--user" in sys.argv[1:]


def package_files(directory: str):
    # from https://stackoverflow.com/a/36693250
    paths = []
    for (path, _, filenames) in os.walk(directory):
        for filename in filenames:
            paths.append(os.path.join('..', path, filename))
    return paths


with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

install_requires = [
    "pyyaml>=5.4.1",
    "tqdm>=4.62.3",
    "torchmetrics>=0.6.0",
    "torch_optimizer==0.1.0",
    "torchvision>=0.9.0",
    "torch>=1.9",
    "yahp>=0.0.14",
    "numpy==1.21.5",
]
extra_deps = {}

extra_deps['base'] = []

extra_deps['dev'] = [
    "custom_inherit==2.3.2",
    'junitparser>=2.1.1',
    'coverage[toml]>=6.1.1',
    'fasteners>=0.16.3',  # run_directory_uploader tests require fasteners
    'pytest>=6.2.0',
    'yapf>=0.13.0',
    'isort>=5.9.3',
    'ipython>=7.29.0',
    'ipykernel>=6.5.0',
    'jupyter>=1.0.0',
    'yamllint>=1.26.2',
    'pytest-timeout>=1.4.2',
    'recommonmark>=0.7.1',
    'sphinx>=4.2.0',
    'sphinx_copybutton>=0.4.0',
    'sphinx_markdown_tables>=0.0.15',
    'sphinx-argparse>=0.3.1',
    'sphinxcontrib.katex>=0.8.6',
    'sphinxext.opengraph>=0.4.2',
    'sphinxemoji>=0.2.0',
    'sphinx_rtd_theme>=1.0.0',
    'testbook>=0.4.2',
    'myst-parser>=0.15.2',
]
extra_deps['logging'] = ['wandb>=0.12.2', 'apache-libcloud>=3.4.1']

extra_deps['nlp'] = [
    'transformers>=4.11.3',
    'datasets>=1.14.0',
]

extra_deps['unet'] = [
    'monai>=0.7.0',
    'scikit-learn>=1.0.1',
]

extra_deps['deepspeed'] = [
    'deepspeed>=0.5.5',
]

extra_deps['all'] = set(dep for deps in extra_deps.values() for dep in deps)

setup(
    name="mosaicml",
    version="0.3.1",
    author="MosaicML",
    author_email="team@mosaicml.com",
    description="composing methods for ML training efficiency",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/mosaicml/composer",
    include_package_data=True,
    package_data={
        "composer": ['py.typed'],
        "": package_files('composer/yamls'),
    },
    packages=setuptools.find_packages(exclude=["tests*"]),
    classifiers=[
        "Programming Language :: Python :: 3",
    ],
    install_requires=install_requires,
    entry_points={
        'console_scripts': ['composer = composer.cli.launcher:main',],
    },
    extras_require=extra_deps,
    dependency_links=['https://developer.download.nvidia.com/compute/redist'],
    python_requires='>=3.7',
    ext_package="composer",
)

# only visible if user installs with verbose -v flag
# Printing to stdout as not to interfere with setup.py CLI flags (e.g. --version)
print("*" * 20, file=sys.stderr)
print(textwrap.dedent("""NOTE: For best performance, we recommend installing Pillow-SIMD
    for accelerated image processing operations. To install:
    \t pip uninstall pillow && pip install pillow-simd"""),
      file=sys.stderr)
print("*" * 20, file=sys.stderr)

setup.py
diff --git a/docker/Makefile b/docker/Makefile
index ecbc390280..f37ff50afc 100644
--- a/docker/Makefile
+++ b/docker/Makefile
@@ -32,11 +32,11 @@ COMPOSER_EXTRA_DEPS ?= base
 # Resulting docker tag for the composer dockerimage
 COMPOSER_TAG ?= mosaicml/composer:$(COMPOSER_EXTRA_DEPS)
 
-.PHONY: build
+.PHONY: build pytorch composer
 
 build: composer
 
-base:
+pytorch:
 	cd pytorch && docker build -t $(PYTORCH_TAG) \
 		--build-arg BASE_IMAGE=$(BASE_IMAGE) \
 		--build-arg CUDA_VERSION_TAG=$(CUDA_VERSION_TAG) \
@@ -45,7 +45,7 @@ base:
 		--build-arg TORCHVISION_VERSION=$(TORCHVISION_VERSION) \
 		.
 
-composer: base
+composer: pytorch
 	cd .. && docker build -t $(COMPOSER_TAG) \
 		--build-arg BASE_IMAGE=$(PYTORCH_TAG) \
 		--build-arg COMPOSER_EXTRA_DEPS=$(COMPOSER_EXTRA_DEPS) \
diff --git a/docker/pytorch/Dockerfile b/docker/pytorch/Dockerfile
index de1fad13b6..66f35b2459 100644
--- a/docker/pytorch/Dockerfile
+++ b/docker/pytorch/Dockerfile
@@ -3,12 +3,18 @@ ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 ARG DEBIAN_FRONTEND=noninteractive
 
+# remove a bad symlink from the base composer image
+# If this file is present after the first command, kaniko
+# won't be able to build the docker image.
+RUN rm -f /usr/local/cuda-11.3/cuda-11.3  
+
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         libgomp1 \
         curl \
         sudo \
         build-essential \
+        git \
         software-properties-common \
         # For PILLOW:
         zlib1g-dev \
@@ -21,25 +27,27 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# # Upgrade NCCL to fix NVLink errors
-# # No Need to install, as the cuda 11.3 image already contains a late enough version
-# RUN DISTRO="$(lsb_release -si | tr '[:upper:]' '[:lower:]')$(lsb_release -sr | sed 's/\.//')" &&  \
-#     apt-get update && \
-#     apt-get install -y --no-install-recommends \
-#         dirmngr \
-#         gpg-agent && \
-#     apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/${DISTRO}/x86_64/7fa2af80.pub && \
-#     add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/${DISTRO}/x86_64/ /" && \
-#     apt-get install -y --no-install-recommends --allow-change-held-packages \
-#         libnccl2 \
-#         libnccl-dev && \
-#     apt-get autoclean && \
-#     apt-get clean && \
-#     rm -rf /var/lib/apt/lists/*
-
-# # Use system installed NCCL per update above, point to library
-# ENV USE_SYSTEM_NCCL=1
-# ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libnccl.so.2.9.6
+##################################################
+# Change the NCCL version to fix NVLink errors
+# This is required to train on Nvidia A100s in GCP
+##################################################
+RUN DISTRO="$(lsb_release -si | tr '[:upper:]' '[:lower:]')$(lsb_release -sr | sed 's/\.//')" &&  \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+        dirmngr \
+        gpg-agent && \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/${DISTRO}/x86_64/7fa2af80.pub && \
+    add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/${DISTRO}/x86_64/ /" && \
+    apt-get install -y --no-install-recommends --allow-change-held-packages --allow-downgrades \
+        libnccl2=2.9.6-1+cuda11.0 \
+        libnccl-dev=2.9.6-1+cuda11.0 && \
+    apt-get autoclean && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Use system installed NCCL per update above, point to library
+ENV USE_SYSTEM_NCCL=1
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libnccl.so.2.9.6
 
 ##############################
 # Install NodeJS (for Pyright)
@@ -64,14 +72,12 @@ RUN add-apt-repository ppa:deadsnakes/ppa && \
     python${PYTHON_VERSION}-dev \
     python${PYTHON_VERSION}-distutils \
     python${PYTHON_VERSION}-venv && \
-    update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
-    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
     apt-get autoclean && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-RUN curl -fsSL https://bootstrap.pypa.io/get-pip.py | python - && \
-    pip install --no-cache-dir --upgrade pip
+RUN curl -fsSL https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} - && \
+    pip${PYTHON_VERSION} install --no-cache-dir --upgrade pip setuptools
 
 #####################
 # Install pillow-simd
@@ -83,8 +89,8 @@ ARG PILLOW_SIMD_VERSION=7.0.0.post3
 # so when pillow_simd is installed, other packages won't later override it
 COPY pillow_stub /tmp/pillow_stub
 
-RUN pip install --no-cache-dir --upgrade /tmp/pillow_stub && \
-    pip install --no-cache-dir --upgrade pillow_simd==${PILLOW_SIMD_VERSION}
+RUN pip${PYTHON_VERSION} install --no-cache-dir --upgrade /tmp/pillow_stub && \
+    pip${PYTHON_VERSION} install --no-cache-dir --upgrade pillow_simd==${PILLOW_SIMD_VERSION}
 
 #################
 # Install Pytorch
@@ -93,7 +99,7 @@ ARG PYTORCH_VERSION
 ARG TORCHVISION_VERSION
 ARG CUDA_VERSION_TAG
 
-RUN pip install --no-cache-dir --find-links https://download.pytorch.org/whl/torch_stable.html \
+RUN pip${PYTHON_VERSION} install --no-cache-dir --find-links https://download.pytorch.org/whl/torch_stable.html \
         torch==${PYTORCH_VERSION}+${CUDA_VERSION_TAG} \
         torchvision==${TORCHVISION_VERSION}+${CUDA_VERSION_TAG}
 
@@ -102,11 +108,41 @@ RUN pip install --no-cache-dir --find-links https://download.pytorch.org/whl/tor
 #########
 RUN rm -rf /tmp/*
 
+
+#######################
+# Set the shell to bash
+#######################
+SHELL ["/bin/bash", "-c"]
+
+
+################################
+# Use the correct python version
+################################
+
+# Set the default python by creating our own folder and hacking the path
+# We don't want to use upgrade-alternatives as that will break system packages
+
+ARG COMPOSER_PYTHON_BIN=/composer-python
+
+RUN mkdir -p ${COMPOSER_PYTHON_BIN} && \
+    ln -s $(which python${PYTHON_VERSION}) ${COMPOSER_PYTHON_BIN}/python && \
+    ln -s $(which python${PYTHON_VERSION}) ${COMPOSER_PYTHON_BIN}/python3 && \
+    ln -s $(which python${PYTHON_VERSION}) ${COMPOSER_PYTHON_BIN}/python${PYTHON_VERSION} && \
+    ln -s $(which pip${PYTHON_VERSION}) ${COMPOSER_PYTHON_BIN}/pip && \
+    ln -s $(which pip${PYTHON_VERSION}) ${COMPOSER_PYTHON_BIN}/pip3 && \
+    ln -s $(which pip${PYTHON_VERSION}) ${COMPOSER_PYTHON_BIN}/pip${PYTHON_VERSION} && \
+    # Include this folder, and the local bin folder, on the path 
+    echo "export PATH=~/.local/bin:$COMPOSER_PYTHON_BIN:$PATH" >> /etc/profile && \
+    echo "export PATH=~/.local/bin:$COMPOSER_PYTHON_BIN:$PATH" >> /etc/bash.bashrc && \
+    echo "export PATH=~/.local/bin:$COMPOSER_PYTHON_BIN:$PATH" >> /etc/zshenv
+
+# Ensure that non-interactive shells load /etc/profile
+ENV BASH_ENV=/etc/profile
+
 #########################
 # Configure non-root user
 #########################
-RUN echo "export PATH=\$PATH:~/.local/bin" >> /etc/skel/.bashrc && \
-    useradd -rm -d /home/mosaicml -s /bin/bash -u 1000 -U -s /bin/bash mosaicml && \
+RUN useradd -rm -d /home/mosaicml -s /bin/bash -u 1000 -U -s /bin/bash mosaicml && \
     usermod -a -G sudo mosaicml && \
     echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
 
diff --git a/setup.py b/setup.py
index 9015603417..6d521f2c8f 100755
--- a/setup.py
+++ b/setup.py
@@ -1,12 +1,16 @@
 # Copyright 2021 MosaicML. All Rights Reserved.
 
 import os
+import site
 import sys
 import textwrap
 
 import setuptools
 from setuptools import setup
 
+# From https://github.com/pypa/pip/issues/7953#issuecomment-645133255
+site.ENABLE_USER_SITE = "--user" in sys.argv[1:]
+
 
 def package_files(directory: str):
     # from https://stackoverflow.com/a/36693250

Add venv into docker image to enable editable `pip install`
When trying to install composer with `pip install -e .` from within the docker image, we are seeing this error:
```
Traceback (most recent call last):
  File "/usr/bin/composer", line 33, in <module>
    sys.exit(load_entry_point('mosaicml', 'console_scripts', 'composer')())
  File "/usr/bin/composer", line 22, in importlib_load_entry_point
    for entry_point in distribution(dist_name).entry_points
  File "/usr/lib/python3.8/importlib/metadata.py", line 445, in distribution
    return Distribution.from_name(distribution_name)
  File "/usr/lib/python3.8/importlib/metadata.py", line 169, in from_name
    raise PackageNotFoundError(name)
importlib.metadata.PackageNotFoundError: mosaicml
```
This seems to be remedied by running the `pip install` from within a virtualenv. Can we bake a virtualenv into the docker image as a workaround?
