cleanup and fix resource handling

Author: shcheklein

src/datachain/catalog/catalog.py

@@ -9,6 +9,7 @@
 import time
 import traceback
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+from contextlib import ExitStack
 from copy import copy
 from dataclasses import dataclass
 from functools import cached_property, reduce
@@ -314,6 +315,16 @@ def download(self, recursive: bool = False, pbar=None) -> None:
         if self.sources:
             self.client.fetch_nodes(self.iternodes(recursive), shared_progress_bar=pbar)
 
+    def close(self) -> None:
+        if self.listing:
+            self.listing.close()
+
+    def __enter__(self) -> "NodeGroup":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.close()
+
 
 def prepare_output_for_cp(
     node_groups: list[NodeGroup],
@@ -1918,38 +1929,40 @@ def cp(
             no_glob,
             client_config=client_config,
         )
+        with ExitStack() as stack:
+            for node_group in node_groups:
+                stack.enter_context(node_group)
+            always_copy_dir_contents, copy_to_filename = prepare_output_for_cp(
+                node_groups, output, force, no_cp
+            )
+            total_size, total_files = collect_nodes_for_cp(node_groups, recursive)
+            if not total_files:
+                return
 
-        always_copy_dir_contents, copy_to_filename = prepare_output_for_cp(
-            node_groups, output, force, no_cp
-        )
-        total_size, total_files = collect_nodes_for_cp(node_groups, recursive)
-        if not total_files:
-            return
-
-        desc_max_len = max(len(output) + 16, 19)
-        bar_format = (
-            "{desc:<"
-            f"{desc_max_len}"
-            "}{percentage:3.0f}%|{bar}| {n_fmt:>5}/{total_fmt:<5} "
-            "[{elapsed}<{remaining}, {rate_fmt:>8}]"
-        )
+            desc_max_len = max(len(output) + 16, 19)
+            bar_format = (
+                "{desc:<"
+                f"{desc_max_len}"
+                "}{percentage:3.0f}%|{bar}| {n_fmt:>5}/{total_fmt:<5} "
+                "[{elapsed}<{remaining}, {rate_fmt:>8}]"
+            )
 
-        if not no_cp:
-            with get_download_bar(bar_format, total_size) as pbar:
-                for node_group in node_groups:
-                    node_group.download(recursive=recursive, pbar=pbar)
+            if not no_cp:
+                with get_download_bar(bar_format, total_size) as pbar:
+                    for node_group in node_groups:
+                        node_group.download(recursive=recursive, pbar=pbar)
 
-        instantiate_node_groups(
-            node_groups,
-            output,
-            bar_format,
-            total_files,
-            force,
-            recursive,
-            no_cp,
-            always_copy_dir_contents,
-            copy_to_filename,
-        )
+            instantiate_node_groups(
+                node_groups,
+                output,
+                bar_format,
+                total_files,
+                force,
+                recursive,
+                no_cp,
+                always_copy_dir_contents,
+                copy_to_filename,
+            )
 
     def du(
         self,

src/datachain/data_storage/metastore.py

@@ -3,6 +3,7 @@
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
+from contextlib import contextmanager, suppress
 from datetime import datetime, timezone
 from functools import cached_property, reduce
 from itertools import groupby
@@ -118,6 +119,16 @@ def close_on_exit(self) -> None:
         differently."""
         self.close()
 
+    @contextmanager
+    def _init_guard(self):
+        """Ensure resources acquired during __init__ are released on failure."""
+        try:
+            yield
+        except Exception:
+            with suppress(Exception):
+                self.close_on_exit()
+            raise
+
     def cleanup_tables(self, temp_table_names: list[str]) -> None:
         """Cleanup temp tables."""
 

src/datachain/data_storage/sqlite.py

@@ -374,11 +374,12 @@ def __init__(
 
         self.db = db or SQLiteDatabaseEngine.from_db_file(db_file)
 
-        self._init_meta_table()
-        self._init_meta_schema_value()
-        self._check_schema_version()
-        self._init_tables()
-        self._init_namespaces_projects()
+        with self._init_guard():
+            self._init_meta_table()
+            self._init_meta_schema_value()
+            self._check_schema_version()
+            self._init_tables()
+            self._init_namespaces_projects()
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:
         """Close connection upon exit from context manager."""

src/datachain/listing.py

@@ -35,6 +35,7 @@ def __init__(
         self.dataset_name = dataset_name  # dataset representing bucket listing
         self.dataset_version = dataset_version  # dataset representing bucket listing
         self.column = column
+        self._closed = False
 
     def clone(self) -> "Listing":
         return self.__class__(
@@ -53,7 +54,13 @@ def __exit__(self, exc_type, exc_value, traceback) -> None:
         self.close()
 
     def close(self) -> None:
-        self.warehouse.close()
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            self.warehouse.close_on_exit()
+        finally:
+            self.metastore.close_on_exit()
 
     @property
     def uri(self):

src/datachain/query/session.py

@@ -1,5 +1,4 @@
 import atexit
-import gc
 import logging
 import os
 import re
@@ -8,6 +7,7 @@
 from collections.abc import Callable
 from typing import TYPE_CHECKING, ClassVar
 from uuid import uuid4
+from weakref import WeakSet
 
 from datachain.catalog import get_catalog
 from datachain.data_storage import JobQueryType, JobStatus
@@ -57,6 +57,7 @@ class Session:
 
     GLOBAL_SESSION_CTX: "Session | None" = None
     SESSION_CONTEXTS: ClassVar[list["Session"]] = []
+    _ALL_SESSIONS: ClassVar[WeakSet["Session"]] = WeakSet()
     ORIGINAL_EXCEPT_HOOK = None
 
     # Job management - class-level to ensure one job per process
@@ -92,6 +93,7 @@ def __init__(
         self.catalog = catalog or get_catalog(
             client_config=client_config, in_memory=in_memory
         )
+        Session._ALL_SESSIONS.add(self)
 
     def __enter__(self):
         # Push the current context onto the stack
@@ -109,6 +111,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
         if Session.SESSION_CONTEXTS:
             Session.SESSION_CONTEXTS.pop()
+        Session._ALL_SESSIONS.discard(self)
 
     def get_or_create_job(self) -> "Job":
         """
@@ -311,6 +314,7 @@ def except_hook(exc_type, exc_value, exc_traceback):
 
     @classmethod
     def cleanup_for_tests(cls):
+        cls._close_all_contexts()
         if cls.GLOBAL_SESSION_CTX is not None:
             cls.GLOBAL_SESSION_CTX.__exit__(None, None, None)
             cls.GLOBAL_SESSION_CTX = None
@@ -333,15 +337,26 @@ def cleanup_for_tests(cls):
 
     @staticmethod
     def _global_cleanup():
+        Session._close_all_contexts()
         if Session.GLOBAL_SESSION_CTX is not None:
             Session.GLOBAL_SESSION_CTX.__exit__(None, None, None)
 
-        for obj in gc.get_objects():  # Get all tracked objects
+        for session in list(Session._ALL_SESSIONS):
             try:
-                if isinstance(obj, Session):
-                    # Cleanup temp dataset for session variables.
-                    obj.__exit__(None, None, None)
+                session.__exit__(None, None, None)
             except ReferenceError:
                 continue  # Object has been finalized already
             except Exception as e:  # noqa: BLE001
                 logger.error(f"Exception while cleaning up session: {e}")  # noqa: G004
+
+    @classmethod
+    def _close_all_contexts(cls) -> None:
+        while cls.SESSION_CONTEXTS:
+            session = cls.SESSION_CONTEXTS.pop()
+            try:
+                session.__exit__(None, None, None)
+            except Exception as exc:  # noqa: BLE001
+                logger.error(
+                    "Exception while closing session context during cleanup: %s",
+                    exc,
+                )

src/datachain/sql/sqlite/base.py

@@ -3,6 +3,7 @@
 import sqlite3
 import warnings
 from collections.abc import Callable, Iterable
+from contextlib import closing
 from datetime import MAXYEAR, MINYEAR, datetime, timezone
 from functools import cache
 from types import MappingProxyType
@@ -111,7 +112,10 @@ def setup():
     compiles(numeric.int_hash_64, "sqlite")(compile_int_hash_64)
     compiles(numeric.bit_hamming_distance, "sqlite")(compile_bit_hamming_distance)
 
-    if load_usearch_extension(sqlite3.connect(":memory:")):
+    with closing(sqlite3.connect(":memory:")) as _usearch_conn:
+        usearch_available = load_usearch_extension(_usearch_conn)
+
+    if usearch_available:
         compiles(array.cosine_distance, "sqlite")(compile_cosine_distance_ext)
         compiles(array.euclidean_distance, "sqlite")(compile_euclidean_distance_ext)
     else:
@@ -145,23 +149,34 @@ def functions_exist(
                 f"Found value of type {type(n).__name__}: {n!r}"
             )
 
+    close_connection = False
     if connection is None:
         connection = sqlite3.connect(":memory:")
+        close_connection = True
 
-    if not names:
-        return True
-    column1 = sa.column("column1", sa.String)
-    func_name_query = column1.not_in(
-        sa.select(sa.column("name", sa.String)).select_from(func.pragma_function_list())
-    )
-    query = (
-        sa.select(func.count() == 0)
-        .select_from(sa.values(column1).data([(n,) for n in names]))
-        .where(func_name_query)
-    )
-    comp = query.compile(dialect=sqlite_dialect)
-    args = (comp.string, comp.params) if comp.params else (comp.string,)
-    return bool(connection.execute(*args).fetchone()[0])
+    try:
+        if not names:
+            return True
+        column1 = sa.column("column1", sa.String)
+        func_name_query = column1.not_in(
+            sa.select(sa.column("name", sa.String)).select_from(
+                func.pragma_function_list()
+            )
+        )
+        query = (
+            sa.select(func.count() == 0)
+            .select_from(sa.values(column1).data([(n,) for n in names]))
+            .where(func_name_query)
+        )
+        comp = query.compile(dialect=sqlite_dialect)
+        if comp.params:
+            result = connection.execute(comp.string, comp.params)
+        else:
+            result = connection.execute(comp.string)
+        return bool(result.fetchone()[0])
+    finally:
+        if close_connection:
+            connection.close()
 
 
 def create_user_defined_sql_functions(connection):

tests/conftest.py

@@ -2,6 +2,7 @@
 import os.path
 import signal
 import subprocess  # nosec B404
+import sys
 import uuid
 from collections.abc import Generator
 from datetime import datetime
@@ -40,6 +41,10 @@
 
 from .utils import DEFAULT_TREE, instantiate_tree, reset_session_job_state
 
+distributed_pythonpath = os.environ.get("DATACHAIN_DISTRIBUTED_PYTHONPATH")
+if distributed_pythonpath and distributed_pythonpath not in sys.path:
+    sys.path.insert(0, distributed_pythonpath)
+
 DEFAULT_DATACHAIN_BIN = "datachain"
 DEFAULT_DATACHAIN_GIT_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
@@ -109,10 +114,11 @@ def monkeypatch_session() -> Generator[MonkeyPatch, None, None]:
 
 
 @pytest.fixture(autouse=True)
-def clean_session() -> None:
+def clean_session() -> Generator[None, None, None]:
     """
-    Make sure we clean leftover session before each test case
+    Clean leftover sessions after each test while storage handles are still open.
     """
+    yield
     Session.cleanup_for_tests()
 
 
@@ -181,11 +187,13 @@ def metastore(monkeypatch):
 
         yield _metastore
 
+        Session.cleanup_for_tests()
         _metastore.cleanup_for_tests()
     else:
         _metastore = SQLiteMetastore(db_file=":memory:")
         yield _metastore
 
+        Session.cleanup_for_tests()
         cleanup_sqlite_db(_metastore.db.clone(), _metastore.default_table_names)
 
     # Close the connection so that the SQLite file is no longer open, to avoid
@@ -256,11 +264,13 @@ def metastore_tmpfile(monkeypatch, tmp_path):
 
         yield _metastore
 
+        Session.cleanup_for_tests()
         _metastore.cleanup_for_tests()
     else:
         _metastore = SQLiteMetastore(db_file=str(tmp_path / "test.db"))
         yield _metastore
 
+        Session.cleanup_for_tests()
         cleanup_sqlite_db(_metastore.db.clone(), _metastore.default_table_names)
 
     # Close the connection so that the SQLite file is no longer open, to avoid
@@ -529,7 +539,10 @@ def cloud_test_catalog(
     metastore,
     warehouse,
 ):
-    return get_cloud_test_catalog(cloud_server, tmp_path, metastore, warehouse)
+    catalog = get_cloud_test_catalog(cloud_server, tmp_path, metastore, warehouse)
+    yield catalog
+
+    reset_session_job_state()
 
 
 @pytest.fixture

tests/func/model/test_yolo.py

@@ -39,7 +39,11 @@ def running_img_masks() -> torch.Tensor:
     mask1_file = os.path.join(os.path.dirname(__file__), "data", "running-mask1.png")
     mask1_np = np.array(Image.open(mask1_file))
 
-    return torch.tensor([mask0_np.astype(np.float32), mask1_np.astype(np.float32)])
+    stacked = np.stack(
+        [mask0_np.astype(np.float32), mask1_np.astype(np.float32)],
+        axis=0,
+    )
+    return torch.from_numpy(stacked)
 
 
 def test_yolo_bbox_from_results_empty(running_img):