Skip to content

Commit

Permalink
test_calculation.py retrieved
Browse files Browse the repository at this point in the history
  • Loading branch information
khsrali committed Jul 16, 2024
1 parent 689ef55 commit 70450fe
Show file tree
Hide file tree
Showing 5 changed files with 173 additions and 7 deletions.
4 changes: 2 additions & 2 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -44,5 +44,5 @@ repos:
additional_dependencies:
- "types-PyYAML"
- "types-requests"
- "pyfirecrest~=2.5.0"
- "aiida-core~=2.5.1.post0"
- "pyfirecrest>=2.5.0" # please change to 2.6.0 when released
- "aiida-core>=2.6.0" # please change to 2.6.2 when released
5 changes: 2 additions & 3 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,16 +1,15 @@
# Changelog

## v0.2.0 - 2024-07-15 (not released yet)
## v0.2.0 - (not released yet)

### Transport plugin
- `dynamic_info()` is added to retrieve machine information without user input.
- Refactor `put` & `get` & `copy` now they mimic behavior `aiida-ssh` transport plugin.
- `put` & `get` & `copy` now support glob patterns.
- Added `dereference` option wherever relevant
- Added `recursive` functionality for `listdir`
- Added `_create_secret_file` to store user secret locally in `~/.firecrest/`
- Added `_validate_temp_directory` to allocate a temporary directory useful for `extract` and `compress` methods on FirecREST server.
- Added `_dynamic_info_direct_size` this is able to get info of direct transfer from the server rather than asking from users. Raise of user inputs fails to make a connection.
- Added `_dynamic_info_direct_size` this is able to get info of direct transfer from the server rather than asking from users. Raise if fails to make a connection.
- Added `_validate_checksum` to check integrity of downloaded/uploaded files.
- Added `_gettreetar` & `_puttreetar` to transfer directories as tar files internally.
- Added `payoff` function to calculate when is gainful to transfer as zip, and when to transfer individually.
Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@ classifiers = [
keywords = ["aiida", "firecrest"]
requires-python = ">=3.9"
dependencies = [
"aiida-core@git+https://github.com/aiidateam/aiida-core.git@954cbdd3",
"aiida-core@git+https://github.com/aiidateam/aiida-core.git@954cbdd",
"click",
"pyfirecrest@git+https://github.com/khsrali/pyfirecrest.git@main#egg=pyfirecrest",
"pyfirecrest@git+https://github.com/eth-cscs/pyfirecrest.git@6cae414",
"pyyaml",
]

Expand Down
22 changes: 22 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ def _firecrest_computer(myfirecrest, tmpdir: Path):
small_file_size_mb=1.0,
temp_directory=str(_temp_directory),
)
computer.store()
return computer


