fixes

pyrocko · Nov 7, 2023 · 2a5cd7b · 2a5cd7b
1 parent b26e1b9
commit 2a5cd7b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/lassie/images/__init__.py b/lassie/images/__init__.py
@@ -86,7 +86,9 @@ async def worker() -> None:
                 "start pre-processing images, queue size %d", self._queue.maxsize
             )
             async for batch in batch_iterator:
+                start_time = datetime_now()
                 images = await self.process_traces(batch.traces)
+                stats.time_per_batch = datetime_now() - start_time
                 if self._queue.empty() and self._processed_images:
                     logger.warning("image queue ran empty, processing is slow")
                 self._processed_images += 1
@@ -98,9 +100,7 @@ async def worker() -> None:
 
         while True:
             stats.queue_size = self._queue.qsize()
-            start_time = datetime_now()
             ret = await self._queue.get()
-            stats.time_per_batch = datetime_now() - start_time
             if ret is None:
                 logger.debug("image function finished")
                 break

diff --git a/lassie/tracers/cake.py b/lassie/tracers/cake.py
@@ -477,7 +477,7 @@ def _interpolate_travel_times(
             traveltimes.append(self._interpolate_traveltimes_sptree(coords))
             PROGRESS.update(status, advance=1)
 
-        PROGRESS.remove_task()
+        PROGRESS.remove_task(status)
 
         return np.asarray(traveltimes).astype(float)