From a65b114149b50e9d16d8699aa07484e451be2834 Mon Sep 17 00:00:00 2001
From: Austin Anderson <angerson@google.com>
Date: Thu, 13 Jul 2023 17:03:33 -0700
Subject: [PATCH] Update runtime containers to Ubuntu 22.04, Python 3.11

---
 .github/workflows/scheduled.yml               | 38 -------------------
 tensorflow_runtime_dockerfiles/cpu.Dockerfile |  4 +-
 tensorflow_runtime_dockerfiles/gpu.Dockerfile |  4 +-
 .../gpu.packages.txt                          |  5 +--
 .../setup.sources.sh                          | 16 +++-----
 5 files changed, 11 insertions(+), 56 deletions(-)
 delete mode 100644 .github/workflows/scheduled.yml

diff --git a/.github/workflows/scheduled.yml b/.github/workflows/scheduled.yml
deleted file mode 100644
index ceeea29..0000000
--- a/.github/workflows/scheduled.yml
+++ /dev/null
@@ -1,38 +0,0 @@
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
-#
-# 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.
-# ============================================================================
-
-name: Continuous job tester
-on:
-  schedule:
-    - cron: '*/5 * * * *'
-  workflow_dispatch:
-
-jobs:
-  dashboards:
-    if: github.repository == 'angerson/build' # Don't do this yet
-    name: A continuous job
-    runs-on: ubuntu-latest
-    permissions:
-      statuses: write
-
-    strategy:
-      matrix:
-        config: [ tensorflow, jax ]
-
-    steps:
-      - uses: myrotvorets/set-commit-status-action@master
-        with:
-          status: ${{ job.status }}
-          sha: ${{ github.sha }}
diff --git a/tensorflow_runtime_dockerfiles/cpu.Dockerfile b/tensorflow_runtime_dockerfiles/cpu.Dockerfile
index 1ead08d..280f9a6 100644
--- a/tensorflow_runtime_dockerfiles/cpu.Dockerfile
+++ b/tensorflow_runtime_dockerfiles/cpu.Dockerfile
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 
-FROM ubuntu:20.04 as base
+FROM ubuntu:22.04 as base
 ENV DEBIAN_FRONTEND=noninteractive
 ENV LANG C.UTF-8
 
@@ -24,7 +24,7 @@ RUN /setup.sources.sh
 RUN /setup.packages.sh /cpu.packages.txt
 
 
-ARG PYTHON_VERSION=python3.8
+ARG PYTHON_VERSION=python3.11
 ARG TENSORFLOW_PACKAGE=tf-nightly
 COPY setup.python.sh /setup.python.sh
 COPY cpu.requirements.txt /cpu.requirements.txt
diff --git a/tensorflow_runtime_dockerfiles/gpu.Dockerfile b/tensorflow_runtime_dockerfiles/gpu.Dockerfile
index 6e3aa17..578d5f0 100644
--- a/tensorflow_runtime_dockerfiles/gpu.Dockerfile
+++ b/tensorflow_runtime_dockerfiles/gpu.Dockerfile
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 
-FROM nvidia/cuda:11.8.0-base-ubuntu20.04 as base
+FROM nvidia/cuda:11.8.0-base-ubuntu22.04 as base
 ENV DEBIAN_FRONTEND=noninteractive
 ENV LANG C.UTF-8
 
@@ -23,7 +23,7 @@ COPY gpu.packages.txt /gpu.packages.txt
 RUN /setup.sources.sh
 RUN /setup.packages.sh /gpu.packages.txt
 
-ARG PYTHON_VERSION=python3.8
+ARG PYTHON_VERSION=python3.11
 ARG TENSORFLOW_PACKAGE=tf-nightly
 COPY setup.python.sh /setup.python.sh
 COPY gpu.requirements.txt /gpu.requirements.txt
diff --git a/tensorflow_runtime_dockerfiles/gpu.packages.txt b/tensorflow_runtime_dockerfiles/gpu.packages.txt
index d0b3691..59d10e3 100644
--- a/tensorflow_runtime_dockerfiles/gpu.packages.txt
+++ b/tensorflow_runtime_dockerfiles/gpu.packages.txt
@@ -13,9 +13,8 @@ libcusparse-11-8
 libcublas-11-8
 # CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
 libcudnn8=8.6.0.163-1+cuda11.8
-# TensorRT: See https://docs.nvidia.com/deeplearning/sdk/tensorrt-install-guide/index.html#maclearn-net-repo-install-rpm
-libnvinfer-plugin8=8.4.3-1+cuda11.6
-libnvinfer8=8.4.3-1+cuda11.6
+libnvinfer-plugin8=8.6.1.6-1+cuda11.8
+libnvinfer8=8.6.1.6-1+cuda11.8
 
 # Other packages
 build-essential
diff --git a/tensorflow_runtime_dockerfiles/setup.sources.sh b/tensorflow_runtime_dockerfiles/setup.sources.sh
index d323121..8ff79b5 100755
--- a/tensorflow_runtime_dockerfiles/setup.sources.sh
+++ b/tensorflow_runtime_dockerfiles/setup.sources.sh
@@ -25,15 +25,9 @@ export DEBIAN_FRONTEND=noninteractive
 
 # Set up shared custom sources
 apt-get update
-apt-get install -y gnupg ca-certificates
+apt-get install -y gnupg ca-certificates wget
 
-# Explicitly request Nvidia repo keys
-# See: https://forums.developer.nvidia.com/t/invalid-public-key-for-cuda-apt-repository/212901/11
-apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
-apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
-
-# Set up custom sources
-cat >/etc/apt/sources.list.d/custom.list <<SOURCES
-# Nvidia CUDA packages: 18.04 has more available than 20.04, and we use those
-deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64 /
-SOURCES
+# Install Nvidia repo keys
+# See: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#network-repo-installation-for-ubuntu
+wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
+dpkg -i cuda-keyring_1.1-1_all.deb
-- 
GitLab