Expand Down Expand Up @@ -105,6 +106,27 @@ def submit(
if script_remote_path and not Path(script_remote_path).exists():
raise FileNotFoundError(f"File {script_remote_path} does not exist")
job_id = random.randint(10000, 99999)

# Filter out lines starting with '#SBATCH'
with open(script_remote_path) as file:
lines = file.readlines()
command = "".join([line for line in lines if not line.strip().startswith("#")])

# Make the dummy files
for line in lines:
if "--error" in line:
error_file = line.split("=")[1].strip()
(Path(script_remote_path).parent / error_file).touch()
elif "--output" in line:
output_file = line.split("=")[1].strip()
(Path(script_remote_path).parent / output_file).touch()

# Execute the job, this is useful for test_calculation.py
if "aiida.in" in command:
# skip blank command like: '/bin/bash'
os.chdir(Path(script_remote_path).parent)
os.system(command)

return {"jobid": job_id}


Expand Down
145 changes: 145 additions & 0 deletions tests/test_calculation.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,145 @@
"""Test for running calculations on a FireCREST computer."""
from pathlib import Path

from aiida import common, engine, manage, orm
from aiida.common.folders import Folder
from aiida.engine.processes.calcjobs.tasks import MAX_ATTEMPTS_OPTION
from aiida.manage.tests.pytest_fixtures import EntryPointManager
from aiida.parsers import Parser
import pytest


@pytest.fixture(name="no_retries")
def _no_retries():
"""Remove calcjob retries, to make failing the test faster."""
# TODO calculation seems to hang on errors still
max_attempts = manage.get_config().get_option(MAX_ATTEMPTS_OPTION)
manage.get_config().set_option(MAX_ATTEMPTS_OPTION, 1)
yield
manage.get_config().set_option(MAX_ATTEMPTS_OPTION, max_attempts)


@pytest.mark.usefixtures("aiida_profile_clean", "no_retries")
def test_calculation_basic(firecrest_computer: orm.Computer):
"""Test running a simple `arithmetic.add` calculation."""
code = orm.InstalledCode(
label="test_code",
description="test code",
default_calc_job_plugin="core.arithmetic.add",
computer=firecrest_computer,
filepath_executable="/bin/sh",
)
code.store()

builder = code.get_builder()
builder.x = orm.Int(1)
builder.y = orm.Int(2)

_, node = engine.run_get_node(builder)
assert node.is_finished_ok


@pytest.mark.usefixtures("aiida_profile_clean", "no_retries")
def test_calculation_file_transfer(
firecrest_computer: orm.Computer, entry_points: EntryPointManager
):
"""Test a calculation, with multiple files copied/uploaded/retrieved."""
# add temporary entry points
entry_points.add(MultiFileCalcjob, "aiida.calculations:testing.multifile")
entry_points.add(NoopParser, "aiida.parsers:testing.noop")

# add a remote file which is used remote_copy_list
firecrest_computer.get_transport()._cwd.joinpath(
firecrest_computer.get_workdir(), "remote_copy.txt"
).touch()

# setup the calculation
code = orm.InstalledCode(
label="test_code",
description="test code",
default_calc_job_plugin="testing.multifile",
computer=firecrest_computer,
filepath_executable="/bin/sh",
)
code.store()
builder = code.get_builder()

node: orm.CalcJobNode
_, node = engine.run_get_node(builder)
assert node.is_finished_ok

if (retrieved := node.get_retrieved_node()) is None:
raise RuntimeError("No retrieved node found")

paths = sorted([str(p) for p in retrieved.base.repository.glob()])
assert paths == [
"_scheduler-stderr.txt",
"_scheduler-stdout.txt",
"folder1",
"folder1/a",
"folder1/a/b.txt",
"folder1/a/c.txt",
"folder2",
"folder2/remote_copy.txt",
"folder2/x",
"folder2/y",
"folder2/y/z",
]


class MultiFileCalcjob(engine.CalcJob):
"""A complex CalcJob that creates/retrieves multiple files."""

@classmethod
def define(cls, spec):
"""Define the process specification."""
super().define(spec)
spec.inputs["metadata"]["options"]["resources"].default = {
"num_machines": 1,
"num_mpiprocs_per_machine": 1,
}
spec.input(
"metadata.options.parser_name", valid_type=str, default="testing.noop"
)
spec.exit_code(400, "ERROR", message="Calculation failed.")

def prepare_for_submission(self, folder: Folder) -> common.CalcInfo:
"""Prepare the calculation job for submission."""
codeinfo = common.CodeInfo()
codeinfo.code_uuid = self.inputs.code.uuid

path = Path(folder.get_abs_path("a")).parent
for subpath in [
"i.txt",
"j.txt",
"folder1/a/b.txt",
"folder1/a/c.txt",
"folder1/a/c.in",
"folder1/c.txt",
"folder2/x",
"folder2/y/z",
]:
path.joinpath(subpath).parent.mkdir(parents=True, exist_ok=True)
path.joinpath(subpath).touch()

calcinfo = common.CalcInfo()
calcinfo.codes_info = [codeinfo]
calcinfo.retrieve_list = [("folder1/*/*.txt", ".", 99), ("folder2", ".", 99)]
comp: orm.Computer = self.inputs.code.computer
calcinfo.remote_copy_list = [
(
comp.uuid,
f"{comp.get_workdir()}/remote_copy.txt",
"folder2/remote_copy.txt",
)
]
# TODO also add remote_symlink_list

return calcinfo


class NoopParser(Parser):
"""Parser that does absolutely nothing!"""

def parse(self, **kwargs):
pass

0 comments on commit 70450fe

Please sign in to comment.