Initial commit for optimizer image

Signed-off-by: Thara Palanivel <[email protected]>

Author: tharapalanivel

build/Dockerfile

@@ -0,0 +1,186 @@
+# Copyright The FMS Model Optimizer Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+## Global Args #################################################################
+ARG BASE_UBI_IMAGE_TAG=latest
+ARG USER=optimizer
+ARG USER_UID=1000
+ARG PYTHON_VERSION=3.11
+ARG WHEEL_VERSION=""
+## Enable Aimstack if requested via ENABLE_AIM set to "true"
+ARG ENABLE_AIM=false
+
+## Base Layer ##################################################################
+FROM registry.access.redhat.com/ubi9/ubi:${BASE_UBI_IMAGE_TAG} AS base
+
+ARG PYTHON_VERSION
+ARG USER
+ARG USER_UID
+
+# Note this works for 3.9, 3.11, 3.12
+RUN dnf remove -y --disableplugin=subscription-manager \
+        subscription-manager \
+    && dnf install -y python${PYTHON_VERSION} procps g++ python${PYTHON_VERSION}-devel \
+    && ln -s /usr/bin/python${PYTHON_VERSION} /bin/python \
+    && python -m ensurepip --upgrade \
+    && python -m pip install --upgrade pip \
+    && python -m pip install --upgrade setuptools \
+    && dnf update -y \
+    && dnf clean all
+
+ENV LANG=C.UTF-8 \
+    LC_ALL=C.UTF-8
+
+RUN useradd -u $USER_UID ${USER} -m -g 0 --system && \
+    chmod g+rx /home/${USER}
+
+## Used as base of the Release stage to removed unrelated the packages and CVEs
+FROM base AS release-base
+
+# Removes the python3.9 code to eliminate possible CVEs.  Also removes dnf
+RUN rpm -e $(dnf repoquery python3-* -q --installed) dnf python3 yum crypto-policies-scripts
+
+
+## CUDA Base ###################################################################
+FROM base AS cuda-base
+
+# Ref: https://docs.nvidia.com/cuda/archive/12.1.0/cuda-toolkit-release-notes/
+ENV CUDA_VERSION=12.1.0 \
+    NV_CUDA_LIB_VERSION=12.1.0-1 \
+    NVIDIA_VISIBLE_DEVICES=all \
+    NVIDIA_DRIVER_CAPABILITIES=compute,utility \
+    NV_CUDA_CUDART_VERSION=12.1.55-1 \
+    NV_CUDA_COMPAT_VERSION=530.30.02-1
+
+RUN dnf config-manager \
+       --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel9/x86_64/cuda-rhel9.repo \
+    && dnf install -y \
+        cuda-cudart-12-1-${NV_CUDA_CUDART_VERSION} \
+        cuda-compat-12-1-${NV_CUDA_COMPAT_VERSION} \
+    && echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf \
+    && echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf \
+    && dnf clean all
+
+ENV CUDA_HOME="/usr/local/cuda" \
+    PATH="/usr/local/nvidia/bin:${CUDA_HOME}/bin:${PATH}" \
+    LD_LIBRARY_PATH="/usr/local/nvidia/lib:/usr/local/nvidia/lib64:$CUDA_HOME/lib64:$CUDA_HOME/extras/CUPTI/lib64:${LD_LIBRARY_PATH}"
+
+## CUDA Development ############################################################
+FROM cuda-base AS cuda-devel
+
+# Ref: https://developer.nvidia.com/nccl/nccl-legacy-downloads
+ENV NV_CUDA_CUDART_DEV_VERSION=12.1.55-1 \
+    NV_NVML_DEV_VERSION=12.1.55-1 \
+    NV_LIBCUBLAS_DEV_VERSION=12.1.0.26-1 \
+    NV_LIBNPP_DEV_VERSION=12.0.2.50-1 \
+    NV_LIBNCCL_DEV_PACKAGE_VERSION=2.18.3-1+cuda12.1
+
+RUN dnf config-manager \
+       --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel9/x86_64/cuda-rhel9.repo \
+    && dnf install -y \
+        cuda-command-line-tools-12-1-${NV_CUDA_LIB_VERSION} \
+        cuda-libraries-devel-12-1-${NV_CUDA_LIB_VERSION} \
+        cuda-minimal-build-12-1-${NV_CUDA_LIB_VERSION} \
+        cuda-cudart-devel-12-1-${NV_CUDA_CUDART_DEV_VERSION} \
+        cuda-nvml-devel-12-1-${NV_NVML_DEV_VERSION} \
+        libcublas-devel-12-1-${NV_LIBCUBLAS_DEV_VERSION} \
+        libnpp-devel-12-1-${NV_LIBNPP_DEV_VERSION} \
+        libnccl-devel-${NV_LIBNCCL_DEV_PACKAGE_VERSION} \
+    && dnf clean all
+
+ENV LIBRARY_PATH="$CUDA_HOME/lib64/stubs"
+
+## Full set of python installations for final image release #########################
+FROM cuda-devel AS python-installations
+
+ARG WHEEL_VERSION
+ARG USER
+ARG USER_UID
+ARG ENABLE_AIM
+
+# consistent arch support anywhere we compile CUDA code
+ENV TORCH_CUDA_ARCH_LIST="8.0;8.6+PTX;8.9"
+
+RUN dnf install -y git && \
+    # perl-Net-SSLeay.x86_64 and server_key.pem are installed with git as dependencies
+    # Twistlock detects it as H severity: Private keys stored in image
+    rm -f /usr/share/doc/perl-Net-SSLeay/examples/server_key.pem && \
+    dnf clean all
+USER ${USER}
+WORKDIR /tmp
+RUN --mount=type=cache,target=/home/${USER}/.cache/pip,uid=${USER_UID} \
+    python -m pip install --user build
+COPY --chown=${USER}:root fms_mo fms_mo
+COPY .git .git
+COPY pyproject.toml pyproject.toml
+
+# Build a wheel if PyPi wheel_version is empty else download the wheel from PyPi
+RUN if [[ -z "${WHEEL_VERSION}" ]]; \
+    then python -m build --wheel --outdir /tmp; \
+    else pip download fms-model-optimizer==${WHEEL_VERSION} --dest /tmp --only-binary=:all: --no-deps; \
+    fi && \
+    ls /tmp/*.whl >/tmp/bdist_name
+
+# Install from the wheel
+RUN --mount=type=cache,target=/home/${USER}/.cache/pip,uid=${USER_UID} \
+    python -m pip install --user wheel && \
+    python -m pip install --user "$(head bdist_name)[opt]" && \
+    python -m pip install --user "$(head bdist_name)" && \
+    python -m pip install --user "$(head bdist_name)[opt]"
+
+
+RUN if [[ "${ENABLE_AIM}" == "true" ]]; then \
+        python -m pip install --user "$(head bdist_name)[aim]"; \
+    fi
+
+# Clean up the wheel module. It's only needed by flash-attn install
+RUN python -m pip uninstall wheel build -y && \
+    # Cleanup the bdist whl file
+    rm $(head bdist_name) /tmp/bdist_name
+
+## Final image ################################################
+FROM release-base AS release
+ARG USER
+ARG PYTHON_VERSION
+ARG ENABLE_AIM
+
+RUN mkdir -p /licenses
+COPY LICENSE /licenses/
+
+RUN mkdir /app && \
+    chown -R $USER:0 /app /tmp && \
+    chmod -R g+rwX /app /tmp
+
+# Need a better way to address these hacks
+RUN if [[ "${ENABLE_AIM}" == "true" ]] ; then \
+        touch /.aim_profile && \
+        chmod -R 777 /.aim_profile; \
+    fi
+RUN mkdir /.cache && \
+    chmod -R 777 /.cache
+
+# Copy scripts and default configs
+COPY build/accelerate_launch.py /app/
+COPY build/utils.py /app/build/
+RUN chmod +x /app/accelerate_launch.py
+
+ENV SET_NUM_PROCESSES_TO_NUM_GPUS="True"
+
+WORKDIR /app
+USER ${USER}
+COPY --from=python-installations /home/${USER}/.local /home/${USER}/.local
+ENV PYTHONPATH="/home/${USER}/.local/lib/python${PYTHON_VERSION}/site-packages"
+
+CMD [ "python", "/app/accelerate_launch.py" ]

build/accelerate_launch.py

@@ -0,0 +1,131 @@
+# Copyright The FMS Model Optimizer Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Script wraps fms_mo to run with accelerate for multi and single GPU cases.
+Read accelerate_launch_args configuration via environment variable `FMS_MO_CONFIG_JSON_PATH`
+for the path to the JSON config file with parameters or `FMS_MO_CONFIG_JSON_ENV_VAR`
+for the encoded config string to parse.
+"""
+
+# Standard
+import os
+import logging
+import subprocess
+import sys
+import traceback
+from pathlib import Path
+
+# Third Party
+from accelerate.commands.launch import launch_command
+
+# Local
+from build.utils import (
+    process_accelerate_launch_args,
+)
+
+from fms_mo.utils.config_utils import get_json_config
+from fms_mo.utils.error_logging import (
+    write_termination_log,
+    USER_ERROR_EXIT_CODE,
+    INTERNAL_ERROR_EXIT_CODE,
+)
+
+ERROR_LOG = "/dev/termination-log"
+
+
+def main():
+    if not os.getenv("TERMINATION_LOG_FILE"):
+        os.environ["TERMINATION_LOG_FILE"] = ERROR_LOG
+
+    ##########
+    #
+    # Parse arguments
+    #
+    ##########
+    try:
+        job_config = get_json_config()
+        if not job_config:
+            raise ValueError(
+                "Must set environment variable 'FMS_MO_CONFIG_JSON_PATH' \
+            or 'FMS_MO_CONFIG_JSON_ENV_VAR'."
+            )
+
+        # Configure log_level of python native logger.
+        # CLI arg takes precedence over env var. And if neither is set, we use default "WARNING"
+        log_level = job_config.get(
+            "log_level"
+        )  # this will be set to either the value found or None
+        if (
+            not log_level
+        ):  # if log level not set by job_config aka by JSON, set it via env var or set default
+            log_level = os.environ.get("LOG_LEVEL", "WARNING") # TPP should be WARNING
+        logging.basicConfig(level=log_level.upper())
+
+        logging.debug("About to parse args") #TPP add
+
+        args = process_accelerate_launch_args(job_config)
+        logging.debug("accelerate launch parsed args: %s", args)
+    except FileNotFoundError as e:
+        logging.error(traceback.format_exc())
+        write_termination_log("Unable to load file: {}".format(e))
+        sys.exit(USER_ERROR_EXIT_CODE)
+    except (TypeError, ValueError, EnvironmentError) as e:
+        logging.error(traceback.format_exc())
+        write_termination_log(
+            f"Exception raised during optimization. This may be a problem with your input: {e}"
+        )
+        sys.exit(USER_ERROR_EXIT_CODE)
+    except Exception as e:  # pylint: disable=broad-except
+        logging.error(traceback.format_exc())
+        write_termination_log(f"Unhandled exception during optimization. {e}")
+        sys.exit(INTERNAL_ERROR_EXIT_CODE)
+
+    ##########
+    #
+    # Launch optimization
+    #
+    ##########
+    output_dir = job_config.get("output_dir")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    try:
+        # checkpoints outputted to tempdir, only final checkpoint copied to output dir
+        launch_command(args)
+    except subprocess.CalledProcessError as e:
+        # If the subprocess throws an exception, the base exception is hidden in the
+        # subprocess call and is difficult to access at this level. However, that is not
+        # an issue because run_quant.py would have already written the exception
+        # message to termination log.
+        logging.error(traceback.format_exc())
+        # The exit code that run_quant.py threw is captured in e.returncode
+
+        return_code = e.returncode
+        if return_code not in [INTERNAL_ERROR_EXIT_CODE, USER_ERROR_EXIT_CODE]:
+            return_code = INTERNAL_ERROR_EXIT_CODE
+            write_termination_log(f"Unhandled exception during optimization. {e}")
+        sys.exit(return_code)
+    except Exception as e:  # pylint: disable=broad-except
+        logging.error(traceback.format_exc())
+        write_termination_log(f"Unhandled exception during optimization. {e}")
+        sys.exit(INTERNAL_ERROR_EXIT_CODE)
+
+    # The .complete file will signal to users that we are finished copying
+    # files over
+    if os.path.exists(output_dir):
+        Path(os.path.join(output_dir, ".complete")).touch()
+
+    return 0
+
+
+if __name__ == "__main__":
+    main()