feat: container build and run enhancements

README.md

@@ -39,7 +39,8 @@ center scale without sacrificing performance or ease of use.
 
 Triton Distributed development and examples are container based.
 
-You can build the Triton Distributed container using the build scripts in `container/`.
+You can build the Triton Distributed container using the build scripts
+in `container/` (or directly with `docker build`).
 
 We provide 3 types of builds:
 
@@ -53,6 +54,38 @@ For example, if you want to build a container for the `VLLM` backend you can run
 
 Please see the instructions in the corresponding example for specific build instructions.
 
+## Running Triton Distributed for Local Testing and Development
+
+You can run the Triton Distributed container using the run scripts in
+`container/` (or directly with `docker run`).
+
+The run script offers a few common workflows:
+
+1. Running a command in a container and exiting.
+
+```
+./container/run.sh -- python3 -c "import tdist.icp.protos.icp_pb2 as icp_proto; print(icp_proto); print(dir(icp_proto));"
+```
+
+2. Starting an interactive shell.
+```
+./container/run.sh -it
+```
+
+3. Mounting the local workspace and Starting an interactive shell.
+
+```
+./container/run.sh -it --mount-workspace
+```
+
+The last command also passes common environment variables ( ```-e
+HF_TOKEN```) and mounts common directories such as ```/tmp:/tmp```,
+```/mnt:/mnt```.
+
+Please see the instructions in the corresponding example for specific
+deployment instructions.
+
+
 <!--
 
 ## Goals

container/Dockerfile

@@ -64,6 +64,9 @@ ENV FRAMEWORK_LD_LIBRARY_PATH=${TENSORRTLLM_FRAMEWORK:+/opt/tritonserver/backend
 ENV LD_LIBRARY_PATH=${FRAMEWORK_LD_LIBRARY_PATH}:${LD_LIBRARY_PATH}
 ENV TENSORRTLLM_BACKEND_COMMIT=$TENSORRTLLM_BACKEND_COMMIT
 
+# TODO set VLLM Version
+# ENV VLLM_VERSION
+
 # Install NATS - pointing toward NATS github instead of binaries.nats.dev due to server instability
 RUN wget https://github.com/nats-io/nats-server/releases/download/v2.10.24/nats-server-v2.10.24-amd64.deb && dpkg -i nats-server-v2.10.24-amd64.deb
 

container/build.sh

@@ -17,6 +17,7 @@
 TAG=
 RUN_PREFIX=
 PLATFORM=linux/amd64
+VERSION=0.1.0
 
 # Frameworks
 #
@@ -104,7 +105,7 @@ get_options() {
             ;;
         --tag)
             if [ "$2" ]; then
-                TAG=$2
+                TAG="--tag $2"
                 shift
             else
 		missing_requirement $1
@@ -121,6 +122,14 @@ get_options() {
 	--no-cache)
 	    NO_CACHE=" --no-cache"
             ;;
+	--cache-from)
+	    if [ "$2" ]; then
+                CACHE_FROM="--cache-from $2"
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
         --)
             shift
             break
@@ -169,7 +178,7 @@ get_options() {
     fi
 
     if [ -z "$TAG" ]; then
-        TAG="triton-distributed:${FRAMEWORK,,}-${BASE_VERSION}"
+        TAG="--tag triton-distributed:${VERSION}-${FRAMEWORK,,}"
     fi
 
     if [ ! -z "$PLATFORM" ]; then
@@ -241,6 +250,8 @@ if [ ! -z ${HF_TOKEN} ]; then
     BUILD_ARGS+=" --build-arg HF_TOKEN=${HF_TOKEN} "
 fi
 
+LATEST_TAG="--tag triton-distributed:latest-${FRAMEWORK,,}"
+
 show_image_options
 
 
@@ -248,7 +259,7 @@ if [ -z "$RUN_PREFIX" ]; then
     set -x
 fi
 
-$RUN_PREFIX docker build -f $DOCKERFILE $PLATFORM $BUILD_ARGS -t $TAG $BUILD_CONTEXT $NO_CACHE
+$RUN_PREFIX docker build -f $DOCKERFILE $PLATFORM $BUILD_ARGS $CACHE_FROM $TAG $LATEST_TAG $BUILD_CONTEXT $NO_CACHE
 
 { set +x; } 2>/dev/null
 

container/run.sh

+#!/bin/bash -e
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+TAG=
+RUN_PREFIX=
+
+# Frameworks
+#
+# Each framework has a corresponding base image.  Additional
+# dependencies are specified in the /container/deps folder and
+# installed within framework specific sections of the Dockerfile.
+
+declare -A FRAMEWORKS=(["STANDARD"]=1 ["TENSORRTLLM"]=2 ["VLLM"]=3)
+DEFAULT_FRAMEWORK=STANDARD
+
+SOURCE_DIR=$(dirname "$(readlink -f "$0")")
+
+IMAGE=
+HF_CACHE=
+DEFAULT_HF_CACHE=${SOURCE_DIR}/.cache/huggingface
+GPUS="all"
+PRIVILEGED=
+VOLUME_MOUNTS=
+MOUNT_WORKSPACE=
+ENVIRONMENT_VARIABLES=
+REMAINING_ARGS=
+INTERACTIVE=
+
+get_options() {
+    while :; do
+        case $1 in
+        -h | -\? | --help)
+            show_help
+            exit
+            ;;
+	--framework)
+            if [ "$2" ]; then
+                FRAMEWORK=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+        --image)
+            if [ "$2" ]; then
+                IMAGE=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+	--name)
+            if [ "$2" ]; then
+                NAME=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+	--hf-cache)
+            if [ "$2" ]; then
+                HF_CACHE=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+
+	--gpus)
+            if [ "$2" ]; then
+                GPUS=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+	--command)
+            if [ "$2" ]; then
+                COMMAND=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+	--privileged)
+            if [ "$2" ]; then
+                PRIVILEGED=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+	-v)
+            if [ "$2" ]; then
+                VOLUME_MOUNTS+=" -v $2 "
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+	-e)
+            if [ "$2" ]; then
+                ENVIRONMENT_VARIABLES+=" -e $2 "
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+	-it)
+	    INTERACTIVE=" -it "
+	    ;;
+	--mount-workspace)
+	    MOUNT_WORKSPACE=TRUE
+	    ;;
+        --dry-run)
+            RUN_PREFIX="echo"
+            echo ""
+            echo "=============================="
+            echo "DRY RUN: COMMANDS PRINTED ONLY"
+            echo "=============================="
+            echo ""
+            ;;
+        --)
+            shift
+            break
+            ;;
+         -?*)
+	    error 'ERROR: Unknown option: ' $1
+            ;;
+	 ?*)
+	    error 'ERROR: Unknown option: ' $1
+            ;;
+        *)
+            break
+            ;;
+        esac
+
+        shift
+    done
+
+    if [ -z "$FRAMEWORK" ]; then
+	FRAMEWORK=$DEFAULT_FRAMEWORK
+    fi
+
+    if [ ! -z "$FRAMEWORK" ]; then
+	FRAMEWORK=${FRAMEWORK^^}
+	if [[ ! -n "${FRAMEWORKS[$FRAMEWORK]}" ]]; then
+	    error 'ERROR: Unknown framework: ' $FRAMEWORK
+	fi
+    fi
+
+    if [ -z "$IMAGE" ]; then
+        IMAGE="triton-distributed:latest-${FRAMEWORK,,}"
+    fi
+
+    if [[ ${GPUS^^} == "NONE" ]]; then
+	GPU_STRING=""
+    else
+	GPU_STRING="--gpus ${GPUS}"
+    fi
+
+    if [[ ${NAME^^} == "" ]]; then
+	NAME_STRING=""
+    else
+	NAME_STRING="--name ${NAME}"
+    fi
+
+    if [ ! -z "$MOUNT_WORKSPACE" ]; then
+	VOLUME_MOUNTS+=" -v ${SOURCE_DIR}/..:/workspace "
+	VOLUME_MOUNTS+=" -v /tmp:/tmp "
+	VOLUME_MOUNTS+=" -v /mnt/:/mnt "
+
+	if [ -z "$HF_CACHE" ]; then
+	    HF_CACHE=$DEFAULT_HF_CACHE
+	fi
+
+	if [ -z ${PRIVILEGED} ]; then
+	    PRIVILEGED="TRUE"
+	fi
+
+	ENVIRONMENT_VARIABLES+=" -e HF_TOKEN"
+
+	if [ ! -d "${SOURCE_DIR}/icp/src/python/tdist/icp/protos" ]; then
+	    $RUN_PREFIX docker run --rm -t -v ${SOURCE_DIR}/..:/workspace -w /workspace $IMAGE /workspace/icp/protos/gen_python.sh > /dev/null 2>&1
+	fi
+	INTERACTIVE=" -it "
+    fi
+
+    if [[ ${HF_CACHE^^} == "NONE" ]]; then
+	HF_CACHE=
+    fi
+
+    if [ ! -z "$HF_CACHE" ]; then
+	mkdir -p $HF_CACHE
+	VOLUME_MOUNTS+=" -v $HF_CACHE:/root/.cache/huggingface"
+    fi
+
+    if [ -z ${PRIVILEGED} ]; then
+	PRIVILEGED="FALSE"
+    fi
+
+    if [[ ${PRIVILEGED^^} == "FALSE" ]]; then
+	PRIVILEGED_STRING=""
+    else
+	PRIVILEGED_STRING="--privileged"
+    fi
+
+    REMAINING_ARGS=("$@")
+}
+
+show_help() {
+    echo "usage: run.sh"
+    echo "  [--image image]"
+    echo "  [--framework framework one of ${!FRAMEWORKS[@]}]"
+    echo "  [--name name for launched container, default NONE] "
+    echo "  [--privileged whether to launch in privileged mode, default FALSE unless mounting workspace]"
+    echo "  [--dry-run print docker commands without running]"
+    echo "  [--hf-cache directory to volume mount as the hf cache, default is NONE unless mounting workspace]"
+    echo "  [--gpus gpus to enable, default is 'all', 'none' disables gpu support]"
+    echo "  [-v add volume mount]"
+    echo "  [-e add environment variable]"
+    echo "  [--mount-workspace set up for local development]"
+    echo "  [-- stop processing and pass remaining args as command to docker run]"
+    exit 0
+}
+
+missing_requirement() {
+    error "ERROR: $1 requires an argument."
+}
+
+error() {
+    printf '%s %s\n' "$1" "$2" >&2
+    exit 1
+}
+
+get_options "$@"
+
+# RUN the image
+
+if [ -z "$RUN_PREFIX" ]; then
+    set -x
+fi
+
+${RUN_PREFIX} docker run ${GPU_STRING} ${INTERACTIVE} --rm --network host --shm-size=10G --ulimit memlock=-1 --ulimit stack=67108864 ${ENVIRONMENT_VARIABLES} ${VOLUME_MOUNTS} -w /workspace --cap-add CAP_SYS_PTRACE --ipc host ${PRIVILEGED_STRING} ${NAME_STRING} ${IMAGE} "${REMAINING_ARGS[@]}"
+
+{ set +x; } 2>/dev/null