Let me try again the task assignment

Author: daquinteroflex

tidy3d/plugins/smatrix/run.py

@@ -18,6 +18,7 @@
 def create_batch(
     modeler: ComponentModelerType,
     path_dir: str = DEFAULT_DATA_DIR,
+    parent_batch_id: str = None,
     file_name: str = "batch.hdf5",
     **kwargs,
 ) -> Batch:
@@ -33,7 +34,15 @@ def create_batch(
         The configured `Batch` object ready for execution.
     """
     filepath = os.path.join(path_dir, file_name)
-    batch = Batch(simulations=modeler.sim_dict, **kwargs)
+
+    if parent_batch_id is not None:
+        parent_task_dict = dict()
+        for key in modeler.sim_dict.keys():
+            parent_task_dict[key] = tuple(parent_batch_id,)
+    else:
+        parent_task_dict = None
+
+    batch = Batch(simulations=modeler.sim_dict, parent_tasks=parent_task_dict, **kwargs)
     batch.to_file(filepath)
     return batch
 

tidy3d/plugins/smatrix/web/api/webapi.py

@@ -177,7 +177,7 @@ def run(
         Monitor progress of each of the running tasks.
     """
     if isinstance(simulation, ComponentModelerType):
-        batch_id = upload(
+        batch_task = upload_modeler(
             simulation=simulation,
             task_name=task_name,
             folder_name=folder_name,
@@ -193,8 +193,9 @@ def run(
         batch = create_batch(
             modeler=simulation,
             folder_name=folder_name,
-            solver_version=solver_version
-            # parent_tasks=[task_id]
+            solver_version=solver_version,
+            parent_batch_id=batch_task.batch_id,
+            group_id=batch_task.group_id
         )
         batch.run()
         # postprocess(
@@ -398,6 +399,152 @@ def upload(
         return task.task_id
 
 
+@wait_for_connection
+def upload_modeler(
+    simulation: SimulationType,
+    task_name: str,
+    folder_name: str = "default",
+    callback_url: Optional[str] = None,
+    verbose: bool = True,
+    progress_callback: Optional[Callable[[float], None]] = None,
+    simulation_type: str = "tidy3d",
+    parent_tasks: Optional[list[str]] = None,
+    source_required: bool = True,
+    solver_version: Optional[str] = None,
+    reduce_simulation: Literal["auto", True, False] = "auto",
+) -> TaskId:
+    """
+    Upload simulation to server, but do not start running :class:`.Simulation`.
+
+    Parameters
+    ----------
+    simulation : Union[:class:`.Simulation`, :class:`.HeatSimulation`, :class:`.EMESimulation`]
+        Simulation to upload to server.
+    task_name : str
+        Name of task.
+    folder_name : str
+        Name of folder to store task on web UI
+    callback_url : str = None
+        Http PUT url to receive simulation finish event. The body content is a json file with
+        fields ``{'id', 'status', 'name', 'workUnit', 'solverVersion'}``.
+    verbose : bool = True
+        If ``True``, will print progressbars and status, otherwise, will run silently.
+    progress_callback : Callable[[float], None] = None
+        Optional callback function called when uploading file with ``bytes_in_chunk`` as argument.
+    simulation_type : str = "tidy3d"
+        Type of simulation being uploaded.
+    parent_tasks : List[str]
+        List of related task ids.
+    source_required: bool = True
+        If ``True``, simulations without sources will raise an error before being uploaded.
+    solver_version: str = None
+        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.
+
+    Returns
+    -------
+    str
+        Unique identifier of task on server.
+
+
+    Notes
+    -----
+
+        Once you've created a ``job`` object using :class:`tidy3d.plugins.smatrix.web.api.container.Job`, you can upload it to our servers with:
+
+        .. code-block:: python
+
+            web.upload(simulation, task_name="task_name", verbose=verbose)
+
+        It will not run until you explicitly tell it to do so with :meth:`tidy3d.plugins.smatrix.web.api.webapi.start`.
+
+    """
+
+    if isinstance(simulation, (ModeSolver, ModeSimulation)):
+        simulation = get_reduced_simulation(simulation, reduce_simulation)
+
+    if isinstance(simulation, ComponentModelerType):
+        port_name_list = [port.name for port in simulation.ports]
+    else:
+        port_name_list = []
+
+    stub = Tidy3dStub(simulation=simulation)
+    stub.validate_pre_upload(source_required=source_required)
+    log.debug("Creating task.")
+
+    task_type = stub.get_type()
+    if (task_type == "COMPONENT_MODELER") or (task_type == "TERMINAL_COMPONENT_MODELER"): # TODO REMOVE
+        task_type = "RF"
+
+    task = SimulationTask.create(
+        task_type, task_name, folder_name, callback_url, simulation_type, parent_tasks, file_type="Gz", port_name_list=port_name_list
+    )
+    print(task)
+    if verbose:
+        console = get_logging_console()
+        if isinstance(simulation, ComponentModelerType):
+            console.log(
+                f"Created modeler task '{task_name}' with batch_id '{task.batch_id}' and task_type '{task_type}'."
+            )
+        else:
+            console.log(
+                f"Created task '{task_name}' with task_id '{task.task_id}' and task_type '{task_type}'."
+            )
+        if task_type in BETA_TASK_TYPES:
+            solver_name = SOLVER_NAME[task_type]
+            console.log(
+                f"Tidy3D's {solver_name} solver is currently in the beta stage. "
+                f"Cost of {solver_name} simulations is subject to change in the future."
+            )
+        if task_type in GUI_SUPPORTED_TASK_TYPES:
+            if isinstance(simulation, ComponentModelerType):
+                url = _get_url(resource_id=task.group_id, resource_type="group")
+                folder_url = _get_folder_url(task.folder_id)
+                console.log(f"View task using web UI at [link={url}]'{url}'[/link].")
+                console.log(f"Task folder: [link={folder_url}]'{task.folder_name}'[/link].")
+            else:
+                url = _get_url(task.task_id)
+                folder_url = _get_folder_url(task.folder_id)
+                console.log(f"View task using web UI at [link={url}]'{url}'[/link].")
+                console.log(f"Task folder: [link={folder_url}]'{task.folder_name}'[/link].")
+
+    remote_sim_file = SIM_FILE_HDF5_GZ
+    if task_type == "MODE_SOLVER":
+        remote_sim_file = MODE_FILE_HDF5_GZ
+    elif isinstance(simulation, ComponentModelerType):
+        remote_sim_file = MODELER_FILE_HDF5_GZ
+
+    if isinstance(simulation, ComponentModelerType):
+        task.upload_simulation(
+            stub=stub,
+            verbose=verbose,
+            progress_callback=progress_callback,
+            remote_sim_file=remote_sim_file,
+        )
+        #  TODO talk momchil
+        # estimate_cost(task_id=task.task_id, solver_version=solver_version, verbose=verbose)
+
+        task.validate_post_upload(parent_tasks=parent_tasks)
+
+        # log the url for the task in the web UI
+        log.debug(f"{Env.current.website_endpoint}/folders/{task.folder_id}/tasks/{task.group_id}")
+        return task.batch_id
+    else:
+        task.upload_simulation(
+            stub=stub,
+            verbose=verbose,
+            progress_callback=progress_callback,
+            remote_sim_file=remote_sim_file,
+        )
+        estimate_cost(task_id=task.task_id, solver_version=solver_version, verbose=verbose)
+
+        task.validate_post_upload(parent_tasks=parent_tasks)
+
+        # log the url for the task in the web UI
+        log.debug(f"{Env.current.website_endpoint}/folders/{task.folder_id}/tasks/{task.task_id}")
+        return task
+
 def get_reduced_simulation(simulation, reduce_simulation):
     """
     Adjust the given simulation object based on the reduce_simulation parameter. Currently only