From 01295487d8fe027d69fbc508dacb6d1ed9b34172 Mon Sep 17 00:00:00 2001
From: Rodrigo Arias <rodrigo.arias@bsc.es>
Date: Mon, 17 Aug 2020 18:50:18 +0200
Subject: [PATCH] Add srun wrapper and use pmi2

---
 bsc/garlic/argv.nix          |  9 +++-
 bsc/garlic/default.nix       |  9 ++--
 bsc/garlic/exp/nbody/bs.nix  | 25 +++++++----
 bsc/garlic/exp/nbody/mpi.nix | 81 ++++++++++++++++++++++++++++++++++++
 bsc/garlic/nbody/default.nix |  8 +++-
 bsc/garlic/sbatch.nix        |  4 +-
 bsc/garlic/srun.nix          | 20 +++++++++
 7 files changed, 138 insertions(+), 18 deletions(-)
 create mode 100644 bsc/garlic/exp/nbody/mpi.nix
 create mode 100644 bsc/garlic/srun.nix

diff --git a/bsc/garlic/argv.nix b/bsc/garlic/argv.nix
index 5978cb3c..701b8d74 100644
--- a/bsc/garlic/argv.nix
+++ b/bsc/garlic/argv.nix
@@ -1,9 +1,11 @@
 {
   stdenv
+, bash
 }:
 
 {
   app
+, env ? ""
 , argv # bash array as string, example: argv=''(-f "file with spaces" -t 10)''
 }:
 
@@ -12,11 +14,14 @@ stdenv.mkDerivation {
   name = "${app.name}-argv";
   preferLocalBuild = true;
   phases = [ "installPhase" ];
-  dontPatchShebangs = true;
   installPhase = ''
     mkdir -p $out/bin
     cat > $out/bin/run <<EOF
-    #!/bin/sh
+    #!${bash}/bin/bash
+    # Requires /nix to use bash
+    
+    ${env}
+
     argv=${argv}
     exec ${app}/bin/run \''${argv[@]}
     EOF
diff --git a/bsc/garlic/default.nix b/bsc/garlic/default.nix
index 4e222dfd..c975b259 100644
--- a/bsc/garlic/default.nix
+++ b/bsc/garlic/default.nix
@@ -11,6 +11,7 @@ let
 
     # Load some helper functions to generate app variants
     inherit (import ./gen.nix) genApps genApp genConfigs;
+    inherit bsc;
 
     mpptest = callPackage ./mpptest { };
 
@@ -24,16 +25,16 @@ let
     };
 
     sbatchWrapper = callPackage ./sbatch.nix { };
-    launcherWrapper = callPackage ./launcher.nix { };
+    srunWrapper = callPackage ./srun.nix { };
+    launchWrapper = callPackage ./launcher.nix { };
     controlWrapper = callPackage ./control.nix { };
     nixsetupWrapper = callPackage ./nix-setup.nix { };
     argvWrapper = callPackage ./argv.nix { };
 
     exp = {
       nbody = {
-        bs = callPackage ./exp/nbody/bs.nix {
-          inherit bsc;
-        };
+        bs = callPackage ./exp/nbody/bs.nix { };
+        mpi = callPackage ./exp/nbody/mpi.nix { };
       };
     };
   };
diff --git a/bsc/garlic/exp/nbody/bs.nix b/bsc/garlic/exp/nbody/bs.nix
index 2ffc8fe1..979b925a 100644
--- a/bsc/garlic/exp/nbody/bs.nix
+++ b/bsc/garlic/exp/nbody/bs.nix
@@ -5,8 +5,9 @@
 , genConfigs
 
 # Wrappers
-, launcherWrapper
+, launchWrapper
 , sbatchWrapper
+, srunWrapper
 , argvWrapper
 , controlWrapper
 , nixsetupWrapper
@@ -28,8 +29,8 @@ let
 
   # Compute the cartesian product of all configurations
   allConfigs = genConfigs config;
-
-  configs = with builtins; filter (c: c.blocksize <= 4096) allConfigs;
+  filteredConfigs = with builtins; filter (c: c.blocksize <= 4096) allConfigs;
+  configs = map (conf: conf // extraConfig) filteredConfigs;
 
   sbatch = conf: app: sbatchWrapper {
     app = app;
@@ -39,6 +40,11 @@ let
     chdirPrefix = "/home/bsc15/bsc15557/bsc-nixpkgs/out";
   };
 
+  srun = app: srunWrapper {
+    app = app;
+    nixPrefix = "/gpfs/projects/bsc15/nix";
+  };
+
   argv = conf: app:
     with conf;
     argvWrapper {
@@ -52,15 +58,16 @@ let
 
   pipeline = conf:
     sbatch conf (
-      nixsetupWrapper (
-        controlWrapper (
-          argv conf (
-            nbodyFn conf))));
+      srun (
+        nixsetupWrapper (
+          controlWrapper (
+            argv conf (
+              nbodyFn conf)))));
 
   # Ideally it should look like this:
   #pipeline = sbatch nixsetup control argv nbodyFn;
 
-  jobs = map (conf: pipeline (conf // extraConfig)) configs;
+  jobs = map pipeline configs;
 
 in
-  launcherWrapper jobs
+  launchWrapper jobs
diff --git a/bsc/garlic/exp/nbody/mpi.nix b/bsc/garlic/exp/nbody/mpi.nix
new file mode 100644
index 00000000..64f12797
--- /dev/null
+++ b/bsc/garlic/exp/nbody/mpi.nix
@@ -0,0 +1,81 @@
+{
+  bsc
+, nbody
+, genApp
+, genConfigs
+
+# Wrappers
+, launchWrapper
+, sbatchWrapper
+, srunWrapper
+, argvWrapper
+, controlWrapper
+, nixsetupWrapper
+}:
+
+let
+  # Set the configuration for the experiment
+  config = {
+    cc = [ bsc.icc ];
+    blocksize = [ 2048 ];
+  };
+
+  extraConfig = {
+    particles = 16384;
+    timesteps = 10;
+    ntasks = 2;
+    mpi = bsc.impi;
+    #mpi = bsc.openmpi;
+    gitBranch = "garlic/mpi+send";
+    gitURL = "ssh://git@bscpm02.bsc.es/garlic/apps/nbody.git";
+  };
+
+  # Compute the cartesian product of all configurations
+  configs = map (conf: conf // extraConfig) (genConfigs config);
+
+  sbatch = conf: app: sbatchWrapper {
+    app = app;
+    nixPrefix = "/gpfs/projects/bsc15/nix";
+    exclusive = false;
+    ntasks = "${toString conf.ntasks}";
+    chdirPrefix = "/home/bsc15/bsc15557/bsc-nixpkgs/out";
+  };
+
+  srun = app: srunWrapper {
+    app = app;
+    nixPrefix = "/gpfs/projects/bsc15/nix";
+  };
+
+  argv = conf: app:
+    with conf;
+    argvWrapper {
+      app = app;
+      argv = ''(-t ${toString timesteps} -p ${toString particles})'';
+      env = ''
+        export I_MPI_THREAD_SPLIT=1
+      '';
+    };
+
+  nbodyFn = conf:
+    with conf;
+    nbody.override { inherit cc mpi blocksize gitBranch gitURL; };
+
+  pipeline = conf:
+#    sbatch conf (
+      srun (
+        nixsetupWrapper (
+          argv conf (
+            nbodyFn conf
+          )
+        )
+      )
+#    )
+    ;
+
+  # Ideally it should look like this:
+  #pipeline = sbatch nixsetup control argv nbodyFn;
+
+  jobs = map pipeline configs;
+
+in
+  launchWrapper jobs
diff --git a/bsc/garlic/nbody/default.nix b/bsc/garlic/nbody/default.nix
index e6d5468b..5f0a3d0d 100644
--- a/bsc/garlic/nbody/default.nix
+++ b/bsc/garlic/nbody/default.nix
@@ -1,22 +1,26 @@
 {
   stdenv
 , cc
+, mpi ? null
 , cflags ? null
 , gitBranch
+, gitURL ? "ssh://git@bscpm02.bsc.es/rarias/nbody.git"
 , blocksize ? 2048
 }:
 
+with stdenv.lib;
 stdenv.mkDerivation rec {
   name = "nbody";
 
   src = builtins.fetchGit {
-    url = "ssh://git@bscpm02.bsc.es/rarias/nbody.git";
+    url = "${gitURL}";
     ref = "${gitBranch}";
   };
 
   buildInputs = [
     cc
-  ];
+  ]
+  ++ optional (mpi != null) [ mpi ];
 
   preBuild = (if cflags != null then ''
     makeFlagsArray+=(CFLAGS="${cflags}")
diff --git a/bsc/garlic/sbatch.nix b/bsc/garlic/sbatch.nix
index dc1603f1..f503e237 100644
--- a/bsc/garlic/sbatch.nix
+++ b/bsc/garlic/sbatch.nix
@@ -10,6 +10,7 @@
 , argv ? ""
 , binary ? "/bin/run"
 , ntasks ? null
+, nodes ? null
 , exclusive ? true # By default we run in exclusive mode
 , qos ? null
 , time ? null
@@ -50,6 +51,7 @@ stdenv.mkDerivation rec {
     #SBATCH --job-name="${name}"
     ''
     + sbatchOpt "ntasks" ntasks
+    + sbatchOpt "nodes" nodes
     + sbatchOpt "chdir" "${chdirPrefix}/$(basename $out)"
     + sbatchOpt "output" output
     + sbatchOpt "error" error
@@ -59,7 +61,7 @@ stdenv.mkDerivation rec {
     + optionalString (extra!=null) extra
     +
     ''
-    srun ${nixPrefix}${app}${binary} ${argv}
+    exec ${nixPrefix}${app}${binary} ${argv}
     EOF
     
     mkdir -p $out/bin
diff --git a/bsc/garlic/srun.nix b/bsc/garlic/srun.nix
new file mode 100644
index 00000000..3ede7375
--- /dev/null
+++ b/bsc/garlic/srun.nix
@@ -0,0 +1,20 @@
+{
+  stdenv
+}:
+{ app , nixPrefix ? "" }:
+
+stdenv.mkDerivation rec {
+  name = "${app.name}-srun";
+  preferLocalBuild = true;
+  phases = [ "installPhase" ];
+  buildInputs = [ app ];
+  dontPatchShebangs = true;
+  installPhase = ''
+    mkdir -p $out/bin
+    cat > $out/bin/run <<EOF
+    #!/bin/sh
+    exec srun --mpi=pmi2 ${nixPrefix}${app}/bin/run
+    EOF
+    chmod +x $out/bin/run
+  '';
+}