From a07ca4d2a068cd6aea1c702eeb35ff71c3f36d1a Mon Sep 17 00:00:00 2001
From: SdgJlbl <sarah.diot-girard@owkin.com>
Date: Tue, 13 Aug 2024 17:05:18 +0200
Subject: [PATCH] fixup! feat: improve Dockerfile

---
 substrafl/remote/register/register.py  | 13 +++++++------
 tests/remote/register/test_register.py | 15 ++++++++-------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/substrafl/remote/register/register.py b/substrafl/remote/register/register.py
index f563f1cd..3152db5e 100644
--- a/substrafl/remote/register/register.py
+++ b/substrafl/remote/register/register.py
@@ -41,12 +41,13 @@
 
 # update image & install Python
 ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update -y && \
-    apt-get install -y software-properties-common && \
-    add-apt-repository -y ppa:deadsnakes/ppa && \
-    apt-get -y upgrade && \
-    apt-get install -y python{python_version} python{python_version}-venv python3-pip && \
-    apt-get clean
+RUN apt-get update -y\
+    && apt-get install -y software-properties-common\
+    && add-apt-repository -y ppa:deadsnakes/ppa\
+    && apt-get -y upgrade\
+    && apt-get install -y python{python_version} python{python_version}-venv python3-pip\
+    && apt-get clean\
+    && rm -rf /var/lib/apt/lists/*
 
 """
 
diff --git a/tests/remote/register/test_register.py b/tests/remote/register/test_register.py
index 6c1cf020..5aca2997 100644
--- a/tests/remote/register/test_register.py
+++ b/tests/remote/register/test_register.py
@@ -58,14 +58,15 @@ def test_get_base_docker_image_gpu():
     expected_dockerfile = """
 FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
 
-# update image
+# update image & install Python
 ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update -y && \
-    apt-get install -y software-properties-common && \
-    add-apt-repository -y ppa:deadsnakes/ppa && \
-    apt-get -y upgrade && \
-    apt-get install -y python{python_version} python{python_version}-venv python3-pip && \
-    apt-get clean
+RUN apt-get update -y\
+    && apt-get install -y software-properties-common\
+    && add-apt-repository -y ppa:deadsnakes/ppa\
+    && apt-get -y upgrade\
+    && apt-get install -y python3.11 python3.11-venv python3-pip\
+    && apt-get clean\
+    && rm -rf /var/lib/apt/lists/*
 
 """
     assert expected_dockerfile == _get_base_docker_image("3.11", use_gpu=True)