Docker images with multi-platform support

.github/actions/build-image/action.yml

@@ -0,0 +1,48 @@
+inputs:
+  dockerhub_username:
+    description: "The DockerHub username."
+    required: true
+
+  dockerhub_token:
+    description: "The DockerHub login token."
+    required: true
+
+  docker_file:
+    description: "The docker file."
+    required: true
+
+  tags:
+    description: "Image tags (csv)."
+    required: true
+
+  platforms:
+    description: "Platforms to build for (csv)."
+    default: "linux/amd64,linux/arm64"
+
+  push_image:
+    description: "Whether to push to DockerHub."
+    default: false
+
+runs:
+  using: "composite"
+  steps:
+    - name: Set up QEMU 🌈
+      uses: docker/setup-qemu-action@v3
+
+    - name: Set up Docker Buildx ✨
+      uses: docker/setup-buildx-action@v3
+
+    - name: Login to Docker Hub 🎪
+      uses: docker/login-action@v3
+      with:
+        username: ${{ inputs.dockerhub_username }}
+        password: ${{ inputs.dockerhub_token }}
+
+    - name: Build and deploy image 🐳
+      uses: docker/build-push-action@v5
+      with:
+        context: docker
+        file: ${{ inputs.docker_file }}
+        platforms: ${{ inputs.platforms }}
+        push: ${{ inputs.push_image }}
+        tags: ${{ inputs.tags }}

.github/workflows/deploy_images.yml

@@ -3,167 +3,185 @@ name: Deploy images
 on:
   workflow_dispatch:
     inputs:
-      skip_push:
-        description: Skip pushing images? (true|false)
-        required: false
-        default: 'false'
-      skip_docs:
-        description: Skip the docs image? (true|false)
-        required: false
-        default: 'false'
+      push_image:
+        description: Push to DockerHub
+        type: boolean
+        default: false
+
+      specific_job:
+        description: Specific job to run
+        type: choice
+        default: all
+        options:
+          - all
+          - py37
+          - py38
+          - py39
+          - py310
+          - py311
+          - py311-cuda
+          - example
+
+      platforms:
+        description: Platforms to build for
+        type: choice
+        default: linux/amd64,linux/arm64
+        options:
+          - linux/amd64,linux/arm64
+          - linux/amd64
+          - linux/arm64
 
 jobs:
   py37:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py37') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_37
-          repository: cmsml/cmsml
-          tags: "3.7"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_37
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.7
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py38:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py38') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_38
-          repository: cmsml/cmsml
-          tags: "3.8"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_38
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.8
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py39:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py39') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_39
-          repository: cmsml/cmsml
-          tags: 3.9,3,latest
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_39
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.9,cmsml/cmsml:3,cmsml/cmsml:latest
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py39_base:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py39') }}
     runs-on: ubuntu-latest
-    if: ${{ github.event.inputs.skip_push != 'true' }}
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_39_base
-          repository: cmsml/cmsml
-          tags: 3.9_base
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_39_base
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.9_base
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py310:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py310') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_310
-          repository: cmsml/cmsml
-          tags: "3.10"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_310
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.10
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py311:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py311') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_311
-          repository: cmsml/cmsml
-          tags: "3.11"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_311
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.11
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
+
   py311-cuda:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py311-cuda') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_311_cuda
-          repository: cmsml/cmsml
-          tags: "3.11-cuda"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-          
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_311_cuda
+          # gpu / cuda only available on x86, not on arm
+          platforms: linux/amd64
+          tags: cmsml/cmsml:3.11-cuda
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
+
   docs:
+    if: ${{ github.event.inputs.push_image == 'true' && (github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'example')) }}
     needs: py39
     runs-on: ubuntu-latest
-    if: ${{ github.event.inputs.skip_docs != 'true' && github.event.inputs.skip_push != 'true' }}
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_docs
-          repository: cmsml/cmsml
-          tags: docs
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_docs
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:docs
+          push_image: true

.github/workflows/lint_and_test.yml

@@ -37,6 +37,7 @@ jobs:
           - {tag: "3.9_base", tf: "2.11.1"}
           - {tag: "3.9_base", tf: "2.12.1"}
           - {tag: "3.9_base", tf: "2.13.0"}
+          - {tag: "3.9_base", tf: "2.16.1"}
     name: test (image=${{ matrix.versions.tag }}, tf=${{ matrix.versions.tf }})
     steps:
       - name: Checkout 🛎️

.gitignore

@@ -19,3 +19,4 @@ build
 tmp
 docs/_build
 .ipynb_checkpoints
+.python-version

README.md

@@ -40,7 +40,7 @@ The documentation of this Python package is hosted on [readthedocs](http://cmsml
 **However**, note that this documentation only covers the API and technical aspects of the package itself.
 Usage examples and further techniques for working with machine learning tools in CMS, alongside a collection of useful guidelines can be found in the [general CMS ML group documentation](https://cms-ml.github.io/documentation).
 
-Click [here](https://github.com/cms-ml/cmsml/issues/new?labels=suggestion&template=feature-suggestion.md&) to submit a feature suggestion!
+Click [here](https://github.com/cms-ml/cmsml/issues/new?labels=suggestion&template=feature-suggestion.md) to submit a feature suggestion!
 
 
 <!-- marker-after-header -->
@@ -53,6 +53,15 @@ Click [here](https://github.com/cms-ml/cmsml/issues/new?labels=suggestion&templa
 
 To use the cmsml package via docker, checkout our [DockerHub](https://hub.docker.com/repository/docker/cmsml/cmsml) which contains tags for several Python versions.
 
+| Image                                                        | Python version |    TF Version    | PyTorch Version | GPU support |
+| :----------------------------------------------------------- | :------------: | :--------------: | :-------------: | :---------: |
+| `cmsml/cmsml:3.7`                                            |      3.7       |      2.11.1      |     1.13.1      |      ✘      |
+| `cmsml/cmsml:3.8`                                            |      3.8       |      2.13.1      | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.9`<br>`cmsml/cmsml:3`<br>`cmsml/cmsml:latest` |      3.9       | latest (~2.16.1) | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.10`                                           |      3.10      | latest (~2.16.1) | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.11`                                           |      3.11      | latest (~2.16.1) | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.11-cuda`                                      |      3.11      | latest (~2.16.1) | latest (~2.3.0) |      ✔︎      |
+
 <!-- marker-after-docker -->
 
 <!-- marker-before-testing -->

cmsml/tensorflow/__init__.py

@@ -6,16 +6,16 @@
 """
 
 __all__ = [
-    "import_tf", "save_frozen_graph", "save_graph", "load_frozen_graph", "load_graph",
-    "write_graph_summary", "load_model", "load_graph_def",
+    "import_tf", "tf_version_check", "tf_keras_version_check", "save_frozen_graph", "save_graph",
+    "load_frozen_graph", "load_graph", "write_graph_summary", "load_model", "load_graph_def",
     "OpsData", "get_graph_ops",
 ]
 
 
 # provisioning imports
 from cmsml.tensorflow.tools import (
-    import_tf, save_frozen_graph, save_graph, load_frozen_graph, load_graph, write_graph_summary,
-    load_model, load_graph_def,
+    import_tf, tf_version_check, tf_keras_version_check, save_frozen_graph, save_graph,
+    load_frozen_graph, load_graph, write_graph_summary, load_model, load_graph_def,
 )
 
 from cmsml.tensorflow.aot import (