add batch id

tests/README.md

@@ -17,7 +17,8 @@ Inside `dir/test_code.py` pytest will run any code in the main file and then cal
 
 # gets evaluated first
 my_two = 2.0
-my_three = add_one(my_two)
+my_three = add_one(my_two) 
+
 
 def add_one(x):
     # calls this function

tidy3d/plugins/adjoint/web.py

@@ -281,7 +281,7 @@ def start(self) -> None:
         To monitor progress of the :class:`Job`, call :meth:`Job.monitor` after started.
         """
         if self.jax_info is not None:
-            upload_jax_info(task_id=self.task_id, jax_info=self.jax_info, verbose=self.verbose)
+            upload_jax_info(task_id=self.task_id(), jax_info=self.jax_info, verbose=self.verbose)
         super().start()
 
 

tidy3d/web/api/autograd/autograd.py

@@ -1167,7 +1167,6 @@ def _run_async_tidy3d_bwd(
     _ = run_kwargs.pop("path_dir", None)
     batch = Batch(simulations=simulations, **batch_init_kwargs)
     td.log.info(f"running {batch.simulation_type} batch with '_run_async_tidy3d_bwd()'")
-
     batch.start()
     batch.monitor()
 

tidy3d/web/api/container.py

@@ -44,9 +44,9 @@
         """Make sure local output directory exists and create it if not."""
 
     @staticmethod
-    def _check_folder(folder_name: str) -> None:
+    def _check_folder(folder_name: str) -> Folder:
         """Make sure ``folder_name`` exists on the web UI and create it if not."""
-        Folder.get(folder_name, create=True)
+        return Folder.get(folder_name, create=True)
 
 
 class Job(WebContainer):
@@ -235,24 +235,24 @@
         self.monitor()
         return self.load(path=path)
 
-    @cached_property
-    def task_id(self) -> TaskId:
+    def task_id(self, batch_id: str = None) -> TaskId:
         """The task ID for this ``Job``. Uploads the ``Job`` if it hasn't already been uploaded."""
         if self.task_id_cached:
             return self.task_id_cached
         self._check_folder(self.folder_name)
-        return self._upload()
+        return self._upload(batch_id)
 
-    def _upload(self) -> TaskId:
+    def _upload(self, batch_id: str) -> TaskId:
         """Upload this job and return the task ID for handling."""
         # upload kwargs with all fields except task_id
         upload_kwargs = {key: getattr(self, key) for key in self._upload_fields}
+        upload_kwargs["batch_id"] = batch_id
         task_id = web.upload(**upload_kwargs)
         return task_id
 
-    def upload(self) -> None:
+    def upload(self, batch_id: str = None) -> None:
         """Upload this ``Job``."""
-        _ = self.task_id
+        _ = self.task_id(batch_id)
 
     def get_info(self) -> TaskInfo:
         """Return information about a :class:`Job`.
@@ -370,7 +370,9 @@
         Cost is calculated assuming the simulation runs for
         the full ``run_time``. If early shut-off is triggered, the cost is adjusted proportionately.
         """
-        return web.estimate_cost(self.task_id, verbose=verbose, solver_version=self.solver_version)
+        return web.estimate_cost(
+            self.task_id(), verbose=verbose, solver_version=self.solver_version
+        )
 
     @staticmethod
     def _check_path_dir(path: str) -> None:
@@ -511,6 +513,18 @@
         description="Name of folder to store member of each batch on web UI.",
     )
 
+    batch_id: str = pd.Field(
+        None,
+        title="Batch Id",
+        description="ID of batch to store member of each batch on web UI.",
+    )
+
+    batch_name: str = pd.Field(
+        "batch",
+        title="Batch Name",
+        description="Name of batch to store member of each batch on web UI.",
+    )
+
     verbose: bool = pd.Field(
         True, title="Verbose", description="Whether to print info messages and progressbars."
     )
@@ -674,8 +688,15 @@
     def upload(self) -> None:
         """Upload a series of tasks associated with this ``Batch`` using multi-threading."""
         self._check_folder(self.folder_name)
+        batch_id = self.create_batch()
         with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
-            futures = [executor.submit(job.upload) for _, job in self.jobs.items()]
+            futures = [
+                executor.submit(
+                    job.upload,
+                    batch_id,
+                )
+                for _, job in self.jobs.items()
+            ]
 
             # progressbar (number of tasks uploaded)
             if self.verbose:
@@ -711,6 +732,10 @@
     def start(self) -> None:
         """Start running all tasks in the :class:`Batch`.
 
+        Parameters:
+            batch_id: str
+                Batch ID to start the tasks in.
+
         Note
         ----
         To monitor the running simulations, can call :meth:`Batch.monitor`.
@@ -1068,3 +1093,8 @@
         """
         if len(path_dir) > 0 and not os.path.exists(path_dir):
             os.makedirs(path_dir, exist_ok=True)
+
+    def create_batch(self) -> str:
+        """Create batch."""
+        folder = self._check_folder(self.folder_name)
+        return OptimizationBatch.create(self.batch_name, folder.folder_id)

