diff --git a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/__snapshots__/test_all_snapshot_ids.ambr b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/__snapshots__/test_all_snapshot_ids.ambr index 5fbe8422e156c..91323bcee7ec0 100644 --- a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/__snapshots__/test_all_snapshot_ids.ambr +++ b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/__snapshots__/test_all_snapshot_ids.ambr @@ -2571,6 +2571,9 @@ "mapped_solid_name": "never_runs_op" } ], + "pools": { + "__set__": [] + }, "tags": {} } ], @@ -33745,6 +33748,9 @@ "name": "simple_graph", "output_def_snaps": [], "output_mapping_snaps": [], + "pools": { + "__set__": [] + }, "tags": {} } ], @@ -33784,7 +33790,7 @@ ''' # --- # name: test_all_snapshot_ids[17] - '2f4a7e8dac16272a9e008a61e766b43bd473eb7e' + 'cfb767a61be6a485474582ec75fc6834cd314c08' # --- # name: test_all_snapshot_ids[18] ''' @@ -34920,6 +34926,9 @@ "mapped_solid_name": "adder_2" } ], + "pools": { + "__set__": [] + }, "tags": {} }, { @@ -35002,6 +35011,9 @@ "mapped_solid_name": "adder_2" } ], + "pools": { + "__set__": [] + }, "tags": {} }, { @@ -35084,6 +35096,9 @@ "mapped_solid_name": "div_2" } ], + "pools": { + "__set__": [] + }, "tags": {} } ], @@ -35165,10 +35180,10 @@ ''' # --- # name: test_all_snapshot_ids[19] - 'f726acb538de96f43dfa0335f7bc79d1b52394c1' + 'b336066bf8dc411b162f9869dd08da0b18057896' # --- # name: test_all_snapshot_ids[1] - '6c2e7891d74bdc0ff647f100a13f4a630081c2ba' + 'fb54440831226acbda6ad87d0d7599c4f3424168' # --- # name: test_all_snapshot_ids[20] ''' @@ -54872,6 +54887,9 @@ "mapped_solid_name": "never_runs_op" } ], + "pools": { + "__set__": [] + }, "tags": {} } ], @@ -55018,7 +55036,7 @@ ''' # --- # name: test_all_snapshot_ids[51] - '709d40d76a7309b11219cb331842fb349e161eb0' + '200c1bc946100875b19e0c6b7f26fbb9edab070b' # --- # name: test_all_snapshot_ids[52] ''' @@ -80111,6 +80129,9 @@ "mapped_solid_name": "plus_one" } ], + "pools": { + "__set__": [] + }, "tags": {} } ], @@ -80254,7 +80275,7 @@ ''' # --- # name: test_all_snapshot_ids[99] - 'cf140114fda778cc12c11ba13f7e26cd0fbccb3c' + '1959e3612cf7d53844928fff25df04128d1b15ff' # --- # name: test_all_snapshot_ids[9] 'cf67ad5622d23845499daf2adba4a11f6b23f9eb' diff --git a/python_modules/dagster/dagster/_core/definitions/graph_definition.py b/python_modules/dagster/dagster/_core/definitions/graph_definition.py index 96a4040b8ca88..5c52dd5de959a 100644 --- a/python_modules/dagster/dagster/_core/definitions/graph_definition.py +++ b/python_modules/dagster/dagster/_core/definitions/graph_definition.py @@ -815,6 +815,13 @@ def tags(self) -> Mapping[str, str]: """The tags associated with the graph.""" return super(GraphDefinition, self).tags + @property + def pools(self) -> Set[str]: + pools = set() + for node_def in self.node_defs: + pools.update(node_def.pools) + return pools + @public def alias(self, name: str) -> "PendingNodeInvocation": """Aliases the graph with a new name. diff --git a/python_modules/dagster/dagster/_core/definitions/node_definition.py b/python_modules/dagster/dagster/_core/definitions/node_definition.py index 65fb36c5a5688..5789991a4ab67 100644 --- a/python_modules/dagster/dagster/_core/definitions/node_definition.py +++ b/python_modules/dagster/dagster/_core/definitions/node_definition.py @@ -7,6 +7,7 @@ Mapping, Optional, Sequence, + Set, Tuple, ) @@ -240,3 +241,7 @@ def get_op_handles(self, parent: "NodeHandle") -> AbstractSet["NodeHandle"]: ... def get_op_output_handles( self, parent: Optional["NodeHandle"] ) -> AbstractSet["NodeOutputHandle"]: ... + + @property + @abstractmethod + def pools(self) -> Set[str]: ... diff --git a/python_modules/dagster/dagster/_core/definitions/op_definition.py b/python_modules/dagster/dagster/_core/definitions/op_definition.py index 2e853b558091c..70f741732b264 100644 --- a/python_modules/dagster/dagster/_core/definitions/op_definition.py +++ b/python_modules/dagster/dagster/_core/definitions/op_definition.py @@ -8,6 +8,7 @@ Mapping, Optional, Sequence, + Set, Tuple, Union, cast, @@ -309,6 +310,11 @@ def pool(self) -> Optional[str]: """Optional[str]: The concurrency group for this op.""" return self._pool + @property + def pools(self) -> Set[str]: + """Optional[str]: The concurrency group for this op.""" + return {self._pool} if self._pool else set() + def is_from_decorator(self) -> bool: from dagster._core.definitions.decorators.op_decorator import DecoratedOpFunction diff --git a/python_modules/dagster/dagster/_core/snap/node.py b/python_modules/dagster/dagster/_core/snap/node.py index 2ae9bd15eb7a1..8801f2a6e4f96 100644 --- a/python_modules/dagster/dagster/_core/snap/node.py +++ b/python_modules/dagster/dagster/_core/snap/node.py @@ -1,5 +1,5 @@ from functools import cached_property -from typing import Mapping, Optional, Sequence, Union +from typing import Mapping, Optional, Sequence, Set, Union import dagster._check as check from dagster._config import ConfigFieldSnap, snap_from_field @@ -158,6 +158,7 @@ class GraphDefSnap: dep_structure_snapshot: DependencyStructureSnapshot input_mapping_snaps: Sequence[InputMappingSnap] output_mapping_snaps: Sequence[OutputMappingSnap] + pools: Set[str] @cached_property def input_def_map(self) -> Mapping[str, InputDefSnap]: @@ -280,6 +281,7 @@ def build_graph_def_snap(graph_def: GraphDefinition) -> GraphDefSnap: dep_structure_snapshot=build_dep_structure_snapshot_from_graph_def(graph_def), input_mapping_snaps=list(map(build_input_mapping_snap, graph_def.input_mappings)), output_mapping_snaps=list(map(build_output_mapping_snap, graph_def.output_mappings)), + pools=graph_def.pools, ) diff --git a/python_modules/dagster/dagster_tests/core_tests/snap_tests/__snapshots__/test_execution_plan.ambr b/python_modules/dagster/dagster_tests/core_tests/snap_tests/__snapshots__/test_execution_plan.ambr index d97d66a36538f..d2405f6c4984e 100644 --- a/python_modules/dagster/dagster_tests/core_tests/snap_tests/__snapshots__/test_execution_plan.ambr +++ b/python_modules/dagster/dagster_tests/core_tests/snap_tests/__snapshots__/test_execution_plan.ambr @@ -310,7 +310,7 @@ }, "step_output_versions": [] }, - "pipeline_snapshot_id": "3057df527127ce93bc08075d3dc3149b5054065a", + "pipeline_snapshot_id": "b57cc7e00dc45fd4927a082265100e665a21f23d", "snapshot_version": 1, "step_keys_to_execute": [ "comp_1.return_one",