add batch id

Author: wei

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 @@ def _check_path_dir(path: str) -> None:
         """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 @@ def run(self, path: str = DEFAULT_DATA_PATH) -> SimulationDataType:
         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 @@ def estimate_cost(self, verbose: bool = True) -> float:
         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 @@ class Batch(WebContainer):
         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 num_jobs(self) -> int:
     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 get_info(self) -> dict[TaskName, TaskInfo]:
     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 @@ def _check_path_dir(path_dir: str) -> None:
         """
         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 @@ def upload(
     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 @@ def upload(
         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 @@ def upload(
     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 @@ def start(
         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

@@ -24,7 +24,7 @@
 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 @@ def create(
         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 @@ def create(
             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 @@ def create(
                 "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 @@ def abort(self):
         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"