add pool as top-level definition arg #26256

prha · 2024-12-03T20:53:24Z

Summary & Motivation

Tags are a bad mechanism for specifying concurrency controls. Adds a top-level argument pool to asset/op definitions to replace the use of op tags to specify concurrency conditions.

This provides more cues for what you can place concurrency limits on (e.g. op defs not graph defs, asset defs not asset specs)

How I Tested These Changes

BK

Changelog

Adds a top-level argument pool to asset/op definitions to replace the use of op tags to specify concurrency conditions.

python_modules/dagster/dagster/_core/definitions/decorators/asset_decorator.py

python_modules/dagster/dagster/_core/definitions/op_definition.py

schrockn

So before we move forward with this I think we need to have a naming session of concurrency group versus pool versus something else. Important decision with long term ramifications.

prha · 2024-12-13T17:47:00Z

add pool tag to display for ops / assets #26804 : 2 dependent PRs (#26805 , #26965 )
RFC: separate configured limit from slot count #26803
add support for run-granularity op concurrency #26458
add new concurrency config to consolidate concurrency settings #26931
add pool graphql fields for runs #26964
hoists the run queue config from the run coordinator up to the instance #26932
thread pool information to graphql layer #26802
thread pool info to remote representation of assets #26801
add pool as top-level definition arg #26256 👈 (View in Graphite)
master

This stack of pull requests is managed by Graphite. Learn more about stacking.

prha · 2025-01-09T21:23:40Z

Discussed offline, but putting here for posterity:

I don't have plans to add PoolDefinition, because that automatically ties pools to a code location, and we specifically want pools to be a concept governing across code locations. If we come up with a richer API at the deployment layer, pools would be a great use case for it.

prha · 2025-01-10T01:18:52Z

Instead of hiding the pool arg in kwarg, just issued a preview_warning which I will remove once this is ready to go in 1.10.0

addressed

OwenKephart

minor comments

OwenKephart · 2025-01-15T00:47:17Z

python_modules/dagster/dagster/_core/definitions/graph_definition.py

@@ -808,6 +808,13 @@ def tags(self) -> Mapping[str, str]:
        """The tags associated with the graph."""
        return super().tags

+    @property
+    def pools(self) -> set[str]:


nit: prefer from collections.abc import Set

is that everyone's preference? the incidence of set is like 4x the count of Set:

(dagster) ➜ dagster git:(prha/asset_op_concurrency_key) s '\bset\[' | wc 423 2458 58005 (dagster) ➜ dagster git:(prha/asset_op_concurrency_key) s '\bSet\[' | wc 99 773 15437

for typing yeah -- I imagine almost all of those lower case sets are at runtime

until recently this was from typing import AbstractSet, from collections.abc import Set is the same thing but is now preferred by our linter -- this marks the return value as immutable, which has some value, but just in general I haven't seen set[x] used as a return type anywhere else in the code so we should be consistent

OwenKephart · 2025-01-15T00:47:23Z

python_modules/dagster/dagster/_core/definitions/node_definition.py

+
+    @property
+    @abstractmethod
+    def pools(self) -> set[str]: ...


OwenKephart · 2025-01-15T00:48:16Z

python_modules/dagster/dagster/_core/definitions/op_definition.py

@@ -286,6 +294,17 @@ def with_retry_policy(self, retry_policy: RetryPolicy) -> "PendingNodeInvocation
        """Creates a copy of this op with the given retry policy."""
        return super().with_retry_policy(retry_policy)

+    @property
+    def pool(self) -> Optional[str]:
+        """Optional[str]: The concurrency group for this op."""


nit: should the docstring say pool instead of group?

OwenKephart · 2025-01-15T00:49:20Z

python_modules/dagster/dagster/_core/execution/plan/active.py

@@ -644,9 +642,9 @@ def concurrency_event_iterator(
            ):
                step = self.get_step_by_key(step_key)
                step_context = plan_context.for_step(step)
-                step_concurrency_key = cast(str, step.tags.get(GLOBAL_CONCURRENCY_TAG))
+                pool = cast(str, step.pool)


nit: pre-existing, but should this be a check.str() instead of a cast?

OwenKephart · 2025-01-15T00:50:02Z

python_modules/dagster/dagster/_core/snap/node.py

@@ -159,6 +159,7 @@ class GraphDefSnap:
    dep_structure_snapshot: DependencyStructureSnapshot
    input_mapping_snaps: Sequence[InputMappingSnap]
    output_mapping_snaps: Sequence[OutputMappingSnap]
+    pools: set[str]


OwenKephart · 2025-01-15T00:50:50Z

python_modules/dagster/dagster/_core/snap/node.py

@@ -176,15 +177,39 @@ def get_output_snap(self, name: str) -> OutputDefSnap:


 @whitelist_for_serdes(storage_name="SolidDefSnap")
-@record
-class OpDefSnap:
+@record_custom


why does this now need to be custom? does the backcompat not work if pool is just pool: Optional[str] = None?

addressed

