Merge pull request #306 from The-Strategy-Unit/299_add_metadata_to_saved_params

add metadata to saved params json file

Author: yiwen-h

docker_run.py

@@ -139,23 +139,30 @@ def _upload_results_json(self, results_file: str, metadata: dict) -> None:
                 overwrite=True,
             )
 
-    def _upload_results_files(self, files: list) -> None:
+    def _upload_results_files(self, files: list, metadata: dict) -> None:
         """Upload the results
 
         once the model has run, upload the files (parquet for model results and json for model params) to blob storage
 
         :param files: list of files to be uploaded
         :type files: list
+        :param metadata: the metadata to attach to the blob
+        :type metadata: dict
 
         """
         container = self._get_container("results")
         for file in files:
             filename = file[8:]
+            if file.endswith(".json"):
+                metadata_to_use = metadata
+            else:
+                metadata_to_use = None
             with open(file, "rb") as f:
                 container.upload_blob(
                     f"aggregated-model-results/{self._app_version}/{filename}",
                     f.read(),
                     overwrite=True,
+                    metadata=metadata_to_use,
                 )
 
     def _upload_full_model_results(self) -> None:
@@ -203,7 +210,7 @@ def finish(
             if not isinstance(v, dict) and not isinstance(v, list)
         }
         self._upload_results_json(results_file, metadata)
-        self._upload_results_files(saved_files)
+        self._upload_results_files(saved_files, metadata)
         if save_full_model_results:
             self._upload_full_model_results()
         self._cleanup()

tests/test_docker_run.py

@@ -229,23 +229,38 @@ def test_RunWithAzureStorage_upload_results_json(mock_run_with_azure_storage, mo
     )
 
 
-def test_RunWithAzureStorage_upload_results_files(
-    mock_run_with_azure_storage, mocker
-):
+def test_RunWithAzureStorage_upload_results_files(mock_run_with_azure_storage, mocker):
     # arrange
     s = mock_run_with_azure_storage
 
     m = mocker.patch("docker_run.RunWithAzureStorage._get_container")
+    metadata = {"k": "v"}
 
     # act
     with patch("builtins.open", mock_open(read_data="data")) as mock_file:
-        s._upload_results_files(["results/filename"])
+        s._upload_results_files(["results/filename", "results/filename.json"], metadata)
 
     # assert
-    mock_file.assert_called_once_with("results/filename", "rb")
+    mock_file.call_args_list == [
+        call("results/filename", "rb"),
+        call("results/filename.json", "rb"),
+    ]
     m.assert_called_once_with("results")
-    m().upload_blob.assert_called_once_with(
-        "aggregated-model-results/dev/filename", "data", overwrite=True
+    m().upload_blob.assert_has_calls(
+        [
+            call(
+                "aggregated-model-results/dev/filename",
+                "data",
+                overwrite=True,
+                metadata=None,
+            ),
+            call(
+                "aggregated-model-results/dev/filename.json",
+                "data",
+                overwrite=True,
+                metadata=metadata,
+            ),
+        ]
     )
 
 
@@ -319,7 +334,7 @@ def test_RunWithAzureStorage_finish_save_full_model_results_false(
 
     # assert
     m1.assert_called_once_with("results_file", metadata)
-    m2.assert_called_once_with(["saved_files"])
+    m2.assert_called_once_with(["saved_files"], metadata)
     m3.assert_not_called()
     m4.assert_called_once()
 
@@ -346,7 +361,7 @@ def test_RunWithAzureStorage_finish_save_full_model_results_true(
 
     # assert
     m1.assert_called_once_with("results_file", metadata)
-    m2.assert_called_once_with(["saved_files"])
+    m2.assert_called_once_with(["saved_files"], metadata)
     m3.assert_called_once()
     m4.assert_called_once()