Skip to content

Commit ce46097

Browse files
Added changes to add --build_variant flag for cpu only build.
1 parent e672c28 commit ce46097

File tree

1 file changed

+19
-4
lines changed

1 file changed

+19
-4
lines changed

build.py

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1239,10 +1239,10 @@ def create_dockerfile_linux(
12391239
12401240
WORKDIR /opt/tritonserver
12411241
COPY --chown=1000:1000 NVIDIA_Deep_Learning_Container_License.pdf .
1242-
RUN find /opt/tritonserver/python -maxdepth 1 -type f -name \\
1243-
"tritonserver-*.whl" | xargs -I {} pip install --upgrade {}[all] && \\
1244-
find /opt/tritonserver/python -maxdepth 1 -type f -name \\
1245-
"tritonfrontend-*.whl" | xargs -I {} pip install --upgrade {}[all]
1242+
RUN find /opt/tritonserver/python -maxdepth 1 -type f -name \
1243+
"tritonserver-*.whl" | xargs -I {} pip install --upgrade {}[$VARIANT] && \
1244+
find /opt/tritonserver/python -maxdepth 1 -type f -name \
1245+
"tritonfrontend-*.whl" | xargs -I {} pip install --upgrade {}[$VARIANT];
12461246
12471247
RUN pip3 install -r python/openai/requirements.txt
12481248
@@ -1295,6 +1295,7 @@ def dockerfile_prepare_container_linux(argmap, backends, enable_gpu, target_mach
12951295
df = """
12961296
ARG TRITON_VERSION
12971297
ARG TRITON_CONTAINER_VERSION
1298+
ARG VARIANT=all
12981299
12991300
ENV TRITON_SERVER_VERSION ${TRITON_VERSION}
13001301
ENV NVIDIA_TRITON_SERVER_VERSION ${TRITON_CONTAINER_VERSION}
@@ -1898,6 +1899,10 @@ def create_docker_build_script(script_name, container_install_dir, container_ci_
18981899
f"--secret id=NVPL_SLIM_URL",
18991900
f"--build-arg BUILD_PUBLIC_VLLM={build_public_vllm}",
19001901
]
1902+
if FLAGS.build_variant:
1903+
finalargs += [
1904+
"--build-arg VARIANT="+(FLAGS.build_variant),
1905+
]
19011906
finalargs += [
19021907
"-t",
19031908
"tritonserver",
@@ -2763,6 +2768,13 @@ def enable_all():
27632768
default=DEFAULT_TRITON_VERSION_MAP["rhel_py_version"],
27642769
help="This flag sets the Python version for RHEL platform of Triton Inference Server to be built. Default: the latest supported version.",
27652770
)
2771+
parser.add_argument(
2772+
"--build_variant",
2773+
required=False,
2774+
type=str,
2775+
default="all",
2776+
help="Can be set to all or cpu,Default value is all."
2777+
)
27662778
parser.add_argument(
27672779
"--build-secret",
27682780
action="append",
@@ -2802,6 +2814,9 @@ def enable_all():
28022814
FLAGS.extra_backend_cmake_arg = []
28032815
if FLAGS.build_secret is None:
28042816
FLAGS.build_secret = []
2817+
if hasattr(FLAGS, 'build_variant') and FLAGS.build_variant not in ["all", "cpu"]:
2818+
raise ValueError(f"Invalid build_variant value: {FLAGS.build_variant}. Expected 'all' or 'cpu'.")
2819+
28052820

28062821
# if --enable-all is specified, then update FLAGS to enable all
28072822
# settings, backends, repo-agents, caches, file systems, endpoints, etc.

0 commit comments

Comments
 (0)