.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
       - id: debug-statements
       - id: check-ast
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.7.2"
+    rev: "v0.7.3"
     hooks:
       - id: ruff
         exclude: docs/source/conf.py|ipynb_filter.py

adaptive_scheduler/_executor.py

@@ -321,19 +321,22 @@ class SlurmExecutor(AdaptiveSchedulerExecutorBase):
     """
 
     # Same as slurm_run, except it has no learners, fnames, dependencies and initializers.
-    # Additionally, the type hints for scheduler arguments are singular instead of tuples.
 
     # slurm_run: Specific to slurm_run
     name: str = "adaptive-scheduler"
     folder: str | Path | None = None  # `slurm_run` defaults to None
     # slurm_run: SLURM scheduler arguments
-    partition: str | None = None
-    nodes: int | None = 1
-    cores_per_node: int | None = 1  # `slurm_run` defaults to `None`
-    num_threads: int = 1
-    exclusive: bool = False
-    executor_type: EXECUTOR_TYPES = "process-pool"
-    extra_scheduler: list[str] | None = None
+    partition: str | tuple[str | Callable[[], str], ...] | None = None
+    nodes: int | tuple[int | None | Callable[[], int | None], ...] | None = 1
+    cores_per_node: int | tuple[int | None | Callable[[], int | None], ...] | None = (
+        1  # `slurm_run` defaults to `None`
+    )
+    num_threads: int | tuple[int | Callable[[], int], ...] = 1
+    exclusive: bool | tuple[bool | Callable[[], bool], ...] = False
+    executor_type: EXECUTOR_TYPES | tuple[EXECUTOR_TYPES | Callable[[], EXECUTOR_TYPES], ...] = (
+        "process-pool"
+    )
+    extra_scheduler: list[str] | tuple[list[str] | Callable[[], list[str]], ...] | None = None
     # slurm_run: Same as RunManager below (except dependencies and initializers)
     goal: GoalTypes | None = None
     check_goal_on_start: bool = True

adaptive_scheduler/_scheduler/local.py

@@ -123,7 +123,7 @@ def start_job(self, name: str, *, index: int | None = None) -> None:
         submit_cmd = f"{self.submit_cmd} {name} {self.batch_fname(name_prefix)}"
         run_submit(submit_cmd, name)
 
-    def extra_scheduler(self, *, index: int | None = None) -> str:
+    def extra_scheduler(self, *, index: int | None = None) -> str:  # noqa: ARG002
         """Get the extra scheduler options."""
         msg = "extra_scheduler is not implemented."
         raise NotImplementedError(msg)

adaptive_scheduler/_scheduler/slurm.py

@@ -5,9 +5,9 @@
 import copy
 import getpass
 import re
+import shutil
 import subprocess
 import textwrap
-from distutils.spawn import find_executable
 from functools import cached_property, lru_cache
 from typing import TYPE_CHECKING, TypeVar
 
@@ -415,7 +415,7 @@ def queue(*, me_only: bool = True) -> dict[str, dict[str, str]]:
         }  # (key -> length) mapping
 
         slurm_format = ",".join(f"{k}:{v}" for k, v in python_format.items())
-        squeue_executable = find_executable("squeue")
+        squeue_executable = shutil.which("squeue")
         assert isinstance(squeue_executable, str)
         cmd = [
             squeue_executable,

adaptive_scheduler/_server_support/multi_run_manager.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING
 
 import ipywidgets as ipw
+from IPython.display import display
 
 from adaptive_scheduler.widgets import _disable_widgets_output_scrollbar, info
 
@@ -128,7 +129,7 @@ def _create_widget(self) -> ipw.VBox:
         children = list(self._info_widgets.values())
         self._tab_widget.children = children
         for i, name in enumerate(self.run_managers.keys()):
-            self._tab_widget.set_title(i, f"RunManager: {name}")
+            self._tab_widget.set_title(i, name)
 
         self._update_all_button = ipw.Button(
             description="Update All",
@@ -158,7 +159,7 @@ def _update_widget(self) -> None:
 
             # Update titles
             for i, name in enumerate(self.run_managers.keys()):
-                self._tab_widget.set_title(i, f"RunManager: {name}")
+                self._tab_widget.set_title(i, name)
 
     def _update_all_callback(self, _: ipw.Button) -> None:
         """Callback function for the Update All button."""
@@ -178,3 +179,7 @@ def info(self) -> ipw.VBox:
     def _repr_html_(self) -> str:
         """HTML representation for Jupyter notebooks."""
         return self.info()
+
+    def display(self) -> None:
+        """Display the widget."""
+        display(self.info())

adaptive_scheduler/scheduler.py

@@ -4,8 +4,8 @@
 
 import os
 import os.path
+import shutil
 import warnings
-from distutils.spawn import find_executable
 
 from adaptive_scheduler._scheduler.base_scheduler import BaseScheduler
 from adaptive_scheduler._scheduler.local import LocalMockScheduler
@@ -38,8 +38,8 @@ def _get_default_scheduler() -> type[BaseScheduler]:
 
     By default it is "SLURM".
     """
