From 5f65f1f18a9ca679d12142017e6503bede71699a Mon Sep 17 00:00:00 2001
From: Sergei Monakhov <sergey.monakhov@gmail.com>
Date: Sun, 15 Oct 2023 12:48:58 +0300
Subject: [PATCH] update images

---
 images/keras/cuda11-8/Dockerfile | 22 +++++++++++++-----
 template/required.j2             | 38 ++++++++++++++++----------------
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/images/keras/cuda11-8/Dockerfile b/images/keras/cuda11-8/Dockerfile
index 0152d96..17512f9 100644
--- a/images/keras/cuda11-8/Dockerfile
+++ b/images/keras/cuda11-8/Dockerfile
@@ -107,13 +107,25 @@ RUN rm -rf /usr/local/bin/pip3 \
 
 RUN ${INTERPRETER} -m pip install -r /tmp/requirements.txt
 
-RUN ${INTERPRETER} -m pip install \
-        tensorflow-gpu==2.14.0 \
+RUN \
+  case ${INTERPRETER} in \
+    python3.9) \
+      CP_FORMAT=cp39 \
+      ;; \
+    python3.10) \
+      CP_FORMAT=cp310 \
+      ;; \
+    python3.11) \
+      CP_FORMAT=cp311 \
+      ;; \
+  esac \
+  && ${INTERPRETER} -m pip install \
+        https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-2.14.0-${CP_FORMAT}-${CP_FORMAT}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl \
         keras==2.14.0 \
         jupyterlab==4.0.7 \
-        --index-url https://download.pytorch.org/whl/cu118 \
-        --extra-index-url https://pypi.org/simple \
-    && ${INTERPRETER} -m pip cache purge
+      --index-url https://download.pytorch.org/whl/cu118 \
+      --extra-index-url https://pypi.org/simple \
+  && ${INTERPRETER} -m pip cache purge
 
 RUN curl https://rclone.org/install.sh -o ~/install.sh \
     && chmod +x ~/install.sh \
diff --git a/template/required.j2 b/template/required.j2
index 719c3fb..efc1256 100644
--- a/template/required.j2
+++ b/template/required.j2
@@ -106,29 +106,29 @@ RUN curl -o ~/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest
 RUN ${INTERPRETER} -m pip install -r /tmp/requirements.txt
 {% endmacro %}
 
-{%- macro convert_interpreter_to_cp(interpreter) -%}
-    {%- if interpreter == "python3.9" -%}
-        cp39
-    {%- elif interpreter == "python3.10" -%}
-        cp310
-    {%- elif interpreter == "python3.11" -%}
-        cp311
-    {%- else -%}
-        unknown
-    {%- endif %}
-{% endmacro %}
-
 {%- macro install_pip_package(packages, INTERPRETER) -%}
 RUN \
-  {%- for package in packages -%}
-    {%- if package.type == "pip" and package.name != "tensorflow" -%}
-      {%- if package.version is defined %}
-        ${INTERPRETER} -m pip install {{ package.name }}=={{ package.version }} \
+  case ${INTERPRETER} in \
+    python3.9) \
+      CP_FORMAT=cp39 \
+      ;; \
+    python3.10) \
+      CP_FORMAT=cp310 \
+      ;; \
+    python3.11) \
+      CP_FORMAT=cp311 \
+      ;; \
+  esac \
+  && ${INTERPRETER} -m pip install \
+  {%- for package in packages %}
+    {%- if package.type == "pip" -%}
+      {%- if package.version is defined and package.name != "tensorflow-gpu" %}
+        {{ package.name }}=={{ package.version }} \
+      {%- elif package.version is defined and package.name == "tensorflow-gpu" %}
+        https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-{{ package.version }}-${CP_FORMAT}-${CP_FORMAT}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl \
       {%- else %}
-        ${INTERPRETER} -m pip install {{ package.name }} \
+        {{ package.name }} \
       {%- endif %}
-    {%- elif package.type == "pip" and package.name == "tensorflow" -%}
-      ${INTERPRETER} -m pip install https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-{{ package.version }}-{{ convert_interpreter_to_cp(INTERPRETER) }}-{{ convert_interpreter_to_cp(INTERPRETER) }}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl \
     {%- endif %}
   {%- endfor %}
       --index-url https://download.pytorch.org/whl/cu118 \