Dawn HPC improvements to pipeline

.gitignore

@@ -30,9 +30,8 @@ loader.*.json
 *.out
 tmp.*
 *.swp
-*test*
 *.png
 
 !ENVS.example
 ENVS.*
-
+/tensorboard

check_and_move_data.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+HEMI=$1
+SPLIT=$2
+DATASET="$3_$HEMI"
+ERROR_FOLDER=network_datasets/${DATASET}/${HEMI}/${SPLIT}.data_errors
+CHECK_LOG=logs/check.${DATASET}.${SPLIT}.log
+
+if [ ! -f $CHECK_LOG ]; then
+    icenet_dataset_check -s $SPLIT dataset_config.${DATASET}.json 2>&1 | tee $CHECK_LOG
+fi
+
+mkdir $ERROR_FOLDER
+
+for FILENAME in $( grep 'WARNING' $CHECK_LOG | sed -r \
+  -e 's/^.+([0-9]{8}\.tfrecord).+$/\1/' \
+   | uniq ); do
+    if [ -f network_datasets/${DATASET}/${HEMI}/${SPLIT}/$FILENAME ]; then
+        echo mv -v network_datasets/${DATASET}/${HEMI}/${SPLIT}/$FILENAME $ERROR_FOLDER;
+    fi
+done
+
+mv -v $CHECK_LOG $ERROR_FOLDER

ensemble/predict.tmpl.yaml

@@ -13,7 +13,12 @@ ensemble:
       - ../../../network_datasets
       - ../../../processed
       - ../../../results
-    mem:          224gb
+    cluster:      short
+    email:        someone@example.com
+    length:       00:30:00
+    mem:          128gb
+    nodes:        1
+    ntasks:       2
 
   pre_process:          []
   post_process:         []
@@ -22,12 +27,7 @@ ensemble:
     templatedir:  ../template
     templates:
     - icenet_predict.sh.j2
-    email:        someone@example.com
     job_file:     icenet_predict.sh
-    cluster:      short
-    nodes:        1
-    ntasks:       8
-    length:       00:30:00
     maxruns:      100
     maxjobs:      10
 

ensemble/template/dawn.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+source $HOME/.bashrc
+
+module purge
+module load default-dawn
+module load dawn-env/2024-04-15 intel-oneapi-ccl intel-oneapi-compilers intel-oneapi-dnn intel-oneapi-dpct intel-oneapi-dpl intel-oneapi-inspector intel-oneapi-mkl intel-oneapi-mpi intel-oneapi-tbb
+

ensemble/train.tmpl.yaml

@@ -15,8 +15,13 @@ ensemble:
       - ../../../network_datasets
       - ../../../processed
       - ../../../results
+    cluster:      long
+    email:        someone@example.com
     gpus:         1
+    length:       4-00:00:00
     mem:          128gb
+    nodes:        1
+    ntasks:       NTASKS
 
   pre_process:
     - name:   execute
@@ -29,12 +34,7 @@ ensemble:
     templatedir:  ../template
     templates:
     - icenet_train.sh.j2
-    email:        someone@example.com
     job_file:     icenet_train.sh
-    cluster:      gpu
-    nodes:        1
-    ntasks:       NTASKS
-    length:       4-00:00:00
     maxruns:      5
     maxjobs:      MAXJOBS
 

environment.dawn.yml

@@ -0,0 +1,11 @@
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - cartopy
+  - eccodes
+  - ffmpeg
+  - hdf5
+  - netcdf4
+  - openh264
+  - xarray

quick_diag.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+
+# set -u -o pipefail
+STRATEGY=${1:-mirrored}
+GPUS=${2:-4}
+
+# srun --gres=gpu:4 --job-name=icenet-test --partition=pvc --nodes=1 --time=01:00:00 --pty bash -i
+LOGNAME="logs/$STRATEGY.$GPUS.`uuidgen`.log"
+
+{
+. ENVS
+conda activate $ICENET_CONDA
+echo "START: `date +%s`"
+icenet_train -b 4 -e 1 -f 1 -n $FILTER_FACTOR -s $STRATEGY --gpus $GPUS -nw --lr 25e-5 -v  exp23_south test_south1 42
+echo "END: `date +%s`"
+} 2>&1 | tee $LOGNAME

