Add docker image build script (#3036)

* Add docker image build script

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

* Add README.md

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

* Add missing .dockerignore

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

---------

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

.dockerignore

@@ -0,0 +1,7 @@
+*
+!src/otx
+!pyproject.toml
+!README.md
+!LICENSE
+!MANIFEST.in
+!docker/download_pretrained_weights.py

docker/Dockerfile.cuda

@@ -0,0 +1,37 @@
+FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime AS base
+
+ARG http_proxy
+ARG https_proxy
+ARG no_proxy
+ARG NON_ROOT_HOME=/home/non-root
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libsm6=2:1.2.3-1 \
+    libxext6=2:1.3.4-0ubuntu1 \
+    ffmpeg=7:4.2.7-0ubuntu0.1 \
+    libfontconfig1=2.13.1-2ubuntu3 \
+    libxrender1=1:0.9.10-1 \
+    libgl1-mesa-glx=21.2.6-0ubuntu0.1~20.04.2 \
+    && rm -rf /var/lib/apt/lists/*
+
+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}
+
+USER non-root
+
+ENV PATH=${PATH}:${NON_ROOT_HOME}/.local/bin
+
+RUN pip install --no-cache-dir src_dir/ && \
+    otx install --do-not-install-torch && \
+    rm -rf src_dir/
+
+FROM base AS cuda
+
+
+FROM base AS cuda_pretrained_ready
+COPY docker/download_pretrained_weights.py download_pretrained_weights.py
+RUN python download_pretrained_weights.py

docker/README.md

@@ -0,0 +1,24 @@
+# How to build cuda and cuda-pretrained-ready Docker images
+
+1. By executing the following commands, it will build two Docker images: `otx:${OTX_VERSION}-cuda` and `otx:${OTX_VERSION}-cuda-pretrained-ready`.
+
+   ```console
+   git clone https://github.com/openvinotoolkit/training_extensions.git
+   cd docker
+   ./build.sh
+   ```
+
+2. After that, you can check whether the images are built correctly such as
+
+   ```console
+   docker image ls | grep otx
+   ```
+
+   Example:
+
+   ```console
+   otx                                           2.0.0-cuda-pretrained-ready                    4f3b5f98f97c   3 minutes ago   14.5GB
+   otx                                           2.0.0-cuda                                     8d14caccb29a   8 minutes ago   10.4GB
+   ```
+
+`otx:${OTX_VERSION}-cuda` is a minimal Docker image where OTX is installed with CUDA supports. On the other hand, `otx:${OTX_VERSION}-cuda-pretrained-ready` includes all the model pre-trained weights that OTX provides in addition to `otx:${OTX_VERSION}-cuda`.

docker/build.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# shellcheck disable=SC2154
+
+OTX_VERSION=$(python -c 'import otx; print(otx.__version__)')
+THIS_DIR=$(dirname "$0")
+
+echo "Build OTX ${OTX_VERSION} CUDA Docker image..."
+docker build \
+    --build-arg http_proxy="${http_proxy}" \
+    --build-arg https_proxy="${https_proxy}" \
+    --build-arg no_proxy="${no_proxy}" \
+    --target cuda \
+    -t "otx:${OTX_VERSION}-cuda" \
+    -f "${THIS_DIR}/Dockerfile.cuda" "${THIS_DIR}"/..
+
+echo "Build OTX ${OTX_VERSION} CUDA pretrained-ready Docker image..."
+docker build \
+    --build-arg http_proxy="${http_proxy}" \
+    --build-arg https_proxy="${https_proxy}" \
+    --build-arg no_proxy="${no_proxy}" \
+    --target cuda_pretrained_ready \
+    -t "otx:${OTX_VERSION}-cuda-pretrained-ready" \
+    -f "${THIS_DIR}/Dockerfile.cuda" "${THIS_DIR}"/..

docker/download_pretrained_weights.py

@@ -0,0 +1,52 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+"""Helper script to download all the model pre-trained weights."""
+
+import logging
+from pathlib import Path
+
+from importlib_resources import files
+from omegaconf import OmegaConf
+from otx.core.utils.instantiators import partial_instantiate_class
+
+logging.basicConfig(
+    level=logging.INFO,
+    filename="download_pretrained_weights.log",
+    filemode="w",
+)
+
+logger = logging.getLogger()
+
+
+def download_all() -> None:
+    """Download pre-trained weights of all models."""
+    recipe_dir = Path(files("otx") / "recipe")
+
+    for config_path in recipe_dir.glob("**/*.yaml"):
+        if "_base_" in str(config_path):
+            msg = f"Skip {config_path} since it is a base config."
+            logger.warning(msg)
+            continue
+        if config_path.name == "openvino_model.yaml":
+            msg = f"Skip {config_path} since it is not a PyTorch config."
+            logger.warning(msg)
+            continue
+        if "anomaly_" in str(config_path) or "otx_dino_v2" in str(config_path) or "h_label_cls" in str(config_path):
+            msg = f"Skip {config_path} since those models show errors on instantiation."
+            logger.warning(msg)
+            continue
+
+        config = OmegaConf.load(config_path)
+        init_model = next(iter(partial_instantiate_class(config.model)))
+        try:
+            model = init_model()
+            msg = f"Downloaded pre-trained model weight of {model!s}"
+            logger.info(msg)
+        except Exception:
+            msg = f"Error on instiating {config_path}"
+            logger.exception(msg)
+
+
+if __name__ == "__main__":
+    download_all()

pyproject.toml

@@ -87,6 +87,9 @@ mmlab = [
     # This causes an error when training the mm model, so fix the version first.
     "oss2==2.17.0",
 ]
+anomaly = [
+    "anomalib==1.0.0",
+]
 
 [project.scripts]
 otx = "otx.cli:main"

src/otx/cli/install.py

@@ -59,10 +59,15 @@ def add_install_parser(subcommands_action: _ActionSubCommands) -> None:
         help="Set Logger level to INFO",
         action="store_true",
     )
+    parser.add_argument(
+        "--do-not-install-torch",
+        help="Do not install PyTorch. Choose this option if you already install PyTorch.",
+        action="store_true",
+    )
     subcommands_action.add_subcommand("install", parser, help="Install OTX requirements.")
 
 
-def otx_install(option: str | None = None, verbose: bool = False) -> int:
+def otx_install(option: str | None = None, verbose: bool = False, do_not_install_torch: bool = False) -> int:
     """Install OTX requirements.
 
     Args:
@@ -92,12 +97,15 @@ def otx_install(option: str | None = None, verbose: bool = False) -> int:
     # 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)
 
-    # Get install args for torch to install it from a specific index-url
     install_args: list[str] = []
-    torch_install_args = get_torch_install_args(torch_requirement)
 
     # Combine torch and other requirements.
-    install_args = other_requirements + torch_install_args
+    install_args = (
+        # Get install args for torch to install it from a specific index-url
+        other_requirements + get_torch_install_args(torch_requirement)
+        if not do_not_install_torch
+        else other_requirements
+    )
 
     # Parse mmX requirements if the task requires mmX packages.
     mmcv_install_args = []

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"]
+        expected_argument = ["help", "option", "verbose", "do_not_install_torch"]
         assert argument_list == expected_argument
 
     def test_install_extra(self, mocker: MockerFixture) -> None: