Using zipped solver data

Author: momchil-flex

CHANGELOG.md

@@ -15,6 +15,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 
+### Changed
+- All solver output is now compressed. However, it is automatically unpacked to the same `simulation_data.hdf5` by default when loading simulation data from the server.
+
+### Fixed
+
+
 ## [2.5.0] - 2023-12-13
 
 ### Added

tests/test_plugins/test_mode_solver.py

@@ -38,7 +38,7 @@ def mock_remote_api(monkeypatch):
     def void(*args, **kwargs):
         return None
 
-    def mock_download(task_id, remote_path, to_file, *args, **kwargs):
+    def mock_download(resource_id, remote_filename, to_file, *args, **kwargs):
         simulation = td.Simulation(
             size=SIM_SIZE,
             grid_spec=td.GridSpec(wavelength=1.0),
@@ -68,6 +68,7 @@ def mock_download(task_id, remote_path, to_file, *args, **kwargs):
     monkeypatch.setattr(httputil, "api_key", lambda: "api_key")
     monkeypatch.setattr(httputil, "get_version", lambda: td.version.__version__)
     monkeypatch.setattr("tidy3d.web.api.mode.upload_file", void)
+    monkeypatch.setattr("tidy3d.web.api.mode.download_gz_file", mock_download)
     monkeypatch.setattr("tidy3d.web.api.mode.download_file", mock_download)
 
     responses.add(

tests/test_web/test_tidy3d_task.py

@@ -89,7 +89,7 @@ def mock_download(*args, **kwargs):
         to_file = kwargs["to_file"]
         sim.to_file(to_file)
 
-    monkeypatch.setattr("tidy3d.web.core.task_core.download_file", mock_download)
+    monkeypatch.setattr("tidy3d.web.core.task_core.download_gz_file", mock_download)
 
     responses.add(
         responses.GET,

tests/test_web/test_webapi.py

@@ -225,6 +225,7 @@ def _mock_download(*args, **kwargs):
         with open(file_path, "w") as f:
             f.write("0.3,5.7")
 
+    monkeypatch.setattr(f"{task_core_path}.download_gz_file", _mock_download)
     monkeypatch.setattr(f"{task_core_path}.download_file", _mock_download)
     download(TASK_ID, str(tmp_path / "web_test_tmp.json"))
     with open(str(tmp_path / "web_test_tmp.json")) as f:
@@ -356,7 +357,7 @@ def mock_download(*args, **kwargs):
     def get_str(*args, **kwargs):
         return sim.json().encode("utf-8")
 
-    monkeypatch.setattr(f"{task_core_path}.download_file", mock_download)
+    monkeypatch.setattr(f"{task_core_path}.download_gz_file", mock_download)
     monkeypatch.setattr(f"{task_core_path}.read_simulation_from_hdf5", get_str)
 
     fname_tmp = str(tmp_path / "web_test_tmp.json")

tests/test_web/test_webapi_heat.py

@@ -3,6 +3,7 @@
 import pytest
 import responses
 from _pytest import monkeypatch
+from botocore.exceptions import ClientError
 
 import tidy3d as td
 from responses import matchers
@@ -170,6 +171,12 @@ def _mock_download(*args, **kwargs):
         with open(file_path, "w") as f:
             f.write("0.3,5.7")
 
+    def _mock_download_gz(*args, **kwargs):
+        raise ClientError(
+            error_response={"Error": {"Message": "File not found"}}, operation_name="download"
+        )
+
+    monkeypatch.setattr(f"{task_core_path}.download_gz_file", _mock_download_gz)
     monkeypatch.setattr(f"{task_core_path}.download_file", _mock_download)
     download(TASK_ID, str(tmp_path / "web_test_tmp.json"))
     with open(str(tmp_path / "web_test_tmp.json"), "r") as f:
@@ -256,7 +263,7 @@ def mock_download(*args, **kwargs):
     def get_str(*args, **kwargs):
         return sim.json().encode("utf-8")
 
-    monkeypatch.setattr(f"{task_core_path}.download_file", mock_download)
+    monkeypatch.setattr(f"{task_core_path}.download_gz_file", mock_download)
     monkeypatch.setattr(f"{task_core_path}.read_simulation_from_hdf5", get_str)
 
     fname_tmp = str(tmp_path / "web_test_tmp.json")

tidy3d/version.py

@@ -1,3 +1,2 @@
 """Defines the front end version of tidy3d"""
-
 __version__ = "2.6.0rc1"

tidy3d/web/api/mode.py

@@ -10,12 +10,15 @@
 import time
 
 import pydantic.v1 as pydantic
+from botocore.exceptions import ClientError
+
 from ...components.simulation import Simulation
 from ...components.data.monitor_data import ModeSolverData
 from ...exceptions import WebError
 from ...log import log, get_logging_console
+from ..core.core_config import get_logger_console
 from ..core.http_util import http
-from ..core.s3utils import download_file, upload_file
+from ..core.s3utils import download_file, download_gz_file, upload_file
 from ..core.task_core import Folder
 from ..core.types import ResourceLifecycle, Submittable
 
@@ -31,6 +34,7 @@
 
 MODESOLVER_LOG = "output/result.log"
 MODESOLVER_RESULT = "output/result.hdf5"
+MODESOLVER_RESULT_GZ = "output/mode_solver_data.hdf5.gz"
 
 
 def run(
@@ -424,18 +428,36 @@ def get_result(
         :class:`.ModeSolverData`
             Mode solver data with the calculated results.
         """
+
+        file = None
         try:
-            download_file(
-                self.solver_id,
-                MODESOLVER_RESULT,
+            file = download_gz_file(
+                resource_id=self.task_id,
+                remote_filename=MODESOLVER_RESULT_GZ,
                 to_file=to_file,
                 verbose=verbose,
                 progress_callback=progress_callback,
             )
-        except Exception:
-            raise WebError(
-                f"Failed to download file '{MODESOLVER_RESULT}' from server. Please confirm that the task was successful."
-            )
+        except ClientError:
+            if verbose:
+                console = get_logger_console()
+                console.log(f"Unable to download '{MODESOLVER_RESULT}'.")
+
+        if not file:
+            try:
+                file = download_file(
+                    resource_id=self.task_id,
+                    remote_filename=MODESOLVER_RESULT,
+                    to_file=to_file,
+                    verbose=verbose,
+                    progress_callback=progress_callback,
+                )
+            except Exception as e:
+                raise WebError(
+                    "Failed to download the simulation data file from the server. "
+                    "Please confirm that the task was successfully run."
+                ) from e
+
         data = ModeSolverData.from_hdf5(to_file)
         data = data.copy(
             update={"monitor": self.mode_solver.to_mode_solver_monitor(name=MODE_MONITOR_NAME)}

tidy3d/web/core/constants.py

@@ -21,6 +21,7 @@
 
 SIMULATION_JSON = "simulation.json"
 SIMULATION_DATA_HDF5 = "output/monitor_data.hdf5"
+SIMULATION_DATA_HDF5_GZ = "output/simulation_data.hdf5.gz"
 RUNNING_INFO = "output/solver_progress.csv"
 SIM_LOG_FILE = "output/tidy3d.log"
 SIM_FILE_HDF5 = "simulation.hdf5"

tidy3d/web/core/s3utils.py

@@ -1,5 +1,8 @@
 """handles filesystem, storage
 """
+import os
+import tempfile
+
 import pathlib
 import urllib
 from datetime import datetime
@@ -12,8 +15,10 @@
 from rich.progress import TextColumn, Progress, BarColumn, DownloadColumn
 from rich.progress import TransferSpeedColumn, TimeRemainingColumn
 from .http_util import http
+from .file_util import extract_gzip_file
 from .environment import Env
 from .core_config import get_logger_console
+from .exceptions import WebError
 
 
 class _UserCredential(BaseModel):
@@ -283,8 +288,8 @@ def download_file(
     ----------
     resource_id : str
         The resource id, e.g. task id.
-    content : str
-        The content of the file
+    remote_filename : str
+        Path to the remote file.
     to_file : str = None
         Local filename to save to, if not specified, use the remote_filename.
     verbose : bool = True
@@ -346,3 +351,58 @@ def _callback(bytes_in_chunk):
             _download(lambda bytes_in_chunk: None)
 
     return to_file
+
+
+def download_gz_file(
+    resource_id: str,
+    remote_filename: str,
+    to_file: str = None,
+    verbose: bool = True,
+    progress_callback: Callable[[float], None] = None,
+) -> pathlib.Path:
+    """Download a ``.gz`` file and unzip it into ``to_file``, unless ``to_file`` itself
+    ends in .gz
+
+    Parameters
+    ----------
+    resource_id : str
+        The resource id, e.g. task id.
+    remote_filename : str
+        Path to the remote file.
+    to_file : str = None
+        Local filename to save to, if not specified, use the remote_filename.
+    verbose : bool = True
+        Whether to display a progressbar for the upload
+    progress_callback : Callable[[float], None] = None
+        User-supplied callback function with ``bytes_in_chunk`` as argument.
+    """
+
+    # If to_file is a gzip extension, just download
+    if to_file.lower().endswith(".gz"):
+        return download_file(
+            resource_id,
+            remote_filename,
+            to_file=to_file,
+            verbose=verbose,
+            progress_callback=progress_callback,
+        )
+
+    # Otherwise, download and unzip
+    # The tempfile is set as ``hdf5.gz`` so that the mock download in the webapi tests works
+    tmp_file, tmp_file_path = tempfile.mkstemp(".hdf5.gz")
+    os.close(tmp_file)
+    try:
+        download_file(
+            resource_id,
+            remote_filename,
+            to_file=tmp_file_path,
+            verbose=verbose,
+            progress_callback=progress_callback,
+        )
+        if os.path.exists(tmp_file_path):
+            extract_gzip_file(tmp_file_path, to_file)
+        else:
+            raise WebError(f"Failed to download and extract '{remote_filename}'.")
+    finally:
+        os.unlink(tmp_file_path)
+    return to_file

tidy3d/web/core/task_core.py

@@ -8,21 +8,21 @@
 from typing import List, Optional, Callable, Tuple
 import pydantic.v1 as pd
 from pydantic.v1 import Extra, Field, parse_obj_as
+from botocore.exceptions import ClientError
 
 from . import http_util
 from .core_config import get_logger_console
 from .exceptions import WebError
 
 from .cache import FOLDER_CACHE
 from .http_util import http
-from .s3utils import download_file, upload_file
+from .s3utils import download_file, download_gz_file, upload_file
 from .stub import TaskStub
 from .types import Queryable, ResourceLifecycle, Submittable
 from .types import Tidy3DResource
 
-
-from .constants import SIM_FILE_HDF5_GZ, SIMULATION_DATA_HDF5, SIM_LOG_FILE
-from .file_util import extract_gzip_file, read_simulation_from_hdf5
+from .constants import SIM_FILE_HDF5_GZ, SIMULATION_DATA_HDF5, SIMULATION_DATA_HDF5_GZ, SIM_LOG_FILE
+from .file_util import read_simulation_from_hdf5
 
 
 class Folder(Tidy3DResource, Queryable, extra=Extra.allow):
@@ -453,7 +453,7 @@ def estimate_cost(self, solver_version=None) -> float:
     def get_sim_data_hdf5(
         self, to_file: str, verbose: bool = True, progress_callback: Callable[[float], None] = None
     ) -> pathlib.Path:
-        """Get output/monitor_data.hdf5 file from Server.
+        """Get simulation data file from Server.
 
         Parameters
         ----------
@@ -472,18 +472,36 @@ def get_sim_data_hdf5(
         if not self.task_id:
             raise WebError("Expected field 'task_id' is unset.")
 
+        file = None
         try:
-            return download_file(
-                self.task_id,
-                SIMULATION_DATA_HDF5,
+            file = download_gz_file(
+                resource_id=self.task_id,
+                remote_filename=SIMULATION_DATA_HDF5_GZ,
                 to_file=to_file,
                 verbose=verbose,
                 progress_callback=progress_callback,
             )
-        except Exception:
-            raise WebError(
-                f"Failed to download file '{SIMULATION_DATA_HDF5}' from server. Please confirm that the task was successful."
-            )
+        except ClientError:
+            if verbose:
+                console = get_logger_console()
+                console.log(f"Unable to download '{SIMULATION_DATA_HDF5_GZ}'.")
+
+        if not file:
+            try:
+                file = download_file(
+                    resource_id=self.task_id,
+                    remote_filename=SIMULATION_DATA_HDF5,
+                    to_file=to_file,
+                    verbose=verbose,
+                    progress_callback=progress_callback,
+                )
+            except Exception as e:
+                raise WebError(
+                    "Failed to download the simulation data file from the server. "
+                    "Please confirm that the task was successfully run."
+                ) from e
+
+        return file
 
     def get_simulation_hdf5(
         self, to_file: str, verbose: bool = True, progress_callback: Callable[[float], None] = None
@@ -507,31 +525,13 @@ def get_simulation_hdf5(
         if not self.task_id:
             raise WebError("Expected field 'task_id' is unset.")
 
-        if to_file.lower().endswith(".gz"):
-            download_file(
-                self.task_id,
-                SIM_FILE_HDF5_GZ,
-                to_file=to_file,
-                verbose=verbose,
-                progress_callback=progress_callback,
-            )
-        else:
-            hdf5_gz_file, hdf5_gz_file_path = tempfile.mkstemp(".hdf5.gz")
-            os.close(hdf5_gz_file)
-            try:
-                download_file(
-                    self.task_id,
-                    SIM_FILE_HDF5_GZ,
-                    to_file=hdf5_gz_file_path,
-                    verbose=verbose,
-                    progress_callback=progress_callback,
-                )
-                if os.path.exists(hdf5_gz_file_path):
-                    extract_gzip_file(hdf5_gz_file_path, to_file)
-                else:
-                    raise WebError("Failed to download simulation.hdf5")
-            finally:
-                os.unlink(hdf5_gz_file_path)
+        return download_gz_file(
+            resource_id=self.task_id,
+            remote_filename=SIM_FILE_HDF5_GZ,
+            to_file=to_file,
+            verbose=verbose,
+            progress_callback=progress_callback,
+        )
 
     def get_running_info(self) -> Tuple[float, float]:
         """Gets the % done and field_decay for a running task.