build: Move dependencies (#11790)

* build: Move dependencies Signed-off-by: oliver könig <[email protected]>
NVIDIA · Jan 21, 2025 · 5b4d091 · 5b4d091
1 parent 1e96841
commit 5b4d091
Show file tree

Hide file tree

Showing 3 changed files with 84 additions and 53 deletions.
diff --git a/.github/workflows/import-test.yml b/.github/workflows/import-test.yml
@@ -5,6 +5,10 @@ on:
     paths:
       - "**"
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
 # Check https://hub.docker.com/r/pytorch/pytorch/tags for latest tags
 jobs:
   test-imports:
@@ -42,7 +46,7 @@ jobs:
         pip install -r requirements/requirements_test.txt
         
         # Install NeMo Domain
-        pip install "${{ steps.nemo-wheel.outputs.DIST_FILE }}[test,${{ matrix.collection }}]"
+        pip install --no-cache-dir "${{ steps.nemo-wheel.outputs.DIST_FILE }}[test,${{ matrix.collection }}]"
     
     - name: Run ${{ matrix.collection }} checks
       run: |

diff --git a/Dockerfile.ci b/Dockerfile.ci
@@ -33,52 +33,19 @@ EOF
 
 WORKDIR /workspace
 
-# Install Mamba Dependancy
-ARG CAUSAL_CONV_TAG=v1.2.2.post1 
-ARG MAMBA_TAG=v2.2.0
-
-RUN <<"EOF" bash -ex
-# Mamba dependancy installation
-MAMBA_FORCE_BUILD=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE pip3 install --no-cache-dir -v git+https://github.com/Dao-AILab/causal-conv1d.git@${CAUSAL_CONV_TAG} git+https://github.com/state-spaces/mamba.git@${MAMBA_TAG}
-EOF
-
-RUN pip install hatchling   # needed to install nemo-run
-ARG NEMO_RUN_TAG=34259bd3e752fef94045a9a019e4aaf62bd11ce2
-RUN pip install nemo_run@git+https://github.com/NVIDIA/NeMo-Run.git@${NEMO_RUN_TAG}
-
-# Install NeMo requirements
-ARG TE_TAG=7d576ed25266a17a7b651f2c12e8498f67e0baea
-ARG MODELOPT_VERSION=0.21.0
-ARG APEX_TAG=810ffae374a2b9cb4b5c5e28eaeca7d7998fca0c
 RUN \
   --mount=type=bind,source=requirements,target=requirements \
   --mount=type=bind,source=tools,target=tools \
+  --mount=type=bind,source=reinstall.sh,target=reinstall.sh \
   --mount=type=bind,source=setup.py,target=setup.py \
   --mount=type=bind,source=nemo/package_info.py,target=nemo/package_info.py \
   --mount=type=bind,source=nemo/__init__.py,target=nemo/__init__.py <<"EOF" bash -ex
-pip install --no-cache-dir --no-build-isolation --extra-index-url https://pypi.nvidia.com \
-"transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@${TE_TAG}" \
-"nvidia-modelopt[torch]~=${MODELOPT_VERSION}" \
-"apex @ git+https://github.com/NVIDIA/apex.git@${APEX_TAG}" \
-"unstructured==0.14.9" \
-"llama-index==0.10.43" \
-"onnxscript @ git+https://github.com/microsoft/onnxscript" \
--r tools/ctc_segmentation/requirements.txt \
-".[all]"
-EOF
-
-ARG MCORE_TAG=4dc8977167d71f86bdec47a60a98e85c4cfa0031
-RUN <<"EOF" bash -ex
-# Megatron-LM installation
-git clone https://github.com/NVIDIA/Megatron-LM.git
-pushd Megatron-LM
-git checkout ${MCORE_TAG} 
-pip install -e .
-export PYTHONPATH="${PYTHONPATH}:/workspace/Megatron-LM"
 
-# Install nvidia-resiliency-ext
-pip install --no-cache-dir "git+https://github.com/NVIDIA/nvidia-resiliency-ext.git@97aad77609d2e25ed38ac5c99f0c13f93c48464e"
+  HEAVY_DEPS=TRUE bash reinstall.sh "dev"
 
+  # Megatron Core installation
+  cp -r /opt/Megatron-LM /workspace
+  export PYTHONPATH="${PYTHONPATH}:/workspace/Megatron-LM"
 EOF
 
 # Copy over NeMo code
@@ -90,11 +57,4 @@ pip install --no-cache-dir --no-build-isolation ".[all]"
 chmod 777 -R /workspace
 EOF
 
-# Install vLLM in virtualenv
-RUN pip install --no-cache-dir --no-build-isolation virtualenv && \
-  virtualenv /opt/venv && \
-  /opt/venv/bin/pip install --no-cache-dir --no-build-isolation \
-      -r /workspace/requirements/requirements_vllm.txt \
-      -r /workspace/requirements/requirements_infer.txt
-
 ENV PYTHONPATH="${PYTHONPATH}:/workspace/Megatron-LM"
diff --git a/reinstall.sh b/reinstall.sh
@@ -1,7 +1,8 @@
 #!/usr/bin/env bash
-set -e
+set -ex
 
 INSTALL_OPTION=${1:-"dev"}
+HEAVY_DEPS=${HEAVY_DEPS:-false}
 
 PIP=pip
 
@@ -16,25 +17,91 @@ ${PIP} uninstall -y nemo_asr
 ${PIP} uninstall -y nemo_nlp
 ${PIP} uninstall -y nemo_tts
 
+export MAMBA_FORCE_BUILD=TRUE
+export CAUSAL_CONV1D_FORCE_BUILD=TRUE
+export TE_TAG=7d576ed25266a17a7b651f2c12e8498f67e0baea
+export NEMO_RUN_TAG=34259bd3e752fef94045a9a019e4aaf62bd11ce2
+export APEX_TAG=810ffae374a2b9cb4b5c5e28eaeca7d7998fca0c
+export CAUSAL_CONV_TAG=v1.2.2.post1
+export MAMBA_TAG=v2.2.0
+export MCORE_TAG=4dc8977167d71f86bdec47a60a98e85c4cfa0031
+export NV_RESILIENCY_EXT_TAG=97aad77609d2e25ed38ac5c99f0c13f93c48464e
+
 if [ -n "${NVIDIA_PYTORCH_VERSION}" ]; then
-  echo 'Installing NeMo in NVIDIA PyTorch container:' "${NVIDIA_PYTORCH_VERSION}" 'so will not install numba'
+  echo "Installing NeMo in NVIDIA PyTorch container: ${NVIDIA_PYTORCH_VERSION}"
+
+  echo "Will not install numba"
+  ${PIP} install --no-build-isolation "apex @ git+https://github.com/NVIDIA/apex.git@${APEX_TAG}"
+
 else
   if [ -n "${CONDA_PREFIX}" ]; then
     NUMBA_VERSION=0.57.1
     echo 'Installing numba=='${NUMBA_VERSION}
     conda install -y -c conda-forge numba==${NUMBA_VERSION}
   fi
+
+  ${PIP} install torch
+  ${PIP} install "apex @ git+https://github.com/NVIDIA/apex.git@${APEX_TAG}"
+
 fi
 
+DEPS=(
+  "nvidia-modelopt[torch]~=0.21.0; sys_platform == 'linux'"
+  "nemo_run@git+https://github.com/NVIDIA/NeMo-Run.git@${NEMO_RUN_TAG}"
+  "git+https://github.com/NVIDIA/Megatron-LM.git@${MCORE_TAG}"
+  "git+https://github.com/NVIDIA/nvidia-resiliency-ext.git@${NV_RESILIENCY_EXT_TAG}"
+  "onnxscript @ git+https://github.com/microsoft/onnxscript"
+)
+
+if [[ "$HEAVY_DEPS" == "TRUE" ]]; then
+  ${PIP} install --no-cache-dir virtualenv &&
+    virtualenv /opt/venv &&
+    /opt/venv/bin/pip install --no-cache-dir --no-build-isolation \
+      -r /workspace/requirements/requirements_vllm.txt \
+      -r /workspace/requirements/requirements_infer.txt
+
+  DEPS+=(
+    "llama-index==0.10.43"
+    "unstructured==0.14.9"
+    "git+https://github.com/Dao-AILab/causal-conv1d.git@${CAUSAL_CONV_TAG}"
+    "git+https://github.com/state-spaces/mamba.git@${MAMBA_TAG}"
+  )
+
+  pip install --no-cache-dir -r tools/ctc_segmentation/requirements.txt
+
+  CURR=$(pwd)
+  cd /opt
+  git clone https://github.com/NVIDIA/Megatron-LM.git &&
+    pushd Megatron-LM &&
+    git checkout ${MCORE_TAG} &&
+    pip install -e . &&
+    popd
+
+  git clone https://github.com/NVIDIA/TransformerEngine.git &&
+    pushd TransformerEngine &&
+    git checkout ${TE_TAG} &&
+    git submodule update --init --recursive -q &&
+    pip install -e . &&
+    popd
+
+  cd "$CURR"
+
+fi
+
+echo 'Installing dependencies of nemo'
+${PIP} install --no-cache-dir --extra-index-url https://pypi.nvidia.com "${DEPS[@]}"
+
 echo 'Installing nemo'
 if [[ "$INSTALL_OPTION" == "dev" ]]; then
-    ${PIP} install --editable ".[all]"
+  ${PIP} install --editable ".[all]"
+
 else
-    rm -rf dist/
-    ${PIP} install build pytest-runner
-    python -m build --no-isolation --wheel
-    DIST_FILE=$(find ./dist -name "*.whl" | head -n 1)
+  rm -rf dist/ &&
+    ${PIP} install build pytest-runner &&
+    python -m build --no-isolation --wheel &&
+    DIST_FILE=$(find ./dist -name "*.whl" | head -n 1) &&
     ${PIP} install "${DIST_FILE}[all]"
+
 fi
 
 echo 'All done!'