fix real cost

Author: tylerflex

tests/test_web/test_webapi.py

@@ -27,6 +27,7 @@
 PROJECT_NAME = "default"
 FLEX_UNIT = 1.0
 
+
 def make_sim():
     """Makes a simulation."""
     return td.Simulation(size=(1, 1, 1), grid_spec=td.GridSpec.auto(wavelength=1.0), run_time=1e-12)
@@ -83,7 +84,7 @@ def mock_get_info(monkeypatch, set_api_key):
             "data": {
                 "taskId": TASK_ID,
                 "createdAt": CREATED_AT,
-                "real_flex_unit": FLEX_UNIT,
+                "realFlexUnit": FLEX_UNIT,
             }
         },
         status=200,
@@ -506,9 +507,10 @@ def test_monitor(mock_monitor):
     monitor(TASK_ID, verbose=True)
     monitor(TASK_ID, verbose=False)
 
+
 @responses.activate
 def test_real_cost(mock_get_info):
-    real_cost(TASK_ID)
+    assert real_cost(TASK_ID) == FLEX_UNIT
 
 
 """ Containers """
@@ -526,6 +528,7 @@ def test_job(mock_webapi, monkeypatch):
     j.get_info()
     # j.download
     j.delete
+    assert j.real_cost() == FLEX_UNIT
 
 
 @pytest.fixture
@@ -543,6 +546,7 @@ def test_batch(mock_webapi, mock_job_status):
     sims = {TASK_NAME: make_sim()}
     b = Batch(simulations=sims, folder_name=PROJECT_NAME)
     batch_data = b.run(path_dir="tests/tmp/")
+    assert b.real_cost() == FLEX_UNIT * len(sims)
 
 
 """ Async """

tidy3d/web/container.py

@@ -170,6 +170,10 @@ def delete(self):
         """Delete server-side data associated with :class:`Job`."""
         web.delete(self.task_id)
 
+    def real_cost(self):
+        """Get the billed cost for the task associated with this job."""
+        return web.real_cost(self.task_id)
+
 
 class BatchData(Tidy3dBaseModel):
     """Holds a collection of :class:`.SimulationData` returned by :class:`.Batch`."""
@@ -512,3 +516,12 @@ def delete(self):
         """Delete server-side data associated with each task in the batch."""
         for _, job in self.jobs.items():
             job.delete()
+
+    def real_cost(self):
+        """Get the sum of billed costs for each task associated with this batch."""
+        real_cost_sum = 0.0
+        for _, job in self.jobs.items():
+            cost_job = job.real_cost()
+            if cost_job is not None:
+                real_cost_sum += cost_job
+        return real_cost_sum or None

tidy3d/web/simulation_task.py

@@ -125,7 +125,7 @@ class SimulationTask(ResourceLifecycle, Submittable, extra=Extra.allow):
     status: Optional[str] = Field(title="status", description="Simulation task status.")
 
     real_flex_unit: float = Field(
-        None, title="real flex units", description="Billed flex units.", alias="realFlexUnit"
+        None, title="real flex units", description="Billed flex units.", alias="realCost"
     )
 
     created_at: Optional[datetime] = Field(
@@ -376,18 +376,6 @@ def estimate_cost(self, solver_version=None, protocol_version=None) -> float:
         )
         return resp
 
-    def real_cost(self) -> float:
-        """Get the billed cost for given task after it has been run.
-
-        Note
-        ----
-            The billed cost may not be immediately available
-            when the task status is set to ``success``,
-            but should be available shortly after.
-        """
-        assert self.task_id
-        return self.get(self.task_id).real_flex_unit
-
     def get_simulation_hdf5(
         self, to_file: str, verbose: bool = True, progress_callback: Callable[[float], None] = None
     ) -> None:

tidy3d/web/webapi.py

@@ -640,13 +640,11 @@ def real_cost(task_id: str) -> float:
         The billed cost may not be immediately available when the task status is set to ``success``,
         but should be available shortly after.
     """
-    task = SimulationTask.get(task_id)
-    if not task:
-        raise ValueError("Task not found.")
-    flex_unit = task.real_flex_unit
+    task_info = get_info(task_id)
+    flex_unit = task_info.realFlexUnit
     if not flex_unit:
         log.warning(
-            f"Billed FlexUnit for task '{task_id}'' is not available. If the task has been "
+            f"Billed FlexUnit for task '{task_id}' is not available. If the task has been "
             "successfully run, it should be available shortly."
         )
     return flex_unit