-    has_pbs = bool(find_executable("qsub")) and bool(find_executable("qstat"))
-    has_slurm = bool(find_executable("sbatch")) and bool(find_executable("squeue"))
+    has_pbs = bool(shutil.which("qsub")) and bool(shutil.which("qstat"))
+    has_slurm = bool(shutil.which("sbatch")) and bool(shutil.which("squeue"))
 
     default = SLURM
     default_msg = f"We set DefaultScheduler to '{default}'."

example.ipynb

@@ -23,25 +23,21 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "\n",
     "import adaptive_scheduler\n",
-    "\n",
+    "import random\n",
     "\n",
     "def h(x, width=0.01, offset=0):\n",
     "    for _ in range(10):  # Burn some CPU time just because\n",
     "        np.linalg.eig(np.random.rand(1000, 1000))\n",
-    "    return x + width**2 / (width**2 + (x - offset) ** 2)\n",
-    "\n",
+    "    return x + width ** 2 / (width ** 2 + (x - offset) ** 2)\n",
     "\n",
     "# Define the sequence/samples we want to run\n",
     "xs = np.linspace(0, 1, 10_000)\n",
     "\n",
     "# ⚠️ Here a `learner` is an `adaptive` concept, read it as `jobs`.\n",
     "# ⚠️ `fnames` are the result locations\n",
     "learners, fnames = adaptive_scheduler.utils.split_sequence_in_sequence_learners(\n",
-    "    h,\n",
-    "    xs,\n",
-    "    n_learners=10,\n",
+    "    h, xs, n_learners=10\n",
     ")\n",
     "\n",
     "run_manager = adaptive_scheduler.slurm_run(\n",
@@ -52,7 +48,7 @@
     "    nodes=1,  # number of nodes per `learner`\n",
     "    cores_per_node=1,  # number of cores on 1 node per `learner`\n",
     "    log_interval=5,  # how often to produce a log message\n",
-    "    save_interval=5,  # how often to save the results\n",
+    "    save_interval=5, # how often to save the results\n",
     ")\n",
     "run_manager.start()"
    ]
@@ -89,18 +85,18 @@
     "from functools import partial\n",
     "\n",
     "import adaptive\n",
-    "\n",
     "import adaptive_scheduler\n",
     "\n",
     "\n",
     "def h(x, width=0.01, offset=0):\n",
     "    import numpy as np\n",
+    "    import random\n",
     "\n",
     "    for _ in range(10):  # Burn some CPU time just because\n",
     "        np.linalg.eig(np.random.rand(1000, 1000))\n",
     "\n",
     "    a = width\n",
