From 23e692f3114a18acec235e890b8c69146b2c1eb0 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Wed, 31 Jan 2024 09:51:41 -0800 Subject: [PATCH] turn off client cleanup in multi-ensemble -> 1 client case --- src/tape/ensemble.py | 10 +++++++++- tests/tape_tests/test_ensemble.py | 2 ++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/tape/ensemble.py b/src/tape/ensemble.py index 8f5809cd..3aa5e442 100644 --- a/src/tape/ensemble.py +++ b/src/tape/ensemble.py @@ -500,9 +500,17 @@ def sample_objects(self, **kwargs): object_subset.set_dirty(True) # make a new ensemble - # TODO: Investigate shared client warning if self.client is not None: new_ens = Ensemble(client=self.client) + + # turn off cleanups -- in the case where multiple ensembles are + # using a client, an individual ensemble should not close the + # client during an __exit__ or __del__ event. This means that + # the client will not be closed without an explicit client.close() + # call, which is unfortunate... not sure of an alternative way + # forward. + self.cleanup_client = False + new_ens.cleanup_client = False else: new_ens = Ensemble(client=False) diff --git a/tests/tape_tests/test_ensemble.py b/tests/tape_tests/test_ensemble.py index 8ba5c375..20c57d0f 100644 --- a/tests/tape_tests/test_ensemble.py +++ b/tests/tape_tests/test_ensemble.py @@ -853,6 +853,8 @@ def test_sample_objects(parquet_ensemble_with_divisions): assert len(ens.object) == prior_obj_len assert len(ens.source) == prior_src_len + ens.client.close() + def test_update_column_map(dask_client): """