fix: Closes threadpool resources upon datastore online_write_batch co…

…mpletion (#2386) Signed-off-by: Pamela Toman <ptoman@paloaltonetworks.com>
feast-dev · Mar 9, 2022 · 7a93152 · 7a93152
1 parent 993b8cc
commit 7a93152
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/sdk/python/feast/infra/online_stores/datastore.py b/sdk/python/feast/infra/online_stores/datastore.py
@@ -159,11 +159,13 @@ def online_write_batch(
         write_batch_size = online_config.write_batch_size
         feast_project = config.project
 
-        pool = ThreadPool(processes=write_concurrency)
-        pool.map(
-            lambda b: self._write_minibatch(client, feast_project, table, b, progress),
-            self._to_minibatches(data, batch_size=write_batch_size),
-        )
+        with ThreadPool(processes=write_concurrency) as pool:
+            pool.map(
+                lambda b: self._write_minibatch(
+                    client, feast_project, table, b, progress
+                ),
+                self._to_minibatches(data, batch_size=write_batch_size),
+            )
 
     @staticmethod
     def _to_minibatches(data: ProtoBatch, batch_size) -> Iterator[ProtoBatch]: