adapt the query_cache_pool

Author: themisvaltinos

sqlmesh/core/loader.py

@@ -1,15 +1,12 @@
 from __future__ import annotations
 
 import abc
-import concurrent.futures
 import glob
 import itertools
 import linecache
-import multiprocessing as mp
 import os
 import re
 import typing as t
-import concurrent
 from collections import Counter, defaultdict
 from dataclasses import dataclass
 from pathlib import Path
@@ -40,6 +37,7 @@
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
 from sqlmesh.utils.metaprogramming import import_python_file
+from sqlmesh.utils.process import create_process_pool_executor
 from sqlmesh.utils.yaml import YAML, load as yaml_load
 
 
@@ -531,8 +529,7 @@ def _load_sql_models(
             )
 
             errors: t.List[str] = []
-            with concurrent.futures.ProcessPoolExecutor(
-                mp_context=mp.get_context("fork"),
+            with create_process_pool_executor(
                 initializer=_init_model_defaults,
                 initargs=(self.config, gateway, defaults, cache),
                 max_workers=c.MAX_FORK_WORKERS,

sqlmesh/core/model/cache.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
 import logging
-import multiprocessing as mp
 import typing as t
-from concurrent.futures import ProcessPoolExecutor
 from pathlib import Path
 
 from sqlglot import exp
@@ -15,6 +13,7 @@
 from sqlmesh.core.model.definition import ExternalModel, Model, SqlModel
 from sqlmesh.utils.cache import FileCache
 from sqlmesh.utils.hashing import crc32
+from sqlmesh.utils.process import PoolExecutor, create_process_pool_executor
 
 from dataclasses import dataclass
 
@@ -135,9 +134,8 @@ def _entry_name(model: SqlModel) -> str:
         return f"{model.name}_{crc32(hash_data)}"
 
 
-def optimized_query_cache_pool(optimized_query_cache: OptimizedQueryCache) -> ProcessPoolExecutor:
-    return ProcessPoolExecutor(
-        mp_context=mp.get_context("fork"),
+def optimized_query_cache_pool(optimized_query_cache: OptimizedQueryCache) -> PoolExecutor:
+    return create_process_pool_executor(
         initializer=_init_optimized_query_cache,
         initargs=(optimized_query_cache,),
         max_workers=c.MAX_FORK_WORKERS,

sqlmesh/core/model/definition.py

@@ -811,8 +811,10 @@ def convert_to_time_column(
         return exp.convert(time)
 
     def set_mapping_schema(self, schema: t.Dict) -> None:
+        # Make a shallow copy to avoid modifying the original in case they're the same
+        temp_schema = schema.copy()
         self.mapping_schema.clear()
-        self.mapping_schema.update(schema)
+        self.mapping_schema.update(temp_schema)
 
     def update_schema(self, schema: MappingSchema) -> None:
         """Updates the schema for this model's dependencies based on the given mapping schema."""

sqlmesh/core/model/schema.py

@@ -28,10 +28,7 @@ def update_model_schemas(
     schema = MappingSchema(normalize=False)
     optimized_query_cache: OptimizedQueryCache = OptimizedQueryCache(context_path / c.CACHE)
 
-    if c.MAX_FORK_WORKERS == 1:
-        _update_model_schemas_sequential(dag, models, schema, optimized_query_cache)
-    else:
-        _update_model_schemas_parallel(dag, models, schema, optimized_query_cache)
+    _update_model_schemas(dag, models, schema, optimized_query_cache)
 
 
 def _update_schema_with_model(schema: MappingSchema, model: Model) -> None:
@@ -49,25 +46,7 @@ def _update_schema_with_model(schema: MappingSchema, model: Model) -> None:
             raise
 
 
-def _update_model_schemas_sequential(
-    dag: DAG[str],
-    models: UniqueKeyDict[str, Model],
-    schema: MappingSchema,
-    optimized_query_cache: OptimizedQueryCache,
-) -> None:
-    for name in dag.sorted:
-        model = models.get(name)
-
-        # External models don't exist in the context, so we need to skip them
-        if not model:
-            continue
-
-        model.update_schema(schema)
-        optimized_query_cache.with_optimized_query(model)
-        _update_schema_with_model(schema, model)
-
-
-def _update_model_schemas_parallel(
+def _update_model_schemas(
     dag: DAG[str],
     models: UniqueKeyDict[str, Model],
     schema: MappingSchema,
@@ -102,17 +81,24 @@ def process_models(completed_model: t.Optional[Model] = None) -> None:
                     )
                 )
 
+    errors: t.List[str] = []
     with optimized_query_cache_pool(optimized_query_cache) as executor:
         process_models()
 
-        while futures:
+        while futures and not errors:
             for future in as_completed(futures):
-                futures.remove(future)
-                fqn, entry_name, data_hash, metadata_hash, mapping_schema = future.result()
-                model = models[fqn]
-                model._data_hash = data_hash
-                model._metadata_hash = metadata_hash
-                model.set_mapping_schema(mapping_schema)
-                optimized_query_cache.with_optimized_query(model, entry_name)
-                _update_schema_with_model(schema, model)
-                process_models(completed_model=model)
+                try:
+                    futures.remove(future)
+                    fqn, entry_name, data_hash, metadata_hash, mapping_schema = future.result()
+                    model = models[fqn]
+                    model._data_hash = data_hash
+                    model._metadata_hash = metadata_hash
+                    model.set_mapping_schema(mapping_schema)
+                    optimized_query_cache.with_optimized_query(model, entry_name)
+                    _update_schema_with_model(schema, model)
+                    process_models(completed_model=model)
+                except Exception as ex:
+                    errors.append(f"{ex}")
+
+    if errors:
+        raise SchemaError(f"Failed to update model schemas\n\n{'\n'.join(errors)}")

sqlmesh/core/snapshot/cache.py

@@ -55,20 +55,19 @@ def get_or_load(
             for snapshot in loaded_snapshots:
                 snapshots[snapshot.snapshot_id] = snapshot
 
-        if c.MAX_FORK_WORKERS != 1:
-            with optimized_query_cache_pool(self._optimized_query_cache) as executor:
-                for key, entry_name in executor.map(
-                    load_optimized_query,
-                    (
-                        (snapshot.model, s_id)
-                        for s_id, snapshot in snapshots.items()
-                        if snapshot.is_model
-                    ),
-                ):
-                    if entry_name:
-                        self._optimized_query_cache.with_optimized_query(
-                            snapshots[key].model, entry_name
-                        )
+        with optimized_query_cache_pool(self._optimized_query_cache) as executor:
+            for key, entry_name in executor.map(
+                load_optimized_query,
+                (
+                    (snapshot.model, s_id)
+                    for s_id, snapshot in snapshots.items()
+                    if snapshot.is_model
+                ),
+            ):
+                if entry_name:
+                    self._optimized_query_cache.with_optimized_query(
+                        snapshots[key].model, entry_name
+                    )
 
         for snapshot in snapshots.values():
             self._update_node_hash_cache(snapshot)

sqlmesh/utils/process.py

@@ -0,0 +1,60 @@
+# mypy: disable-error-code=no-untyped-def
+
+from concurrent.futures import Future, ProcessPoolExecutor
+import typing as t
+import multiprocessing as mp
+from sqlmesh.core import constants as c
+
+
+class SynchronousPoolExecutor:
+    """A mock implementation of the ProcessPoolExecutor for synchronous use.
+
+    This executor runs functions synchronously in the same process, avoiding the issues
+    with forking in test environments or when forking isn't possible (non-posix).
+    """
+
+    def __init__(self, max_workers=None, mp_context=None, initializer=None, initargs=()):
+        if initializer is not None:
+            try:
+                initializer(*initargs)
+            except BaseException as ex:
+                raise RuntimeError(f"Exception in initializer: {ex}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        return True
+
+    def submit(self, fn, *args, **kwargs):
+        """Execute the function synchronously and return a Future with the result."""
+        future = Future()
+        try:
+            result = fn(*args, **kwargs)
+            future.set_result(result)
+        except Exception as e:
+            future.set_exception(e)
+        return future
+
+    def map(self, fn, *iterables, timeout=None, chunksize=1):
+        """Synchronous implementation of ProcessPoolExecutor.map.
+
+        This executes the function for each set of inputs from the iterables in the
+        current process using Python's built-in map, rather than distributing work.
+        """
+        return map(fn, *iterables)
+
+
+PoolExecutor = t.Union[SynchronousPoolExecutor, ProcessPoolExecutor]
+
+
+def create_process_pool_executor(
+    initializer: t.Callable, initargs: t.Tuple, max_workers: t.Optional[int] = c.MAX_FORK_WORKERS
+) -> PoolExecutor:
+    executor = SynchronousPoolExecutor if max_workers == 1 else ProcessPoolExecutor
+    return executor(
+        mp_context=mp.get_context("fork"),
+        initializer=initializer,
+        initargs=initargs,
+        max_workers=max_workers,
+    )

tests/conftest.py

@@ -506,15 +506,3 @@ def _make_function(table_name: str, random_id: str) -> exp.Table:
         return temp_table
 
     return _make_function
-
-
-@pytest.fixture(autouse=True)
-def patch_process_pool_executor(mocker: MockerFixture, request):
-    """Patch ProcessPoolExecutor with MockProcessPoolExecutor in all tests except test_forking.py."""
-    # Skip mocking for test_forking.py
-    if request.node.fspath.basename == "test_forking.py":
-        return
-
-    from tests.mock_executor import MockProcessPoolExecutor
-
-    mocker.patch("concurrent.futures.ProcessPoolExecutor", MockProcessPoolExecutor)

tests/mock_executor.py

@@ -10,7 +10,7 @@
 
 def test_parallel_load(assert_exp_eq, mocker):
     mocker.patch("sqlmesh.core.constants.MAX_FORK_WORKERS", 2)
-    spy = mocker.spy(schema, "_update_model_schemas_parallel")
+    spy = mocker.spy(schema, "_update_model_schemas")
     context = Context(paths="examples/sushi")
 
     if hasattr(os, "fork"):