Make sure TensorFlow also works on offline machines

eessi/testsuite/tests/apps/tensorflow/src/mnist_setup.py

@@ -1,9 +1,20 @@
+import os
+
 import tensorflow as tf
 import numpy as np
 
+if os.environ.get('EESSI_TEST_SUITE_DISABLE_DOWNLOAD') == 'True':
+    eessi_test_suite_download = False
+else:
+    eessi_test_suite_download = True
+
 
 def mnist_dataset(batch_size, test_batch_size):
-    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
+    if eessi_test_suite_download:
+        (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
+    else:
+        tensorflow_data = os.environ['RFM_TENSORFLOW_DATA']
+        (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(path=tensorflow_data)
     # The `x` arrays are in uint8 and have values in the [0, 255] range.
     # You need to convert them to float32 with values in the [0, 1] range.
     x_train = x_train / np.float32(255)

eessi/testsuite/tests/apps/tensorflow/tensorflow.py

@@ -3,6 +3,7 @@
 The test itself is based on an official multi-worker with Keras tutoral at
 https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
 """
+import os
 
 import reframe as rfm
 from reframe.core.builtins import deferrable, parameter, run_after, sanity_function, performance_function
@@ -97,3 +98,21 @@ def set_thread_count_args(self):
         self.executable_opts += ['--intra-op-parallelism', '%s' % self.num_cpus_per_task]
         self.executable_opts += ['--inter-op-parallelism', '1']
         utils.log(f'executable_opts set to {self.executable_opts}')
+
+    @run_after('setup')
+    def set_up_offline_run(self):
+        """
+        Set environments variables to run offline or skip the test
+        """
+        if 'offline' in self.current_partition.features:
+            resourcesdir = self.current_system.resourcesdir
+            data = os.path.join(resourcesdir, self.module_name, 'datasets/mnist.npz')
+            if os.path.exists(data):
+                self.env_vars['EESSI_TEST_SUITE_DISABLE_DOWNLOAD'] = 'True'
+                self.env_vars['RFM_TENSORFLOW_DATA'] = data
+            else:
+                msg = f'{self.current_partition.name} does not have internet access. '
+                msg += f'And could not find {data}. '
+                msg += 'You can download the file running tf.keras.datasets.mnist.load_data() '
+                msg += f'with {self.module_name} on a system with internet access.'
+                self.skip_if(self.module_name == self.module_name, msg)