Skip to content

Commit

Permalink
Added new test
Browse files Browse the repository at this point in the history
  • Loading branch information
Cmurilochem committed Feb 22, 2024
1 parent 335521f commit f294a85
Showing 1 changed file with 64 additions and 35 deletions.
99 changes: 64 additions & 35 deletions n3fit/src/n3fit/tests/test_hyperopt.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,35 +82,6 @@ def test_restart_from_pickle(tmp_path):
assert restart_json[i]['result'] == direct_json[i]['result']


def start_mongo_database(tmp_path):
"""Creates MongoDB database and returns the Popen object."""
db_command = ["mongod", "--dbpath", f"{tmp_path}/hyperopt-db"]
directory_path = f"{tmp_path}/hyperopt-db"
try:
# create database directory
sp.run(["mkdir", "-p", directory_path], check=True)
# launch database
process = sp.Popen(db_command, cwd=tmp_path)
return process
except (sp.CalledProcessError, OSError) as err:
msg = f"Error creating directory or executing {db_command}: {err}"
raise EnvironmentError(msg) from err


def stop_mongod_command(process):
"""Stops the MongoDB database."""
# directory_path = f"{tmp_path}/hyperopt"
try:
# stop mongod command
process.terminate()
process.wait()
# remove database files
# sp.run(f"rm -r {directory_path} && rm -r {tmp_path}/65*", check=True)
except (sp.CalledProcessError, OSError) as err:
msg = f"Error stopping the MongoDB process or removing database files: {err}"
raise EnvironmentError(msg) from err


def test_parallel_hyperopt(tmp_path):
"""Ensure that the parallel implementation of hyperopt with MongoDB works as expected."""
# Prepare the run
Expand Down Expand Up @@ -138,9 +109,6 @@ def test_parallel_hyperopt(tmp_path):
end_time = time.time()
sequential_run_time = end_time - start_time

# Generate on-the-fly a real MongoDB database
process = start_mongo_database(tmp_path)

# Run hyperopt in parallel
start_time = time.time()
sp.run(
Expand All @@ -153,9 +121,6 @@ def test_parallel_hyperopt(tmp_path):
end_time = time.time()
parallel_run_time = end_time - start_time

# Stop mongod command
stop_mongod_command(process)

# Read up generated json files
sequential_json_path = f"{output_sequential}/nnfit/replica_{REPLICA}/tries.json"
sequential_json = load_data(sequential_json_path)
Expand All @@ -175,3 +140,67 @@ def test_parallel_hyperopt(tmp_path):
# Note: cannot check that they share exactly the same history
# as the hyperopt algorithm depends on the results from previous runs
# which is obviously different between parallel and sequential runs


def clean_up_database(tmp_path):
"""Stops the MongoDB database."""
directory_path = f"{tmp_path}/hyperopt-db"
try:
sp.run(f"rm -r {directory_path} {tmp_path}/65*", shell=True, check=True)
except (sp.CalledProcessError, OSError) as err:
msg = f"Error cleaning up database: {err}"
raise EnvironmentError(msg) from err


def test_restart_from_tar(tmp_path):
"""Ensure that our parallel hyperopt restart works as expected"""
# Prepare the run
quickcard = f"hyper-{QUICKNAME}.yml"
quickpath = REGRESSION_FOLDER / quickcard

# Set up some options
n_mongo_workers = 3
n_trials_stop = 3
n_trials_total = 6
output = tmp_path / "output"

# cp runcard to tmp folder
shutil.copy(quickpath, tmp_path)
# run some trials for the first time
sp.run(
f"{EXE} {quickpath} {REPLICA} --hyperopt {n_trials_stop} "
f"--parallel-hyperopt --num-mongo-workers {n_mongo_workers} "
f"-o {output}".split(),
cwd=tmp_path,
check=True,
)
json_path = f"{output}/nnfit/replica_{REPLICA}/tries.json"
initial_json = load_data(json_path)

# check if the calculation went well
assert len(initial_json) == n_trials_stop

# just in case, remove old database files to ensure that the restart occurs via tar file
clean_up_database(tmp_path)

# restart and calculate more trials
sp.run(
f"{EXE} {quickpath} {REPLICA} --hyperopt {n_trials_total} "
f"--parallel-hyperopt --num-mongo-workers {n_mongo_workers} "
f"-o {output}".split(),
cwd=tmp_path,
check=True,
)
final_json = load_data(json_path)

# check if the run completed all trials
assert len(final_json) == n_trials_total

print(initial_json)

for i in range(n_trials_stop):
# check that the files share exactly the same hyperopt history until the restart
assert initial_json[i]['misc'] == final_json[i]['misc']
assert initial_json[i]['state'] == final_json[i]['state']
assert initial_json[i]['tid'] == final_json[i]['tid']
# assert initial_json[i]['result'] == final_json[i]['result']

0 comments on commit f294a85

Please sign in to comment.