|
| 1 | +ARG BASE_IMAGE="loongforge/loongforge_kunlun:py310_torch25" |
| 2 | +FROM ${BASE_IMAGE} |
| 3 | + |
| 4 | +# ========================================================= |
| 5 | +# Build Arguments |
| 6 | +# ========================================================= |
| 7 | +ARG XPYTORCH_URL_ARG="" |
| 8 | +ARG KUNLUN_OPS_URL_ARG="" |
| 9 | +ARG COCOPOD_URL_ARG="" |
| 10 | +ARG XSPEEDGATE_URL_ARG="" |
| 11 | + |
| 12 | +# ========================================================= |
| 13 | +# Base System Configuration + APT Source Replacement |
| 14 | +# ========================================================= |
| 15 | +COPY ./LoongForge /workspace/LoongForge |
| 16 | +SHELL ["/bin/bash", "-c"] |
| 17 | +ENV PATH=/opt/loongforge_kunlun/bin:${PATH} |
| 18 | + |
| 19 | +RUN wget https://ai-data.bj.bcebos.com/open-source/xcudart.tar.gz && tar -xzvf xcudart.tar.gz -C /opt/loongforge_kunlun/ && rm -rf xcudart* |
| 20 | +ENV LD_LIBRARY_PATH=/usr/local/cuda-11.7/:/opt/loongforge_kunlun/xcudart/lib/:${LD_LIBRARY_PATH} |
| 21 | + |
| 22 | +# ======================== Install XPU Requirements ======================== |
| 23 | +RUN set -exuo pipefail && \ |
| 24 | + download_xpytorch() { \ |
| 25 | + local XPYTORCH_URL=$1; \ |
| 26 | + local XPYTORCH_FILE=$(basename "${XPYTORCH_URL}"); \ |
| 27 | + echo "download xpytorch package: ${XPYTORCH_URL}" && \ |
| 28 | + wget "${XPYTORCH_URL}" && \ |
| 29 | + bash "${XPYTORCH_FILE}" --noexec --target xpytorch_unpack && cd xpytorch_unpack && \ |
| 30 | + sed -i 's/CONDA_PREFIX/VIRTUAL_ENV/g' setup.sh && \ |
| 31 | + python -m ensurepip --upgrade && \ |
| 32 | + python -m pip install --upgrade pip && \ |
| 33 | + export VIRTUAL_ENV=/opt/loongforge_kunlun && bash setup.sh && \ |
| 34 | + cd .. && rm -rf "${XPYTORCH_FILE}" xpytorch_unpack; \ |
| 35 | + }; \ |
| 36 | + source /opt/loongforge_kunlun/bin/activate && \ |
| 37 | + if [ -n "${XPYTORCH_URL_ARG}" ]; then \ |
| 38 | + XPYTORCH_URL="${XPYTORCH_URL_ARG}"; \ |
| 39 | + echo "use specified xpytorch version: ${XPYTORCH_URL}"; \ |
| 40 | + else \ |
| 41 | + XPYTORCH_URL=$(python3 -c "import json; print(json.load(open('/workspace/LoongForge/docker/xpu_packages.json'))['torch25']['xpytorch']['url'])"); \ |
| 42 | + echo "use default xpytorch version from json: ${XPYTORCH_URL}"; \ |
| 43 | + fi && \ |
| 44 | + download_xpytorch "${XPYTORCH_URL}" |
| 45 | + |
| 46 | +# ======================== Fix xpytorch_import_hook ======================== |
| 47 | +RUN source /opt/loongforge_kunlun/bin/activate && bash /workspace/LoongForge/docker/fix_xpytorch_hook.sh |
| 48 | + |
| 49 | +# ======================== Install XPU Packages ======================== |
| 50 | +RUN source /opt/loongforge_kunlun/bin/activate && \ |
| 51 | + if [ -n "${COCOPOD_URL_ARG}" ]; then \ |
| 52 | + COCOPOD_URL="${COCOPOD_URL_ARG}"; \ |
| 53 | + else \ |
| 54 | + COCOPOD_URL=$(python3 -c "import json; print(json.load(open('/workspace/LoongForge/docker/xpu_packages.json'))['torch25']['cocopod']['url'])"); \ |
| 55 | + fi && \ |
| 56 | + if [ -n "${XSPEEDGATE_URL_ARG}" ]; then \ |
| 57 | + XSPEEDGATE_URL="${XSPEEDGATE_URL_ARG}"; \ |
| 58 | + else \ |
| 59 | + XSPEEDGATE_URL=$(python3 -c "import json; print(json.load(open('/workspace/LoongForge/docker/xpu_packages.json'))['torch25']['xspeedgate']['url'])"); \ |
| 60 | + fi && \ |
| 61 | + if [ -n "${KUNLUN_OPS_URL_ARG}" ]; then \ |
| 62 | + KUNLUN_OPS_URL="${KUNLUN_OPS_URL_ARG}"; \ |
| 63 | + else \ |
| 64 | + KUNLUN_OPS_URL=$(python3 -c "import json; print(json.load(open('/workspace/LoongForge/docker/xpu_packages.json'))['torch25']['kunlun_ops']['url'])"); \ |
| 65 | + fi && \ |
| 66 | + if [ "${COCOPOD_URL}" = "__SET_WITH_COCOPOD_URL_ARG__" ] || [ "${XSPEEDGATE_URL}" = "__SET_WITH_XSPEEDGATE_URL_ARG__" ]; then \ |
| 67 | + echo "Set COCOPOD_URL_ARG and XSPEEDGATE_URL_ARG when building the XPU image." >&2; \ |
| 68 | + exit 1; \ |
| 69 | + fi && \ |
| 70 | + pip install "${COCOPOD_URL}" "${XSPEEDGATE_URL}" "${KUNLUN_OPS_URL}" |
| 71 | + |
| 72 | +# ======================== Install Base Requirements ======================== |
| 73 | +RUN set -exuo pipefail && \ |
| 74 | + source /opt/loongforge_kunlun/bin/activate && \ |
| 75 | + # for deepspeed cpu adam |
| 76 | + DS_BUILD_CPU_ADAM=1 BUILD_UTILS=1 pip install --no-build-isolation deepspeed==0.18.5 -U && \ |
| 77 | + echo "alias ll='ls -alF'" > /etc/profile.d/alias.sh; \ |
| 78 | + echo "source /etc/profile.d/alias.sh" >> ~/.bashrc && \ |
| 79 | + apt-get update && \ |
| 80 | + apt-get install -y ffmpeg=7:6.1.1-3ubuntu5 && \ |
| 81 | + apt-get install -y libomp-dev && \ |
| 82 | + cd /workspace/LoongForge/ && \ |
| 83 | + pip install uv && \ |
| 84 | + uv pip install --no-config --no-cache -e ".[xpu]" |
| 85 | + |
| 86 | +# ========================================================= |
| 87 | +# Lerobot |
| 88 | +# ========================================================= |
| 89 | +RUN set -exuo pipefail && \ |
| 90 | + install_lerobot() { \ |
| 91 | + echo "Installing lerobot..." && \ |
| 92 | + pip install "dmuon @ git+https://github.com/X-Square-Robot/dmuon.git" && \ |
| 93 | + cd /workspace/LoongForge/ && \ |
| 94 | + uv pip install --no-config --system --break-system-packages --no-cache -e ".[vla]" && \ |
| 95 | + cd /workspace && \ |
| 96 | + if [ -d "lerobot" ]; then \ |
| 97 | + rm -rf lerobot; \ |
| 98 | + fi && \ |
| 99 | + git clone --no-checkout https://github.com/huggingface/lerobot.git && \ |
| 100 | + cd lerobot && \ |
| 101 | + git checkout tags/v0.5.0 && \ |
| 102 | + if [ -f "requirements-ubuntu.txt" ]; then \ |
| 103 | + sed -i '/evdev/d' requirements-ubuntu.txt; \ |
| 104 | + fi && \ |
| 105 | + pip install func_timeout pytest-xdist prettytable==3.11.0 huggingface-hub==0.34.0 && \ |
| 106 | + sed -i 's/^class Backtrackable\[T\]:$/from typing import Iterable, TypeVar, Generic\nT = TypeVar('\''T'\'')\nclass Backtrackable(Generic[T]):/' ./src/lerobot/datasets/utils.py && \ |
| 107 | + sed -i 's/^def deserialize_json_into_object\[T: JsonLike\](fpath: Path, obj: T) -> T:$/from typing import TypeVar\nT = TypeVar('\''T'\'', bound=JsonLike)\ndef deserialize_json_into_object(fpath: Path, obj: T) -> T:/' ./src/lerobot/utils/io_utils.py && \ |
| 108 | + sed -i '/^@dataclass$/{N; s/^@dataclass\nclass DataProcessorPipeline\[TInput, TOutput\](HubMixin):$/from typing import TypeVar, Generic\nTInput = TypeVar('\''TInput'\'')\nTOutput = TypeVar('\''TOutput'\'')\n@dataclass\nclass DataProcessorPipeline(Generic[TInput, TOutput], HubMixin):/}' ./src/lerobot/processor/pipeline.py && \ |
| 109 | + sed -i 's/^\([[:space:]]*\)from typing import TYPE_CHECKING, Literal, TypedDict, Unpack$/\1from typing import TYPE_CHECKING, Literal, TypedDict\n\1from typing_extensions import Unpack/' ./src/lerobot/policies/pi0_fast/modeling_pi0_fast.py && \ |
| 110 | + sed -i 's/^\([[:space:]]*\)from typing import TYPE_CHECKING, Literal, TypedDict, Unpack$/\1from typing import TYPE_CHECKING, Literal, TypedDict\n\1from typing_extensions import Unpack/' ./src/lerobot/policies/pi05/modeling_pi05.py && \ |
| 111 | + sed -i 's/^\([[:space:]]*\)from typing import TYPE_CHECKING, Literal, TypedDict, Unpack$/\1from typing import TYPE_CHECKING, Literal, TypedDict\n\1from typing_extensions import Unpack/' ./src/lerobot/policies/pi0/modeling_pi0.py && \ |
| 112 | + sed -i 's/^\([[:space:]]*\)from typing import TypedDict, Unpack$/\1from typing import TypedDict\n\1from typing_extensions import Unpack/' ./src/lerobot/policies/smolvla/modeling_smolvla.py && \ |
| 113 | + sed -i 's/^\([[:space:]]*\)from typing import TypedDict, TypeVar, Unpack$/\1from typing import TypedDict, TypeVar\n\1from typing_extensions import Unpack/' ./src/lerobot/policies/pretrained.py && \ |
| 114 | + sed -i 's/^\([[:space:]]*\)from typing import Any, TypedDict, Unpack$/\1from typing import Any, TypedDict\n\1from typing_extensions import Unpack/' ./src/lerobot/policies/factory.py && \ |
| 115 | + sed -i '/^[[:space:]]*type NameOrID = str | int$/{N; s/^\([[:space:]]*\)type NameOrID = str | int\n[[:space:]]*type Value = int | float$/from typing import TypeAlias, Union\n\1NameOrID: TypeAlias = Union[str, int]\n\1Value: TypeAlias = Union[int, float]/}' ./src/lerobot/motors/motors_bus.py && \ |
| 116 | + sed -i -e '/"torch[>=<]/d' \ |
| 117 | + -e 's/"torchvision[^"]*",/"torchvision==0.20.1",/' \ |
| 118 | + -e '/python-can/d' \ |
| 119 | + -e '/lerobot\[can-dep\]/d' \ |
| 120 | + -e '/lerobot\[damiao\]/d' \ |
| 121 | + -e '/unitree_g1 = \[/,/]/ { /matplotlib/d; /pin/d; /meshcat/d; /casadi/d; }' \ |
| 122 | + -e '/lerobot-setup-can/d' \ |
| 123 | + -e '/\[tool.setuptools.package-data\]/d' \ |
| 124 | + -e '/lerobot = \["envs\/\*\.json"\]/d' \ |
| 125 | + -e '/"OT_VALUE"/d' \ |
| 126 | + -e 's/requires-python = ">=3\.12"/requires-python = ">=3\.10"/g' \ |
| 127 | + -e '/huggingface-hub>=1\.0\.0,<2\.0\.0/d' \ |
| 128 | + pyproject.toml && \ |
| 129 | + pip install --no-build-isolation -e ".[pi,peft-dep,scipy-dep]" && \ |
| 130 | + sed -i 's|tokenizer_name="google/paligemma-3b-pt-224",|tokenizer_name=config.tokenizer_name,|g' ./src/lerobot/policies/pi05/processor_pi05.py && \ |
| 131 | + pip install "numpy==1.26.4" && \ |
| 132 | + pip install "scipy==1.15.3" && \ |
| 133 | + echo "Lerobot installation completed"; \ |
| 134 | + }; \ |
| 135 | + source /opt/loongforge_kunlun/bin/activate && \ |
| 136 | + install_lerobot; |
| 137 | + |
| 138 | + |
| 139 | +# ========================================================= |
| 140 | +# Megatron Symlink |
| 141 | +# Scripts reference /workspace/Loong-Megatron; the actual location is |
| 142 | +# /workspace/LoongForge/third_party/Loong-Megatron (git submodule). |
| 143 | +# ========================================================= |
| 144 | +RUN ln -sf /workspace/LoongForge/third_party/Loong-Megatron /workspace/Loong-Megatron |
| 145 | + |
| 146 | +# ========================================================= |
| 147 | +# Jupyter Configuration |
| 148 | +# ========================================================= |
| 149 | +RUN mkdir -p /root/.jupyter && \ |
| 150 | + touch /root/.jupyter/enterpoint.sh && \ |
| 151 | + chmod +x /root/.jupyter/enterpoint.sh |
| 152 | + |
| 153 | +# ========================================================= |
| 154 | +# Clean Up Unused Files (Consistent with install.sh) |
| 155 | +# ========================================================= |
| 156 | +RUN rm -rf /tmp/* ~/.bash_history && \ |
| 157 | + rm -rf /workspace/LoongForge/.git && \ |
| 158 | + rm -rf /workspace/LoongForge/third_party/Loong-Megatron/.git && \ |
| 159 | + rm -rf /workspace/LoongForge/build.sh && \ |
| 160 | + rm -rf /workspace/LoongForge/ci.yml && \ |
| 161 | + rm -rf /workspace/LoongForge/docker/ci && \ |
| 162 | + rm -rf /workspace/*.tar.gz && \ |
| 163 | + rm -f /workspace/README.md /workspace/license.txt && \ |
| 164 | + rm -rf /workspace/docker-examples /workspace/tutorials /workspace/lerobot |
| 165 | + |
| 166 | +WORKDIR /workspace |
0 commit comments