github.com/apache/beam/sdks/v2@v2.48.2/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile (about) 1 # Licensed to the Apache Software Foundation (ASF) under one 2 # or more contributor license agreements. See the NOTICE file 3 # distributed with this work for additional information 4 # regarding copyright ownership. The ASF licenses this file 5 # to you under the Apache License, Version 2.0 (the 6 # "License"); you may not use this file except in compliance 7 # with the License. You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 17 ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:22.05-py3 18 19 FROM ${BUILD_IMAGE} 20 21 ENV PATH="/usr/src/tensorrt/bin:${PATH}" 22 23 WORKDIR /workspace 24 25 RUN pip install --no-cache-dir apache-beam[gcp]==2.42.0 26 COPY --from=apache/beam_python3.8_sdk:2.42.0 /opt/apache/beam /opt/apache/beam 27 28 RUN pip install --upgrade pip \ 29 && pip install torch>=1.7.1 \ 30 && pip install torchvision>=0.8.2 \ 31 && pip install pillow>=8.0.0 \ 32 && pip install transformers>=4.18.0 \ 33 && pip install cuda-python 34 35 ENTRYPOINT [ "/opt/apache/beam/boot" ]