From 9fb9987841772efda2cb81731b62fa15e5c3632c Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 2 Sep 2026 18:30:56 +0000 Subject: [PATCH] mason: add --min-instances to `mason deploy` for manual horizontal scaling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires `--min-instances` through to Databricks Apps compute scaling. Apps doesn't autoscale yet, so a fixed replica count is set by pinning compute_min_instances == compute_max_instances (the API rejects setting one bound alone). Instance count is create-only: `apps create` sets it, and re-deploying an existing app only proceeds when the requested count matches the running one — otherwise it fails with a clear "delete and recreate to re-scale" message rather than a doomed update-API call. Verified e2e on e2-dogfood: create at 2 instances (app RUNNING, min==max==2), matching-count re-deploy is a no-op, and a mismatched re-scale is rejected before any deploy. Co-authored-by: Isaac --- integrations/mason/README.md | 2 +- .../mason/src/databricks_mason/deploy.py | 55 +++++++- .../mason/tests/unit_tests/deploy_test.py | 117 ++++++++++++++++++ 3 files changed, 172 insertions(+), 2 deletions(-) diff --git a/integrations/mason/README.md b/integrations/mason/README.md index d5dd5d77b..d437d02da 100644 --- a/integrations/mason/README.md +++ b/integrations/mason/README.md @@ -99,7 +99,7 @@ mason [-p ] [-o text|json] list [--source PATH] deploy --source PATH [--with-memory-store N] [--with-session-store N] [--actor-id ID] - [--with-traces C.S] [--create-stores] + [--with-traces C.S] [--create-stores] [--min-instances N] deployments list | get | logs | start | stop | delete ``` diff --git a/integrations/mason/src/databricks_mason/deploy.py b/integrations/mason/src/databricks_mason/deploy.py index cf5ce02b0..606b17abd 100644 --- a/integrations/mason/src/databricks_mason/deploy.py +++ b/integrations/mason/src/databricks_mason/deploy.py @@ -45,6 +45,23 @@ def _deployment_exists(name: str, profile: Optional[str]) -> bool: return _databricks(["apps", "get", name], profile, capture=True, check=False).returncode == 0 +def _compute_flags(instances: Optional[int]) -> list[str]: + """`databricks apps` compute-scaling flags for a fixed instance count (empty if unset). + + Apps doesn't autoscale, so manual replicas are a fixed count — and the API rejects setting one + bound without the other ("both compute_min_instances and compute_max_instances must be + provided"), so pin min == max to the requested count. + """ + if instances is None: + return [] + return [ + "--compute-min-instances", + str(instances), + "--compute-max-instances", + str(instances), + ] + + def _app_service_principal(name: str, profile: Optional[str]) -> Optional[str]: """The app's service principal client id (its Postgres role identity), or None if unavailable.""" result = _databricks(["apps", "get", name, "-o", "json"], profile, capture=True, check=False) @@ -56,6 +73,17 @@ def _app_service_principal(name: str, profile: Optional[str]) -> Optional[str]: return None +def _app_instances(name: str, profile: Optional[str]) -> Optional[int]: + """The existing app's fixed instance count (compute_min_instances), or None if unreadable.""" + result = _databricks(["apps", "get", name, "-o", "json"], profile, capture=True, check=False) + if result.returncode != 0: + return None + try: + return json.loads(result.stdout).get("compute_min_instances") + except json.JSONDecodeError: + return None + + def _app_compute_state(name: str, profile: Optional[str]) -> Optional[str]: """The app's compute state (e.g. RUNNING), or None if it can't be read.""" result = _databricks(["apps", "get", name, "-o", "json"], profile, capture=True, check=False) @@ -347,6 +375,14 @@ def _grant_store_access( default=None, help="Workspace destination for the synced source (defaults to a per-user path).", ) +@click.option( + "--min-instances", + type=click.IntRange(min=1), + default=None, + help="Number of app instances to run (manual horizontal scaling). More instances add capacity " + "and redundancy; Apps does not autoscale yet, so this is a fixed count. Only settable when the " + "app is first created — re-scaling an existing app requires deleting and recreating it.", +) @click.pass_obj def deploy( obj, @@ -360,6 +396,7 @@ def deploy( create_stores, pip_index_url, workspace_path, + min_instances, ) -> None: """Deploy an agent: provision its stores, wire them in, and roll out the deployment.""" _validate_deployment_name(name) @@ -389,6 +426,8 @@ def deploy( for env in _PIP_INDEX_ENVS: env_updates[env] = pip_index_url provisioned["Package index"] = pip_index_url + if min_instances is not None: + provisioned["Instances"] = str(min_instances) # 2. Patch the app.yaml manifest with the store identifiers. scaffolded = False @@ -397,10 +436,24 @@ def deploy( # 3. Roll out the deployment (Databricks Apps runtime). if not _deployment_exists(name, obj.profile): - _databricks(["apps", "create", name], obj.profile) + _databricks(["apps", "create", name] + _compute_flags(min_instances), obj.profile) # `apps create` returns before the app's compute is up, but `apps deploy` requires it to be # RUNNING — so wait for it, or the first deploy races and fails ("not in RUNNING state"). _wait_for_running(name, obj.profile) + elif min_instances is not None: + # Instance count can only be set at create time. Neither Apps update path re-scales a + # fixed-count app in place: the sync API rejects the change outright, and the async one + # ("apps create-update") refuses because a fixed-count app counts as "scalable" and it won't + # let scaling be downgraded. So if the existing app already runs the requested count this is + # a no-op; otherwise fail clearly rather than passing a doomed flag downstream. + current = _app_instances(name, obj.profile) + if current is not None and current != min_instances: + raise AgentCliError( + f"App '{name}' already exists with {current} instance(s); its instance count can't " + f"be changed to {min_instances} on deploy (Apps only sets it at creation).", + hint=f"To re-scale, delete and recreate it: `mason deployments delete {name}` then " + "re-run deploy with --min-instances.", + ) ws_path = workspace_path or f"/Workspace/Users/{client.current_user}/mason_deployments/{name}" # Don't ship uv.lock: it pins exact package URLs from whatever index the developer's machine # resolved against (often an internal proxy). The Apps build must resolve against its own diff --git a/integrations/mason/tests/unit_tests/deploy_test.py b/integrations/mason/tests/unit_tests/deploy_test.py index 87a58ed5b..6448f7552 100644 --- a/integrations/mason/tests/unit_tests/deploy_test.py +++ b/integrations/mason/tests/unit_tests/deploy_test.py @@ -184,6 +184,123 @@ def test_wait_for_running_times_out(monkeypatch): pass +def test_create_passes_min_instances_to_apps_create(tmp_path: pathlib.Path, monkeypatch): + # For a new app, --min-instances is set at creation time via `apps create`. + src = tmp_path / "app" + src.mkdir() + (src / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + + calls: list[list[str]] = [] + monkeypatch.setattr(deploy_mod, "_deployment_exists", lambda a, p: False) + monkeypatch.setattr(deploy_mod, "_wait_for_running", lambda name, profile: None) + monkeypatch.setattr(deploy_mod, "_app_service_principal", lambda name, p: None) + monkeypatch.setattr( + deploy_mod, + "_databricks", + lambda args, profile, **kw: calls.append(args) + or types.SimpleNamespace(returncode=0, stdout="", stderr=""), + ) + + result = CliRunner().invoke( + deploy_mod.deploy, + ["myapp", "--source", str(src), "--min-instances", "3"], + obj=_FakeCtx(), + ) + + assert result.exit_code == 0, result.output + # Apps has no autoscaling: a fixed count pins both min and max to the same value. + assert [ + "apps", + "create", + "myapp", + "--compute-min-instances", + "3", + "--compute-max-instances", + "3", + ] in calls + # scaling is only set on create here; no separate `apps update` + assert not any(a[:2] == ["apps", "update"] for a in calls) + + +def test_existing_app_with_matching_instances_is_noop(tmp_path: pathlib.Path, monkeypatch): + # Instance count is create-only: re-deploying an existing app at its current count just proceeds + # (no `apps update`, which the API rejects for instance-count changes). + src = tmp_path / "app" + src.mkdir() + (src / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + + calls: list[list[str]] = [] + monkeypatch.setattr(deploy_mod, "_deployment_exists", lambda a, p: True) + monkeypatch.setattr(deploy_mod, "_app_instances", lambda name, p: 2) + monkeypatch.setattr(deploy_mod, "_app_service_principal", lambda name, p: None) + monkeypatch.setattr( + deploy_mod, + "_databricks", + lambda args, profile, **kw: calls.append(args) + or types.SimpleNamespace(returncode=0, stdout="", stderr=""), + ) + + result = CliRunner().invoke( + deploy_mod.deploy, + ["myapp", "--source", str(src), "--min-instances", "2"], + obj=_FakeCtx(), + ) + + assert result.exit_code == 0, result.output + assert not any(a[:2] == ["apps", "update"] for a in calls) + assert any(a[:2] == ["apps", "deploy"] for a in calls) + + +def test_existing_app_rescale_is_rejected(tmp_path: pathlib.Path, monkeypatch): + # Asking for a different count than the running app fails clearly (Apps can't re-scale in place). + src = tmp_path / "app" + src.mkdir() + (src / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + + calls: list[list[str]] = [] + monkeypatch.setattr(deploy_mod, "_deployment_exists", lambda a, p: True) + monkeypatch.setattr(deploy_mod, "_app_instances", lambda name, p: 2) + monkeypatch.setattr( + deploy_mod, + "_databricks", + lambda args, profile, **kw: calls.append(args) + or types.SimpleNamespace(returncode=0, stdout="", stderr=""), + ) + + result = CliRunner().invoke( + deploy_mod.deploy, + ["myapp", "--source", str(src), "--min-instances", "5"], + obj=_FakeCtx(), + ) + + assert result.exit_code != 0 + assert not any(a[:2] == ["apps", "deploy"] for a in calls), ( + "must not deploy on a doomed re-scale" + ) + + +def test_no_instance_lookup_without_min_instances(tmp_path: pathlib.Path, monkeypatch): + # Without --min-instances an existing-app deploy must not touch the app's compute config. + src = tmp_path / "app" + src.mkdir() + (src / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + + calls: list[list[str]] = [] + monkeypatch.setattr(deploy_mod, "_deployment_exists", lambda a, p: True) + monkeypatch.setattr(deploy_mod, "_app_service_principal", lambda name, p: None) + monkeypatch.setattr( + deploy_mod, + "_databricks", + lambda args, profile, **kw: calls.append(args) + or types.SimpleNamespace(returncode=0, stdout="", stderr=""), + ) + + result = CliRunner().invoke(deploy_mod.deploy, ["myapp", "--source", str(src)], obj=_FakeCtx()) + + assert result.exit_code == 0, result.output + assert not any(a[:2] == ["apps", "update"] for a in calls) + + def test_deploy_injects_shared_actor_for_managed_stores(tmp_path: pathlib.Path, monkeypatch): src = tmp_path / "app" src.mkdir()