Repository: GreenScheduler/cats

Issue:
scheduler: add sbatch

Discussion:

Solution patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a395542..ffd5104 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -2,14 +2,14 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v5.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-yaml
       - id: check-added-large-files
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.3.2
+    rev: v0.9.9
     hooks:
       - id: ruff
         args: [ --fix ]
diff --git a/cats/__init__.py b/cats/__init__.py
index 7210aa7..b571c28 100644
--- a/cats/__init__.py
+++ b/cats/__init__.py
@@ -20,7 +20,7 @@
 
 # To add a scheduler, add a date format here
 # and create a scheduler_<new>(...) function
-SCHEDULER_DATE_FORMAT = {"at": "%Y%m%d%H%M"}
+SCHEDULER_DATE_FORMAT = {"at": "%Y%m%d%H%M", "sbatch": "%Y-%m-%dT%H:%M"}
 
 
 def indent_lines(lines, spaces):
@@ -117,8 +117,8 @@ def positive_integer(string):
         "-s",
         "--scheduler",
         type=str,
-        help="Pass command using `-c` to scheduler. Currently, the only supported scheduler is at",
-        choices=["at"],
+        help="Pass command using `-c` to scheduler.",
+        choices=["at", "sbatch"],
     )
     parser.add_argument(
         "-a",
@@ -263,9 +263,7 @@ def print_banner(disable_colour):
         print(CATS_ASCII_BANNER_COLOUR)
 
 
-def schedule_at(
-    output: CATSOutput, args: list[str], at_command: str = "at"
-) -> Optional[str]:
+def schedule_at(output: CATSOutput, args: list[str]) -> Optional[str]:
     """Schedule job with optimal start time using at(1)
 
     :return: Error as a string, or None if successful
@@ -274,7 +272,7 @@ def schedule_at(
     try:
         subprocess.check_output(
             (
-                at_command,
+                "at",
                 "-t",
                 output.carbonIntensityOptimal.start.strftime(
                     SCHEDULER_DATE_FORMAT["at"]
@@ -289,6 +287,29 @@ def schedule_at(
         return f"Scheduling with at failed with code {e.returncode}, see output below:\n{e.output}"
 
 
+def schedule_sbatch(output: CATSOutput, args: list[str]) -> Optional[str]:
+    """Schedule job with optimal start time using sbatch(1)
+
+    :return: Error as a string, or None if successful
+    """
+    try:
+        subprocess.check_output(
+            [
+                "sbatch",
+                "--begin",
+                output.carbonIntensityOptimal.start.strftime(
+                    SCHEDULER_DATE_FORMAT["sbatch"]
+                ),
+                *args,
+            ]
+        )
+        return None
+    except FileNotFoundError:
+        return "No sbatch command found in PATH, ensure slurm is configured correctly"
+    except subprocess.CalledProcessError as e:  # pragma: no cover
+        return f"Scheduling with sbatch failed with code {e.returncode}, see output below:\n{e.output}"
+
+
 def main(arguments=None) -> int:
     parser = parse_arguments()
     args = parser.parse_args(arguments)
@@ -299,7 +320,7 @@ def main(arguments=None) -> int:
 
     if args.command and not args.scheduler:
         print(
-            "cats: To run a command with the -c or --command option, you must\n"
+            "cats: To run a command or sbatch script with the -c or --command option, you must\n"
             "      specify the scheduler with the -s or --scheduler option"
         )
         return 1
@@ -361,8 +382,14 @@ def main(arguments=None) -> int:
         print(output.to_json(dateformat, sort_keys=True, indent=2))
     else:
         print(output)
-    if args.command and args.scheduler == "at":
-        if err := schedule_at(output, args.command.split()):
+    if args.command:
+        if args.scheduler == "at":
+            err = schedule_at(output, args.command.split())
+        elif args.scheduler == "sbatch":
+            err = schedule_sbatch(output, args.command.split())
+        else:  # pragma: no cover - we already check for valid scheduler in parse_arguments
+            err = f"Scheduler {args.scheduler} not in supported schedulers: {SCHEDULER_DATE_FORMAT.keys()}"
+        if err:
             print(err)
             return 1
     return 0
diff --git a/docs/source/use-with-schedulers.rst b/docs/source/use-with-schedulers.rst
index 8fae568..968ff1a 100644
--- a/docs/source/use-with-schedulers.rst
+++ b/docs/source/use-with-schedulers.rst
@@ -19,6 +19,16 @@ duration, in the 'OX1' postcode that would look like:
 
    $ cats -d 5 --loc OX1 --scheduler at --command 'ls'
 
+Using CATS with the ``sbatch`` scheduler
+----------------------------------------
+
+CATS provides a wrapper around ``sbatch`` to enable deferring job execution
+till the optimal start time:
+
+.. code-block:: console
+
+   $ cats -d <job_duration> --loc <postcode> --scheduler sbatch --command ./script.sh
+
 
 Demonstration
 ^^^^^^^^^^^^^


Test patch:
diff --git a/tests/test_main.py b/tests/test_main.py
index f14b9c8..50dd6aa 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -3,8 +3,18 @@
 from datetime import datetime, timedelta
 from unittest.mock import patch
 
-from cats import SCHEDULER_DATE_FORMAT, CATSOutput, main, schedule_at
+import pytest
+
+from cats import (
+    SCHEDULER_DATE_FORMAT,
+    CATSOutput,
+    main,
+    print_banner,
+    schedule_at,
+    schedule_sbatch,
+)
 from cats.CI_api_interface import API_interfaces, InvalidLocationError
+from cats.constants import CATS_ASCII_BANNER_COLOUR, CATS_ASCII_BANNER_NO_COLOUR
 from cats.forecast import CarbonIntensityAverageEstimate
 
 API = API_interfaces["carbonintensity.org.uk"]
@@ -21,6 +31,52 @@
 )
 
 
+@pytest.mark.parametrize("disable_colour", [False, True])
+def test_print_banner(disable_colour, capsys):
+    print_banner(disable_colour)
+    expected_output = (
+        CATS_ASCII_BANNER_NO_COLOUR if disable_colour else CATS_ASCII_BANNER_COLOUR
+    )
+    assert capsys.readouterr().out.strip() == expected_output.strip()
+
+
+def test_schedule_sbatch_success(fp):
+    fp.register_subprocess(
+        [
+            "sbatch",
+            "--begin",
+            OUTPUT.carbonIntensityOptimal.start.strftime(
+                SCHEDULER_DATE_FORMAT["sbatch"]
+            ),
+            "./script.sh",
+        ],
+        stdout=b"Submitted batch job 123456",
+    )
+    schedule_sbatch(OUTPUT, ["./script.sh"])
+
+
+def test_schedule_sbatch_failure():
+    assert schedule_sbatch(OUTPUT, ["./script.sh"])
+
+
+@pytest.mark.parametrize(
+    "exc,err",
+    [
+        (
+            FileNotFoundError,
+            "No sbatch command found in PATH, ensure slurm is configured correctly",
+        ),
+        (
+            subprocess.CalledProcessError(1, "sbatch"),
+            "Scheduling with sbatch failed with code 1, see output below:\nNone",
+        ),
+    ],
+)
+def test_schedule_sbatch_side_effects(exc, err):
+    with patch("subprocess.check_output", side_effect=exc):
+        assert schedule_sbatch(OUTPUT, ["./script.sh"]) == err
+
+
 def test_schedule_at_success(fp):
     fp.register_subprocess(["ls"], stdout=b"foobar.txt")
     fp.register_subprocess(
@@ -40,17 +96,19 @@ def test_schedule_at_success(fp):
     )
 
 
-def test_schedule_at_missing():
-    assert (
-        schedule_at(OUTPUT, ["ls"], at_command="at_imaginary")
-        == "No at command found in PATH, please install one"
-    )
-
-
-def test_schedule_at_failure():
-    assert schedule_at(OUTPUT, ["ls"], at_command="/usr/bin/false").startswith(
-        "Scheduling with at failed with code 1, see output below:"
-    )
+@pytest.mark.parametrize(
+    "exc,err",
+    [
+        (FileNotFoundError, "No at command found in PATH, please install one"),
+        (
+            subprocess.CalledProcessError(1, "at"),
+            "Scheduling with at failed with code 1, see output below:\nNone",
+        ),
+    ],
+)
+def test_schedule_at_side_effects(exc, err):
+    with patch("subprocess.check_output", side_effect=exc):
+        assert schedule_at(OUTPUT, ["ls"]) == err
 
 
 def raiseLocationError():