-
Notifications
You must be signed in to change notification settings - Fork 1.7k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Rewrite execution of microbatch models to avoid blocking the main thread #11332
Merged
QMalcolm
merged 36 commits into
main
from
qmalcolm--11243-stop-microbatch-from-blocking-main-thread
Mar 3, 2025
+454
−387
Merged
Changes from 34 commits
Commits
Show all changes
36 commits
Select commit
Hold shift + click to select a range
1da252a
Add comments to help guide the work of pushing microbatch batch orche…
QMalcolm 90d53c2
Stub out new `MicrobatchModelRunner` and `MicrobatchBatchRunner`.
QMalcolm b149ec4
Make `handle_job_queue` of `RunTask` call invoke the `MicrobatchModel…
QMalcolm 85b4f23
Straight copy `handle_microbatch_model` into `MicrobatchModelRunner`
QMalcolm d23095d
Convert `handle_microbatch_model` into runner `execute` type method
QMalcolm 5279140
Begin getting batches in `MicrobatchModelRunner`
QMalcolm e45c2b1
Begin getting `_has_relation` in execute method of `MicrobatchModelRu…
QMalcolm c4d1b1c
Begin initializing empty `RunResult` for `MicrobatchModelRunner`
QMalcolm 374a9bd
Make `RunTask` and `ThreadPool` available in `MicrobatchModelRunner`
QMalcolm a2331df
Make `merge_batch_results` to `MicrobatchModelRunner`
QMalcolm ce9333c
Begin instantiating `MicrobatchBatchRunner` during `_submit_batch`
QMalcolm 8397264
Move `should_run_in_parallel` to `MicrobatchBatchRunner`
QMalcolm 046e08a
Move `describe_batch` to `MicrobatchBatchRunner`
QMalcolm 86a8439
Move batch print eventing methods to `MicrobatchBatchRunner`
QMalcolm da1b08c
Move `before_execute` and `after_execute` to `MicrobatchBatchRunner`
QMalcolm 4bc33fb
Move `describe_node` to `MicrobatchBatchRunner` and `MicrobatchModelR…
QMalcolm df3a3f1
Implement `compile` for `MicrobatchBatchRunner` and `MicrobatchModelR…
QMalcolm 1a8fbf1
Implement `on_skip` for `MicrobatchBatchRunner`
QMalcolm 590fbb0
Move batch execution logic from old runner into `MicrobatchBatchRunner`
QMalcolm 1c54050
Make `build_jinja_context_for_batch` a static method
QMalcolm 5c5558e
Ensure first batch is run as full refresh when relevant
QMalcolm e60ceb6
Update `BuildTask` to use `MicrobatchModelRunner` directly
QMalcolm ce0d4c1
Fix mocking of `should_run_in_parallel` in microbatch tests
QMalcolm 7ec1ce5
Fix unit tests for `build_jinja_context_for_batch`
QMalcolm 35bc7ce
Fix unit tests for `test_should_run_in_parallel`
QMalcolm a4f62e3
Allow for cancellation of microbatch runs when batches are running se…
QMalcolm 94fc86c
Enable graceful shutdown of microbatch model executing batches concur…
QMalcolm c07e071
Remove unecessary tracking/printing at end of `MicrobatchModelRunner.…
QMalcolm 1ba67a9
Remove `describe_node` method from `MicrobatchBatchRunner`
QMalcolm 20a1446
Rename print line methods in `MicrobatchBatchRunner` to simplify impl…
QMalcolm c81010e
Abstract initialization of `MicrobatchBuilder` into utility method
QMalcolm c8b78f1
Remove special skip logic in `MicrobatchModelRunner.execute`
QMalcolm 303cdf7
Remove TODO statement which is no longer relevant
QMalcolm 486f351
Add changie doc
QMalcolm a32a7ae
Use setters for `parent_task` and `pool` of `MicrobatchModelRunner`
QMalcolm ef461da
Add comments around the necessity of `_parent_task` and `_pool` on `M…
QMalcolm File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,6 @@ | ||
kind: Fixes | ||
body: Fix microbatch execution to not block main thread nor hang | ||
time: 2025-03-03T13:14:40.432874-06:00 | ||
custom: | ||
Author: QMalcolm | ||
Issue: 11243 11306 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
from __future__ import annotations | ||
|
||
from multiprocessing.pool import ThreadPool | ||
|
||
|
||
class DbtThreadPool(ThreadPool): | ||
"""A ThreadPool that tracks whether or not it's been closed""" | ||
|
||
def __init__(self, *args, **kwargs): | ||
super().__init__(*args, **kwargs) | ||
self.closed = False | ||
|
||
def close(self): | ||
self.closed = True | ||
super().close() | ||
|
||
def is_closed(self): | ||
return self.closed |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We created
DbtThreadPool
so that we can have visibility on whether.close()
has been called on the pool. This class is now used instead ofThreadPool
.