test run blocking op concurrency

Author: prha

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

@@ -82,6 +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",
         inst_data: Optional[ConfigurableClassData] = None,
     ):
         self._inst_data: Optional[ConfigurableClassData] = check.opt_inst_param(
@@ -128,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._logger = logging.getLogger("dagster.run_coordinator.queued_run_coordinator")
         super().__init__()
 
@@ -144,6 +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,
         )
 
     @property
@@ -264,6 +266,17 @@ def config_type(cls) -> UserConfigSchema:
                     ),
                 }
             ),
+            "concurrency_group_granularity": Field(
+                str,
+                is_required=False,
+                default_value="op",
+                description=(
+                    "Determines the granularity at which concurrency limits are applied. If set to "
+                    "'op', dequeues runs as long as any op in the run can make progress.  If set to "
+                    "'run', dequeues runs as long as all of the concurrency groups assigned to the run "
+                    "have free slots available."
+                ),
+            ),
         }
 
     @classmethod
@@ -280,6 +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"),
         )
 
     def submit_run(self, context: SubmitRunContext) -> DagsterRun:

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

@@ -3,31 +3,36 @@
 from dagster._core.storage.tags import GLOBAL_CONCURRENCY_TAG
 
 
-@asset(op_tags={GLOBAL_CONCURRENCY_TAG: "foo"}, key_prefix=["prefix"])
+@asset(concurrency_group="foo", key_prefix=["prefix"])
 def foo_limited_asset():
     return 1
 
 
-@asset(op_tags={GLOBAL_CONCURRENCY_TAG: "bar"}, key_prefix=["prefix"])
+@asset(concurrency_group="bar", key_prefix=["prefix"])
 def bar_limited_asset():
     return 1
 
 
 @asset(
-    op_tags={GLOBAL_CONCURRENCY_TAG: "baz"},
+    concurrency_group="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"])
+def baz_limited_asset():
+    return 1
+
+
 concurrency_limited_asset_job = define_asset_job(
     "concurrency_limited_asset_job",
-    [foo_limited_asset, bar_limited_asset, baz_limited_asset_depends_on_foo],
+    [foo_limited_asset, bar_limited_asset, baz_limited_asset, baz_limited_asset_depends_on_foo],
 ).resolve(
     asset_graph=AssetGraph.from_assets(
-        [foo_limited_asset, bar_limited_asset, baz_limited_asset_depends_on_foo]
+        [foo_limited_asset, bar_limited_asset, baz_limited_asset_depends_on_foo, baz_limited_asset]
     )
 )
 

python_modules/dagster/dagster_tests/daemon_tests/test_queued_run_coordinator_daemon.py

@@ -1179,6 +1179,60 @@ def test_concurrency_buffer_with_default_slot(
         assert set(self.get_run_ids(instance.run_launcher.queue())) == {run_id_1}
         caplog.text.count(f"Run {run_id_2} is blocked by global concurrency limits") == 1  # pyright: ignore[reportUnusedExpression]
 
+    @pytest.mark.parametrize(
+        "run_coordinator_config",
+        [
+            {
+                "block_op_concurrency_limited_runs": {
+                    "enabled": True,
+                },
+                "concurrency_group_granularity": "run",
+            },
+        ],
+    )
+    def test_concurrency_run_granularity(
+        self,
+        concurrency_limited_workspace_context,
+        daemon,
+        instance,
+    ):
+        run_id_1, run_id_2 = [make_new_run_id() for _ in range(2)]
+        workspace = concurrency_limited_workspace_context.create_request_context()
+        # concurrency_limited_asset_job
+        remote_job = self.get_concurrency_job(workspace)
+        foo_key = AssetKey(["prefix", "foo_limited_asset"])
+        bar_key = AssetKey(["prefix", "bar_limited_asset"])
+        baz_key = AssetKey(["prefix", "baz_limited_asset"])
+        downstream_baz_key = AssetKey(["prefix", "baz_limited_asset_depends_on_foo"])
+
+        # first submit a run that occupies the baz slot
+        self.submit_run(
+            instance, remote_job, workspace, run_id=run_id_1, asset_selection=set([baz_key])
+        )
+        list(daemon.run_iteration(concurrency_limited_workspace_context))
+        assert set(self.get_run_ids(instance.run_launcher.queue())) == set([run_id_1])
+
+        # submit a run that has 2 root nodes, respectively with foo, bar concurrency groups
+        # also with a downstream extra baz concurrency group asset
+        self.submit_run(
+            instance,
+            remote_job,
+            workspace,
+            run_id=run_id_2,
+            asset_selection=set([foo_key, bar_key, downstream_baz_key]),
+        )
+
+        # even though the root nodes have slots available, the second run is not launched
+        list(daemon.run_iteration(concurrency_limited_workspace_context))
+        assert set(self.get_run_ids(instance.run_launcher.queue())) == set([run_id_1])
+
+        instance.event_log_storage.set_concurrency_slots("baz", 2)
+
+        # all group slots available, run launched, even though there is only one slot open for baz
+        # and the run has two baz-grouped nodes
+        list(daemon.run_iteration(concurrency_limited_workspace_context))
+        assert set(self.get_run_ids(instance.run_launcher.queue())) == set([run_id_1, run_id_2])
+
 
 class TestQueuedRunCoordinatorDaemon(QueuedRunCoordinatorDaemonTests):
     @pytest.fixture