Fix deployment level concurrency

README.md

@@ -112,6 +112,86 @@ Note: Starting the server and running work pool is unnecessary if local Mac Pref
 6. Start up system services that hosts the 2 sets of visualization
 
 
+## Configuring deployment concurrency
+
+Deploy recipes support two independent concurrency controls:
+
+- `concurrency_group` limits the total number of simultaneous runs shared by
+  multiple deployments in the same work pool. Echodataflow implements each group
+  as a Prefect work queue with a concurrency limit.
+- `deployment_concurrency` limits simultaneous runs of one deployment, regardless
+  of whether that deployment belongs to a concurrency group.
+
+### Shared concurrency across deployments
+
+Define groups at the top level of the deploy recipe, then assign flows to them by
+name:
+
+```yaml
+concurrency_groups:
+  acoustic_ingestion:
+    limit: 3
+
+flows:
+  ingest_NASC:
+    concurrency_group: acoustic_ingestion
+
+  ingest_MVBS:
+    concurrency_group: acoustic_ingestion
+```
+
+In this example, `ingest_NASC` and `ingest_MVBS` can use at most three running
+slots in total. Runs beyond the shared limit wait in the group's work queue. A
+concurrency group must remain within one work pool.
+
+Flows without `concurrency_group` use the work pool's default queue.
+
+### Per-deployment concurrency
+
+Use `deployment_concurrency` within a flow to limit only that deployment:
+
+```yaml
+flows:
+  ingest_NASC:
+    deployment_concurrency:
+      limit: 1
+      collision_strategy: CANCEL_NEW
+```
+
+Supported fields are:
+
+- `limit` (required): maximum number of concurrent runs for the deployment.
+- `collision_strategy` (optional): `ENQUEUE` or `CANCEL_NEW`; defaults to
+  `ENQUEUE`. `ENQUEUE` makes a new run wait for a slot, while `CANCEL_NEW`
+  cancels it when the limit is full.
+- `grace_period_seconds` (optional): time allowed for run infrastructure to start
+  before its concurrency slot is released. The value must be between 60 and
+  86,400 seconds.
+
+### Combining both controls
+
+The controls can be used together:
+
+```yaml
+concurrency_groups:
+  acoustic_ingestion:
+    limit: 3
+
+flows:
+  ingest_NASC:
+    concurrency_group: acoustic_ingestion
+    deployment_concurrency:
+      limit: 1
+      collision_strategy: CANCEL_NEW
+
+  ingest_MVBS:
+    concurrency_group: acoustic_ingestion
+```
+
+Here, the two deployments share three work-queue slots, while `ingest_NASC` may
+occupy only one of those slots. A run must satisfy both limits before it can run.
+
+
 ## Running Local Prefect and auto mounting services on macOS (launchd)
 
 To run a local Prefect server and worker as background services on macOS, you can

src/echodataflow/deployment/core.py

@@ -10,6 +10,7 @@
 }
 ALLOWED_FLOW_DEPLOY_KEYS = {
     "concurrency_group",
+    "deployment_concurrency",
     "deployment_name",
     "flow",
     "interval",
@@ -20,6 +21,11 @@
     "work_pool_name",
 }
 ALLOWED_CONCURRENCY_GROUP_KEYS = {"limit"}
+ALLOWED_DEPLOYMENT_CONCURRENCY_KEYS = {
+    "limit",
+    "collision_strategy",
+    "grace_period_seconds",
+}
 ALLOWED_TASK_RUNNER_KEYS = {"type", "cluster_kwargs"}
 ALLOWED_DASK_CLUSTER_KEYS = {
     "memory_limit",

src/echodataflow/deployment/deployment_engine.py

@@ -20,6 +20,7 @@
     ALLOWED_CONCURRENCY_GROUP_KEYS,
     ALLOWED_DASK_CLUSTER_KEYS,
     ALLOWED_DEPLOY_KEYS,
+    ALLOWED_DEPLOYMENT_CONCURRENCY_KEYS,
     ALLOWED_FLOW_DEPLOY_KEYS,
     ALLOWED_GIT_SOURCE_KEYS,
     ALLOWED_SOURCE_KEYS,
@@ -40,6 +41,7 @@ class DeploymentSpec:
     entrypoint: str  # source-relative entrypoint for the actual deployed flow
     parameters: dict[str, Any]  # parameters passed directly to the deployed flow
     concurrency_group: str | None = None
+    deployment_concurrency: dict[str, Any] | None = None
     task_runner: dict[str, Any] | None = None
     cron: str | None = None  # precomputed cron schedule, when interval mode is used
     work_pool_name: str | None = (
@@ -293,6 +295,37 @@ def validate_task_runner_config(value: Any, *, path: str) -> None:
         raise ValueError(f"{cluster_path}.processes must be a boolean")
 
 
+def validate_deployment_concurrency_config(value: Any, *, path: str) -> None:
+    """Validate deployment-scoped flow-run concurrency settings."""
+    if not isinstance(value, dict):
+        raise ValueError(f"{path} must be a mapping")
+    _reject_unknown_keys(
+        value,
+        allowed=ALLOWED_DEPLOYMENT_CONCURRENCY_KEYS,
+        path=path,
+    )
+    if "limit" not in value:
+        raise ValueError(f"{path}.limit is required")
+    _validate_positive_integer(value["limit"], path=f"{path}.limit")
+
+    collision_strategy = value.get("collision_strategy", "ENQUEUE")
+    if collision_strategy not in {"ENQUEUE", "CANCEL_NEW"}:
+        raise ValueError(
+            f"{path}.collision_strategy must be 'ENQUEUE' or 'CANCEL_NEW'"
+        )
+
+    grace_period_seconds = value.get("grace_period_seconds")
+    if grace_period_seconds is not None:
+        _validate_positive_integer(
+            grace_period_seconds,
+            path=f"{path}.grace_period_seconds",
+        )
+        if not 60 <= grace_period_seconds <= 86400:
+            raise ValueError(
+                f"{path}.grace_period_seconds must be between 60 and 86400"
+            )
+
+
 def validate_deploy_config(deploy_cfg: Any) -> None:
     """Reject unknown fields throughout a deployment specification."""
     if not isinstance(deploy_cfg, dict):
@@ -348,6 +381,13 @@ def validate_deploy_config(deploy_cfg: Any) -> None:
                     f"{concurrency_group!r}"
                 )
 
+        deployment_concurrency = deploy_meta.get("deployment_concurrency")
+        if deployment_concurrency is not None:
+            validate_deployment_concurrency_config(
+                deployment_concurrency,
+                path=f"{flow_path}.deployment_concurrency",
+            )
+
         task_runner = deploy_meta.get("task_runner")
         if task_runner is not None:
             validate_task_runner_config(task_runner, path=f"{flow_path}.task_runner")
@@ -574,6 +614,7 @@ def build_deploy_specs(
                 entrypoint=flow_info["entrypoint"],
                 parameters=deployment_parameters,
                 concurrency_group=deploy_meta.get("concurrency_group"),
+                deployment_concurrency=deploy_meta.get("deployment_concurrency"),
                 task_runner=deploy_meta.get("task_runner"),
                 cron=cron,
                 work_pool_name=deploy_meta.get("work_pool_name"),
@@ -616,6 +657,13 @@ def create_deployments(
         if spec.concurrency_group is not None:
             deployment_kwargs["work_queue_name"] = spec.concurrency_group
 
+        if spec.deployment_concurrency is not None:
+            from prefect.client.schemas.objects import ConcurrencyLimitConfig
+
+            deployment_kwargs["concurrency_limit"] = ConcurrencyLimitConfig(
+                **spec.deployment_concurrency
+            )
+
         # The worker reloads the flow entrypoint, so runner settings must be
         # present in its runtime environment instead of only on this Flow object
         if spec.task_runner is not None:

tests/deployment/conftest.py

@@ -23,6 +23,11 @@ class FakeRunnerDeployment:
     pass
 
 
+class FakeConcurrencyLimitConfig:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+
 class FakePrefectFlowGeneric:
     @classmethod
     def __class_getitem__(cls, _item):
@@ -57,18 +62,27 @@ def fake_deploy(*deployments, **kwargs):
         events_mod = types.ModuleType("prefect.events")
         events_mod.DeploymentEventTrigger = FakeTrigger
 
+        client_mod = types.ModuleType("prefect.client")
+        schemas_mod = types.ModuleType("prefect.client.schemas")
+        objects_mod = types.ModuleType("prefect.client.schemas.objects")
+        objects_mod.ConcurrencyLimitConfig = FakeConcurrencyLimitConfig
+
         monkeypatch.setitem(sys.modules, "prefect", prefect_mod)
         monkeypatch.setitem(sys.modules, "prefect.deployments", deployments_mod)
         monkeypatch.setitem(sys.modules, "prefect.deployments.runner", runner_mod)
         monkeypatch.setitem(sys.modules, "prefect.flows", flows_mod)
         monkeypatch.setitem(sys.modules, "prefect.variables", variables_mod)
         monkeypatch.setitem(sys.modules, "prefect.events", events_mod)
+        monkeypatch.setitem(sys.modules, "prefect.client", client_mod)
+        monkeypatch.setitem(sys.modules, "prefect.client.schemas", schemas_mod)
+        monkeypatch.setitem(sys.modules, "prefect.client.schemas.objects", objects_mod)
 
         return {
             "FakeVariable": FakeVariable,
             "FakeTrigger": FakeTrigger,
             "FakeRunnerDeployment": FakeRunnerDeployment,
+            "FakeConcurrencyLimitConfig": FakeConcurrencyLimitConfig,
             "FakePrefectFlowGeneric": FakePrefectFlowGeneric,
         }
 
-    return _install
+    return _install

tests/deployment/test_deploy_engine.py

@@ -162,6 +162,10 @@ def test_build_deploy_specs_preserves_runner_and_concurrency_group(install_prefe
             "flows": {
                 "raw2Sv_postprocessing": {
                     "concurrency_group": "postprocessing",
+                    "deployment_concurrency": {
+                        "limit": 1,
+                        "collision_strategy": "CANCEL_NEW",
+                    },
                     "task_runner": runner_config,
                 }
             },
@@ -175,6 +179,10 @@ def test_build_deploy_specs_preserves_runner_and_concurrency_group(install_prefe
     )
 
     assert specs[0].concurrency_group == "postprocessing"
+    assert specs[0].deployment_concurrency == {
+        "limit": 1,
+        "collision_strategy": "CANCEL_NEW",
+    }
     assert specs[0].task_runner == runner_config
 
 
@@ -218,6 +226,78 @@ def from_source(self, **kwargs):
     assert standalone == []
 
 
+def test_deployment_concurrency_is_independent_of_concurrency_group(
+    install_prefect_stubs,
+):
+    install_prefect_stubs()
+    engine = importlib.import_module("echodataflow.deployment.deployment_engine")
+    calls = {}
+
+    class SourcedFlow:
+        def to_deployment(self, **kwargs):
+            calls["deployment"] = kwargs
+            return kwargs
+
+    class RegisteredFlow:
+        def from_source(self, **kwargs):
+            return SourcedFlow()
+
+    deployment_concurrency = {
+        "limit": 1,
+        "collision_strategy": "CANCEL_NEW",
+        "grace_period_seconds": 120,
+    }
+    spec = engine.DeploymentSpec(
+        flow_key="ingest_NASC",
+        deployment_name="ingest-NASC",
+        flow_obj=RegisteredFlow(),
+        entrypoint="echodataflow/flows/flows_integration.py:flow_ingest_NASC",
+        parameters={},
+        deployment_concurrency=deployment_concurrency,
+    )
+
+    engine.create_deployments(
+        specs=[spec],
+        source="local-source",
+        default_work_pool_name="local",
+    )
+
+    assert "work_queue_name" not in calls["deployment"]
+    limit_config = calls["deployment"]["concurrency_limit"]
+    assert limit_config.limit == 1
+    assert limit_config.collision_strategy == "CANCEL_NEW"
+    assert limit_config.grace_period_seconds == 120
+
+
+@pytest.mark.parametrize(
+    ("deployment_concurrency", "expected_message"),
+    [
+        ({}, "limit is required"),
+        ({"limit": 0}, "limit must be a positive integer"),
+        (
+            {"limit": 1, "collision_strategy": "DROP_OLD"},
+            "collision_strategy must be 'ENQUEUE' or 'CANCEL_NEW'",
+        ),
+        (
+            {"limit": 1, "grace_period_seconds": 30},
+            "grace_period_seconds must be between 60 and 86400",
+        ),
+    ],
+)
+def test_validate_deploy_config_rejects_invalid_deployment_concurrency(
+    install_prefect_stubs,
+    deployment_concurrency,
+    expected_message,
+):
+    install_prefect_stubs()
+    engine = importlib.import_module("echodataflow.deployment.deployment_engine")
+
+    with pytest.raises(ValueError, match=expected_message):
+        engine.validate_deploy_config(
+            {"flows": {"ingest_NASC": {"deployment_concurrency": deployment_concurrency}}}
+        )
+
+
 @pytest.mark.parametrize(
     ("task_runner", "expected_message"),
     [
@@ -273,6 +353,11 @@ def test_validate_deploy_config_accepts_every_allowed_key(install_prefect_stubs)
         "flows": {
             "scheduled": {
                 "concurrency_group": "postprocessing",
+                "deployment_concurrency": {
+                    "limit": 1,
+                    "collision_strategy": "CANCEL_NEW",
+                    "grace_period_seconds": 120,
+                },
                 "deployment_name": "scheduled-deployment",
                 "flow": "actual_flow_name",
                 "interval": 10,
@@ -310,6 +395,7 @@ def test_validate_deploy_config_accepts_every_allowed_key(install_prefect_stubs)
     }
     assert core.ALLOWED_FLOW_DEPLOY_KEYS == {
         "concurrency_group",
+        "deployment_concurrency",
         "deployment_name",
         "flow",
         "interval",
@@ -320,6 +406,11 @@ def test_validate_deploy_config_accepts_every_allowed_key(install_prefect_stubs)
         "work_pool_name",
     }
     assert core.ALLOWED_CONCURRENCY_GROUP_KEYS == {"limit"}
+    assert core.ALLOWED_DEPLOYMENT_CONCURRENCY_KEYS == {
+        "limit",
+        "collision_strategy",
+        "grace_period_seconds",
+    }
     assert core.ALLOWED_TASK_RUNNER_KEYS == {"type", "cluster_kwargs"}
     assert core.ALLOWED_DASK_CLUSTER_KEYS == {
         "memory_limit",