run_data.sh

@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/bin/bash -l
 
 . ENVS
 
@@ -17,11 +17,6 @@ BATCH_SIZE=${2:-2}
 WORKERS=${3:-8}
 
 if [ ! -f loader.${DATANAME}_${HEMI}.json ]; then
-    [ ! -z "$PROC_ARGS_ERA5" ] && icenet_process_era5 -v -l $LAG \
-        $PROC_ARGS_ERA5 \
-        -ns $TRAIN_START -ne $TRAIN_END -vs $VAL_START -ve $VAL_END -ts $TEST_START -te $TEST_END \
-        ${DATANAME}_${HEMI} $HEMI
-
     [ ! -z "$PROC_ARGS_ORAS5" ] && icenet_process_oras5 -v -l $LAG \
         $PROC_ARGS_ORAS5 \
         -ns $TRAIN_START -ne $TRAIN_END -vs $VAL_START -ve $VAL_END -ts $TEST_START -te $TEST_END \
@@ -32,6 +27,11 @@ if [ ! -f loader.${DATANAME}_${HEMI}.json ]; then
         -ns $TRAIN_START -ne $TRAIN_END -vs $VAL_START -ve $VAL_END -ts $TEST_START -te $TEST_END \
         ${DATANAME}_${HEMI} $HEMI
 
+    [ ! -z "$PROC_ARGS_ERA5" ] && icenet_process_era5 -v -l $LAG \
+        $PROC_ARGS_ERA5 \
+        -ns $TRAIN_START -ne $TRAIN_END -vs $VAL_START -ve $VAL_END -ts $TEST_START -te $TEST_END \
+        ${DATANAME}_${HEMI} $HEMI
+
     icenet_process_metadata ${DATANAME}_${HEMI} $HEMI
 else
     echo "Skipping preprocessing as loader.${DATANAME}_${HEMI}.json already exists..."

run_predict_ensemble.sh

@@ -19,6 +19,7 @@ ENSEMBLE_SEEDS_DEFAULT=42,46,45
 while getopts ":b:df:i:lm:p:r:x" opt; do
   case "$opt" in
     b)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}arg_batch=$OPTARG ";;
+    c)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}cluster=$OPTARG ";;
     d)  ENSEMBLE_TARGET="dummy";;
     f)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}arg_filter_factor=$OPTARG ";;
     i)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}arg_ident=$OPTARG ";;

run_train_ensemble.sh

@@ -29,6 +29,7 @@ while getopts ":b:c:de:f:g:j:l:m:n:p:q:r:s:t:" opt; do
     l)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}arg_preload=$OPTARG ";;
     m)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}mem=$OPTARG ";;
     n)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}nodelist=$OPTARG ";;
+    o)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}nodes=$OPTARG ";;
     p)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}arg_prep=$OPTARG ";;
     q)  ENSEMBLE_ARGS="${ENSEMBLE_ARGS}arg_queue=$OPTARG ";;
     r)  ENSEMBLE_RUNS=$OPTARG ;; # Ensemble member run seed values

scripts/horovod_north.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+#SBATCH --job-name=hnorth1
+#SBATCH --partition=pvc
+#SBATCH --nodes=2
+#SBATCH --ntasks=16
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:4
+#SBATCH --cpus-per-task=12          # split from 96 cores
+#SBATCH --time=1-00:00:00             # job length
+#SBATCH --output=logs/train.north.%j.out
+#SBATCH --error=logs/train.north.%j.err
+
+source $HOME/.bashrc
+
+module purge
+module load default-dawn
+module load dawn-env/2024-04-15 intel-oneapi-ccl intel-oneapi-compilers intel-oneapi-dnn intel-oneapi-dpct intel-oneapi-dpl intel-oneapi-inspector intel-oneapi-mkl intel-oneapi-mpi intel-oneapi-tbb
+
+conda activate icenet
+
+mpirun -np 16 icenet_train_horovod --device-type XPU -v --early-stopping 5 -wp test -wu jambyr --shuffle-train -e 100 -b 4 -n 1.44 dataset_config.full_train_north.json hv_north1 42

scripts/horovod_slurm_sbatch.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+#SBATCH --job-name=icy_test
+#SBATCH --partition=pvc
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=4
+#SBATCH --gres=gpu:4
+#SBATCH --cpus-per-task=24          # split from 96 cores
+#SBATCH --time=08:00:00             # job length
+#SBATCH --output=train.%j.out
+#SBATCH --error=train.%j.err
+
+source $HOME/.bashrc
+
+module purge
+module load default-dawn
+module load dawn-env/2024-04-15 intel-oneapi-ccl intel-oneapi-compilers intel-oneapi-dnn intel-oneapi-dpct intel-oneapi-dpl intel-oneapi-inspector intel-oneapi-mkl intel-oneapi-mpi intel-oneapi-tbb
+
+conda activate icenet
+
+mpirun -np 8 python scripts/horovod_test.py

scripts/horovod_small_test.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+#SBATCH --job-name=hsm1
+#SBATCH --partition=pvc
+#SBATCH --nodes=1
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=2
+#SBATCH --gres=gpu:4
+#SBATCH --exclusive
+#SBATCH --cpus-per-task=24          # split from 96 cores
+#SBATCH --time=12:00:00             # job length
+#SBATCH --output=logs/train.small_north_test.%j.out
+#SBATCH --error=logs/train.small_north_test.%j.err
+
+source $HOME/.bashrc
+
+module purge
+module load default-dawn
+module load dawn-env/2024-04-15 intel-oneapi-ccl intel-oneapi-compilers intel-oneapi-dnn intel-oneapi-dpct intel-oneapi-dpl intel-oneapi-inspector intel-oneapi-mkl intel-oneapi-mpi intel-oneapi-tbb
+
+conda activate icenet
+
+mpirun -np 2 icenet_train_horovod --device-type XPU -v --early-stopping 5 -wp test -wu jambyr --shuffle-train -e 3 -b 4 -n 1.44 dataset_config.full_train_north.json hv_small_test1 42

scripts/horovod_south.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+#SBATCH --job-name=hsouth1
+#SBATCH --partition=pvc
+#SBATCH --nodes=2
+#SBATCH --ntasks=16
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:4
+#SBATCH --cpus-per-task=12          # split from 96 cores
+#SBATCH --time=1-00:00:00             # job length
+#SBATCH --output=logs/train.south.%j.out
+#SBATCH --error=logs/train.south.%j.err
+
+source $HOME/.bashrc
+
+module purge
+module load default-dawn
+module load dawn-env/2024-04-15 intel-oneapi-ccl intel-oneapi-compilers intel-oneapi-dnn intel-oneapi-dpct intel-oneapi-dpl intel-oneapi-inspector intel-oneapi-mkl intel-oneapi-mpi intel-oneapi-tbb
+
+conda activate icenet
+
+mpirun -np 16 icenet_train_horovod --device-type XPU -v --early-stopping 5 -wp test -wu jambyr --shuffle-train -e 100 -b 4 -n 1.44 dataset_config.full_train_south.json hv_south1 42

scripts/horovod_test.py

@@ -0,0 +1,61 @@
+import logging
+logging.basicConfig(level=logging.DEBUG)
+import tensorflow as tf
+import horovod.tensorflow.keras as hvd
+from tensorflow.keras.optimizers import Adam
+
+hvd.init()
+
+# https://www.tensorflow.org/guide/keras/distributed_training
+
+# Create a MirroredStrategy.
+gpus = tf.config.list_physical_devices('XPU')
+print("XPU count is {}".format(len(gpus)))
+gpu_ids = []
+for gpu in gpus:
+    tf.config.experimental.set_memory_growth(gpu, True)
+if gpus:
+    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'XPU')
+
+from icenet.data.dataset import IceNetDataSet
+import icenet.model.losses as losses
+import icenet.model.metrics as metrics
+import icenet.model.networks.tensorflow as models
+
+batch_size = 1
+dataset = IceNetDataSet("dataset_config.full_train_north.json", batch_size=batch_size, shuffling=True)
+input_shape = (*dataset.shape, dataset.num_channels)
+
+loss = losses.WeightedMSE()
+metrics_list = [
+    metrics.WeightedBinaryAccuracy(),
+    metrics.WeightedMAE(),
+    metrics.WeightedRMSE(),
+    losses.WeightedMSE()
+]
+network = models.unet_batchnorm(
+    custom_optimizer=hvd.DistributedOptimizer(Adam(0.001)),
+    experimental_run_tf_function=False,
+    input_shape=input_shape,
+    loss=loss,
+    metrics=metrics_list,
+    filter_size=3,
+    n_filters_factor=1.44,
+    n_forecast_days=dataset.n_forecast_days,
+)
+
+network.summary()
+train_ds, val_ds, test_ds = dataset.get_split_datasets(ratio=1.0)
+
+model_history = network.fit(
+        #strategy.experimental_distribute_dataset(train_ds),
+        train_ds,
+        epochs=100,
+        steps_per_epoch=dataset.counts["train"] // (batch_size * hvd.size()),
+        verbose=1 if hvd.rank() == 0 else 0,
+        callbacks=[
+            hvd.callbacks.BroadcastGlobalVariablesCallback(0),    
+        ],
+        validation_data=val_ds)
+        #validation_data=strategy.experimental_distribute_dataset(val_ds),
+        #max_queue_size=10)

scripts/tensorflow_resnet_dist.py

@@ -0,0 +1,65 @@
+def main():
+    import logging
+    logging.basicConfig(level=logging.DEBUG)
+    import tensorflow as tf, numpy as np
+    # tf.config.run_functions_eagerly(True)
+    from icenet.model.train import train_model
+    from icenet.data.dataset import IceNetDataSet
+    import icenet.model.losses as losses
+    import icenet.model.metrics as metrics
+    import icenet.model.models as models
+    from tensorflow.keras.models import Model
+    from tensorflow.keras.layers import Conv2D, BatchNormalization, UpSampling2D, \
+        concatenate, MaxPooling2D, Input
+    from tensorflow.keras.optimizers import Adam
+    # Create a MirroredStrategy.
+    gpus = tf.config.list_physical_devices('XPU')
+    print("XPU count is {}".format(len(gpus)))
+    gpu_ids = []
+    for gpu in gpus:
+        print("Setting memory growth for XPU: {}".format(gpu))
+        tf.config.experimental.set_memory_growth(gpu, True)
+        gpu_ids.append(gpu.name[-5:])
+    strategy = tf.distribute.MirroredStrategy(gpu_ids)
+    print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
+    #dataset = IceNetDataSet("dataset_config.exp23_south.json", batch_size=32, shuffling=False)
+    #input_shape = (*dataset.shape, dataset.num_channels)
+    (mnist_images, mnist_labels), _ = \
+        tf.keras.datasets.mnist.load_data(path='mnist.npz')
+    dataset = tf.data.Dataset.from_tensor_slices(
+        ((mnist_images[..., tf.newaxis] / 255.0).astype(np.float32),
+        tf.cast(mnist_labels, tf.int32)))
+    dataset = dataset.repeat().shuffle(1000).batch(1000)
+    with strategy.scope():
+        #loss = losses.WeightedMSE()
+        #metrics_list = [
+        #    losses.WeightedMSE()
+        #]
+        model = tf.keras.Sequential([
+            tf.keras.layers.Input(shape=(28, 28, 1,)),
+            tf.keras.layers.Conv2D(32, [3, 3], activation='relu'),
+            tf.keras.layers.Conv2D(64, [3, 3], activation='relu'),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Dropout(0.25),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(128, activation='relu'),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(10, activation='softmax')
+        ])
+        opt = tf.optimizers.Adam(0.001)
+        model.compile(loss=tf.losses.SparseCategoricalCrossentropy(),
+                        optimizer=opt,
+                        metrics=['accuracy'],
+                        experimental_run_tf_function=False)
+    model.summary()
+    tboard_callback = tf.keras.callbacks.TensorBoard(log_dir = "tensorboard",
+                                                     histogram_freq = 1,
+                                                     profile_batch = 5)
+    #train_ds, val_ds, test_ds = dataset.get_split_datasets(ratio=1.0)
+    #model_history = network.fit(
+    #        train_ds,
+    #        epochs=5,
+    #        verbose=2,
+    #        callbacks=list(),
+    #        validation_data=val_ds)
+    model.fit(dataset, epochs=12, verbose=1, steps_per_epoch=50, callbacks=[tboard_callback])