Skip to content

Commit f6dee1d

Browse files
Added changes to add --build_variant flag for cpu only build.
1 parent 4f98ebf commit f6dee1d

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
@@ -1246,10 +1246,10 @@ def create_dockerfile_linux(
12461246
12471247
WORKDIR /opt/tritonserver
12481248
COPY --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
12541254
RUN 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 = """
12981298
ARG TRITON_VERSION
12991299
ARG TRITON_CONTAINER_VERSION
1300+
ARG VARIANT=all
13001301
13011302
ENV TRITON_SERVER_VERSION ${TRITON_VERSION}
13021303
ENV 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

Comments
 (0)