Revisit Docker image build script (#3536)

* Revisit

Signed-off-by: Kim, Vinnam <vinnam.kim@intel.com>

* Fix test

Signed-off-by: Kim, Vinnam <vinnam.kim@intel.com>

---------

Signed-off-by: Kim, Vinnam <vinnam.kim@intel.com>

docker/Dockerfile.cuda

@@ -1,9 +1,9 @@
-FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime@sha256:971fbeae82c0a5a7a970a264a8b8ce1c3426aa79df7111004ad2bc2640f7d89c AS base
-
 ARG http_proxy
 ARG https_proxy
 ARG no_proxy
-ARG NON_ROOT_HOME=/home/non-root
+
+
+FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime@sha256:971fbeae82c0a5a7a970a264a8b8ce1c3426aa79df7111004ad2bc2640f7d89c AS build_base
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     libsm6=2:1.2.3-1 \
@@ -14,27 +14,64 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     libgl1-mesa-glx=21.2.6-0ubuntu0.1~20.04.2 \
     && rm -rf /var/lib/apt/lists/*
 
+ARG NON_ROOT_HOME=/home/non-root
+
 RUN useradd -l -u 10001 non-root \
     && mkdir -p ${NON_ROOT_HOME}
 
-WORKDIR ${NON_ROOT_HOME}
-COPY . src_dir
 RUN chown -R non-root:non-root ${NON_ROOT_HOME}
 
+ENV PATH=${PATH}:${NON_ROOT_HOME}/.local/bin
+
+
+FROM build_base AS install_packages
+
+ARG OTX_INSTALL_PATH=/opt/otx
+
+RUN mkdir -p ${OTX_INSTALL_PATH} && chown 10001 ${OTX_INSTALL_PATH}
+
+WORKDIR /workspace
+COPY --chown=10001 . src_dir
+
 USER non-root
 
-ENV PATH=${PATH}:${NON_ROOT_HOME}/.local/bin
+RUN pip install --user --no-cache-dir --require-hashes --no-deps -r src_dir/.ci/requirements/piptools/requirements.txt && \
+    pip-compile --generate-hashes -o /tmp/requirements.txt src_dir/pyproject.toml
 
-RUN pip install --no-cache-dir --require-hashes --no-deps -r src_dir/.ci/requirements/piptools/requirements.txt && \
-    pip-compile --generate-hashes -o /tmp/requirements.txt src_dir/pyproject.toml && \
-    pip install --no-cache-dir --no-deps -e src_dir/ && \
-    pip install --no-cache-dir --require-hashes --no-deps -r /tmp/requirements.txt && \
-    otx install --do-not-install-torch && \
+RUN pip install --no-cache-dir --no-deps --target ${OTX_INSTALL_PATH} src_dir/
+
+ENV PYTHONPATH=${OTX_INSTALL_PATH}:${PYTHONPATH:+:${PYTHONPATH}}
+ENV PATH=${OTX_INSTALL_PATH}/bin:${PATH:+:${PATH}}
+
+RUN pip install --user --no-cache-dir --require-hashes --no-deps -r /tmp/requirements.txt && \
+    otx install --do-not-install-torch --user && \
     rm /tmp/requirements.txt
 
-FROM base AS cuda
 
+FROM install_packages AS download_pretrained_weights
+
+WORKDIR ${NON_ROOT_HOME}
 
-FROM base AS cuda_pretrained_ready
 COPY docker/download_pretrained_weights.py download_pretrained_weights.py
 RUN python download_pretrained_weights.py
+
+
+FROM build_base AS cuda
+
+ARG USER_SITE=/home/non-root/.local/lib/python3.10/site-packages
+ARG OTX_INSTALL_PATH=/opt/otx
+
+ENV PYTHONPATH=${OTX_INSTALL_PATH}:${PYTHONPATH:+:${PYTHONPATH}}
+ENV PATH=${OTX_INSTALL_PATH}/bin:${PATH:+:${PATH}}
+
+COPY --chown=10001 --from=install_packages ${USER_SITE} ${USER_SITE}
+COPY --chown=10001 --from=install_packages ${OTX_INSTALL_PATH} ${OTX_INSTALL_PATH}
+
+USER non-root
+WORKDIR ${NON_ROOT_HOME}
+
+
+FROM cuda AS cuda_pretrained_ready
+ARG TORCH_CACHE=/home/non-root/.cache/torch
+
+COPY --chown=10001 --from=download_pretrained_weights ${TORCH_CACHE} ${TORCH_CACHE}

src/otx/cli/install.py

@@ -61,7 +61,12 @@ def add_install_parser(subcommands_action: _ActionSubCommands) -> None:
     )
     parser.add_argument(
         "--do-not-install-torch",
-        help="Do not install PyTorch. Choose this option if you already install PyTorch.",
+        help="Do not install PyTorch. Choose this option if you already installed PyTorch.",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--user",
+        help="Install packages in the user site directory, e.g., `pip install --user ...`",
         action="store_true",
     )
 
@@ -72,12 +77,16 @@ def otx_install(
     option: str | None = None,
     verbose: bool = False,
     do_not_install_torch: bool = False,
+    user: bool = False,
 ) -> int:
     """Install OTX requirements.
 
     Args:
         option (str): Optional-dependency to install requirements for.
         verbose (bool): Set pip logger level to INFO
+        do_not_install_torch (bool): If true, skip PyTorch installation.
+        user (bool): If true, install packages in the user site directory,
+            e.g., `pip install --user ...`
 
     Raises:
         ValueError: When the task is not supported.
@@ -102,7 +111,7 @@ def otx_install(
     # This is done to parse the correct version of torch (cpu/cuda) and mmcv (mmcv/mmcv-full).
     torch_requirement, mmcv_requirements, other_requirements = parse_requirements(requirements)
 
-    install_args: list[str] = []
+    install_args: list[str] = ["--user"] if user else []
 
     # Combine torch and other requirements.
     install_args = (
@@ -113,7 +122,7 @@ def otx_install(
     )
 
     # Parse mmX requirements if the task requires mmX packages.
-    mmcv_install_args = []
+    mmcv_install_args = ["--user"] if user else []
     if mmcv_requirements:
         mmcv_install_args = get_mmcv_install_args(torch_requirement, mmcv_requirements)
         install_args += ["openmim"]
@@ -146,7 +155,9 @@ def otx_install(
 
         # TODO(harimkang): Remove this reinstalling after resolving conflict with anomalib==1.0.1
         # https://github.com/openvinotoolkit/training_extensions/actions/runs/8531851027/job/23372146228?pr=3258#step:5:2587
-        status_code = create_command("install").main(["jsonargparse==4.27.7"])
+        install_args = ["--user"] if user else []
+        install_args += ["jsonargparse==4.27.7"]
+        status_code = create_command("install").main(install_args)
         if status_code != 0:
             msg = "Cannot install jsonargparse==4.27.7"
             raise RuntimeError(msg)

tests/unit/cli/test_install.py

@@ -28,7 +28,7 @@ def test_add_install_parser(self) -> None:
         assert parser_subcommands.choices.get("install") is not None
         install_parser = parser_subcommands.choices.get("install")
         argument_list = [action.dest for action in install_parser._actions]
-        expected_argument = ["help", "option", "verbose", "do_not_install_torch"]
+        expected_argument = ["help", "option", "verbose", "do_not_install_torch", "user"]
         assert argument_list == expected_argument
 
     def test_install_extra(self, mocker: MockerFixture) -> None: