From 1641f583eb7ceaed1ead406d467f6caefcc66c8a Mon Sep 17 00:00:00 2001
From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com>
Date: Wed, 6 Nov 2024 14:56:45 -0500
Subject: [PATCH] Add `--email` to `flepimop batch`

Added an email option for slurm jobs for notifying users of job status
via email if desired.
---
 flepimop/gempyor_pkg/src/gempyor/batch.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/flepimop/gempyor_pkg/src/gempyor/batch.py b/flepimop/gempyor_pkg/src/gempyor/batch.py
index 3401ba31b..57b8f21b6 100644
--- a/flepimop/gempyor_pkg/src/gempyor/batch.py
+++ b/flepimop/gempyor_pkg/src/gempyor/batch.py
@@ -700,6 +700,13 @@ def _resolve_batch_system(
             default=None,
             help="Unique prefix for this run.",
         ),
+        click.Option(
+            ["--email"],
+            "email",
+            type=str,
+            default=None,
+            help="Optionally an email that can be notified on job begin and end.",
+        ),
     ]
     + list(verbosity_options.values()),
 )
@@ -741,6 +748,13 @@ def _click_batch(ctx: click.Context = mock_context, **kwargs) -> None:
             "The `flepimop batch` CLI only supports batch submission to slurm."
         )
     logger.info("Constructing a job to submit to %s", batch_system)
+    if batch_system != "slurm" and kwargs["email"] is not None:
+        logger.warning(
+            "The email option, given '%s', is only used when "
+            "the batch system is slurm, but is instead %s.",
+            kwargs["email"],
+            batch_system,
+        )
 
     # Job size
     iterations_per_slot = (
@@ -837,6 +851,9 @@ def _click_batch(ctx: click.Context = mock_context, **kwargs) -> None:
     }
     if kwargs["partition"] is not None:
         options["partition"] = kwargs["partition"]
+    if kwargs["email"] is not None:
+        options["mail-type"] = "BEGIN,END"
+        options["mail-user"] = kwargs["email"]
     _sbatch_template(
         "emcee_inference.sbatch.j2",
         None,