From 1bfe812bce46bb72bf403c0914a41b338be7a9a7 Mon Sep 17 00:00:00 2001 From: Dominik Krupke Date: Sat, 3 Aug 2024 12:19:23 +0200 Subject: [PATCH] CHORES: Lots of cleaning up --- .pre-commit-config.yaml | 9 ++-- src/slurminade/__init__.py | 28 +++++------ src/slurminade/bundling.py | 1 + src/slurminade/dispatcher.py | 8 +-- src/slurminade/function.py | 9 ++-- src/slurminade/function_map.py | 9 ++-- src/slurminade/guard.py | 1 + tests/test_create_command.py | 31 ++++++------ tests/test_create_command_with_noise.py | 32 ++++++------ tests/test_dispatch_guard.py | 46 +++++++++-------- tests/test_local.py | 46 +++++++++-------- tests/test_local_function.py | 4 +- tests/test_node_setup.py | 7 ++- tests/test_subprocess.py | 66 ++++++++++++------------- 14 files changed, 148 insertions(+), 149 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e0e8192..5045dd4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -53,11 +53,14 @@ repos: args: [--prose-wrap=always] exclude: "^tests" - - repo: https://github.com/charliermarsh/ruff-pre-commit - rev: v0.0.269 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.5.1 hooks: - id: ruff - args: ["--fix", "--show-fixes"] + types_or: [python, pyi, jupyter] + #args: [--fix,--ignore, T201, --ignore, E402] + - id: ruff-format + types_or: [python, pyi, jupyter] - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.3.0 diff --git a/src/slurminade/__init__.py b/src/slurminade/__init__.py index 784beb9..a3bfaf7 100644 --- a/src/slurminade/__init__.py +++ b/src/slurminade/__init__.py @@ -55,25 +55,25 @@ def clean_up(): """ # flake8: noqa F401 -from .function import slurmify, shell -from .conf import update_default_configuration, set_default_configuration -from .guard import ( - set_dispatch_limit, - allow_recursive_distribution, - disable_warning_on_repeated_flushes, -) -from .bundling import JobBundling, Batch +from .bundling import Batch, JobBundling +from .conf import set_default_configuration, update_default_configuration from .dispatcher import ( - srun, - sbatch, - join, SlurmDispatcher, - set_dispatcher, - get_dispatcher, - TestDispatcher, SubprocessDispatcher, + TestDispatcher, + get_dispatcher, + join, + sbatch, + set_dispatcher, + srun, ) +from .function import shell, slurmify from .function_map import set_entry_point +from .guard import ( + allow_recursive_distribution, + disable_warning_on_repeated_flushes, + set_dispatch_limit, +) from .node_setup import node_setup __all__ = [ diff --git a/src/slurminade/bundling.py b/src/slurminade/bundling.py index d91b8a3..3b85cd8 100644 --- a/src/slurminade/bundling.py +++ b/src/slurminade/bundling.py @@ -1,6 +1,7 @@ """ Contains code for bundling function calls together. """ + import logging import typing from collections import defaultdict diff --git a/src/slurminade/dispatcher.py b/src/slurminade/dispatcher.py index 0b43dc8..a055c4d 100644 --- a/src/slurminade/dispatcher.py +++ b/src/slurminade/dispatcher.py @@ -4,6 +4,7 @@ This allows to change the behaviour of the distribution, e.g., we use it for batch: Batch simply wraps the dispatcher by a buffered version. """ + import abc import logging import os @@ -21,13 +22,12 @@ from .function_call import FunctionCall from .function_map import FunctionMap, get_entry_point from .guard import dispatch_guard +from .job_reference import JobReference from .options import SlurmOptions # MAX_ARG_STRLEN on a Linux system with PAGE_SIZE 4096 is 131072 DEFAULT_MAX_ARG_LENGTH = 100000 -from .job_reference import JobReference - class Dispatcher(abc.ABC): """ @@ -443,7 +443,7 @@ def get_dispatcher() -> Dispatcher: to allow compatibility. :return: The dispatcher. """ - global __dispatcher + global __dispatcher # noqa: PLW0603 if __dispatcher is None: try: __dispatcher = SlurmDispatcher() @@ -460,7 +460,7 @@ def set_dispatcher(dispatcher: Dispatcher) -> None: :param dispatcher: The dispatcher to be used. :return: None """ - global __dispatcher + global __dispatcher # noqa: PLW0603 __dispatcher = dispatcher assert dispatcher == get_dispatcher() diff --git a/src/slurminade/function.py b/src/slurminade/function.py index 941776d..bee029a 100644 --- a/src/slurminade/function.py +++ b/src/slurminade/function.py @@ -118,13 +118,12 @@ def __call__(self, *args, **kwargs): """ if self.call_policy == CallPolicy.LOCALLY: return self.run_locally(*args, **kwargs) - elif self.call_policy == CallPolicy.DISTRIBUTED: + if self.call_policy == CallPolicy.DISTRIBUTED: return self.distribute(*args, **kwargs) - elif self.call_policy == CallPolicy.DISTRIBUTED_BLOCKING: + if self.call_policy == CallPolicy.DISTRIBUTED_BLOCKING: return self.distribute_and_wait(*args, **kwargs) - else: - msg = "Unknown call policy." - raise RuntimeError(msg) + msg = "Unknown call policy." + raise RuntimeError(msg) def get_entry_point(self) -> Path: """ diff --git a/src/slurminade/function_map.py b/src/slurminade/function_map.py index 4627991..bc1ef69 100644 --- a/src/slurminade/function_map.py +++ b/src/slurminade/function_map.py @@ -7,8 +7,8 @@ import logging import pathlib import typing -from typing import Optional from pathlib import Path +from typing import Optional from .execute_cmds import call_slurminade_to_get_function_ids @@ -24,8 +24,8 @@ class FunctionMap: # The slurm node just executes the file content of a script, so the file name is lost. # slurminade will set this value in the beginning to reconstruct it. entry_point: typing.Optional[str] = None - _data = {} - _ids: Optional[set] = set() + _data: typing.ClassVar[dict] = {} + _ids: typing.ClassVar[Optional[set]] = set() @staticmethod def get_id(func: typing.Callable) -> str: @@ -113,7 +113,8 @@ def check_id(func_id: str, entry_point: Path) -> bool: FunctionMap._ids = call_slurminade_to_get_function_ids(entry_point) except Exception as e: logging.getLogger("slurminade").warning( - "Cannot verify function ids before submitting to slurm: %s. This is not critical, things will just be more difficult to debug in case you make an error.", e + "Cannot verify function ids before submitting to slurm: %s. This is not critical, things will just be more difficult to debug in case you make an error.", + e, ) FunctionMap._ids = None return True diff --git a/src/slurminade/guard.py b/src/slurminade/guard.py index b894958..88436e1 100644 --- a/src/slurminade/guard.py +++ b/src/slurminade/guard.py @@ -10,6 +10,7 @@ ``allow_recursive_distribution``, ``set_dispatch_limit(None)``, and ``disable_warning_for_multiple_flushes``. """ + import logging import typing diff --git a/tests/test_create_command.py b/tests/test_create_command.py index 6004e98..e5adf1b 100644 --- a/tests/test_create_command.py +++ b/tests/test_create_command.py @@ -1,4 +1,3 @@ -import unittest from pathlib import Path import slurminade @@ -8,22 +7,20 @@ @slurminade.slurmify() def f(s): - with open(test_file_path, "w") as file: + with test_file_path.open("w") as file: file.write(s) -class TestCreateCommand(unittest.TestCase): - def test_dispatch_with_temp_file(self): - slurminade.set_entry_point(__file__) - if test_file_path.exists(): - test_file_path.unlink() - dispatcher = slurminade.SubprocessDispatcher() - dispatcher.max_arg_length = 1 - slurminade.set_dispatcher(dispatcher) - s = "test" - f.distribute(s) - assert test_file_path.is_file() - with open(test_file_path) as file: - assert file.readline() == s - if test_file_path.exists(): # delete the file - test_file_path.unlink() +def test_create_command(): + slurminade.set_entry_point(__file__) + if test_file_path.exists(): + test_file_path.unlink() + dispatcher = slurminade.SubprocessDispatcher() + dispatcher.max_arg_length = 1 + slurminade.set_dispatcher(dispatcher) + s = "test" + f.distribute(s) + assert test_file_path.is_file() + with test_file_path.open() as file: + assert file.readline() == s + test_file_path.unlink(missing_ok=True) diff --git a/tests/test_create_command_with_noise.py b/tests/test_create_command_with_noise.py index 1ebc759..87ee0fc 100644 --- a/tests/test_create_command_with_noise.py +++ b/tests/test_create_command_with_noise.py @@ -1,4 +1,3 @@ -import unittest from pathlib import Path import slurminade @@ -6,24 +5,23 @@ test_file_path = Path("./f_test_file.txt") print("NOISE123") + @slurminade.slurmify() def f(s): - with open(test_file_path, "w") as file: + with test_file_path.open("w") as file: file.write(s) -class TestCreateCommand(unittest.TestCase): - def test_dispatch_with_temp_file(self): - slurminade.set_entry_point(__file__) - if test_file_path.exists(): - test_file_path.unlink() - dispatcher = slurminade.SubprocessDispatcher() - dispatcher.max_arg_length = 1 - slurminade.set_dispatcher(dispatcher) - s = "test" - f.distribute(s) - assert test_file_path.is_file() - with open(test_file_path) as file: - assert file.readline() == s - if test_file_path.exists(): # delete the file - test_file_path.unlink() +def test_create_command_with_noise(): + slurminade.set_entry_point(__file__) + if test_file_path.exists(): + test_file_path.unlink() + dispatcher = slurminade.SubprocessDispatcher() + dispatcher.max_arg_length = 1 + slurminade.set_dispatcher(dispatcher) + s = "test" + f.distribute(s) + assert test_file_path.is_file() + with test_file_path.open() as file: + assert file.readline() == s + test_file_path.unlink(missing_ok=True) diff --git a/tests/test_dispatch_guard.py b/tests/test_dispatch_guard.py index d5aae36..83f32a4 100644 --- a/tests/test_dispatch_guard.py +++ b/tests/test_dispatch_guard.py @@ -1,4 +1,4 @@ -import unittest +import pytest import slurminade from slurminade.guard import TooManyDispatchesError, _DispatchGuard, set_dispatch_limit @@ -9,27 +9,31 @@ def f(): pass -class TestDispatchGuard(unittest.TestCase): - def test_simple(self): - slurminade.set_entry_point(__file__) - dg = _DispatchGuard(3) +def test_dispatch_guard_simple(): + slurminade.set_entry_point(__file__) + dg = _DispatchGuard(3) + dg() + dg() + dg() + with pytest.raises(TooManyDispatchesError): dg() - dg() - dg() - self.assertRaises(TooManyDispatchesError, dg) - def test_dispatch_limit(self): - slurminade.set_entry_point(__file__) - set_dispatch_limit(3) - f.distribute() + +def test_dispatch_guard_dispatch_limit(): + slurminade.set_entry_point(__file__) + set_dispatch_limit(3) + f.distribute() + f.distribute() + f.distribute() + with pytest.raises(TooManyDispatchesError): f.distribute() + + +def test_dispatch_guard_dispatch_limit_batch(): + slurminade.set_entry_point(__file__) + set_dispatch_limit(2) + with slurminade.JobBundling(max_size=2): + for _ in range(4): + f.distribute() + with pytest.raises(TooManyDispatchesError): f.distribute() - self.assertRaises(TooManyDispatchesError, f.distribute) - - def test_dispatch_limit_batch(self): - slurminade.set_entry_point(__file__) - set_dispatch_limit(2) - with slurminade.JobBundling(max_size=2): - for _ in range(4): - f.distribute() - self.assertRaises(TooManyDispatchesError, f.distribute) diff --git a/tests/test_local.py b/tests/test_local.py index 63dd9ef..17add23 100644 --- a/tests/test_local.py +++ b/tests/test_local.py @@ -1,5 +1,4 @@ -import os.path -import unittest +from pathlib import Path import slurminade from slurminade.function import SlurmFunction @@ -10,37 +9,36 @@ @slurminade.slurmify() def f(): - with open(f_file, "w") as file: + with Path(f_file).open("w") as file: file.write("test") def delete_f(): - if os.path.exists(f_file): - os.remove(f_file) + Path(f_file).unlink(missing_ok=True) @slurminade.slurmify() def g(x, y): - with open(g_file, "w") as file: + with Path(g_file).open("w") as file: file.write(f"{x}:{y}") def delete_g(): - if os.path.exists(g_file): - os.remove(g_file) - - -class TestLocal(unittest.TestCase): - def test_1(self): - delete_f() - SlurmFunction.call(f.func_id) - assert os.path.exists(f_file) - delete_f() - - def test_2(self): - delete_g() - SlurmFunction.call(g.func_id, x="a", y=2) - assert os.path.exists(g_file) - with open(g_file) as file: - assert file.readline() == "a:2" - delete_g() + Path(g_file).unlink(missing_ok=True) + + +def test_local_1(): + delete_f() + SlurmFunction.call(f.func_id) + assert Path(f_file).exists() + delete_f() + + +def test_local_2(): + delete_g() + SlurmFunction.call(g.func_id, x="a", y=2) + + assert Path(g_file).exists() + with Path(g_file).open() as file: + assert file.readline() == "a:2" + delete_g() diff --git a/tests/test_local_function.py b/tests/test_local_function.py index 0066090..ba4982d 100644 --- a/tests/test_local_function.py +++ b/tests/test_local_function.py @@ -1,4 +1,4 @@ -from pytest import raises +import pytest import slurminade @@ -10,5 +10,5 @@ def f(): slurminade.set_entry_point(__file__) - with raises(KeyError): + with pytest.raises(KeyError): f.distribute() diff --git a/tests/test_node_setup.py b/tests/test_node_setup.py index f413709..3fbacf6 100644 --- a/tests/test_node_setup.py +++ b/tests/test_node_setup.py @@ -7,7 +7,7 @@ @slurminade.node_setup def f(): - with open(test_file_path, "w") as file: + with test_file_path.open("w") as file: file.write("node_setup") @@ -18,13 +18,12 @@ def nil(): def test_node_setup(): slurminade.set_entry_point(__file__) - if test_file_path.exists(): - test_file_path.unlink() + test_file_path.unlink(missing_ok=True) dispatcher = slurminade.SubprocessDispatcher() slurminade.set_dispatcher(dispatcher) slurminade.set_dispatch_limit(100) nil.distribute() - with open(test_file_path) as file: + with test_file_path.open() as file: assert file.readline() == "node_setup" if test_file_path.exists(): # delete the file test_file_path.unlink() diff --git a/tests/test_subprocess.py b/tests/test_subprocess.py index 4f72484..4854bd4 100644 --- a/tests/test_subprocess.py +++ b/tests/test_subprocess.py @@ -1,5 +1,5 @@ -import os.path -import unittest +import os +from pathlib import Path import slurminade @@ -9,53 +9,51 @@ @slurminade.slurmify() def f(): - with open(f_file, "w") as file: + with Path(f_file).open("w") as file: file.write("test") def delete_f(): - if os.path.exists(f_file): - os.remove(f_file) + Path(f_file).unlink(missing_ok=True) @slurminade.slurmify() def g(x, y): - with open(g_file, "w") as file: + with Path(g_file).open("w") as file: file.write(f"{x}:{y}") def delete_g(): - if os.path.exists(g_file): - os.remove(g_file) + Path(g_file).unlink(missing_ok=True) def get_file_name(): filename = os.getenv("PYTEST_CURRENT_TEST") # doesn't seem to work always filename = filename.split("::")[0] if filename else __file__ # workaround - if not os.path.exists(filename): # sometimes the test folder gets duplicated. + if not Path(filename).exists(): # sometimes the test folder gets duplicated. filename = "/".join(filename.split("/")[1:]) - return os.path.abspath(filename) - - -class TestSubprocess(unittest.TestCase): - def test_1(self): - slurminade.set_dispatcher(slurminade.SubprocessDispatcher()) - slurminade.set_dispatch_limit(100) - slurminade.set_entry_point(get_file_name()) - - delete_f() - f.distribute() - assert os.path.exists(f_file) - delete_f() - - def test_2(self): - slurminade.set_dispatcher(slurminade.SubprocessDispatcher()) - slurminade.set_entry_point(get_file_name()) - slurminade.set_dispatch_limit(100) - - delete_g() - g.distribute(x="a", y=2) - assert os.path.exists(g_file) - with open(g_file) as file: - assert file.readline() == "a:2" - delete_g() + return Path(filename).resolve() + + +def test_subprocess_1(): + slurminade.set_dispatcher(slurminade.SubprocessDispatcher()) + slurminade.set_dispatch_limit(100) + slurminade.set_entry_point(get_file_name()) + + delete_f() + f.distribute() + assert Path(f_file).exists() + delete_f() + + +def test_subprocess_2(): + slurminade.set_dispatcher(slurminade.SubprocessDispatcher()) + slurminade.set_entry_point(get_file_name()) + slurminade.set_dispatch_limit(100) + + delete_g() + g.distribute(x="a", y=2) + assert Path(g_file).exists() + with Path(g_file).open() as file: + assert file.readline() == "a:2" + delete_g()