OwenKephart

approving but I do think the return types should be Set[str] instead of set[str]

OwenKephart · 2025-01-16T23:52:18Z

python_modules/dagster/dagster/_core/snap/node.py

@@ -177,14 +178,15 @@ def get_output_snap(self, name: str) -> OutputDefSnap:

 @whitelist_for_serdes(storage_name="SolidDefSnap")
 @record
-class OpDefSnap:
+class OpDefSnap(IHaveNew):


why IHaveNew here?

prha force-pushed the prha/asset_op_concurrency_key branch from 3e46d69 to 0c1183c Compare December 3, 2024 23:26

graphite-app bot reviewed Dec 3, 2024

View reviewed changes

python_modules/dagster/dagster/_core/definitions/decorators/asset_decorator.py Outdated Show resolved Hide resolved

prha force-pushed the prha/asset_op_concurrency_key branch from 0c1183c to 041872f Compare December 3, 2024 23:28

graphite-app bot reviewed Dec 3, 2024

View reviewed changes

python_modules/dagster/dagster/_core/definitions/op_definition.py Outdated Show resolved Hide resolved

prha force-pushed the prha/asset_op_concurrency_key branch 6 times, most recently from addf80e to d288e73 Compare December 4, 2024 19:30

prha requested review from schrockn and deepyaman December 5, 2024 20:50

prha force-pushed the prha/asset_op_concurrency_key branch 2 times, most recently from e51874e to 571ead4 Compare December 11, 2024 22:01

schrockn previously requested changes Dec 11, 2024

View reviewed changes

prha mentioned this pull request Dec 12, 2024

add support for run-granularity op concurrency #26458

Open

prha removed the request for review from deepyaman December 13, 2024 17:45

prha force-pushed the prha/asset_op_concurrency_key branch from 571ead4 to 238ba49 Compare December 13, 2024 17:46

prha force-pushed the prha/asset_op_concurrency_key branch from 238ba49 to a36d157 Compare January 3, 2025 05:18

prha force-pushed the prha/asset_op_concurrency_key branch 5 times, most recently from f52347d to ce97f66 Compare January 8, 2025 19:51

prha force-pushed the prha/asset_op_concurrency_key branch from ba49ba7 to 80d098b Compare January 10, 2025 01:17

prha requested review from schrockn and OwenKephart January 10, 2025 01:19

OwenKephart previously requested changes Jan 15, 2025

View reviewed changes

prha force-pushed the prha/asset_op_concurrency_key branch 2 times, most recently from a120802 to 58ac167 Compare January 16, 2025 00:18

prha requested a review from OwenKephart January 16, 2025 06:26

prha mentioned this pull request Jan 16, 2025

Update docs to reflect pools concept for concurrency #27171

Open

prha force-pushed the prha/asset_op_concurrency_key branch from 58ac167 to 20be47d Compare January 16, 2025 20:22

OwenKephart approved these changes Jan 16, 2025

View reviewed changes

OwenKephart reviewed Jan 16, 2025

View reviewed changes

prha force-pushed the prha/asset_op_concurrency_key branch 3 times, most recently from 197d7c1 to b7bfa67 Compare January 17, 2025 02:54

prha added 11 commits January 17, 2025 14:26

add concurrency key to the op definition

08cf1fc

thread concurrency key into execution step

4995d4b

group

b590aea

definitions

55715de

concurrency_group => pool

f4d6680

snap op node

ad3e307

fix error message, descriptions

5d79968

fix tests

a129020

handle graph snapshots

929a677

coerce tag-based concurrency keys into pools

80dadc7

update types

9d20984

prha force-pushed the prha/asset_op_concurrency_key branch from b7bfa67 to 9d20984 Compare January 17, 2025 22:29

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add pool as top-level definition arg #26256

add pool as top-level definition arg #26256

prha commented Dec 3, 2024 •

edited

Loading

schrockn left a comment

prha commented Dec 13, 2024 •

edited

Loading

prha commented Jan 9, 2025

prha commented Jan 10, 2025

OwenKephart left a comment

OwenKephart Jan 15, 2025

prha Jan 16, 2025

OwenKephart Jan 16, 2025

OwenKephart Jan 15, 2025

OwenKephart Jan 15, 2025

OwenKephart Jan 15, 2025

OwenKephart Jan 15, 2025

OwenKephart Jan 15, 2025

OwenKephart left a comment

OwenKephart Jan 16, 2025

add pool as top-level definition arg #26256

Are you sure you want to change the base?

add pool as top-level definition arg #26256

Conversation

prha commented Dec 3, 2024 • edited Loading

Summary & Motivation

How I Tested These Changes

Changelog

schrockn left a comment

Choose a reason for hiding this comment

prha commented Dec 13, 2024 • edited Loading

prha commented Jan 9, 2025

prha commented Jan 10, 2025

OwenKephart left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

OwenKephart left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

prha commented Dec 3, 2024 •

edited

Loading

prha commented Dec 13, 2024 •

edited

Loading