group => pool

Author: prha

python_modules/dagster/dagster/_core/op_concurrency_limits_counter.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Dict, Mapping, Optional, Sequence, Set
 
 from dagster._core.instance import DagsterInstance
-from dagster._core.run_coordinator.queued_run_coordinator import ConcurrencyGranularity
+from dagster._core.run_coordinator.queued_run_coordinator import PoolGranularity
 from dagster._core.snap.execution_plan_snapshot import ExecutionPlanSnapshot
 from dagster._core.storage.dagster_run import (
     IN_PROGRESS_RUN_STATUSES,
@@ -58,14 +58,14 @@ def __init__(
         runs: Sequence[DagsterRun],
         in_progress_run_records: Sequence[RunRecord],
         slot_count_offset: int = 0,
-        concurrency_group_granularity: ConcurrencyGranularity = ConcurrencyGranularity.OP,
+        pool_granularity: PoolGranularity = PoolGranularity.OP,
     ):
         self._root_concurrency_keys_by_run = {}
         self._concurrency_info_by_key: Dict[str, "ConcurrencyKeyInfo"] = {}
         self._launched_concurrency_key_counts = defaultdict(int)
         self._in_progress_concurrency_key_counts = defaultdict(int)
         self._slot_count_offset = slot_count_offset
-        self._concurrency_group_granularity = concurrency_group_granularity
+        self._pool_granularity = pool_granularity
         self._in_progress_run_ids: Set[str] = set(
             [record.dagster_run.run_id for record in in_progress_run_records]
         )
@@ -91,7 +91,7 @@ def _fetch_concurrency_info(self, instance: DagsterInstance, queued_runs: Sequen
                 # if using op granularity, consider only the root keys
                 run_concurrency_keys = (
                     run.run_op_concurrency.root_key_counts.keys()
-                    if self._concurrency_group_granularity == ConcurrencyGranularity.OP
+                    if self._pool_granularity == PoolGranularity.OP
                     else run.run_op_concurrency.all_keys or []
                 )
                 all_concurrency_keys.update(run_concurrency_keys)
@@ -115,7 +115,7 @@ def _should_allocate_slots_for_in_progress_run(self, record: RunRecord):
         if status not in IN_PROGRESS_RUN_STATUSES:
             return False
 
-        if self._concurrency_group_granularity == ConcurrencyGranularity.RUN:
+        if self._pool_granularity == PoolGranularity.RUN:
             return True
 
         if status == DagsterRunStatus.STARTING:
@@ -132,11 +132,11 @@ def _slot_counts_for_run(self, run: DagsterRun) -> Mapping[str, int]:
         if not run.run_op_concurrency:
             return {}
 
-        if self._concurrency_group_granularity == ConcurrencyGranularity.OP:
+        if self._pool_granularity == PoolGranularity.OP:
             return {**run.run_op_concurrency.root_key_counts}
 
         else:
-            assert self._concurrency_group_granularity == ConcurrencyGranularity.RUN
+            assert self._pool_granularity == PoolGranularity.RUN
             return {concurrency_key: 1 for concurrency_key in run.run_op_concurrency.all_keys or []}
 
     def _process_in_progress_runs(self, in_progress_records: Sequence[RunRecord]):
@@ -154,14 +154,14 @@ def is_blocked(self, run: DagsterRun) -> bool:
             return False
 
         if (
-            self._concurrency_group_granularity == ConcurrencyGranularity.OP
+            self._pool_granularity == PoolGranularity.OP
             and run.run_op_concurrency.has_unconstrained_root_nodes
         ):
             # if the granularity is at the op level and there exists a root node that is not
             # concurrency blocked, we should dequeue.
             return False
 
-        if self._concurrency_group_granularity == ConcurrencyGranularity.OP:
+        if self._pool_granularity == PoolGranularity.OP:
             # we just need to check all of the root concurrency keys, instead of all the concurrency keys
             # in the run
             for concurrency_key in run.run_op_concurrency.root_key_counts.keys():
@@ -189,7 +189,7 @@ def is_blocked(self, run: DagsterRun) -> bool:
             return True
 
         else:
-            assert self._concurrency_group_granularity == ConcurrencyGranularity.RUN
+            assert self._pool_granularity == PoolGranularity.RUN
 
             # if the granularity is at the run level, we should check if any of the concurrency
             # keys are blocked
@@ -228,7 +228,7 @@ def get_blocked_run_debug_info(self, run: DagsterRun) -> Mapping:
                 continue
 
             log_info[concurrency_key] = {
-                "granularity": self._concurrency_group_granularity.value,
+                "granularity": self._pool_granularity.value,
                 "slot_count": concurrency_info.slot_count,
                 "pending_step_count": len(concurrency_info.pending_steps),
                 "pending_step_run_ids": list(

python_modules/dagster/dagster/_core/run_coordinator/queued_run_coordinator.py

@@ -20,7 +20,7 @@
 from dagster._serdes import ConfigurableClass, ConfigurableClassData
 
 
-class ConcurrencyGranularity(Enum):
+class PoolGranularity(Enum):
     OP = "op"
     RUN = "run"
 
@@ -35,7 +35,7 @@ class RunQueueConfig(
             ("user_code_failure_retry_delay", int),
             ("should_block_op_concurrency_limited_runs", bool),
             ("op_concurrency_slot_buffer", int),
-            ("concurrency_group_granularity", ConcurrencyGranularity),
+            ("pool_granularity", PoolGranularity),
         ],
     )
 ):
@@ -47,7 +47,7 @@ def __new__(
         user_code_failure_retry_delay: int = 60,
         should_block_op_concurrency_limited_runs: bool = False,
         op_concurrency_slot_buffer: int = 0,
-        concurrency_group_granularity: ConcurrencyGranularity = ConcurrencyGranularity.OP,
+        pool_granularity: PoolGranularity = PoolGranularity.OP,
     ):
         return super(RunQueueConfig, cls).__new__(
             cls,
@@ -60,9 +60,9 @@ def __new__(
             ),
             check.int_param(op_concurrency_slot_buffer, "op_concurrency_slot_buffer"),
             check.inst_param(
-                concurrency_group_granularity,
-                "concurrency_group_granularity",
-                ConcurrencyGranularity,
+                pool_granularity,
+                "pool_granularity",
+                PoolGranularity,
             ),
         )
 
@@ -82,7 +82,7 @@ def __init__(
         max_user_code_failure_retries: Optional[int] = None,
         user_code_failure_retry_delay: Optional[int] = None,
         block_op_concurrency_limited_runs: Optional[Mapping[str, Any]] = None,
-        concurrency_group_granularity: str = "op",
+        pool_granularity: str = "op",
         inst_data: Optional[ConfigurableClassData] = None,
     ):
         self._inst_data: Optional[ConfigurableClassData] = check.opt_inst_param(
@@ -129,7 +129,7 @@ def __init__(
                 "op_concurrency_slot_buffer can only be set if block_op_concurrency_limited_runs "
                 "is enabled",
             )
-        self._concurrency_group_granularity = ConcurrencyGranularity(concurrency_group_granularity)
+        self._pool_granularity = PoolGranularity(pool_granularity)
         self._logger = logging.getLogger("dagster.run_coordinator.queued_run_coordinator")
         super().__init__()
 
@@ -145,7 +145,7 @@ def get_run_queue_config(self) -> RunQueueConfig:
             user_code_failure_retry_delay=self._user_code_failure_retry_delay,
             should_block_op_concurrency_limited_runs=self._should_block_op_concurrency_limited_runs,
             op_concurrency_slot_buffer=self._op_concurrency_slot_buffer,
-            concurrency_group_granularity=self._concurrency_group_granularity,
+            pool_granularity=self._pool_granularity,
         )
 
     @property
@@ -266,7 +266,7 @@ def config_type(cls) -> UserConfigSchema:
                     ),
                 }
             ),