tidy3d/web/api/webapi.py

@@ -201,6 +201,7 @@
     source_required: bool = True,
     solver_version: Optional[str] = None,
     reduce_simulation: Literal["auto", True, False] = "auto",
+    batch_id: str = None,
 ) -> TaskId:
     """
     Upload simulation to server, but do not start running :class:`.Simulation`.
@@ -230,7 +231,8 @@
         target solver version.
     reduce_simulation: Literal["auto", True, False] = "auto"
         Whether to reduce structures in the simulation to the simulation domain only. Note: currently only implemented for the mode solver.
-
+    batch_id: str = None
+        Batch id to upload tasks to.
     Returns
     -------
     str
@@ -260,7 +262,14 @@
     task_type = stub.get_type()
 
     task = SimulationTask.create(
-        task_type, task_name, folder_name, callback_url, simulation_type, parent_tasks, "Gz"
+        task_type,
+        task_name,
+        folder_name,
+        callback_url,
+        simulation_type,
+        parent_tasks,
+        "Gz",
+        batch_id,
     )
     if verbose:
         console = get_logging_console()
@@ -385,6 +394,8 @@
         worker group
     pay_type: Union[PayType, str] = PayType.AUTO
         Which method to pay the simulation
+    batch_id: str
+        batch id
     Note
     ----
     To monitor progress, can call :meth:`monitor` after starting simulation.

tidy3d/web/core/task_core.py

@@ -1,30 +1,30 @@
 """Tidy3d webapi types."""
 
 from __future__ import annotations
 
 import os
 import pathlib
 import tempfile
 from datetime import datetime
 from typing import Callable, Optional, Union
 
 import pydantic.v1 as pd
 from botocore.exceptions import ClientError
 from pydantic.v1 import Extra, Field, parse_obj_as
 
 import tidy3d as td
 
 from . import http_util
 from .cache import FOLDER_CACHE
 from .constants import SIM_ERROR_FILE, SIM_FILE_HDF5_GZ, SIM_LOG_FILE, SIMULATION_DATA_HDF5_GZ
 from .core_config import get_logger_console
 from .environment import Env
 from .exceptions import WebError, WebNotFoundError
 from .file_util import read_simulation_from_hdf5
 from .http_util import http
 from .s3utils import download_file, download_gz_file, upload_file
 from .stub import TaskStub
-from .types import PayType, Queryable, ResourceLifecycle, Submittable, Tidy3DResource
+from .types import OptimizationType, PayType, Queryable, ResourceLifecycle, Submittable, Tidy3DResource
 
 
 class Folder(Tidy3DResource, Queryable, extra=Extra.allow):
@@ -208,6 +208,7 @@
         simulation_type: str = "tidy3d",
         parent_tasks: Optional[list[str]] = None,
         file_type: str = "Gz",
+        batch_id: str = None,
     ) -> SimulationTask:
         """Create a new task on the server.
 
@@ -228,7 +229,8 @@
             List of related task ids.
         file_type: str
             the simulation file type Json, Hdf5, Gz
-
+        batch_id: str
+            batch id
         Returns
         -------
         :class:`SimulationTask`
@@ -250,6 +252,7 @@
                 "simulationType": simulation_type,
                 "parentTasks": parent_tasks,
                 "fileType": file_type,
+                "batchId": batch_id,
             },
         )
         return SimulationTask(**resp, taskType=task_type, folder_name=folder_name)
@@ -671,3 +674,48 @@
         return http.put(
             "tidy3d/tasks/abort", json={"taskType": self.task_type, "taskId": self.task_id}
         )
+
+
+class OptimizationBatch(ResourceLifecycle, extra=Extra.allow):
+    """OptimizationBatch."""
+
+    batch_id: Optional[str] = Field(
+        ...,
+        title="batch_id",
+        description="Batch ID number, set when the batch is uploaded, leave as None.",
+        alias="optimizationId",
+    )
+    batch_name: Optional[str] = Field(
+        None,
+        title="batch_name",
+        description="The name of batch, leave as None.",
+        alias="optimizationName",
+    )
+
+    @classmethod
+    def create(cls, batch_name: str, folder_id: str) -> str:
+        """Create batch from server.
+
+        Parameters
+        ----------
+        batch_name: str = None
+            batch name.
+        folder_id: str = None
+            folder id.
+
+        Returns
+        -------
+        batch id: str
+
+        """
+        if folder_id is None:
+            raise WebError("folder_id can't be None.")
+
+        resp = http.post(
+            "tidy3d/optimization",
+            json={"type": OptimizationType.BATCH.value, "name": batch_name, "folderId": folder_id},
+        )
+        return resp["optimizationId"]
+
+    def delete(self):
+        return http.delete("tidy3d/optimization", json={"batch_id": self.batch_id})

tidy3d/web/core/types.py

@@ -67,3 +67,7 @@ def _missing_(cls, value: object) -> PayType:
             if key in cls.__members__:
                 return cls.__members__[key]
         return super()._missing_(value)
+
+
+class OptimizationType(str, Enum):
+    BATCH = "BATCH"