-    "    return x + a**2 / (a**2 + (x - offset) ** 2)\n",
+    "    return x + a ** 2 / (a ** 2 + (x - offset) ** 2)\n",
     "\n",
     "\n",
     "offsets = [i / 10 - 0.5 for i in range(5)]\n",
@@ -270,16 +266,16 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "from adaptive import SequenceLearner\n",
     "\n",
-    "from adaptive_scheduler.utils import split\n",
+    "from adaptive import SequenceLearner\n",
+    "from adaptive_scheduler.utils import split, combo_to_fname\n",
     "\n",
     "\n",
     "def g(xyz):\n",
     "    x, y, z = xyz\n",
     "    for _ in range(5):  # Burn some CPU time just because\n",
     "        np.linalg.eig(np.random.rand(1000, 1000))\n",
-    "    return x**2 + y**2 + z**2\n",
+    "    return x ** 2 + y ** 2 + z ** 2\n",
     "\n",
     "\n",
     "xs = np.linspace(0, 10, 11)\n",
@@ -306,17 +302,11 @@
     "\n",
     "\n",
     "scheduler = adaptive_scheduler.scheduler.DefaultScheduler(\n",
-    "    cores=10,\n",
-    "    executor_type=\"ipyparallel\",\n",
+    "    cores=10, executor_type=\"ipyparallel\",\n",
     ")  # PBS or SLURM\n",
     "\n",
     "run_manager2 = adaptive_scheduler.server_support.RunManager(\n",
-    "    scheduler,\n",
-    "    learners,\n",
-    "    fnames,\n",
-    "    goal=goal,\n",
-    "    log_interval=30,\n",
-    "    save_interval=30,\n",
+    "    scheduler, learners, fnames, goal=goal, log_interval=30, save_interval=30,\n",
     ")\n",
     "run_manager2.start()"
    ]
@@ -353,19 +343,19 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
+    "\n",
     "from adaptive import SequenceLearner\n",
+    "from adaptive_scheduler.utils import split, combo2fname\n",
     "from adaptive.utils import named_product\n",
     "\n",
-    "from adaptive_scheduler.utils import combo2fname\n",
-    "\n",
     "\n",
     "def g(combo):\n",
     "    x, y, z = combo[\"x\"], combo[\"y\"], combo[\"z\"]\n",
     "\n",
     "    for _ in range(5):  # Burn some CPU time just because\n",
     "        np.linalg.eig(np.random.rand(1000, 1000))\n",
     "\n",
-    "    return x**2 + y**2 + z**2\n",
+    "    return x ** 2 + y ** 2 + z ** 2\n",
     "\n",
     "\n",
     "combos = named_product(x=np.linspace(0, 10), y=np.linspace(-1, 1), z=np.linspace(-3, 3))\n",
@@ -374,15 +364,15 @@
     "\n",
     "# We could run this as 1 job with N nodes, but we can also split it up in multiple jobs.\n",
     "# This is desireable when you don't want to run a single job with 300 nodes for example.\n",
-    "# Note that\n",
+    "# Note that \n",
     "# `adaptive_scheduler.utils.split_sequence_in_sequence_learners(g, combos, 100, \"data\")`\n",
     "# does the same!\n",
     "\n",
     "njobs = 100\n",
     "split_combos = list(split(combos, njobs))\n",
     "\n",
     "print(\n",
-    "    f\"Length of split_combos: {len(split_combos)} and length of split_combos[0]: {len(split_combos[0])}.\",\n",
+    "    f\"Length of split_combos: {len(split_combos)} and length of split_combos[0]: {len(split_combos[0])}.\"\n",
     ")\n",
     "\n",
     "learners = [SequenceLearner(g, combos_part) for combos_part in split_combos]\n",
