Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 1 addition & 44 deletions src/nhp/docker/run.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
"""Run the model inside of the docker container."""

import gzip
import json
import logging
import os
Expand All @@ -17,7 +16,7 @@

from nhp.docker.config import Config
from nhp.model.params import load_params
from nhp.model.results import generate_results_json, save_results_files
from nhp.model.results import save_results_files
from nhp.model.run import noop_progress_callback
Comment thread
tomjemmett marked this conversation as resolved.


Expand Down Expand Up @@ -165,39 +164,6 @@ def _get_data(self, year: str, dataset: str) -> None:
file_client = fs_client.get_file_client(filename)
local_file.write(file_client.download_file().readall())

def _upload_results_json(
self, results: dict[str, pd.DataFrame], metadata: dict[str, Any], variants: list[str]
) -> None:
"""Upload the results.

Once the model has run, upload the results to blob storage.

Args:
results: Dictionary containing the results dataframes.
metadata: The metadata to attach to the blob.
variants: A list of the variants that were run.
"""
container = self._get_container(self._config.RESULTS_STORAGE_ACCOUNT, "results")

logging.info("Generating results json file")
results_file = generate_results_json(results, self.params, variants)

results_json_gz_path = f"prod/{self._app_version}/{results_file}.json.gz"
logging.info("Uploading results json file to blob storage: %s", results_json_gz_path)
with open(f"results/{results_file}.json", "rb") as file:
container.upload_blob(
results_json_gz_path,
gzip.compress(file.read()),
metadata={k: str(v) for k, v in metadata.items()},
overwrite=True,
)

logging.info("Updating table storage with results json path: %s", results_json_gz_path)
self._update_table_storage(
results_json_gz_path=results_json_gz_path,
)
logging.info("Results json file uploaded and table storage updated successfully.")

def _upload_results_files(
self,
file_path: str,
Expand Down Expand Up @@ -317,15 +283,6 @@ def finish(
self._upload_results_files(
file_path, results, {"model_run_id": str(self._model_run_id)}, variants
)
# ---
# see issue #286, this should be removed once we no longer need the results json file
metadata = {
k: v
for k, v in self.params.items()
if not isinstance(v, dict) and not isinstance(v, list)
}
metadata.update(additional_metadata)
self._upload_results_json(results, metadata, variants)
## ---
if save_full_model_results:
self._upload_full_model_results()
Expand Down
68 changes: 0 additions & 68 deletions src/nhp/model/results.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,74 +100,6 @@ def _combine_model_results(
}


def generate_results_json(
results: dict[str, pd.DataFrame],
params: dict,
variants: list[str],
) -> str:
"""Generate the results in the json format and save."""

def agg_to_dict(res):
results_df = res.set_index("model_run")
return (
pd.concat(
[
results_df.loc[0]
.set_index([i for i in results_df.columns if i != "value"])
.rename(columns={"value": "baseline"}),
results_df.loc[results_df.index != 0]
.groupby([i for i in results_df.columns if i != "value"])
.agg(list)
.rename(columns={"value": "model_runs"}),
],
axis=1,
)
.reset_index()
.to_dict(orient="records")
)

dict_results = {
k: agg_to_dict(v) for k, v in results.items() if k != "step_counts" if len(v) > 0
}

dict_results["step_counts"] = (
results["step_counts"]
.groupby(
[
"pod",
"change_factor",
"strategy",
"sitetret",
"activity_type",
"measure",
]
)[["value"]]
.agg(list)
.reset_index()
.to_dict("records")
)

for i in dict_results["step_counts"]:
i["model_runs"] = i.pop("value")
if i["change_factor"] == "baseline":
i["model_runs"] = i["model_runs"][0:1]
if i["strategy"] == "-":
i.pop("strategy")

filename = f"{params['dataset']}/{params['scenario']}-{params['create_datetime']}"
os.makedirs(f"results/{params['dataset']}", exist_ok=True)
with open(f"results/{filename}.json", "w", encoding="utf-8") as file:
json.dump(
{
"params": params,
"population_variants": variants,
"results": dict_results,
},
file,
)
return filename


def save_results_files(results: dict, params: dict, variants: list[str]) -> list:
"""Save aggregated and combined results as parquet, and params as JSON.

Expand Down
41 changes: 0 additions & 41 deletions tests/unit/nhp/docker/test_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -282,39 +282,6 @@ def test_RunWithAzureStorage_get_data_skips_missing_directory(mock_run_with_azur
mock_file.assert_not_called()


@pytest.mark.unit
def test_RunWithAzureStorage_upload_results_json(mock_run_with_azure_storage, mocker):
# arrange
s = mock_run_with_azure_storage

m_get_container = mocker.patch("nhp.docker.run.RunWithAzureStorage._get_container")
m_update_table_storage = mocker.patch(
"nhp.docker.run.RunWithAzureStorage._update_table_storage"
)
m_generate_results_json = mocker.patch(
"nhp.docker.run.generate_results_json", return_value="filename"
)
mocker.patch("gzip.compress", return_value="gzdata")
metadata = {"k": "v", "count": 1}

# act
with patch("builtins.open", mock_open(read_data="data")) as mock_file:
s._upload_results_json("filename", metadata, "variants")

# assert
mock_file.assert_called_once_with("results/filename.json", "rb")
m_get_container.assert_called_once_with("results-sa", "results")
m_get_container().upload_blob.assert_called_once_with(
"prod/dev/filename.json.gz",
"gzdata",
metadata={"k": "v", "count": "1"},
overwrite=True,
)
m_generate_results_json.assert_called_once_with("filename", {"dataset": "test"}, "variants")

m_update_table_storage.assert_called_once_with(results_json_gz_path="prod/dev/filename.json.gz")


@pytest.mark.unit
def test_RunWithAzureStorage_upload_results_files(mock_run_with_azure_storage, mocker):
# arrange
Expand Down Expand Up @@ -435,8 +402,6 @@ def test_RunWithAzureStorage_finish_save_full_model_results_false(
m3 = mocker.patch("nhp.docker.run.RunWithAzureStorage._upload_full_model_results")
m4 = mocker.patch("nhp.docker.run.RunWithAzureStorage._cleanup")

m5 = mocker.patch("nhp.docker.run.RunWithAzureStorage._upload_results_json")

metadata = {
"id": "1",
"dataset": "synthetic",
Expand Down Expand Up @@ -486,8 +451,6 @@ def test_RunWithAzureStorage_finish_save_full_model_results_false(
m3.assert_not_called()
m4.assert_called_once_with()

m5.assert_called_once_with("results", metadata_expected, "variants")


@pytest.mark.unit
def test_RunWithAzureStorage_finish_save_full_model_results_true(
Expand All @@ -500,8 +463,6 @@ def test_RunWithAzureStorage_finish_save_full_model_results_true(
m3 = mocker.patch("nhp.docker.run.RunWithAzureStorage._upload_full_model_results")
m4 = mocker.patch("nhp.docker.run.RunWithAzureStorage._cleanup")

m5 = mocker.patch("nhp.docker.run.RunWithAzureStorage._upload_results_json")

metadata = {
"id": "1",
"dataset": "synthetic",
Expand Down Expand Up @@ -549,8 +510,6 @@ def test_RunWithAzureStorage_finish_save_full_model_results_true(
m3.assert_called_once()
m4.assert_called_once_with()

m5.assert_called_once_with("results", metadata_expected, "variants")


@pytest.mark.unit
def test_RunWithAzureStorage_error(mock_run_with_azure_storage, mocker):
Expand Down
121 changes: 0 additions & 121 deletions tests/unit/nhp/model/test_results.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
_save_parquet_file,
_save_variants_file,
combine_results,
generate_results_json,
save_results_files,
)

Expand Down Expand Up @@ -144,126 +143,6 @@ def test_combine_model_results(mocker):
assert cfd_mock.call_count == 10


@pytest.mark.unit
def test_generate_results_json(mocker):
# arrange
results = {
"default": pd.DataFrame(
{
"a": [i for i in [0, 1] for _ in range(5)],
"model_run": list(range(5)) * 2,
"value": range(10),
}
),
"a": pd.DataFrame(
{
"a": [i for i in [0, 1] for _ in list(range(5)) * 2],
"b": [i for i in [0, 1] for _ in list(range(5))] * 2,
"model_run": list(range(5)) * 4,
"value": list(range(20)),
}
),
"step_counts": pd.DataFrame(
{
"pod": ["a1"] * 4 * 5,
"change_factor": ["baseline", "a", "b", "c", "c"] * 4,
"strategy": ["-", "-", "-", "a", "b"] * 4,
"sitetret": ["s"] * 4 * 5,
"activity_type": ["a"] * 4 * 5,
"measure": ["x"] * 4 * 5,
"value": range(20),
}
),
}

os_m = mocker.patch("os.makedirs")
jd_m = mocker.patch("json.dump")

json_content = {
"default": [
{"a": 0, "baseline": 0, "model_runs": [1, 2, 3, 4]},
{"a": 1, "baseline": 5, "model_runs": [6, 7, 8, 9]},
],
"a": [
{"a": 0, "b": 0, "baseline": 0, "model_runs": [1, 2, 3, 4]},
{"a": 0, "b": 1, "baseline": 5, "model_runs": [6, 7, 8, 9]},
{"a": 1, "b": 0, "baseline": 10, "model_runs": [11, 12, 13, 14]},
{"a": 1, "b": 1, "baseline": 15, "model_runs": [16, 17, 18, 19]},
],
"step_counts": [
{
"pod": "a1",
"change_factor": "a",
"sitetret": "s",
"activity_type": "a",
"measure": "x",
"model_runs": [1, 6, 11, 16],
},
{
"pod": "a1",
"change_factor": "b",
"sitetret": "s",
"activity_type": "a",
"measure": "x",
"model_runs": [2, 7, 12, 17],
},
{
"pod": "a1",
"change_factor": "baseline",
"sitetret": "s",
"activity_type": "a",
"measure": "x",
"model_runs": [0],
},
{
"pod": "a1",
"change_factor": "c",
"strategy": "a",
"sitetret": "s",
"activity_type": "a",
"measure": "x",
"model_runs": [3, 8, 13, 18],
},
{
"pod": "a1",
"change_factor": "c",
"strategy": "b",
"sitetret": "s",
"activity_type": "a",
"measure": "x",
"model_runs": [4, 9, 14, 19],
},
],
}

params = {
"dataset": "synthetic",
"scenario": "test",
"create_datetime": "create_datetime",
}

expected = "synthetic/test-create_datetime"

# act
with patch("builtins.open", mock_open()) as mock_file:
actual = generate_results_json(results, params, [1, 2, 3]) # ty: ignore

# assert
assert actual == expected
mock_file.assert_called_once_with(
"results/synthetic/test-create_datetime.json", "w", encoding="utf-8"
)
os_m.assert_called_once_with("results/synthetic", exist_ok=True)
jd_m.assert_called_once_with(
{
"params": params,
"population_variants": [1, 2, 3],
"results": json_content,
},
mock_file(),
)


@pytest.mark.unit
def test_combine_results(mocker):
# arrange
Expand Down
Loading