From 7e7efd5511a0a13b4fea9e4e523d5053032e2f27 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Thu, 14 Nov 2024 13:49:46 -0500 Subject: [PATCH] Infer cluster from FQDN * Added `gempyor.info._infer_cluster_from_fqdn` internal utility. * `gempyor.info.get_cluster_info` now accepts `None` for `name`. * `gempyor.batch._click_submit` accepts `None` for `--cluster` when using `--slurm`. * Minor testing fixes. --- flepimop/gempyor_pkg/src/gempyor/batch.py | 2 - flepimop/gempyor_pkg/src/gempyor/info.py | 32 ++++++++++++++-- .../tests/batch/test__click_submit.py | 32 ++-------------- .../info/test__infer_cluster_from_fqdn.py | 37 +++++++++++++++++++ .../tests/info/test_get_cluster_info.py | 20 +++++++++- .../templates/test_cluster_setup_bash.py | 15 ++++---- 6 files changed, 96 insertions(+), 42 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/info/test__infer_cluster_from_fqdn.py diff --git a/flepimop/gempyor_pkg/src/gempyor/batch.py b/flepimop/gempyor_pkg/src/gempyor/batch.py index 7df377582..b2d56c99f 100644 --- a/flepimop/gempyor_pkg/src/gempyor/batch.py +++ b/flepimop/gempyor_pkg/src/gempyor/batch.py @@ -786,8 +786,6 @@ def _submit_scenario_job( ) elif batch_system == BatchSystem.SLURM: # Cluster info - if kwargs["cluster"] is None: - raise ValueError("When submitting a batch job to slurm a cluster is required.") cluster = get_cluster_info(kwargs["cluster"]) if verbosity is not None: logger.info( diff --git a/flepimop/gempyor_pkg/src/gempyor/info.py b/flepimop/gempyor_pkg/src/gempyor/info.py index 1048f3f8a..00d094c09 100644 --- a/flepimop/gempyor_pkg/src/gempyor/info.py +++ b/flepimop/gempyor_pkg/src/gempyor/info.py @@ -3,7 +3,9 @@ import os from pathlib import Path -from typing import TypeVar +import re +from socket import getfqdn +from typing import Pattern, TypeVar from pydantic import BaseModel import yaml @@ -29,6 +31,12 @@ class Cluster(BaseModel): T = TypeVar("T", bound=BaseModel) +_CLUSTER_FQDN_REGEXES: tuple[tuple[str, Pattern], ...] = ( + ("longleaf", re.compile(r"^longleaf\-login[0-9]+\.its\.unc\.edu$")), + ("rockfish", re.compile(r"^login[0-9]+\.cm\.cluster$")), +) + + def _get_info( category: str, name: str, model: type[T], flepi_path: os.PathLike | None ) -> T: @@ -61,17 +69,35 @@ def _get_info( return model.model_validate(yaml.safe_load(info.read_text())) -def get_cluster_info(name: str, flepi_path: os.PathLike | None = None) -> Cluster: +def get_cluster_info(name: str | None, flepi_path: os.PathLike | None = None) -> Cluster: """ Get cluster specific info. Args: name: The name of the cluster to pull information for. Currently only 'longleaf' - and 'rockfish' are supported. + and 'rockfish' are supported or `None` to infer from the FQDN. flepi_path: Either a path like determine the directory to look for the info directory in or `None` to use the `FLEPI_PATH` environment variable. Returns An object containing the information about the `name` cluster. """ + name = _infer_cluster_from_fqdn() if name is None else name return _get_info("cluster", name, Cluster, flepi_path) + + +def _infer_cluster_from_fqdn() -> str: + """ + Infer the cluster name from the FQDN. + + Returns: + The name of the cluster inferred from the FQDN. + + Raises: + ValueError: If the value of `socket.getfqdn()` does not match an expected regex. + """ + fqdn = getfqdn() + for cluster, regex in _CLUSTER_FQDN_REGEXES: + if regex.match(fqdn): + return cluster + raise ValueError(f"The fqdn, '{fqdn}', does not match any of the expected clusters.") diff --git a/flepimop/gempyor_pkg/tests/batch/test__click_submit.py b/flepimop/gempyor_pkg/tests/batch/test__click_submit.py index 31582c9f6..993c0a8c1 100644 --- a/flepimop/gempyor_pkg/tests/batch/test__click_submit.py +++ b/flepimop/gempyor_pkg/tests/batch/test__click_submit.py @@ -25,6 +25,10 @@ def test_batch_system_aws_not_implemented_error(tmp_path: Path) -> None: "1", "--blocks", "1", + "--flepi-path", + str(tmp_path.absolute()), + "--project-path", + str(tmp_path.absolute()), str(config_file.absolute()), ], ) @@ -35,31 +39,3 @@ def test_batch_system_aws_not_implemented_error(tmp_path: Path) -> None: str(result.exception) == "The `flepimop submit` CLI does not support batch submission to AWS yet." ) - - -def test_cluster_required_for_slurm_value_error(tmp_path: Path) -> None: - config_file = tmp_path / "config.yml" - with config_file.open(mode="w") as f: - yaml.dump({"name": "foobar", "inference": {"method": "emcee"}}, f) - - runner = CliRunner() - result = runner.invoke( - _click_submit, - [ - "--slurm", - "--simulations", - "1", - "--jobs", - "1", - "--blocks", - "1", - str(config_file.absolute()), - ], - ) - - assert result.exit_code == 1 - assert isinstance(result.exception, ValueError) - assert ( - str(result.exception) - == "When submitting a batch job to slurm a cluster is required." - ) diff --git a/flepimop/gempyor_pkg/tests/info/test__infer_cluster_from_fqdn.py b/flepimop/gempyor_pkg/tests/info/test__infer_cluster_from_fqdn.py new file mode 100644 index 000000000..a26705503 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/info/test__infer_cluster_from_fqdn.py @@ -0,0 +1,37 @@ +from unittest.mock import patch + +import pytest + +from gempyor.info import _infer_cluster_from_fqdn + + +@pytest.mark.parametrize("fqdn", ("new.cluster.com", "unsupported.cluster")) +def test_no_matching_fqdn_found_value_error(fqdn: str) -> None: + def socket_fqdn_wraps() -> str: + return fqdn + + with patch("gempyor.info.getfqdn", wraps=socket_fqdn_wraps) as socket_fqdn_patch: + with pytest.raises( + ValueError, + match=f"^The fqdn, '{fqdn}', does not match any of the expected clusters.$", + ): + _infer_cluster_from_fqdn() + socket_fqdn_patch.assert_called_once() + + +@pytest.mark.parametrize( + ("fqdn", "expected"), + ( + ("login01.cm.cluster", "rockfish"), + ("login3.cm.cluster", "rockfish"), + ("longleaf-login1.its.unc.edu", "longleaf"), + ("longleaf-login07.its.unc.edu", "longleaf"), + ), +) +def test_exact_results_for_select_values(fqdn: str, expected: str) -> None: + def socket_fqdn_wraps() -> str: + return fqdn + + with patch("gempyor.info.getfqdn", wraps=socket_fqdn_wraps) as socket_fqdn_patch: + assert _infer_cluster_from_fqdn() == expected + socket_fqdn_patch.assert_called_once() diff --git a/flepimop/gempyor_pkg/tests/info/test_get_cluster_info.py b/flepimop/gempyor_pkg/tests/info/test_get_cluster_info.py index b8ab64436..66fcd9836 100644 --- a/flepimop/gempyor_pkg/tests/info/test_get_cluster_info.py +++ b/flepimop/gempyor_pkg/tests/info/test_get_cluster_info.py @@ -1,4 +1,5 @@ import os +from unittest.mock import patch import pytest @@ -10,6 +11,23 @@ os.getenv("FLEPI_PATH") is None, reason="The $FLEPI_PATH environment variable is not set.", ) -def test_output_validation(name: str) -> None: +def test_exact_results_given_cluster_name(name: str) -> None: cluster = get_cluster_info(name) assert isinstance(cluster, Cluster) + + +@pytest.mark.parametrize("name", ("longleaf", "rockfish")) +@pytest.mark.skipif( + os.getenv("FLEPI_PATH") is None, + reason="The $FLEPI_PATH environment variable is not set.", +) +def test_exact_results_when_inferred_from_fqdn(name: str) -> None: + def infer_cluster_from_fqdn_wraps() -> str: + return name + + with patch( + "gempyor.info._infer_cluster_from_fqdn", wraps=infer_cluster_from_fqdn_wraps + ) as infer_cluster_from_fqdn_patch: + cluster = get_cluster_info(None) + assert isinstance(cluster, Cluster) + infer_cluster_from_fqdn_patch.assert_called_once() diff --git a/flepimop/gempyor_pkg/tests/templates/test_cluster_setup_bash.py b/flepimop/gempyor_pkg/tests/templates/test_cluster_setup_bash.py index 89048444a..90e1e4cc2 100644 --- a/flepimop/gempyor_pkg/tests/templates/test_cluster_setup_bash.py +++ b/flepimop/gempyor_pkg/tests/templates/test_cluster_setup_bash.py @@ -1,3 +1,4 @@ +import os from typing import Any import pytest @@ -6,15 +7,13 @@ from gempyor.info import Cluster, Module, PathExport, get_cluster_info -@pytest.mark.parametrize( - "cluster", - ( - None, - get_cluster_info("longleaf").model_dump(), - get_cluster_info("rockfish").model_dump(), - ), +@pytest.mark.parametrize("cluster", (None, "longleaf", "rockfish")) +@pytest.mark.skipif( + os.getenv("FLEPI_PATH") is None, + reason="The $FLEPI_PATH environment variable is not set.", ) -def test_output_validation(cluster: dict[str, Any]) -> None: +def test_output_validation(cluster: str | None) -> None: + cluster = cluster if cluster is None else get_cluster_info(cluster).model_dump() rendered_template = _render_template("cluster_setup.bash.j2", {"cluster": cluster}) lines = rendered_template.split("\n") assert "module purge" in lines