diff --git a/aiida_hyperqueue/cli/__init__.py b/aiida_hyperqueue/cli/__init__.py index 76dd0bd..61ba344 100644 --- a/aiida_hyperqueue/cli/__init__.py +++ b/aiida_hyperqueue/cli/__init__.py @@ -1,8 +1,5 @@ # -*- coding: utf-8 -*- -from aiida.cmdline.params import options as core_options -from aiida.cmdline.params import types as core_types - -from .root import cmd_root -from .install import cmd_install -from .server import cmd_info, cmd_start, cmd_stop -from .alloc import cmd_list, cmd_add, cmd_remove +from .root import cmd_root # noqa: F401 +from .install import cmd_install # noqa: F401 +from .server import cmd_info, cmd_start, cmd_stop # noqa: F401 +from .alloc import cmd_list, cmd_add, cmd_remove # noqa: F401 diff --git a/aiida_hyperqueue/cli/alloc.py b/aiida_hyperqueue/cli/alloc.py index 607e655..c810d06 100644 --- a/aiida_hyperqueue/cli/alloc.py +++ b/aiida_hyperqueue/cli/alloc.py @@ -1,3 +1,4 @@ +# -*- coding: utf-8 -*- import click from aiida.cmdline.params import options, arguments @@ -5,6 +6,7 @@ from .root import cmd_root + @cmd_root.group("alloc") def alloc_group(): """Commands to configure HQ allocations.""" diff --git a/aiida_hyperqueue/cli/install.py b/aiida_hyperqueue/cli/install.py index 592eeb1..0e049fb 100644 --- a/aiida_hyperqueue/cli/install.py +++ b/aiida_hyperqueue/cli/install.py @@ -30,14 +30,16 @@ "--hq-version", type=str, default="0.19.0", help="the hq version will be installed." ) # TODO: should also support different arch binary?? -def cmd_install(computer: orm.Computer, remote_bin_dir: Path, hq_version: str, write_bashrc: bool): +def cmd_install( + computer: orm.Computer, remote_bin_dir: Path, hq_version: str, write_bashrc: bool +): """Install the hq binary to the computer through the transport""" # The minimal hq version we support is 0.13.0, check the minor version try: - _, minor, _ = hq_version.split('.') + _, minor, _ = hq_version.split(".") except ValueError as e: - echo.echo_critical(f"Cannot parse the version {hq_version}: {e}") + echo.echo_critical(f"Cannot parse the version {hq_version}: {e}") else: if int(minor) < 13: # `--no-hyper-threading` replace `--cpus=no-ht` from 0.13.0 @@ -78,9 +80,13 @@ def cmd_install(computer: orm.Computer, remote_bin_dir: Path, hq_version: str, w # TODO: try not override if the binary exist, put has overwrite=True as default with computer.get_transport() as transport: # Get the abs path of remote bin dir - retval, stdout, stderr = transport.exec_command_wait(f"echo {str(remote_bin_dir)}") - if retval !=0: - echo.echo_critical(f"Not able to parse remote bin dir {remote_bin_dir}, exit_code={retval}") + retval, stdout, stderr = transport.exec_command_wait( + f"echo {str(remote_bin_dir)}" + ) + if retval != 0: + echo.echo_critical( + f"Not able to parse remote bin dir {remote_bin_dir}, exit_code={retval}" + ) else: remote_bin_dir = Path(stdout.strip()) diff --git a/aiida_hyperqueue/cli/server.py b/aiida_hyperqueue/cli/server.py index 768a476..7f2fae3 100644 --- a/aiida_hyperqueue/cli/server.py +++ b/aiida_hyperqueue/cli/server.py @@ -1,3 +1,4 @@ +# -*- coding: utf-8 -*- import click from aiida.cmdline.utils import echo @@ -5,6 +6,7 @@ from .root import cmd_root from .params import arguments + @cmd_root.group("server") def server_group(): """Commands for interacting with the HQ server.""" @@ -12,7 +14,13 @@ def server_group(): @server_group.command("start") @arguments.COMPUTER() -@click.option("-d", "--domain", required=False, type=click.STRING, help="domain that will attached to the `hostname` of remote.") +@click.option( + "-d", + "--domain", + required=False, + type=click.STRING, + help="domain that will attached to the `hostname` of remote.", +) def cmd_start(computer, domain: str): """Start the HyperQueue server.""" @@ -32,16 +40,20 @@ def cmd_start(computer, domain: str): start_command_lst = ["nohup", "hq", "server", "start"] if domain is not None: - retval, stdout, stderr = transport.exec_command_wait( - "hostname" - ) + retval, stdout, stderr = transport.exec_command_wait("hostname") if retval != 0: echo.echo_critical(f"unable to get the hostname: {stderr}") else: hostname = stdout.strip() start_command_lst.extend(["--host", f"{hostname}.{domain}"]) - start_command_lst.extend(["1>$HOME/.hq-stdout", "2>$HOME/.hq-stderr", "&",]) + start_command_lst.extend( + [ + "1>$HOME/.hq-stdout", + "2>$HOME/.hq-stderr", + "&", + ] + ) start_command = " ".join(start_command_lst) echo.echo_debug(f"Run start command {start_command} on the remote") @@ -58,6 +70,7 @@ def cmd_start(computer, domain: str): echo.echo_success("HQ server started!") + @server_group.command("stop") @arguments.COMPUTER() def cmd_stop(computer): @@ -73,15 +86,14 @@ def cmd_stop(computer): echo.echo_info("Stop the hq server will close all allocs.") with computer.get_transport() as transport: - retval, _, stderr = transport.exec_command_wait( - "hq server stop" - ) + retval, _, stderr = transport.exec_command_wait("hq server stop") if retval != 0: echo.echo_critical(f"unable to stop the server: {stderr}") echo.echo_success("HQ server stopped!") + @server_group.command("restart") @arguments.COMPUTER() # TODO: how to pass domain to restart??? diff --git a/aiida_hyperqueue/scheduler.py b/aiida_hyperqueue/scheduler.py index 7972006..2ccb1bf 100644 --- a/aiida_hyperqueue/scheduler.py +++ b/aiida_hyperqueue/scheduler.py @@ -7,14 +7,12 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### -"""Plugin for the HyperQueue meta scheduler. -""" +"""Plugin for the HyperQueue meta scheduler.""" import re -from typing import Union, Optional +import typing as t from aiida.common.extendeddicts import AttributeDict -from aiida.common.exceptions import FeatureNotAvailable from aiida.schedulers import Scheduler, SchedulerError from aiida.schedulers.datastructures import JobInfo, JobState, JobResource, JobTemplate @@ -85,8 +83,7 @@ def get_tot_num_mpiprocs(self): class HyperQueueScheduler(Scheduler): - """Support for the HyperQueue scheduler (https://it4innovations.github.io/hyperqueue/stable/). - """ + """Support for the HyperQueue scheduler (https://it4innovations.github.io/hyperqueue/stable/).""" _logger = Scheduler._logger.getChild("hyperqueue") @@ -200,7 +197,7 @@ def _parse_submit_output(self, retval: int, stdout: str, stderr: str) -> str: ) def _get_joblist_command( - self, jobs: Optional[list] = None, user: Optional[str] = None + self, jobs: t.Optional[list] = None, user: t.Optional[str] = None ) -> str: """Return the ``hq`` command for listing the active jobs. @@ -247,8 +244,7 @@ def _parse_joblist_output(self, retval: int, stdout: str, stderr: str) -> list: return job_info_list def _get_kill_command(self, jobid): - """Return the command to kill the job with specified jobid. - """ + """Return the command to kill the job with specified jobid.""" submit_command = f"hq job cancel {jobid}" self.logger.info(f"killing job {jobid}") diff --git a/tests/test_cli_alloc.py b/tests/test_cli_alloc.py index ebb2a5b..2730bf8 100644 --- a/tests/test_cli_alloc.py +++ b/tests/test_cli_alloc.py @@ -1,12 +1,15 @@ +# -*- coding: utf-8 -*- import pytest from click.testing import CliRunner from aiida.transports.transport import Transport as TransportClass + @pytest.fixture def runner(): return CliRunner() + original_exec_command_wait = TransportClass.exec_command_wait # TODO: Not yet implemented, seems hard to just use hq_env diff --git a/tests/test_cli_install.py b/tests/test_cli_install.py index 4251a3d..6ed4a87 100644 --- a/tests/test_cli_install.py +++ b/tests/test_cli_install.py @@ -1,21 +1,30 @@ +# -*- coding: utf-8 -*- import pytest from click.testing import CliRunner from aiida_hyperqueue.cli import cmd_install + @pytest.fixture def runner(): return CliRunner() + def test_install(runner, tmp_path, aiida_computer_ssh): aiida_computer_ssh(label="localhost-hq") version = "0.19.0" - result = runner.invoke(cmd_install, ["-p", f"{str(tmp_path.resolve())}", "--hq-version", version, "--no-write-bashrc", "localhost-hq"]) + result = runner.invoke( + cmd_install, + [ + "-p", + f"{str(tmp_path.resolve())}", + "--hq-version", + version, + "--no-write-bashrc", + "localhost-hq", + ], + ) assert result.exit_code == 0 assert f"hq version {version}" in result.output - - - - diff --git a/tests/test_cli_server.py b/tests/test_cli_server.py index 7a47116..b2d7c1d 100644 --- a/tests/test_cli_server.py +++ b/tests/test_cli_server.py @@ -1,7 +1,6 @@ +# -*- coding: utf-8 -*- import pytest import time -from functools import partial -from pathlib import Path from click.testing import CliRunner from aiida.transports.transport import Transport as TransportClass @@ -9,65 +8,83 @@ from .conftest import HqEnv, get_hq_binary + @pytest.fixture def runner(): return CliRunner() + original_exec_command_wait = TransportClass.exec_command_wait + @pytest.fixture def hq_env_mock_exec_command_wait(hq_env: HqEnv): - def _mock_exec_command_wait(obj, command: str, **kwargs): """mock `exec_command_wait` method of Transport Abstract class. Since it is an object method, its first parameter is self. - + This function is for running the `hq` command through the temperory hq_env fixture. If the command start with "hq" we regard it running the hq command and pass it to the hq_env to run, otherwith use the original behavior. """ if command.startswith("hq"): - process = hq_env.command(command.split(' ')[1:], wait=False, ignore_stderr=True) + process = hq_env.command( + command.split(" ")[1:], wait=False, ignore_stderr=True + ) stdout = process.communicate()[0].decode() stderr = "" retval = process.returncode else: retval, stdout, stderr = original_exec_command_wait(obj, command, **kwargs) - + return retval, stdout, stderr return _mock_exec_command_wait + @pytest.fixture def server_dir_mock_exec_command_wait(tmp_path): - def _mock_exec_command_wait(obj, command: str, **kwargs): """mock `exec_command_wait` method of Transport Abstract class. Since it is an object method, its first parameter is self. - + Different from `hq_env_mock_exec_command_wait` this one run with changing the server_dir and did not handle the clean up of stop the hq. """ hq = get_hq_binary() - cmd_list = command.split(' ') - if command.startswith('hq'): # `hq` - command = ' '.join([f'{hq}', '--server-dir', f'{str(tmp_path.resolve())}'] + cmd_list[1:]) - elif command.startswith('nohup hq'): # `nohup hq` - command = ' '.join([f'nohup {hq}', '--server-dir', f'{str(tmp_path.resolve())}'] + cmd_list[2:]) + cmd_list = command.split(" ") + if command.startswith("hq"): # `hq` + command = " ".join( + [f"{hq}", "--server-dir", f"{str(tmp_path.resolve())}"] + cmd_list[1:] + ) + elif command.startswith("nohup hq"): # `nohup hq` + command = " ".join( + [f"nohup {hq}", "--server-dir", f"{str(tmp_path.resolve())}"] + + cmd_list[2:] + ) retval, stdout, stderr = original_exec_command_wait(obj, command, **kwargs) - + return retval, stdout, stderr return _mock_exec_command_wait -def test_server_info_direct(runner: CliRunner, aiida_computer_local, monkeypatch: pytest.MonkeyPatch, hq_env_mock_exec_command_wait, hq_env: HqEnv): + +def test_server_info_direct( + runner: CliRunner, + aiida_computer_local, + monkeypatch: pytest.MonkeyPatch, + hq_env_mock_exec_command_wait, + hq_env: HqEnv, +): """Test server info, test against local transport""" - aiida_computer_local(label='localhost') + aiida_computer_local(label="localhost") - monkeypatch.setattr(TransportClass, "exec_command_wait", hq_env_mock_exec_command_wait) + monkeypatch.setattr( + TransportClass, "exec_command_wait", hq_env_mock_exec_command_wait + ) result = runner.invoke(cmd_info, "localhost") assert result.exit_code == 1 @@ -79,11 +96,20 @@ def test_server_info_direct(runner: CliRunner, aiida_computer_local, monkeypatch assert result.exit_code == 0 assert "Start date" in result.output -def test_server_info_ssh(runner: CliRunner, aiida_computer_ssh, monkeypatch: pytest.MonkeyPatch, hq_env_mock_exec_command_wait, hq_env: HqEnv): + +def test_server_info_ssh( + runner: CliRunner, + aiida_computer_ssh, + monkeypatch: pytest.MonkeyPatch, + hq_env_mock_exec_command_wait, + hq_env: HqEnv, +): """Test server info, test against ssh transport""" - aiida_computer_ssh(label='localhost-hq') + aiida_computer_ssh(label="localhost-hq") - monkeypatch.setattr(TransportClass, "exec_command_wait", hq_env_mock_exec_command_wait) + monkeypatch.setattr( + TransportClass, "exec_command_wait", hq_env_mock_exec_command_wait + ) result = runner.invoke(cmd_info, "localhost-hq") assert result.exit_code == 1 @@ -95,16 +121,26 @@ def test_server_info_ssh(runner: CliRunner, aiida_computer_ssh, monkeypatch: pyt assert result.exit_code == 0 assert "Start date" in result.output -def test_server_start_info_stop_circle(runner: CliRunner, aiida_computer_ssh, monkeypatch: pytest.MonkeyPatch, hq_env_mock_exec_command_wait, server_dir_mock_exec_command_wait, tmp_path): + +def test_server_start_info_stop_circle( + runner: CliRunner, + aiida_computer_ssh, + monkeypatch: pytest.MonkeyPatch, + hq_env_mock_exec_command_wait, + server_dir_mock_exec_command_wait, + tmp_path, +): """Test server start, info, stop cile - + NOTE: this test is run hq command using the system process instead of using the hq_env which handle the drop operation to clean the files and process. Therefor, it may cause problem that the process are left not cleaned. """ - aiida_computer_ssh(label='localhost-hq') + aiida_computer_ssh(label="localhost-hq") - monkeypatch.setattr(TransportClass, "exec_command_wait", server_dir_mock_exec_command_wait) + monkeypatch.setattr( + TransportClass, "exec_command_wait", server_dir_mock_exec_command_wait + ) # Before start, server info will fail result = runner.invoke(cmd_info, "localhost-hq") @@ -120,7 +156,7 @@ def test_server_start_info_stop_circle(runner: CliRunner, aiida_computer_ssh, mo # Wait a bit until the server is start time.sleep(2) - # Check the hq server is started + # Check the hq server is started result = runner.invoke(cmd_info, "localhost-hq") assert result.exit_code == 0 @@ -139,4 +175,3 @@ def test_server_start_info_stop_circle(runner: CliRunner, aiida_computer_ssh, mo assert result.exit_code == 1 assert "Critical: cannot obtain HyperQueue server information" in result.output - diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index dfdeb60..5900f07 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -162,14 +162,16 @@ def test_get_and_parse_joblist(hq_env: HqEnv): """Test whether _parse_joblist_output can parse the hq job list output""" scheduler = HyperQueueScheduler() - joblist_command = scheduler._get_joblist_command().split(' ')[1:] + joblist_command = scheduler._get_joblist_command().split(" ")[1:] hq_env.start_server() # waiting and cancel hq_env.command(["submit", "--", "bash", "-c", "echo '1 canceled'"]) hq_env.command(["submit", "--", "bash", "-c", "echo '2 waiting'"]) - hq_env.command(["submit", "--", "bash", "-c", "echoooo '3 failed'"]) # This job will failed + hq_env.command( + ["submit", "--", "bash", "-c", "echoooo '3 failed'"] + ) # This job will failed r = hq_env.command(["job", "cancel", "1"]) assert "Job 1 canceled" in r @@ -178,7 +180,7 @@ def test_get_and_parse_joblist(hq_env: HqEnv): # Test parse waiting job_info_list = scheduler._parse_joblist_output(0, joblist, "") - + assert len(job_info_list) == 2 for job_info in job_info_list: @@ -191,7 +193,6 @@ def test_get_and_parse_joblist(hq_env: HqEnv): wait_for_job_state(hq_env, 2, "FINISHED") wait_for_job_state(hq_env, 3, "FAILED") - # Test parse running hq_env.command(["submit", "--", "sleep", "1"]) @@ -205,4 +206,3 @@ def test_get_and_parse_joblist(hq_env: HqEnv): assert len(job_info_list) == 1 assert job_info_list[0].job_state == JobState.RUNNING assert job_info_list[0].title == "sleep" -