"""Web API for mode solver"""
from __future__ import annotations
import os
import tempfile
import time
from datetime import datetime
from typing import TYPE_CHECKING
from botocore.exceptions import ClientError
from joblib import Parallel, delayed
from pydantic import Field
from rich.progress import Progress
from tidy3d.components.data.monitor_data import ModeSolverData
from tidy3d.components.eme.simulation import EMESimulation
from tidy3d.components.medium import AbstractCustomMedium
from tidy3d.components.simulation import Simulation
from tidy3d.config import config
from tidy3d.exceptions import SetupError, WebError, format_chained_exception_message
from tidy3d.log import get_logging_console, log
from tidy3d.plugins.mode.mode_solver import MODE_MONITOR_NAME, ModeSolver
from tidy3d.version import __version__
from tidy3d.web.api.run_options import log_deprecated_run_args, resolve_pay_type
from tidy3d.web.core.core_config import get_logger_console
from tidy3d.web.core.http_util import http
from tidy3d.web.core.s3utils import download_file, download_gz_file, upload_file
from tidy3d.web.core.task_core import Folder
from tidy3d.web.core.types import ResourceLifecycle, Submittable
if TYPE_CHECKING:
import pathlib
from collections.abc import Callable
from os import PathLike
from typing import Literal
import requests
from rich.progress import TaskID
from tidy3d.web.core.types import PayType
SIMULATION_JSON = "simulation.json"
SIM_FILE_HDF5_GZ = "simulation.hdf5.gz"
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"
MODESOLVER_RESULT_GZ = "output/mode_solver_data.hdf5.gz"
DEFAULT_NUM_WORKERS = 10
DEFAULT_MAX_RETRIES = 3
DEFAULT_RETRY_DELAY = 10 # in seconds
[docs]
def run(
mode_solver: ModeSolver,
task_name: str = "Untitled",
mode_solver_name: str = "mode_solver",
folder_name: str = "Mode Solver",
results_file: PathLike = "mode_solver.hdf5",
verbose: bool = True,
progress_callback_upload: Callable[[float], None] | None = None,
progress_callback_download: Callable[[float], None] | None = None,
reduce_simulation: Literal["auto", True, False] = "auto",
pay_type: PayType | str | None = None,
) -> ModeSolverData:
"""Submits a :class:`.ModeSolver` to server, starts running, monitors progress, downloads,
and loads results as a :class:`.ModeSolverData` object.
Parameters
----------
mode_solver : :class:`.ModeSolver`
Mode solver to upload to server.
task_name : str = "Untitled"
Name of task.
mode_solver_name: str = "mode_solver"
The name of the mode solver to create the in task.
folder_name : str = "Mode Solver"
Name of folder to store task on web UI.
results_file : PathLike = "mode_solver.hdf5"
Path to download results file (.hdf5).
verbose : bool = True
If ``True``, will print status, otherwise, will run silently.
progress_callback_upload : Callable[[float], None] = None
Optional callback function called when uploading file with ``bytes_in_chunk`` as argument.
progress_callback_download : Callable[[float], None] = None
Optional callback function called when downloading file with ``bytes_in_chunk`` as argument.
reduce_simulation : Literal["auto", True, False] = "auto"
Restrict simulation to mode solver region. If "auto", then simulation is automatically
restricted if it contains custom mediums.
pay_type : Optional[Union[PayType, str]] = None
Deprecated direct argument for internal use only. Internal workflows should set
``td.config.run.pay_type`` instead; external users should leave unset.
Returns
-------
:class:`.ModeSolverData`
Mode solver data with the calculated results.
"""
log_deprecated_run_args(pay_type=pay_type)
log_level = "DEBUG" if verbose else "INFO"
if verbose:
console = get_logging_console()
if reduce_simulation == "auto":
sim_mediums = mode_solver.simulation.scene.mediums
contains_custom = any(isinstance(med, AbstractCustomMedium) for med in sim_mediums)
reduce_simulation = contains_custom
if reduce_simulation:
log.warning(
"The associated 'Simulation' object contains custom mediums. It will be "
"automatically restricted to the mode solver plane to reduce data for uploading. "
"To force uploading the original 'Simulation' object use 'reduce_simulation=False'."
" Setting 'reduce_simulation=True' will force simulation reduction in all cases and"
" silence this warning."
)
if reduce_simulation:
mode_solver = mode_solver.reduced_simulation_copy
task = ModeSolverTask.create(mode_solver, task_name, mode_solver_name, folder_name)
if verbose:
console.log(
f"Mode solver created with task_id='{task.task_id}', solver_id='{task.solver_id}'."
)
task.upload(verbose=verbose, progress_callback=progress_callback_upload)
task.submit(pay_type=pay_type)
# Wait for task to finish
prev_status = "draft"
status = task.status
while status not in ("success", "error", "diverged", "deleted"):
if status != prev_status:
log.log(log_level, f"Mode solver status: {status}")
if verbose:
console.log(f"Mode solver status: {status}")
prev_status = status
time.sleep(0.5)
status = task.get_info().status
if status == "error":
raise WebError("Error running mode solver.")
log.log(log_level, f"Mode solver status: {status}")
if verbose:
console.log(f"Mode solver status: {status}")
if status != "success":
# Our cache discards None, so the user is able to re-run
return None
return task.get_result(
to_file=results_file, verbose=verbose, progress_callback=progress_callback_download
)
[docs]
def run_batch(
mode_solvers: list[ModeSolver],
task_name: str = "BatchModeSolver",
folder_name: str = "BatchModeSolvers",
results_files: list[str] | None = None,
verbose: bool = True,
max_workers: int = DEFAULT_NUM_WORKERS,
max_retries: int = DEFAULT_MAX_RETRIES,
retry_delay: float = DEFAULT_RETRY_DELAY,
progress_callback_upload: Callable[[float], None] | None = None,
progress_callback_download: Callable[[float], None] | None = None,
) -> list[ModeSolverData]:
"""
Submits a batch of ModeSolver to the server concurrently, manages progress, and retrieves results.
Parameters
----------
mode_solvers : list[ModeSolver]
List of mode solvers to be submitted to the server.
task_name : str
Base name for tasks. Each task in the batch will have a unique index appended to this base name.
folder_name : str
Name of the folder where tasks are stored on the server's web UI.
results_files : list[str], optional
List of file paths where the results for each ModeSolver should be downloaded. If None, a default path based on the folder name and index is used.
verbose : bool
If True, displays a progress bar. If False, runs silently.
max_workers : int
Maximum number of concurrent workers to use for processing the batch of simulations.
max_retries : int
Maximum number of retries for each simulation in case of failure before giving up.
retry_delay : int
Delay in seconds between retries when a simulation fails.
progress_callback_upload : Callable[[float], None], optional
Optional callback function called when uploading file with ``bytes_in_chunk`` as argument.
progress_callback_download : Callable[[float], None], optional
Optional callback function called when downloading file with ``bytes_in_chunk`` as argument.
Returns
-------
list[ModeSolverData]
A list of ModeSolverData objects containing the results from each simulation in the batch. ``None`` is placed in the list for simulations that fail after all retries.
"""
console = get_logging_console()
if not all(isinstance(x, ModeSolver) for x in mode_solvers):
console.log(
"Validation Error: All items in `mode_solvers` must be instances of `ModeSolver`."
)
return []
num_mode_solvers = len(mode_solvers)
if results_files is None:
results_files = [f"mode_solver_batch_results_{i}.hdf5" for i in range(num_mode_solvers)]
def handle_mode_solver(index: int, progress: Progress, pbar: TaskID | None) -> Parallel | None:
retries = 0
while retries <= max_retries:
try:
result = run(
mode_solver=mode_solvers[index],
task_name=f"{task_name}_{index}",
mode_solver_name=f"mode_solver_batch_{index}",
folder_name=folder_name,
results_file=results_files[index],
verbose=False,
progress_callback_upload=progress_callback_upload,
progress_callback_download=progress_callback_download,
)
if verbose:
progress.update(pbar, advance=1)
return result
except Exception as e:
console.log(f"Error in mode solver {index}: {e!s}")
if retries < max_retries:
time.sleep(retry_delay)
retries += 1
else:
console.log(f"The {index}-th mode solver failed after {max_retries} tries.")
if verbose:
progress.update(pbar, advance=1)
return None
return None
if verbose:
console.log(f"[cyan]Running a batch of [deep_pink4]{num_mode_solvers} mode solvers.\n")
# Create the common folder before running the parallel computation
_ = Folder.create(folder_name=folder_name)
with Progress(console=console) as progress:
pbar = progress.add_task("Status:", total=num_mode_solvers)
results = Parallel(n_jobs=max_workers, backend="threading")(
delayed(handle_mode_solver)(i, progress, pbar) for i in range(num_mode_solvers)
)
# Make sure the progress bar is complete
progress.update(pbar, completed=num_mode_solvers, refresh=True)
console.log("[green]A batch of `ModeSolver` tasks completed successfully!")
else:
results = Parallel(n_jobs=max_workers, backend="threading")(
delayed(handle_mode_solver)(i, None, None) for i in range(num_mode_solvers)
)
return results
[docs]
class ModeSolverTask(ResourceLifecycle, Submittable, extra="allow"):
"""Interface for managing the running of a :class:`.ModeSolver` task on server."""
task_id: str | None = Field(
None,
title="task_id",
description="Task ID number, set when the task is created, leave as None.",
alias="refId",
)
solver_id: str | None = Field(
None,
title="solver",
description="Solver ID number, set when the task is created, leave as None.",
alias="id",
)
real_flex_unit: float | None = Field(
None, title="real FlexCredits", description="Billed FlexCredits.", alias="charge"
)
created_at: datetime | None = Field(
title="created_at", description="Time at which this task was created.", alias="createdAt"
)
status: str | None = Field(
None,
title="status",
description="Mode solver task status.",
)
file_type: str | None = Field(
None,
title="file_type",
description="File type used to upload the mode solver.",
alias="fileType",
)
mode_solver: ModeSolver | None = Field(
None,
title="mode_solver",
description="Mode solver being run by this task.",
)
[docs]
@classmethod
def create(
cls,
mode_solver: ModeSolver,
task_name: str = "Untitled",
mode_solver_name: str = "mode_solver",
folder_name: str = "Mode Solver",
) -> ModeSolverTask:
"""Create a new mode solver task on the server.
Parameters
----------
mode_solver: :class".ModeSolver"
The object that will be uploaded to server in the submitting phase.
task_name: str = "Untitled"
The name of the task.
mode_solver_name: str = "mode_solver"
The name of the mode solver to create the in task.
folder_name: str = "Mode Solver"
The name of the folder to store the task.
Returns
-------
:class:`ModeSolverTask`
:class:`ModeSolverTask` object containing information about the created task.
"""
folder = Folder.get(folder_name, create=True)
mode_solver.validate_pre_upload()
if isinstance(mode_solver.simulation, Simulation):
mode_solver.simulation.validate_pre_upload(source_required=False)
elif isinstance(mode_solver.simulation, EMESimulation):
# TODO: replace this with native web api support
raise SetupError(
"'EMESimulation' is not yet supported in the "
"remote mode solver web api. Please instead call 'ModeSolver.to_fdtd_mode_solver' "
"before using the web api; this replaces the 'EMESimulation' with a 'Simulation' "
"that can be used in the remote mode solver. "
"Alternatively, you can add a 'ModeSolverMonitor' to the 'EMESimulation' "
"and use the EME solver web api."
)
# mode_solver.simulation.validate_pre_upload()
else:
raise SetupError("Simulation type not supported in the remote mode solver web api.")
response_body = {
"projectId": folder.folder_id,
"taskName": task_name,
"protocolVersion": __version__,
"modeSolverName": mode_solver_name,
"fileType": "Gz",
"source": "Python",
}
resp = http.post(MODESOLVER_API, response_body)
# TODO: actually fix the root cause later.
# sometimes POST for mode returns None and crashes the log.info call.
# For now this catches it and raises a better error that we can debug.
if resp is None:
raise WebError(
"'ModeSolver' POST request returned 'None'. If you received this error, please "
"raise an issue on the Tidy3D front end GitHub repository referencing the following"
f"response body '{response_body}'."
)
log.info(
"Mode solver created with task_id='%s', solver_id='%s'.", resp["refId"], resp["id"]
)
return ModeSolverTask(**resp, mode_solver=mode_solver)
[docs]
@classmethod
def get(
cls,
task_id: str,
solver_id: str,
to_file: PathLike = "mode_solver.hdf5",
sim_file: PathLike = "simulation.hdf5",
verbose: bool = True,
progress_callback: Callable[[float], None] | None = None,
) -> ModeSolverTask:
"""Get mode solver task from the server by id.
Parameters
----------
task_id: str
Unique identifier of the task on server.
solver_id: str
Unique identifier of the mode solver in the task.
to_file: PathLike = "mode_solver.hdf5"
File to store the mode solver downloaded from the task.
sim_file: PathLike = "simulation.hdf5"
File to store the simulation downloaded from the task.
verbose: bool = True
Whether to display progress bars.
progress_callback : Callable[[float], None] = None
Optional callback function called while downloading the data.
Returns
-------
:class:`ModeSolverTask`
:class:`ModeSolverTask` object containing information about the task.
"""
resp = http.get(f"{MODESOLVER_API}/{task_id}/{solver_id}")
task = ModeSolverTask(**resp)
mode_solver = task.get_modesolver(to_file, sim_file, verbose, progress_callback)
return task.model_copy(update={"mode_solver": mode_solver})
[docs]
def get_info(self) -> ModeSolverTask:
"""Get the current state of this task on the server.
Returns
-------
:class:`ModeSolverTask`
:class:`ModeSolverTask` object containing information about the task, without the mode
solver.
"""
resp = http.get(f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}")
return ModeSolverTask(**resp, mode_solver=self.mode_solver)
[docs]
def upload(
self, verbose: bool = True, progress_callback: Callable[[float], None] | None = None
) -> None:
"""Upload this task's 'mode_solver' to the server.
Parameters
----------
verbose: bool = True
Whether to display progress bars.
progress_callback : Callable[[float], None] = None
Optional callback function called while uploading the data.
"""
mode_solver = self.mode_solver.copy()
sim = mode_solver.simulation
# Upload simulation.hdf5.gz for GUI display
file, file_name = tempfile.mkstemp(".hdf5.gz")
os.close(file)
try:
sim.to_hdf5_gz(file_name)
upload_file(
self.task_id,
file_name,
SIM_FILE_HDF5_GZ,
verbose=verbose,
progress_callback=progress_callback,
)
finally:
os.unlink(file_name)
# Upload a single HDF5 file with the full data
file, file_name = tempfile.mkstemp(".hdf5.gz")
os.close(file)
try:
mode_solver.to_hdf5_gz(file_name)
upload_file(
self.solver_id,
file_name,
MODESOLVER_GZ,
verbose=verbose,
progress_callback=progress_callback,
)
finally:
os.unlink(file_name)
[docs]
def submit(
self,
pay_type: PayType | str | None = None,
) -> None:
"""Start the execution of this task.
The mode solver must be uploaded to the server with the :meth:`ModeSolverTask.upload` method
before this step.
"""
http.post(
f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}/run",
{
"enableCaching": config.web.enable_caching,
"payType": resolve_pay_type(pay_type).value,
},
)
[docs]
def delete(self) -> None:
"""Delete the mode solver and its corresponding task from the server."""
# Delete mode solver
http.delete(f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}")
# Delete parent task
http.delete(f"tidy3d/tasks/{self.task_id}")
[docs]
def abort(self) -> requests.Response:
"""Abort the mode solver and its corresponding task from the server."""
return http.put(
"tidy3d/tasks/abort", json={"taskType": "MODE_SOLVER", "taskId": self.solver_id}
)
[docs]
def get_modesolver(
self,
to_file: PathLike = "mode_solver.hdf5",
sim_file: PathLike = "simulation.hdf5",
verbose: bool = True,
progress_callback: Callable[[float], None] | None = None,
) -> ModeSolver:
"""Get mode solver associated with this task from the server.
Parameters
----------
to_file: PathLike = "mode_solver.hdf5"
File to store the mode solver downloaded from the task.
sim_file: PathLike = "simulation.hdf5"
File to store the simulation downloaded from the task, if any.
verbose: bool = True
Whether to display progress bars.
progress_callback : Callable[[float], None] = None
Optional callback function called while downloading the data.
Returns
-------
:class:`.ModeSolver`
:class:`.ModeSolver` object associated with this task.
"""
if self.file_type == "Gz":
file, file_path = tempfile.mkstemp(".hdf5.gz")
os.close(file)
try:
download_file(
self.solver_id,
MODESOLVER_GZ,
to_file=file_path,
verbose=verbose,
progress_callback=progress_callback,
)
mode_solver = ModeSolver.from_hdf5_gz(file_path)
finally:
os.unlink(file_path)
elif self.file_type == "Hdf5":
file, file_path = tempfile.mkstemp(".hdf5")
os.close(file)
try:
download_file(
self.solver_id,
MODESOLVER_HDF5,
to_file=file_path,
verbose=verbose,
progress_callback=progress_callback,
)
mode_solver = ModeSolver.from_hdf5(file_path)
finally:
os.unlink(file_path)
else:
file, file_path = tempfile.mkstemp(".json")
os.close(file)
try:
download_file(
self.solver_id,
MODESOLVER_JSON,
to_file=file_path,
verbose=verbose,
progress_callback=progress_callback,
)
mode_solver_dict = ModeSolver.dict_from_json(file_path)
finally:
os.unlink(file_path)
download_file(
self.task_id,
SIMULATION_JSON,
to_file=sim_file,
verbose=verbose,
progress_callback=progress_callback,
)
mode_solver_dict["simulation"] = Simulation.from_json(sim_file)
mode_solver = ModeSolver.model_validate(mode_solver_dict)
# Store requested mode solver file
mode_solver.to_file(to_file)
return mode_solver
[docs]
def get_result(
self,
to_file: PathLike = "mode_solver_data.hdf5",
verbose: bool = True,
progress_callback: Callable[[float], None] | None = None,
) -> ModeSolverData:
"""Get mode solver results for this task from the server.
Parameters
----------
to_file: PathLike = "mode_solver_data.hdf5"
File to store the mode solver downloaded from the task.
verbose: bool = True
Whether to display progress bars.
progress_callback : Callable[[float], None] = None
Optional callback function called while downloading the data.
Returns
-------
:class:`.ModeSolverData`
Mode solver data with the calculated results.
"""
file = None
try:
file = download_gz_file(
resource_id=self.solver_id,
remote_filename=MODESOLVER_RESULT_GZ,
to_file=to_file,
verbose=verbose,
progress_callback=progress_callback,
)
except ClientError:
if verbose:
console = get_logger_console()
console.log(f"Unable to download '{MODESOLVER_RESULT_GZ}'.")
if not file:
try:
file = download_file(
resource_id=self.solver_id,
remote_filename=MODESOLVER_RESULT,
to_file=to_file,
verbose=verbose,
progress_callback=progress_callback,
)
except Exception as e:
raise WebError(
format_chained_exception_message(
"Failed to download the simulation data file from the server. Please "
"confirm that the task was successfully run.",
e,
)
) from e
data = ModeSolverData.from_hdf5(to_file)
data = data.copy(
update={"monitor": self.mode_solver.to_mode_solver_monitor(name=MODE_MONITOR_NAME)}
)
self.mode_solver._cached_properties["data_raw"] = data
# Perform symmetry expansion
self.mode_solver._cached_properties.pop("data", None)
return self.mode_solver.data
[docs]
def get_log(
self,
to_file: PathLike = "mode_solver.log",
verbose: bool = True,
progress_callback: Callable[[float], None] | None = None,
) -> pathlib.Path:
"""Get execution log for this task from the server.
Parameters
----------
to_file: PathLike = "mode_solver.log"
File to store the mode solver downloaded from the task.
verbose: bool = True
Whether to display progress bars.
progress_callback : Callable[[float], None] = None
Optional callback function called while downloading the data.
Returns
-------
path: pathlib.Path
Path to saved file.
"""
return download_file(
self.solver_id,
MODESOLVER_LOG,
to_file=to_file,
verbose=verbose,
progress_callback=progress_callback,
)
