@@ -1246,10 +1246,10 @@ def create_dockerfile_linux(
12461246
12471247WORKDIR /opt/tritonserver
12481248COPY --chown=1000:1000 NVIDIA_Deep_Learning_Container_License.pdf .
1249- RUN find /opt/tritonserver/python -maxdepth 1 -type f -name \\
1250- "tritonserver-*.whl" | xargs -I {} pip install --upgrade {}[all ] && \ \
1251- find /opt/tritonserver/python -maxdepth 1 -type f -name \\
1252- "tritonfrontend-*.whl" | xargs -I {} pip install --upgrade {}[all]
1249+ RUN find /opt/tritonserver/python -maxdepth 1 -type f -name \
1250+ "tritonserver-*.whl" | xargs -I {} pip install --upgrade {}[$VARIANT ] && \
1251+ find /opt/tritonserver/python -maxdepth 1 -type f -name \
1252+ "tritonfrontend-*.whl" | xargs -I {} pip install --upgrade {}[$VARIANT];
12531253
12541254RUN pip3 install -r python/openai/requirements.txt
12551255
@@ -1297,6 +1297,7 @@ def dockerfile_prepare_container_linux(argmap, backends, enable_gpu, target_mach
12971297 df = """
12981298ARG TRITON_VERSION
12991299ARG TRITON_CONTAINER_VERSION
1300+ ARG VARIANT=all
13001301
13011302ENV TRITON_SERVER_VERSION ${TRITON_VERSION}
13021303ENV NVIDIA_TRITON_SERVER_VERSION ${TRITON_CONTAINER_VERSION}
@@ -1911,6 +1912,10 @@ def create_docker_build_script(script_name, container_install_dir, container_ci_
19111912 f"--secret id=NVPL_SLIM_URL" ,
19121913 f"--build-arg BUILD_PUBLIC_VLLM={ build_public_vllm } " ,
19131914 ]
1915+ if FLAGS .build_variant :
1916+ finalargs += [
1917+ "--build-arg VARIANT=" + (FLAGS .build_variant ),
1918+ ]
19141919 finalargs += [
19151920 "-t" ,
19161921 "tritonserver" ,
@@ -2776,6 +2781,13 @@ def enable_all():
27762781 default = DEFAULT_TRITON_VERSION_MAP ["rhel_py_version" ],
27772782 help = "This flag sets the Python version for RHEL platform of Triton Inference Server to be built. Default: the latest supported version." ,
27782783 )
2784+ parser .add_argument (
2785+ "--build_variant" ,
2786+ required = False ,
2787+ type = str ,
2788+ default = "all" ,
2789+ help = "Can be set to all or cpu,Default value is all."
2790+ )
27792791 parser .add_argument (
27802792 "--build-secret" ,
27812793 action = "append" ,
@@ -2815,6 +2827,9 @@ def enable_all():
28152827 FLAGS .extra_backend_cmake_arg = []
28162828 if FLAGS .build_secret is None :
28172829 FLAGS .build_secret = []
2830+ if hasattr (FLAGS , 'build_variant' ) and FLAGS .build_variant not in ["all" , "cpu" ]:
2831+ raise ValueError (f"Invalid build_variant value: { FLAGS .build_variant } . Expected 'all' or 'cpu'." )
2832+
28182833
28192834 FLAGS .boost_url = os .getenv (
28202835 "TRITON_BOOST_URL" ,
0 commit comments