Feat: more cores for loading #4427
Conversation
Makefile
Outdated
| @@ -64,7 +64,7 @@ engine-up: engine-clickhouse-up engine-mssql-up engine-mysql-up engine-postgres- | |||
| engine-down: engine-clickhouse-down engine-mssql-down engine-mysql-down engine-postgres-down engine-spark-down engine-trino-down | |||
|
|
|||
| fast-test: | |||
| pytest -n auto -m "fast and not cicdonly" && pytest -m "isolated" | |||
| pytest -n auto -m "fast and not cicdonly" && pytest -m "isolated" && pytest -m "isolated2" | |||
There was a problem hiding this comment.
This label name is not very descriptive. Can we be more precise
sqlmesh/core/loader.py
Outdated
| for row in YAML().load(file.read()) | ||
| ] | ||
|
|
||
| cache.put(external_models, path) |
There was a problem hiding this comment.
I don't like the fact that we now need to manually manage the cache everywhere by splitting the load from the get. Any way to preserve the previous declarative API?
There was a problem hiding this comment.
i'm not sure. this is the simplest because the loading all happens in separate processes and must all be complete separate, we only communicate through the file system
sqlmesh/core/loader.py
Outdated
| models[model.fqn] = model | ||
|
|
||
| if paths: | ||
| defaults = dict( |
There was a problem hiding this comment.
This overlaps with the defaults that is in the dict itself. Can we have a more descriptive name for this?
There was a problem hiding this comment.
Renamed to model_loading_defaults for clarity, couldn't think of a better alternative
sqlmesh/core/loader.py
Outdated
| _selected_gateway = selected_gateway | ||
|
|
||
|
|
||
| def load_sql_models(path: Path) -> t.Tuple[Path, list[Model]]: |
There was a problem hiding this comment.
I don't see the 1st item of the return value being used anywhere. Is it used?
There was a problem hiding this comment.
no this was leftover from before I added the futures_to_paths dictionary. The path is already accessible so it’s no longer needed. I removed it and updated _load_sql_models to use the simplified return value—loaded = future.result() instead of unpacking with _.
sqlmesh/core/loader.py
Outdated
| self._track_file(seed_path) | ||
| if errors: | ||
| error_string = "\n".join(errors) | ||
| raise ConfigError(f"Failed to load models\n\n{error_string}") |
There was a problem hiding this comment.
[Nit] I believe Failed to load models\n\n is redundant
pytest.ini
Outdated
| @@ -8,6 +8,7 @@ markers = | |||
| remote: test that involves interacting with a remote DB | |||
| cicdonly: test that only runs on CI/CD | |||
| isolated: tests that need to run sequentially usually because they use fork | |||
| isolated2: tests that need to run isolated because they interfere | |||
There was a problem hiding this comment.
Why can't we use isolated here, btw?
There was a problem hiding this comment.
I've renamed the marker to registry_isolation and updated all references to it, to indicate this test needs isolation because of an issue with the registry. This particular test test_duplicate_python_model_names_raise_error added here: #3945 is successful when it is run in isolation, but it breaks if it is grouped with the rest of the tests, which is why they originally had it with isolated. But having it as before with isolated which are forking tests, since it’s designed to raise an error, leads to these tests failing as well
| except ConfigError: | ||
| from sqlmesh.core.console import get_console | ||
|
|
||
| get_console().log_warning( |
There was a problem hiding this comment.
Does this respect --ignore-warnings when this runs in a worker process?
There was a problem hiding this comment.
good catch, I tested it and your suspicion is correct it wasn't working because the console wasn’t shared across processes. I added a console parameter to _init_model_defaults to pass self._console to worker processes when creating the process pool, so that each worker used set_console(console) to match the parent’s console and preserves ignore_warnings along with all console settings
| pass | ||
|
|
||
| @abc.abstractmethod | ||
| def get(self, path: Path) -> t.List[Model]: |
There was a problem hiding this comment.
This interface is only needed in one place but impacts all other places. Can we keep both interfaces (get + put AND get_or_load_models) and revert unrelated changes?
There was a problem hiding this comment.
reverted to use in all other places except for sqlmodels the previous interface
| future.set_exception(e) | ||
| return future | ||
|
|
||
| def map(self, fn, *iterables, timeout=None, chunksize=1): |
There was a problem hiding this comment.
Don't you need a shutdown method too?
There was a problem hiding this comment.
since shutdown was called in __exit__ and since the synchronous executor runs everything in the main process, there are no resources to release or cleanup and futures are complete from the start this is why I didn't add it. but I revise to have it to keep the api similar to python’s ProcessPoolExecutor
This update adds a
ProcessPoolExecutorfor parallel loading of a project's models. It also adds a mock executor for single-process scenarios, such as when the system doesn’t support fork. Also, it refactorsoptimized_query_cache_poolto use this unified execution logic, eliminating the need for separate sequential and parallel code paths.