diff --git a/tests/test_plugins/test_mode_solver.py b/tests/test_plugins/test_mode_solver.py index b7015e3b3d..7644c57889 100644 --- a/tests/test_plugins/test_mode_solver.py +++ b/tests/test_plugins/test_mode_solver.py @@ -11,6 +11,7 @@ from tidy3d.plugins.mode.solver import compute_modes from tidy3d import ScalarFieldDataArray from tidy3d.web.environment import Env +from tidy3d.version import __version__ WAVEGUIDE = td.Structure(geometry=td.Box(size=(100, 0.5, 0.5)), medium=td.Medium(permittivity=4.0)) @@ -44,7 +45,7 @@ def mock_download(task_id, remote_path, to_file, *args, **kwargs): sources=[SRC], ) mode_spec = td.ModeSpec( - num_modes=1, + num_modes=3, target_neff=2.0, filter_pol="tm", precision="double", @@ -80,7 +81,9 @@ def mock_download(task_id, remote_path, to_file, *args, **kwargs): "projectId": PROJECT_ID, "taskName": TASK_NAME, "modeSolverName": MODESOLVER_NAME, - "fileType": "Json", + "fileType": "Gz", + "source": "Python", + "protocolVersion": __version__, } ) ], @@ -91,33 +94,7 @@ def mock_download(task_id, remote_path, to_file, *args, **kwargs): "status": "draft", "createdAt": "2023-05-19T16:47:57.190Z", "charge": 0, - "fileType": "Json", - } - }, - status=200, - ) - - responses.add( - responses.POST, - f"{Env.current.web_api_endpoint}/tidy3d/modesolver/py", - match=[ - responses.matchers.json_params_matcher( - { - "projectId": PROJECT_ID, - "taskName": TASK_NAME, - "modeSolverName": MODESOLVER_NAME, - "fileType": "Hdf5", - } - ) - ], - json={ - "data": { - "refId": TASK_ID, - "id": SOLVER_ID, - "status": "draft", - "createdAt": "2023-05-19T16:47:57.190Z", - "charge": 0, - "fileType": "Hdf5", + "fileType": "Gz", } }, status=200, @@ -149,7 +126,7 @@ def mock_download(task_id, remote_path, to_file, *args, **kwargs): "status": "queued", "createdAt": "2023-05-19T16:47:57.190Z", "charge": 0, - "fileType": "Json", + "fileType": "Gz", } }, status=200, diff --git a/tests/test_web/test_tidy3d_task.py b/tests/test_web/test_tidy3d_task.py index 49979f7989..bfeaab1615 100644 --- a/tests/test_web/test_tidy3d_task.py +++ b/tests/test_web/test_tidy3d_task.py @@ -148,6 +148,7 @@ def test_create(set_api_key): { "taskName": "test task", "callbackUrl": None, + "fileType": "Gz", "simulationType": "tidy3d", "parentTasks": None, } @@ -186,6 +187,7 @@ def test_submit(set_api_key): { "taskName": task_name, "callbackUrl": None, + "fileType": "Gz", "simulationType": "tidy3d", "parentTasks": None, } diff --git a/tests/test_web/test_webapi.py b/tests/test_web/test_webapi.py index 859e248c49..e3858ce5ab 100644 --- a/tests/test_web/test_webapi.py +++ b/tests/test_web/test_webapi.py @@ -6,7 +6,6 @@ import os import tidy3d as td -import tidy3d.web as web from responses import matchers @@ -16,7 +15,6 @@ from tidy3d.web.webapi import download_log, estimate_cost, get_info, get_run_info, get_tasks from tidy3d.web.webapi import load, load_simulation, start, upload, monitor, real_cost from tidy3d.web.container import Job, Batch -from tidy3d.web.task import TaskInfo from tidy3d.web.asynchronous import run_async from tidy3d.__main__ import main @@ -80,6 +78,7 @@ def mock_upload(monkeypatch, set_api_key): "callbackUrl": None, "simulationType": "tidy3d", "parentTasks": None, + "fileType": "Gz", } ) ], @@ -97,6 +96,7 @@ def mock_download(*args, **kwargs): pass monkeypatch.setattr("tidy3d.web.simulation_task.upload_string", mock_download) + monkeypatch.setattr("tidy3d.web.simulation_task.upload_file", mock_download) @pytest.fixture diff --git a/tidy3d/plugins/mode/web.py b/tidy3d/plugins/mode/web.py index 16847a8f20..d6cbb268fb 100644 --- a/tidy3d/plugins/mode/web.py +++ b/tidy3d/plugins/mode/web.py @@ -16,16 +16,20 @@ from ...components.data.monitor_data import ModeSolverData from ...exceptions import WebError from ...log import log +from ...web.file_util import compress_file_to_gzip, extract_gz_file from ...web.http_management import http from ...web.s3utils import download_file, upload_file, upload_string -from ...web.simulation_task import Folder, SIMULATION_JSON +from ...web.simulation_task import Folder, SIMULATION_JSON, SIM_FILE_HDF5_GZ from ...web.types import ResourceLifecycle, Submittable from .mode_solver import ModeSolver, MODE_MONITOR_NAME +from ...version import __version__ MODESOLVER_API = "tidy3d/modesolver/py" MODESOLVER_JSON = "mode_solver.json" MODESOLVER_HDF5 = "mode_solver.hdf5" +MODESOLVER_GZ = "mode_solver.hdf5.gz" + MODESOLVER_LOG = "output/result.log" MODESOLVER_RESULT = "output/result.hdf5" @@ -94,7 +98,7 @@ def run( status = task.get_info().status if status == "error": - raise WebError("Error runnig mode solver.") + raise WebError("Error running mode solver.") log.log(log_level, f"Mode solver status: {status}") if verbose: @@ -154,6 +158,7 @@ class ModeSolverTask(ResourceLifecycle, Submittable, extra=pydantic.Extra.allow) ) # pylint: disable=arguments-differ + # pylint: disable=protected-access @classmethod def create( cls, @@ -195,8 +200,10 @@ def create( { "projectId": folder.folder_id, "taskName": task_name, + "protocolVersion": __version__, "modeSolverName": mode_solver_name, - "fileType": "Hdf5" if len(mode_solver.simulation.custom_datasets) > 0 else "Json", + "fileType": "Gz", + "source": "Python", }, ) log.info( @@ -259,6 +266,7 @@ def get_info(self) -> ModeSolverTask: resp = http.get(f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}") return ModeSolverTask(**resp, mode_solver=self.mode_solver) + # pylint: disable=protected-access def upload( self, verbose: bool = True, progress_callback: Callable[[float], None] = None ) -> None: @@ -273,46 +281,60 @@ def upload( """ mode_solver = self.mode_solver.copy() - # Upload simulation as json for GUI display + sim = mode_solver.simulation + + file, file_name = tempfile.mkstemp() + gz_file, gz_file_name = tempfile.mkstemp() + os.close(file) + os.close(gz_file) + + # Upload simulation as json for download_json upload_string( self.task_id, - mode_solver.simulation._json_string, # pylint: disable=protected-access + sim._json_string, # pylint: disable=protected-access SIMULATION_JSON, verbose=verbose, progress_callback=progress_callback, ) - if self.file_type == "Hdf5": - # Upload a single HDF5 file with the full data - file, file_name = tempfile.mkstemp() - os.close(file) - mode_solver.to_hdf5(file_name) - - try: - upload_file( - self.solver_id, - file_name, - MODESOLVER_HDF5, - verbose=verbose, - progress_callback=progress_callback, - ) - finally: - os.unlink(file_name) - else: - # Send only mode solver, without simulation - mode_solver_spec = mode_solver.dict() + sim.to_hdf5(file_name) + try: + # Upload simulation.hdf5.gz for GUI display + # compress .hdf5 to .hdf5.gz + compress_file_to_gzip(file_name, gz_file_name) - # Upload mode solver without simulation: 'construct' skips all validation - mode_solver_spec["simulation"] = None - mode_solver = ModeSolver.construct(**mode_solver_spec) - upload_string( + upload_file( + self.task_id, + gz_file_name, + SIM_FILE_HDF5_GZ, + verbose=verbose, + progress_callback=progress_callback, + ) + finally: + os.unlink(file_name) + os.unlink(gz_file_name) + + # Upload a single HDF5 file with the full data + file, file_name = tempfile.mkstemp() + gz_file, gz_file_name = tempfile.mkstemp() + os.close(file) + os.close(gz_file) + mode_solver.to_hdf5(file_name) + + try: + # compress .hdf5 to .hdf5.gz + compress_file_to_gzip(file_name, gz_file_name) + + upload_file( self.solver_id, - mode_solver._json_string, # pylint: disable=protected-access - MODESOLVER_JSON, + gz_file_name, + MODESOLVER_GZ, verbose=verbose, progress_callback=progress_callback, - extra_arguments={"type": "ms"}, ) + finally: + os.unlink(file_name) + os.unlink(gz_file_name) # pylint: disable=arguments-differ def submit(self): @@ -362,7 +384,21 @@ def get_modesolver( stored in the same path as 'to_file', but with '.hdf5' extension, and neither 'to_file' or 'sim_file' will be created. """ - if self.file_type == "Hdf5": + if self.file_type == "Gz": + to_gz = pathlib.Path(to_file).with_suffix(".hdf5.gz") + to_hdf5 = pathlib.Path(to_file).with_suffix(".hdf5") + download_file( + self.task_id, + self.mode_solver_path + MODESOLVER_GZ, + to_file=to_gz, + verbose=verbose, + progress_callback=progress_callback, + ) + extract_gz_file(to_gz, to_hdf5) + to_file = str(to_hdf5) + mode_solver = ModeSolver.from_hdf5(to_hdf5) + + elif self.file_type == "Hdf5": to_hdf5 = pathlib.Path(to_file).with_suffix(".hdf5") download_file( self.task_id, diff --git a/tidy3d/web/__init__.py b/tidy3d/web/__init__.py index f9b7d8b7a1..dc8f927004 100644 --- a/tidy3d/web/__init__.py +++ b/tidy3d/web/__init__.py @@ -3,7 +3,7 @@ from .cli.migrate import migrate from .webapi import run, upload, get_info, start, monitor, delete, download, load, estimate_cost -from .webapi import get_tasks, delete_old, download_json, download_log, load_simulation, real_cost +from .webapi import get_tasks, delete_old, download_log, download_json, load_simulation, real_cost from .container import Job, Batch, BatchData from .cli import tidy3d_cli from .cli.app import configure_fn as configure diff --git a/tidy3d/web/file_util.py b/tidy3d/web/file_util.py new file mode 100644 index 0000000000..92e6ff5fa2 --- /dev/null +++ b/tidy3d/web/file_util.py @@ -0,0 +1,29 @@ +"""compress and extract file""" + +import gzip + + +def compress_file_to_gzip(input_file, output_gz_file): + """ + Compresses a file using gzip. + + Args: + input_file (str): The path of the input file. + output_gz_file (str): The path of the output gzip file. + """ + with open(input_file, "rb") as file_in: + with gzip.open(output_gz_file, "wb") as file_out: + file_out.writelines(file_in) + + +def extract_gz_file(input_gz_file, output_file): + """ + Extract the GZ file + + Args: + input_gz_file (str): The path of the gzip input file. + output_file (str): The path of the output file. + """ + with gzip.open(input_gz_file, "rb") as f_in: + with open(output_file, "wb") as f_out: + f_out.write(f_in.read()) diff --git a/tidy3d/web/s3utils.py b/tidy3d/web/s3utils.py index b0b19a31b2..11b7a9166f 100644 --- a/tidy3d/web/s3utils.py +++ b/tidy3d/web/s3utils.py @@ -311,6 +311,9 @@ def _upload(_callback: Callable) -> None: Key=token.get_s3_key(), Callback=_callback, Config=_s3_config, + ExtraArgs={"ContentEncoding": "gzip"} + if token.get_s3_key().endswith(".gz") + else None, ) if progress_callback is not None: diff --git a/tidy3d/web/simulation_task.py b/tidy3d/web/simulation_task.py index c11f0ba531..28d27191a7 100644 --- a/tidy3d/web/simulation_task.py +++ b/tidy3d/web/simulation_task.py @@ -17,12 +17,14 @@ from .s3utils import download_file, upload_file, upload_string from .types import Queryable, ResourceLifecycle, Submittable from .types import Tidy3DResource +from .file_util import compress_file_to_gzip SIMULATION_JSON = "simulation.json" SIMULATION_HDF5 = "output/monitor_data.hdf5" RUNNING_INFO = "output/solver_progress.csv" SIM_LOG_FILE = "output/tidy3d.log" SIM_FILE_HDF5 = "simulation.hdf5" +SIM_FILE_HDF5_GZ = "simulation.hdf5.gz" class Folder(Tidy3DResource, Queryable, extra=Extra.allow): @@ -197,6 +199,7 @@ def create( callback_url: str = None, simulation_type: str = "tidy3d", parent_tasks: List[str] = None, + file_type: str = "Gz", ) -> SimulationTask: """Create a new task on the server. @@ -217,6 +220,8 @@ def create( Type of simulation being uploaded. parent_tasks : List[str] List of related task ids. + file_type: str + the simulation file type Json, Hdf5, Gz Returns ------- @@ -232,6 +237,7 @@ def create( "callbackUrl": callback_url, "simulationType": simulation_type, "parentTasks": parent_tasks, + "fileType": file_type, }, ) @@ -316,7 +322,7 @@ def get_simulation_json( Path to saved file. """ assert self.task_id - return download_file( + download_file( self.task_id, SIMULATION_JSON, to_file=to_file, @@ -338,6 +344,7 @@ def upload_simulation( """ assert self.task_id assert self.simulation + # upload simulation.json for download_json simulation.json upload_string( self.task_id, self.simulation._json_string, # pylint: disable=protected-access @@ -345,21 +352,25 @@ def upload_simulation( verbose=verbose, progress_callback=progress_callback, ) - if len(self.simulation.custom_datasets) > 0: - # Also upload hdf5 containing all data. - file, file_name = tempfile.mkstemp() - os.close(file) - self.simulation.to_hdf5(file_name) - try: - upload_file( - self.task_id, - file_name, - SIM_FILE_HDF5, - verbose=verbose, - progress_callback=progress_callback, - ) - finally: - os.unlink(file_name) + # Also upload hdf5 containing all data. + file, file_name = tempfile.mkstemp() + gz_file, gz_file_name = tempfile.mkstemp() + os.close(file) + os.close(gz_file) + self.simulation.to_hdf5(file_name) + try: + # compress .hdf5 to .hdf5.gz + compress_file_to_gzip(file_name, gz_file_name) + upload_file( + self.task_id, + gz_file_name, + SIM_FILE_HDF5_GZ, + verbose=verbose, + progress_callback=progress_callback, + ) + finally: + os.unlink(file_name) + os.unlink(gz_file_name) def upload_file( self, @@ -411,6 +422,7 @@ def submit( worker group """ if self.simulation: + # upload simulation.json for download_json() upload_string( self.task_id, self.simulation._json_string, # pylint: disable=protected-access @@ -418,6 +430,27 @@ def submit( verbose=False, ) + # Also upload hdf5 containing all data. + file, file_name = tempfile.mkstemp() + gz_file, gz_file_name = tempfile.mkstemp() + os.close(file) + os.close(gz_file) + self.simulation.to_hdf5(file_name) + try: + # compress .hdf5 to .hdf5.gz + compress_file_to_gzip(file_name, gz_file_name) + + upload_file( + self.task_id, + gz_file_name, + SIM_FILE_HDF5_GZ, + verbose=False, + progress_callback=None, + ) + finally: + os.unlink(file_name) + os.unlink(gz_file_name) + if solver_version: protocol_version = None else: diff --git a/tidy3d/web/webapi.py b/tidy3d/web/webapi.py index 4307081a4d..49e2e5019f 100644 --- a/tidy3d/web/webapi.py +++ b/tidy3d/web/webapi.py @@ -182,7 +182,7 @@ def upload( # pylint:disable=too-many-locals,too-many-arguments log.debug("Creating task.") task = SimulationTask.create( - simulation, task_name, folder_name, callback_url, simulation_type, parent_tasks + simulation, task_name, folder_name, callback_url, simulation_type, parent_tasks, "Gz" ) if verbose: console = Console() @@ -388,7 +388,8 @@ def monitor_preprocess() -> None: # verbose case, update progressbar console.log("running solver") console.log( - "To cancel the simulation, use 'web.delete(task_id)' or delete the task in the web" + "To cancel the simulation, use 'web.abort(task_id)' or 'web.delete(task_id)'" + " or abort/delete the task in the web" " UI. Terminating the Python script will not stop the job running on the cloud." ) with Progress(console=console) as progress: