Update DaskTaskRunner for compatibility with the updated engine
#13555
Conversation
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
desertaxle
force-pushed
the
udpated-dask-task-runner-2
branch
from
6600f22
to
e298937
Compare
desertaxle
force-pushed
the
udpated-dask-task-runner-2
branch
from
e298937
to
aff963e
Compare
desertaxle
commented
May 24, 2024
Comment on lines
+328
to
+401
| from prefect.utilities.engine import ( | ||
| collect_task_run_inputs_sync, | ||
| resolve_inputs_sync, | ||
| ) | ||
|
|
||
| # We need to resolve some futures to map over their data, collect the upstream | ||
| # links beforehand to retain relationship tracking. | ||
| task_inputs = { | ||
| k: collect_task_run_inputs_sync(v, max_depth=0) | ||
| for k, v in parameters.items() | ||
| } | ||
|
|
||
| # Resolve the top-level parameters in order to get mappable data of a known length. | ||
| # Nested parameters will be resolved in each mapped child where their relationships | ||
| # will also be tracked. | ||
| parameters = resolve_inputs_sync(parameters, max_depth=0) | ||
|
|
||
| # Ensure that any parameters in kwargs are expanded before this check | ||
| parameters = explode_variadic_parameter(task.fn, parameters) | ||
|
|
||
| iterable_parameters = {} | ||
| static_parameters = {} | ||
| annotated_parameters = {} | ||
| for key, val in parameters.items(): | ||
| if isinstance(val, (allow_failure, quote)): | ||
| # Unwrap annotated parameters to determine if they are iterable | ||
| annotated_parameters[key] = val | ||
| val = val.unwrap() | ||
|
|
||
| if isinstance(val, unmapped): | ||
| static_parameters[key] = val.value | ||
| elif isiterable(val): | ||
| iterable_parameters[key] = list(val) | ||
| else: | ||
| static_parameters[key] = val | ||
|
|
||
| if not len(iterable_parameters): | ||
| raise MappingMissingIterable( | ||
| "No iterable parameters were received. Parameters for map must " | ||
| f"include at least one iterable. Parameters: {parameters}" | ||
| ) | ||
|
|
||
| iterable_parameter_lengths = { | ||
| key: len(val) for key, val in iterable_parameters.items() | ||
| } | ||
| lengths = set(iterable_parameter_lengths.values()) | ||
| if len(lengths) > 1: | ||
| raise MappingLengthMismatch( | ||
| "Received iterable parameters with different lengths. Parameters for map" | ||
| f" must all be the same length. Got lengths: {iterable_parameter_lengths}" | ||
| ) | ||
|
|
||
| map_length = list(lengths)[0] | ||
|
|
||
| futures = [] | ||
| for i in range(map_length): | ||
| call_parameters = { | ||
| key: value[i] for key, value in iterable_parameters.items() | ||
| } | ||
| call_parameters.update( | ||
| {key: value for key, value in static_parameters.items()} | ||
| ) | ||
|
|
||
| # Add default values for parameters; these are skipped earlier since they should | ||
| # not be mapped over | ||
| for key, value in get_parameter_defaults(task.fn).items(): | ||
| call_parameters.setdefault(key, value) | ||
|
|
||
| # Re-apply annotations to each key again | ||
| for key, annotation in annotated_parameters.items(): | ||
| call_parameters[key] = annotation.rewrap(call_parameters[key]) | ||
|
|
||
| # Collapse any previously exploded kwargs | ||
| call_parameters = collapse_variadic_parameters(task.fn, call_parameters) |
There was a problem hiding this comment.
This is identical to a chunk of logic in the ThreadPoolTaskRunner. If I need to repeat this in the RayTaskRunner too, I'll move it up to the TaskRunner base class.
desertaxle
requested review from
zzstoatzz,
chrisguidry and
a team
as code owners
cicdw
reviewed
May 25, 2024
| parameters: Dict[str, Any], | ||
| wait_for: Iterable[PrefectFuture], | ||
| dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None, | ||
| ) -> PrefectDaskFuture: |
There was a problem hiding this comment.
Every future type that we implement for each task runner will have the same interface, right? It's just there will be some special handling based on the underlying system?
There was a problem hiding this comment.
Yes, Each specific future will have a different implementation for waiting for a future and getting a future result, but the .wait, .result, and .state interfaces will all be the same.
| """ | ||
| self.__dict__.update(data) | ||
| self._client = distributed.get_client() | ||
| def __exit__(self, *args): |
There was a problem hiding this comment.
What let you drop the serialization requirement?
There was a problem hiding this comment.
Previously, a PrefectFuture carried an instance of the task runner that submitted the run around with it. We don't have to do that anymore because we directly wrap a future, so we shouldn't be pickling task runners anymore.
cicdw
approved these changes
May 27, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Updates the task runner to implement the new
TaskRunnerinterface. The new task runner delegates most responsibilities to thePrefectDistributedClient, but the task runner handles wrapping and unwrappingPrefectDaskFutures to ensure work is efficiently scheduled on a Dask cluster.Example
Checklist
<link to issue>"maintenance,fix,feature,enhancement,docs.For documentation changes:
netlify.tomlfor files that are removed or renamed.For new functions or classes in the Python SDK:
mkdocs.ymlnavigation.