From f749954f1df550d54214749b3e4c7f130d23f5ea Mon Sep 17 00:00:00 2001
From: Scott Davidson <scott@stackhpc.com>
Date: Fri, 22 Sep 2023 15:31:25 +0100
Subject: [PATCH] Add pytorch-benchmark image

---
 .github/workflows/publish-benchmark-images.yaml |  1 +
 images/pytorch-benchmarks/Dockerfile            | 13 +++++++++++++
 images/pytorch-benchmarks/run-benchmark.sh      |  3 +++
 3 files changed, 17 insertions(+)
 create mode 100644 images/pytorch-benchmarks/Dockerfile
 create mode 100644 images/pytorch-benchmarks/run-benchmark.sh

diff --git a/.github/workflows/publish-benchmark-images.yaml b/.github/workflows/publish-benchmark-images.yaml
index a5a1bcf..20fbede 100644
--- a/.github/workflows/publish-benchmark-images.yaml
+++ b/.github/workflows/publish-benchmark-images.yaml
@@ -21,6 +21,7 @@ jobs:
           - component: mpi-benchmarks
           - component: openfoam
           - component: perftest
+          - component: pytorch-benchmarks
     steps:
       - name: Check out the repository
         uses: actions/checkout@v2
diff --git a/images/pytorch-benchmarks/Dockerfile b/images/pytorch-benchmarks/Dockerfile
new file mode 100644
index 0000000..72fe8c4
--- /dev/null
+++ b/images/pytorch-benchmarks/Dockerfile
@@ -0,0 +1,13 @@
+# TODO: Try newer versions of base image - 23.08 fails with
+# 'CUDA on host is too old' type error but something in between
+# should work. Latest 23.08 might also work on Ubuntu-22.04.
+FROM nvcr.io/nvidia/pytorch:22.12-py3
+
+RUN pip install -U pip
+RUN pip install torchvision torchaudio
+RUN git clone https://github.com/pytorch/benchmark
+WORKDIR /workspace/benchmark
+# Add other benchmarks here?
+RUN python install.py alexnet resnet50 llama
+
+COPY run-benchmark.sh /usr/local/bin/
\ No newline at end of file
diff --git a/images/pytorch-benchmarks/run-benchmark.sh b/images/pytorch-benchmarks/run-benchmark.sh
new file mode 100644
index 0000000..11cfcaa
--- /dev/null
+++ b/images/pytorch-benchmarks/run-benchmark.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+# Env vars should be set in container spec
+python run.py $BENCHMARK_NAME -t $BENCHMARK_TYPE -d $BENCHMARK_DEVICE --bs $INPUT_DATA_BATCH_SIZE
\ No newline at end of file