-            "concurrency_group_granularity": Field(
+            "pool_granularity": Field(
                 str,
                 is_required=False,
                 default_value="op",
@@ -293,7 +293,7 @@ def from_config_value(
             max_user_code_failure_retries=config_value.get("max_user_code_failure_retries"),
             user_code_failure_retry_delay=config_value.get("user_code_failure_retry_delay"),
             block_op_concurrency_limited_runs=config_value.get("block_op_concurrency_limited_runs"),
-            concurrency_group_granularity=config_value.get("concurrency_group_granularity", "op"),
+            pool_granularity=config_value.get("pool_granularity", "op"),
         )
 
     def submit_run(self, context: SubmitRunContext) -> DagsterRun:

python_modules/dagster/dagster/_daemon/run_coordinator/queued_run_coordinator_daemon.py

@@ -269,7 +269,7 @@ def _get_runs_to_dequeue(
                         batch,
                         in_progress_run_records,
                         run_queue_config.op_concurrency_slot_buffer,
-                        run_queue_config.concurrency_group_granularity,
+                        run_queue_config.pool_granularity,
                     )
                 except:
                     self._logger.exception("Failed to initialize op concurrency counter")

python_modules/dagster/dagster_tests/daemon_tests/test_locations/concurrency_limited_workspace.py

@@ -3,26 +3,26 @@
 from dagster._core.storage.tags import GLOBAL_CONCURRENCY_TAG
 
 
-@asset(concurrency_group="foo", key_prefix=["prefix"])
+@asset(pool="foo", key_prefix=["prefix"])
 def foo_limited_asset():
     return 1
 
 
-@asset(concurrency_group="bar", key_prefix=["prefix"])
+@asset(pool="bar", key_prefix=["prefix"])
 def bar_limited_asset():
     return 1
 
 
 @asset(
-    concurrency_group="baz",
+    pool="baz",
     key_prefix=["prefix"],
     ins={"foo_limited_asset": AssetIn(key_prefix="prefix")},
 )
 def baz_limited_asset_depends_on_foo(foo_limited_asset):
     return 1
 
 
-@asset(concurrency_group="baz", key_prefix=["prefix"])
+@asset(pool="baz", key_prefix=["prefix"])
 def baz_limited_asset():
     return 1
 

python_modules/dagster/dagster_tests/daemon_tests/test_queued_run_coordinator_daemon.py

@@ -1186,7 +1186,7 @@ def test_concurrency_buffer_with_default_slot(
                 "block_op_concurrency_limited_runs": {
                     "enabled": True,
                 },
-                "concurrency_group_granularity": "run",
+                "pool_granularity": "run",
             },
         ],
     )