scheduler: better and more tests (#57)

* scheduler: better tests

Co-authored-by: Alexander Goscinski <[email protected]>

---------

Co-authored-by: Alexander Goscinski <[email protected]>

.firecrest-demo-config.json

@@ -7,5 +7,6 @@
     "temp_directory": "",
     "small_file_size_mb": 5.0,
     "workdir": "",
-    "api_version": "1.16.0"
+    "api_version": "1.16.0",
+    "builder_metadata_options_custom_scheduler_commands": []
 }

aiida_firecrest/scheduler.py

@@ -221,14 +221,17 @@ def get_jobs(
             try:
                 for page_iter in itertools.count():
                     results += transport._client.poll_active(
-                        transport._machine, jobs, page_number=page_iter
+                        transport._machine,
+                        jobs,
+                        page_number=page_iter,
+                        page_size=self._DEFAULT_PAGE_SIZE,
                     )
                     if len(results) < self._DEFAULT_PAGE_SIZE * (page_iter + 1):
                         break
             except FirecrestException as exc:
-                # firecrest returns error if the job is completed
                 # TODO: check what type of error is returned and handle it properly
-                if "Invalid job id specified" not in str(exc):
+                if "Invalid job id" not in str(exc):
+                    # firecrest returns error if the job is completed, while aiida expect a silent return
                     raise SchedulerError(str(exc)) from exc
         job_list = []
         for raw_result in results:

aiida_firecrest/transport.py

@@ -401,6 +401,8 @@ def __init__(
         # aiida-core/src/aiida/orm/utils/remote:clean_remote()
         self._is_open = True
 
+        self.checksum_check = False
+
     def __str__(self) -> str:
         """Return the name of the plugin."""
         return self.__class__.__name__
@@ -723,7 +725,8 @@ def getfile(
                 down_obj = self._client.external_download(self._machine, str(remote))
                 down_obj.finish_download(local)
 
-        self._validate_checksum(local, remote)
+        if self.checksum_check:
+            self._validate_checksum(local, remote)
 
     def _validate_checksum(
         self, localpath: str | Path, remotepath: str | FcPath
@@ -965,7 +968,8 @@ def putfile(
                 )
                 up_obj.finish_upload()
 
-        self._validate_checksum(localpath, str(remote))
+        if self.checksum_check:
+            self._validate_checksum(localpath, str(remote))
 
     def payoff(self, path: str | FcPath | Path) -> bool:
         """

tests/conftest.py

@@ -1,25 +1,27 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import hashlib
 import itertools
 import json
 import os
 from pathlib import Path
 import shutil
 import stat
-from typing import Any, Callable
+from typing import Any, Callable, ClassVar
+from unittest.mock import MagicMock
 from urllib.parse import urlparse
 
 from aiida import orm
 import firecrest
-import firecrest.path
 import pytest
 import requests
 
 
-class Values:
-    _DEFAULT_PAGE_SIZE: int = 25
+class Slurm:
+    """Save the submitted job ids for testing purposes."""
+
+    all_jobs: ClassVar[list] = []
 
 
 @pytest.fixture
@@ -75,7 +77,12 @@ def submit(
             raise DeprecationWarning("local_file is not supported")
 
         if script_remote_path and not Path(script_remote_path).exists():
-            raise FileNotFoundError(f"File {script_remote_path} does not exist")
+            # Firecrest raises FirecrestException instead of FileNotFoundError
+            mock_response = MagicMock()
+            mock_response.status_code = 999  # I don't really know
+            mock_response.json.return_value = {"error": "Mock error message"}
+            raise firecrest.FirecrestException(mock_response)
+
         job_id = next(self.job_id_generator)
 
         # Filter out lines starting with '#SBATCH'
@@ -98,9 +105,18 @@ def submit(
             os.chdir(Path(script_remote_path).parent)
             os.system(command)
 
+        Slurm.all_jobs.append(job_id)
+
         return {"jobid": job_id}
 
-    def poll_active(self, machine: str, jobs: list[str], page_number: int = 0):
+    def cancel(self, machine: str, job_id: str):
+        job_id = int(job_id)
+        if job_id in Slurm.all_jobs:
+            Slurm.all_jobs.remove(job_id)
+
+    def poll_active(
+        self, machine: str, jobs: list[str], page_number: int = 0, page_size: int = 25
+    ):
         response = []
         # 12 satets are defined in firecrest
         states = [
@@ -118,6 +134,11 @@ def poll_active(self, machine: str, jobs: list[str], page_number: int = 0):
             "COMPLETING",
         ]
         for i in range(len(jobs)):
+            if int(jobs[i]) not in Slurm.all_jobs:
+                mock_response = MagicMock()
+                mock_response.status_code = 999  # I don't really know
+                mock_response.json.return_value = {"error": "Invalid job id"}
+                raise firecrest.FirecrestException([mock_response])
             response.append(
                 {
                     "job_data_err": "",
@@ -139,11 +160,7 @@ def poll_active(self, machine: str, jobs: list[str], page_number: int = 0):
                 }
             )
 
-        return response[
-            page_number
-            * Values._DEFAULT_PAGE_SIZE : (page_number + 1)
-            * Values._DEFAULT_PAGE_SIZE
-        ]
+        return response[page_number * page_size : (page_number + 1) * page_size]
 
     def whoami(self, machine: str):
         assert machine == "MACHINE_NAME"
@@ -373,7 +390,22 @@ def __init__(self, *args, **kwargs):
 
 @dataclass
 class ComputerFirecrestConfig:
-    """Configuration of a computer using FirecREST as transport plugin."""
+    """Configuration of a computer using FirecREST as transport plugin.
+
+    :param url: The URL of the FirecREST server.
+    :param token_uri: The URI to receive  tokens.
+    :param client_id: The client ID for the client credentials.
+    :param client_secret: The client secret for the client credentials.
+    :param compute_resource: The name of the compute resource. This is the name of the machine.
+    :param temp_directory: A temporary directory on the machine for transient zip files.
+    :param workdir: The aiida working directory on the machine.
+    :param api_version: The version of the FirecREST API.
+    :param builder_metadata_options_custom_scheduler_commands: A list of custom
+           scheduler commands when submitting a job, for example
+           ["#SBATCH --account=mr32",
+            "#SBATCH --constraint=mc",
+            "#SBATCH --mem=10K"].
+    :param small_file_size_mb: The maximum file size for direct upload & download."""
 
     url: str
     token_uri: str
@@ -384,6 +416,9 @@ class ComputerFirecrestConfig:
     workdir: str
     api_version: str
     small_file_size_mb: float = 1.0
+    builder_metadata_options_custom_scheduler_commands: list[str] = field(
+        default_factory=list
+    )
 
 
 class RequestTelemetry:
@@ -514,4 +549,5 @@ def firecrest_config(
             small_file_size_mb=1.0,
             temp_directory=str(_temp_directory),
             api_version="2",
+            builder_metadata_options_custom_scheduler_commands=[],
         )

tests/test_calculation.py

@@ -20,8 +20,9 @@ def _no_retries():
     manage.get_config().set_option(MAX_ATTEMPTS_OPTION, max_attempts)
 
 
+@pytest.mark.timeout(180)
 @pytest.mark.usefixtures("aiida_profile_clean", "no_retries")
-def test_calculation_basic(firecrest_computer: orm.Computer):
+def test_calculation_basic(firecrest_computer: orm.Computer, firecrest_config):
     """Test running a simple `arithmetic.add` calculation."""
     code = orm.InstalledCode(
         label="test_code",
@@ -35,6 +36,10 @@ def test_calculation_basic(firecrest_computer: orm.Computer):
     builder = code.get_builder()
     builder.x = orm.Int(1)
     builder.y = orm.Int(2)
+    custom_scheduler_commands = "\n".join(
+        firecrest_config.builder_metadata_options_custom_scheduler_commands
+    )
+    builder.metadata.options.custom_scheduler_commands = custom_scheduler_commands
 
     _, node = engine.run_get_node(builder)
     assert node.is_finished_ok

tests/test_scheduler.py

@@ -1,47 +1,125 @@
 from pathlib import Path
+import textwrap
+from time import sleep
 
 from aiida import orm
+from aiida.schedulers import SchedulerError
 from aiida.schedulers.datastructures import CodeRunMode, JobTemplate
 import pytest
 
 from aiida_firecrest.scheduler import FirecrestScheduler
-from conftest import Values
 
 
 @pytest.mark.usefixtures("aiida_profile_clean")
-def test_submit_job(firecrest_computer: orm.Computer, tmp_path: Path):
+def test_submit_job(firecrest_computer: orm.Computer, firecrest_config, tmpdir: Path):
+    """Test submitting a job to the scheduler.
+    Note: this test relies on a functional transport.put() method."""
+
     transport = firecrest_computer.get_transport()
     scheduler = FirecrestScheduler()
     scheduler.set_transport(transport)
 
-    with pytest.raises(FileNotFoundError):
-        scheduler.submit_job(transport.getcwd(), "unknown.sh")
+    # raise error if file not found
+    with pytest.raises(SchedulerError):
+        scheduler.submit_job(firecrest_config.workdir, "unknown.sh")
+
+    custom_scheduler_commands = "\n    ".join(
+        firecrest_config.builder_metadata_options_custom_scheduler_commands
+    )
+
+    shell_script = f"""
+    #!/bin/bash
+    #SBATCH --no-requeue
+    #SBATCH --job-name="aiida-1928"
+    #SBATCH --get-user-env
+    #SBATCH --output=_scheduler-stdout.txt
+    #SBATCH --error=_scheduler-stderr.txt
+    #SBATCH --nodes=1
+    #SBATCH --ntasks-per-node=1
+    {custom_scheduler_commands}
+
+    echo 'hello world'
+    """
+
+    dedented_script = textwrap.dedent(shell_script).strip()
+    Path(tmpdir / "job.sh").write_text(dedented_script)
+    remote_ = transport._cwd.joinpath(firecrest_config.workdir, "job.sh")
+    transport.put(tmpdir / "job.sh", remote_)
 
-    _script = Path(tmp_path / "job.sh")
-    _script.write_text("#!/bin/bash\n\necho 'hello world'")
+    job_id = scheduler.submit_job(firecrest_config.workdir, "job.sh")
 
-    job_id = scheduler.submit_job(transport.getcwd(), _script)
-    # this is how aiida expects the job_id to be returned
     assert isinstance(job_id, str)
 
 
+@pytest.mark.timeout(180)
 @pytest.mark.usefixtures("aiida_profile_clean")
-def test_get_jobs(firecrest_computer: orm.Computer):
+def test_get_and_kill_jobs(
+    firecrest_computer: orm.Computer, firecrest_config, tmpdir: Path
+):
+    """Test getting and killing jobs from the scheduler.
+    We test the two together for performance reasons, as this test might run against
+      a real server and we don't want to leave parasitic jobs behind.
+      also less billing for the user.
+    Note: this test relies on a functional transport.put() method.
+    """
+    import time
+
     transport = firecrest_computer.get_transport()
     scheduler = FirecrestScheduler()
     scheduler.set_transport(transport)
 
-    # test pagaination
-    scheduler._DEFAULT_PAGE_SIZE = 2
-    Values._DEFAULT_PAGE_SIZE = 2
+    # verify that no error is raised in the case of an invalid job id 000
+    scheduler.get_jobs(["000"])
+
+    custom_scheduler_commands = "\n    ".join(
+        firecrest_config.builder_metadata_options_custom_scheduler_commands
+    )
+    shell_script = f"""
+    #!/bin/bash
+    #SBATCH --no-requeue
+    #SBATCH --job-name="aiida-1929"
+    #SBATCH --get-user-env
+    #SBATCH --output=_scheduler-stdout.txt
+    #SBATCH --error=_scheduler-stderr.txt
+    #SBATCH --nodes=1
+    #SBATCH --ntasks-per-node=1
+    {custom_scheduler_commands}
+
+    sleep 180
+    """
 
-    joblist = ["111", "222", "333", "444", "555"]
+    joblist = []
+    dedented_script = textwrap.dedent(shell_script).strip()
+    Path(tmpdir / "job.sh").write_text(dedented_script)
+    remote_ = transport._cwd.joinpath(firecrest_config.workdir, "job.sh")
+    transport.put(tmpdir / "job.sh", remote_)
+
+    for _ in range(5):
+        joblist.append(scheduler.submit_job(firecrest_config.workdir, "job.sh"))
+
+    # test pagaination is working
+    scheduler._DEFAULT_PAGE_SIZE = 2
     result = scheduler.get_jobs(joblist)
     assert len(result) == 5
     for i in range(5):
-        assert result[i].job_id == str(joblist[i])
+        assert result[i].job_id in joblist
         # TODO: one could check states as well
 
+    # test kill jobs
+    for jobid in joblist:
+        scheduler.kill_job(jobid)
+
+    # sometimes it takes time for the server to actually kill the jobs
+    timeout_kill = 5  # seconds
+    start_time = time.time()
+    while time.time() - start_time < timeout_kill:
+        result = scheduler.get_jobs(joblist)
+        if not len(result):
+            break
+        sleep(0.5)
+
+    assert not len(result)
+
 
 def test_write_script_full():
     # to avoid false positive (overwriting on existing file),
@@ -68,9 +146,9 @@ def test_write_script_full():
     #SBATCH --mem=1
     test_command
     """
-    expectaion_flat = "\n".join(line.strip() for line in expectaion.splitlines()).strip(
-        "\n"
-    )
+    expectation_flat = "\n".join(
+        line.strip() for line in expectaion.splitlines()
+    ).strip("\n")
     scheduler = FirecrestScheduler()
     template = JobTemplate(
         {
@@ -98,7 +176,7 @@ def test_write_script_full():
         }
     )
     try:
-        assert scheduler.get_submit_script(template).rstrip() == expectaion_flat
+        assert scheduler.get_submit_script(template).rstrip() == expectation_flat
     except AssertionError:
         print(scheduler.get_submit_script(template).rstrip())
         print(expectaion)
@@ -116,9 +194,9 @@ def test_write_script_minimal():
     #SBATCH --ntasks-per-node=1
     """
 
-    expectaion_flat = "\n".join(line.strip() for line in expectaion.splitlines()).strip(
-        "\n"
-    )
+    expectation_flat = "\n".join(
+        line.strip() for line in expectaion.splitlines()
+    ).strip("\n")
     scheduler = FirecrestScheduler()
     template = JobTemplate(
         {
@@ -131,7 +209,7 @@ def test_write_script_minimal():
     )
 
     try:
-        assert scheduler.get_submit_script(template).rstrip() == expectaion_flat
+        assert scheduler.get_submit_script(template).rstrip() == expectation_flat
     except AssertionError:
         print(scheduler.get_submit_script(template).rstrip())
         print(expectaion)