@@ -403,16 +393,17 @@
    "outputs": [],
    "source": [
     "from functools import partial\n",
-    "\n",
     "import adaptive_scheduler\n",
-    "from adaptive_scheduler.scheduler import SLURM, DefaultScheduler\n",
+    "from adaptive_scheduler.scheduler import DefaultScheduler, PBS, SLURM\n",
     "\n",
     "\n",
     "def goal(learner):\n",
     "    return learner.done()  # the standard goal for a SequenceLearner\n",
     "\n",
     "\n",
-    "extra_scheduler = [\"--exclusive\", \"--time=24:00:00\"] if DefaultScheduler is SLURM else []\n",
+    "extra_scheduler = (\n",
+    "    [\"--exclusive\", \"--time=24:00:00\"] if DefaultScheduler is SLURM else []\n",
+    ")\n",
     "\n",
     "scheduler = adaptive_scheduler.scheduler.DefaultScheduler(\n",
     "    cores=10,\n",
@@ -468,8 +459,7 @@
    "source": [
     "run_manager3.load_learners()  # load the data into the learners\n",
     "result = sum(\n",
-    "    [l.result() for l in learners],\n",
-    "    [],\n",
+    "    [l.result() for l in learners], []\n",
     ")  # combine all learner's result into 1 list"
    ]
   }

tests/conftest.py

@@ -23,13 +23,13 @@
     import zmq.asyncio
 
 
-@pytest.fixture
+@pytest.fixture()
 def mock_scheduler(tmp_path: Path) -> MockScheduler:
     """Fixture for creating a MockScheduler instance."""
     return MockScheduler(log_folder=str(tmp_path), cores=8)
 
 
-@pytest.fixture
+@pytest.fixture()
 def db_manager(
     mock_scheduler: MockScheduler,
     learners: list[adaptive.Learner1D]
@@ -99,14 +99,14 @@ def fnames(
     raise NotImplementedError(msg)
 
 
-@pytest.fixture
+@pytest.fixture()
 def socket(db_manager: DatabaseManager) -> zmq.asyncio.Socket:
     """Fixture for creating a ZMQ socket."""
     with get_socket(db_manager) as socket:
         yield socket
 
 
-@pytest.fixture
+@pytest.fixture()
 def job_manager(
     db_manager: DatabaseManager,
     mock_scheduler: MockScheduler,
@@ -116,7 +116,7 @@ def job_manager(
     return JobManager(job_names, db_manager, mock_scheduler, interval=0.05)
 
 
-@pytest.fixture
+@pytest.fixture()
 def _mock_slurm_partitions_output() -> Generator[None, None, None]:
     """Mock `slurm_partitions` function."""
     mock_output = "hb120v2-low\nhb60-high\nnc24-low*\nnd40v2-mpi\n"
@@ -125,7 +125,7 @@ def _mock_slurm_partitions_output() -> Generator[None, None, None]:
         yield
 
 
-@pytest.fixture
+@pytest.fixture()
 def _mock_slurm_partitions() -> Generator[None, None, None]:
     """Mock `slurm_partitions` function."""
     with (
@@ -141,7 +141,7 @@ def _mock_slurm_partitions() -> Generator[None, None, None]:
         yield
 
 
-@pytest.fixture
+@pytest.fixture()
 def _mock_slurm_queue() -> Generator[None, None, None]:
     """Mock `SLURM.queue` function."""
     with patch(

tests/test_client_support.py

@@ -27,7 +27,7 @@ def client(zmq_url: str) -> zmq.Socket:
     return client
 
 
-@pytest.mark.asyncio
+@pytest.mark.asyncio()
 async def test_get_learner(zmq_url: str) -> None:
     """Test `get_learner` function."""
     with tempfile.NamedTemporaryFile() as tmpfile:
@@ -94,7 +94,7 @@ async def test_get_learner(zmq_url: str) -> None:
                 mock_log.exception.assert_called_with("got an exception")
 
 
-@pytest.mark.asyncio
+@pytest.mark.asyncio()
 async def test_tell_done(zmq_url: str) -> None:
     """Test `tell_done` function."""
     fname = "test_learner_file.pkl"