From 02b2232df8c79148aea8547d58475c2458bf411b Mon Sep 17 00:00:00 2001
From: Toby Boyd <tobyboyd@google.com>
Date: Fri, 29 Sep 2023 14:23:56 -0700
Subject: [PATCH] Fixes nightly builds moving to python 3.9. - Removed
 tensorflow-datasets, it triggers trying to install etils[array_types] which
 doesn't exist. I am not sure when it existed. I tried 1.4.1 and 1.5.0
 (latest). - ubuntu_ci has python3.9 and 3.10. - Did a successful build and
 colab test with TF 2.14.0, Reverb 2.13.0, and tfp-nightly (need because the
 last release is not compatibility with TF 2.14.0. - Had to turn off
 `bandits.agents.linear_bandit_agent_test` test due to assertion error. Sees 2
 somethings in the event log and expects 1. Given our issues turning this off
 is prudent. P1 bug opened on that team.

PiperOrigin-RevId: 569588372
Change-Id: Id403aa14226d586a72628120b4a59d9a9759da0a
---
 broken_tests.txt       | 3 ++-
 broken_tests_gpu.txt   | 5 +++--
 setup.py               | 1 -
 tools/docker/ubuntu_ci | 9 ++++++---
 4 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/broken_tests.txt b/broken_tests.txt
index e558b2a36..76ebbace0 100644
--- a/broken_tests.txt
+++ b/broken_tests.txt
@@ -2,4 +2,5 @@ agents.cql.cql_sac_agent_test  # b/239448722
 environments.suite_atari_test # b/201537188
 examples.cql_sac.kumar20.dataset.file_utils_test  # b/199583100
 examples.cql_sac.kumar20.dataset.dataset_utils_test  # b/199583100
-examples.cql_sac.kumar20.cql_sac_train_eval_test  # b/199583100
\ No newline at end of file
+examples.cql_sac.kumar20.cql_sac_train_eval_test  # b/199583100
+bandits.agents.linear_bandit_agent_test  # b/302734159
\ No newline at end of file
diff --git a/broken_tests_gpu.txt b/broken_tests_gpu.txt
index fd4d6873c..594fd82e8 100644
--- a/broken_tests_gpu.txt
+++ b/broken_tests_gpu.txt
@@ -6,5 +6,6 @@ examples.cql_sac.kumar20.cql_sac_train_eval_test  # b/199583100
 replay_buffers.episodic_replay_buffer_driver_test  # b/199587793
 replay_buffers.episodic_replay_buffer_test  # b/199587793
 networks.q_rnn_network_test # b/237573967
-networks.actor_distribution_rnn_network_test # b/237573967
-agents.ddpg.critic_rnn_network_test # b/237573967
\ No newline at end of file
+networks.actor_distribution_rnn_network_test  # b/237573967
+agents.ddpg.critic_rnn_network_test  # b/237573967
+bandits.agents.linear_bandit_agent_test  # b/302734159
\ No newline at end of file
diff --git a/setup.py b/setup.py
index f245fe978..8357d6443 100644
--- a/setup.py
+++ b/setup.py
@@ -212,7 +212,6 @@ def get_test_packages():
       'opencv-python >= 3.4.1.15',
       'pybullet',
       'scipy >= 1.1.0',
-      'tensorflow_datasets',
   ]
   return test_packages
 
diff --git a/tools/docker/ubuntu_ci b/tools/docker/ubuntu_ci
index 415476558..9356e179c 100644
--- a/tools/docker/ubuntu_ci
+++ b/tools/docker/ubuntu_ci
@@ -12,7 +12,7 @@ FROM $base_image
 LABEL maintainer="tobyboyd@google.com"
 
 # Supports setting up a single version of python.
-ARG python_version="python3"
+ARG python_version="python3.9"
 ARG APT_COMMAND="apt-get -o Acquire::Retries=3 -y"
 
 # Stops tzdata from asking about timezones and blocking install on user input.
@@ -42,8 +42,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
       xvfb \
       freeglut3-dev \
       ffmpeg \
-      python3-opencv \
-      python3-dev \
+      python3.9-dev \
+      python3.10-dev \
+      # python >= 3.8 needs distutils for packaging.
+      python3.9-distutils \
+      python3.10-distutils \
       gfortran \
       libopenblas-dev \
       liblapack-dev