durable fixes (#1724)

<!-- This is an auto-generated description by cubic. -->
## Summary by cubic
Fixes durable disk handling and private-pool fallback, and simplifies
the `beta9` database CLI by nesting commands under `db`, removing unused
flags, and cleaning output. Also ensures build containers always resolve
the worker address.

- **Bug Fixes**
- Sync durable disks during cleanup only when a durable disk is mounted;
default the durable-disks host path to `<storagePath>/durable-disks`
(else `/var/lib/beta9/durable-disks`).
- Always resolve the worker address for build containers; remove the
`ClipVersion2` TTL refresh/early-return; handle durable-disk requests
before private-pool fallback.

- **New Features**
- Group database commands under `db` (`db postgres`, `db redis`); drop
the unused `--size` flag from `db create/scale` and trim unneeded fields
from output; errors use the active CLI name.

<sup>Written for commit f19b3c7.
Summary will update on new commits.</sup>

<a
href="https://cubic.dev/pr/beam-cloud/beta9/pull/1724?utm_source=github"
target="_blank" rel="noopener noreferrer"
data-no-image-dialog="true"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://www.cubic.dev/buttons/review-in-cubic-light.svg"><img
alt="Review in cubic"
src="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"></picture></a>

<!-- End of auto-generated description by cubic. -->

Author: luke-lombardi

pkg/abstractions/image/builder.go

@@ -100,11 +100,6 @@ func (b *Builder) startBuildContainer(ctx context.Context, build *Build) error {
 		return err
 	}
 
-	if b.config.ImageService.ClipVersion == uint32(types.ClipVersion2) {
-		go b.refreshBuildContainerTTL(ctx, build.containerID)
-		return nil
-	}
-
 	hostname, err := b.containerRepo.GetWorkerAddress(ctx, build.containerID)
 	if err != nil {
 		build.log(true, "Failed to connect to build container.\n")

pkg/common/config.default.yaml

@@ -180,7 +180,8 @@ worker:
           - --overlay2=none
           - --file-access=shared
       # Optional per-pool durable disk and cache overrides for nodes with different storage layouts.
-      # durableDisksPath: /var/lib/beta9/durable-disks
+      # Defaults to <storagePath>/durable-disks when storagePath is set, otherwise /var/lib/beta9/durable-disks.
+      # durableDisksPath: /mnt/nvme/beta9/durable-disks
       # cache:
       #   disk:
       #     hostPath: /mnt/nvme/beta9/cache
@@ -210,7 +211,8 @@ worker:
           - --overlay2=none
           - --file-access=shared
       # Optional per-pool durable disk and cache overrides.
-      # durableDisksPath: /var/lib/beta9/durable-disks
+      # Defaults to <storagePath>/durable-disks when storagePath is set, otherwise /var/lib/beta9/durable-disks.
+      # durableDisksPath: /mnt/build-cache/beta9/durable-disks
       # cache:
       #   disk:
       #     hostPath: /mnt/build-cache/beta9/cache

pkg/scheduler/pool.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"time"
@@ -121,6 +122,9 @@ func workerDurableDisksHostPath(poolConfig types.WorkerPoolConfig) string {
 	if poolConfig.DurableDisksPath != "" {
 		return poolConfig.DurableDisksPath
 	}
+	if poolConfig.StoragePath != "" {
+		return filepath.Join(poolConfig.StoragePath, "durable-disks")
+	}
 	return types.DefaultDurableDisksPath
 }
 

pkg/scheduler/pool_test.go

@@ -83,7 +83,11 @@ func TestWorkerPodCommandUsesInitReaper(t *testing.T) {
 
 func TestWorkerDurableDisksHostPathUsesPoolOverride(t *testing.T) {
 	assert.Equal(t, types.DefaultDurableDisksPath, workerDurableDisksHostPath(types.WorkerPoolConfig{}))
+	assert.Equal(t, "/mnt/nvme/beta9/storage/durable-disks", workerDurableDisksHostPath(types.WorkerPoolConfig{
+		StoragePath: "/mnt/nvme/beta9/storage",
+	}))
 	assert.Equal(t, "/mnt/nvme/beta9/disks", workerDurableDisksHostPath(types.WorkerPoolConfig{
+		StoragePath:      "/mnt/nvme/beta9/storage",
 		DurableDisksPath: "/mnt/nvme/beta9/disks",
 	}))
 }

pkg/scheduler/private_pool_fallback.go

@@ -35,7 +35,12 @@ func (a *schedulingAttempt) tryPrivatePoolFallback() bool {
 }
 
 func (a *schedulingAttempt) privatePoolFallbackRequest() (*types.ContainerRequest, string, bool) {
-	if a == nil || a.request == nil || a.request.HasDurableDiskMount() {
+	if a == nil || a.request == nil {
+		return nil, "", false
+	}
+	if a.request.HasDurableDiskMount() {
+		// Durable disk fallback is handled before scheduling so snapshot
+		// availability is checked before the pool selector is cleared.
 		return nil, "", false
 	}
 

pkg/worker/lifecycle.go

@@ -148,19 +148,21 @@ func (s *Worker) finalizeContainer(containerId string, request *types.ContainerR
 }
 
 func (s *Worker) clearContainer(containerId string, request *types.ContainerRequest, exitCode int) {
+	if request != nil && request.HasDurableDiskMount() {
+		if err := s.syncDurableDiskMounts(request); err != nil {
+			log.Error().Str("container_id", containerId).Err(err).Msg("failed to sync durable disks during container cleanup")
+		}
+	}
+
 	s.setContainerExitCode(containerId, exitCode)
 
-	// Set container exit code on instance before any slower cleanup work.
+	// Keep the local instance state consistent with the reported exit code.
 	instance, exists := s.containerInstances.Get(containerId)
 	if exists {
 		instance.ExitCode = exitCode
 		s.containerInstances.Set(containerId, instance)
 	}
 
-	if err := s.syncDurableDiskMounts(request); err != nil {
-		log.Error().Str("container_id", containerId).Err(err).Msg("failed to sync durable disks during container cleanup")
-	}
-
 	s.containerLock.Lock()
 
 	// De-allocate GPU devices so they are available for new containers

sdk/src/beta9/cli/database.py

@@ -45,12 +45,17 @@ def common(**_):
     pass
 
 
-@common.group(name="postgres", help="Create and manage Postgres services.")
+@common.group(name="db", help="Create and manage database services.")
+def db():
+    pass
+
+
+@db.group(name="postgres", help="Create and manage Postgres services.")
 def postgres():
     pass
 
 
-@common.group(name="redis", help="Create and manage Redis services.")
+@db.group(name="redis", help="Create and manage Redis services.")
 def redis():
     pass
 
@@ -172,12 +177,21 @@ def _deployments_by_name(service: ServiceClient, name: str):
     return [deployment for deployment in res.deployments if deployment.name == name]
 
 
+def _cli_name() -> str:
+    ctx = click.get_current_context(silent=True)
+    if ctx is not None and ctx.command_path:
+        return ctx.command_path.split()[0]
+    return "beta9"
+
+
 def _ensure_database_name_available(service: ServiceClient, product: DatabaseProduct, name: str) -> None:
     if _deployment_by_name(service, name) is None:
         return
+    cli_name = _cli_name()
     raise click.ClickException(
         f"{product.kind} service {name!r} already exists. "
-        f"Use `beam {product.kind} credentials {name}`, `beam {product.kind} status {name}`, "
+        f"Use `{cli_name} db {product.kind} credentials {name}`, "
+        f"`{cli_name} db {product.kind} status {name}`, "
         f"or delete it before creating a replacement."
     )
 
@@ -247,9 +261,7 @@ def _print_result(format: str, payload: Dict[str, object]) -> None:
         "name",
         "kind",
         "deployment_id",
-        "version",
         "disk",
-        "disk_size",
         "host",
         "username",
         "database",
@@ -521,7 +533,7 @@ def _deploy_database_service(
         cpu=cpu,
         memory=memory,
     )
-    stub_id, deployment_id, version = _deploy_database_stub(service, db_service)
+    stub_id, deployment_id, _ = _deploy_database_stub(service, db_service)
     host = _tcp_host_for_stub(service, stub_id, deployment_id)
     if product.kind == "postgres":
         connection_url = _postgres_url(username, password, host, database)
@@ -535,9 +547,7 @@ def _deploy_database_service(
             "name": name,
             "kind": product.kind,
             "deployment_id": deployment_id,
-            "version": version,
             "disk": db_service.disks[0].name,
-            "disk_size": db_service.disks[0].size,
             "host": host,
             "username": username,
             "connection_string": connection_url,
@@ -570,7 +580,6 @@ def _create_options(func):
         help="Minimum database replicas to keep warm. Use 1 to keep the service warm.",
     )(func)
     func = click.option("--pool", type=click.STRING, default=None, help="Run on a private pool.")(func)
-    func = click.option("--size", type=click.STRING, default=None, help="Durable disk size.")(func)
     func = click.option("--password-stdin", is_flag=True, help="Read password from stdin.")(func)
     func = click.option("--password-from-env", type=click.STRING, default="", help="Read password from an environment variable.")(func)
     func = click.option("--password", type=click.STRING, default="", help="Database password. Generated if omitted.")(func)
@@ -591,7 +600,6 @@ def create_postgres(
     password: str,
     password_from_env: str,
     password_stdin: bool,
-    size: str,
     pool: Optional[str],
     min_replicas: int,
     format: str,
@@ -605,7 +613,7 @@ def create_postgres(
         username=username or name.replace("-", "_"),
         database=database or name.replace("-", "_") or POSTGRES.default_database,
         password=_password(password, password_from_env, password_stdin),
-        size=size or POSTGRES.default_size,
+        size=POSTGRES.default_size,
         pool=pool,
         min_replicas=min_replicas,
         format=format,
@@ -625,7 +633,6 @@ def create_redis(
     password: str,
     password_from_env: str,
     password_stdin: bool,
-    size: str,
     pool: Optional[str],
     min_replicas: int,
     format: str,
@@ -639,7 +646,7 @@ def create_redis(
         username=username or "default",
         database="",
         password=_password(password, password_from_env, password_stdin),
-        size=size or REDIS.default_size,
+        size=REDIS.default_size,
         pool=pool,
         min_replicas=min_replicas,
         format=format,
@@ -795,7 +802,6 @@ def _status(service: ServiceClient, product: DatabaseProduct, name: str, format:
         "kind": product.kind,
         "deployment_id": deployment.id,
         "active": deployment.active,
-        "version": deployment.version,
         "connection_string_secret": _secret_names(product, name)["url"],
     }
     if format == "json":
@@ -873,7 +879,6 @@ def redis_delete(service: ServiceClient, name: str):
 
 def _scale_options(func):
     func = click.option("--format", type=click.Choice(("table", "json")), default="table")(func)
-    func = click.option("--size", type=click.STRING, default=None, help="Durable disk size to use when redeploying.")(func)
     func = click.option("--pool", type=click.STRING, default=None, help="Pool to use when redeploying with new resources.")(func)
     func = click.option("--memory", type=click.STRING, default=None, help="Memory to allocate, for example 1024 or 2Gi.")(func)
     func = click.option("--cpu", type=click.FLOAT, default=None, help="CPU cores to allocate, for example 0.5 or 2.")(func)
@@ -893,10 +898,9 @@ def postgres_scale(
     cpu: Optional[float],
     memory: Optional[str],
     pool: Optional[str],
-    size: Optional[str],
     format: str,
 ):
-    _scale_database(service, POSTGRES, name, always_on, serverless, cpu, memory, pool, size, format)
+    _scale_database(service, POSTGRES, name, always_on, serverless, cpu, memory, pool, format)
 
 
 @redis.command(name="scale", help="Scale a Redis service.")
@@ -911,10 +915,9 @@ def redis_scale(
     cpu: Optional[float],
     memory: Optional[str],
     pool: Optional[str],
-    size: Optional[str],
     format: str,
 ):
-    _scale_database(service, REDIS, name, always_on, serverless, cpu, memory, pool, size, format)
+    _scale_database(service, REDIS, name, always_on, serverless, cpu, memory, pool, format)
 
 
 def _scale_mode_containers(always_on: bool, serverless: bool) -> int:
@@ -932,7 +935,6 @@ def _scale_database(
     cpu: Optional[float],
     memory: Optional[str],
     pool: Optional[str],
-    size: Optional[str],
     format: str,
 ) -> None:
     containers = _scale_mode_containers(always_on, serverless)
@@ -946,7 +948,7 @@ def _scale_database(
             username=_get_secret_value(service, secret_names["username"]),
             database=_get_secret_value(service, secret_names["database"]) if product.kind == "postgres" else "",
             password=_get_secret_value(service, secret_names["password"]),
-            size=size or product.default_size,
+            size=product.default_size,
             pool=pool,
             min_replicas=containers,
             format=format,

sdk/tests/test_cli.py

@@ -0,0 +1,41 @@
+import click
+import pytest
+
+from beta9.cli import database as database_cli
+from beta9.cli.main import load_cli
+
+
+def test_database_commands_are_nested_under_db():
+    cli = load_cli(check_config=False)
+
+    assert cli.common_group.get_command(None, "postgres") is None
+    assert cli.common_group.get_command(None, "redis") is None
+
+    db = cli.common_group.get_command(None, "db")
+    assert db is not None
+    assert sorted(db.commands) == ["postgres", "redis"]
+    assert "create" in db.commands["postgres"].commands
+    assert "create" in db.commands["redis"].commands
+
+
+def test_database_commands_do_not_expose_unenforced_disk_size():
+    cli = load_cli(check_config=False)
+    db = cli.common_group.get_command(None, "db")
+
+    for product in ("postgres", "redis"):
+        for command in ("create", "scale"):
+            options = {param.name for param in db.commands[product].commands[command].params}
+            assert "size" not in options
+
+
+def test_database_exists_error_uses_active_cli_name(monkeypatch):
+    monkeypatch.setattr(database_cli, "_deployment_by_name", lambda service, name: object())
+
+    with click.Context(click.Command("create"), info_name="beta9"):
+        with pytest.raises(click.ClickException) as exc:
+            database_cli._ensure_database_name_available(None, database_cli.REDIS, "myredis")
+
+    message = str(exc.value)
+    assert "beta9 db redis credentials myredis" in message
+    assert "beta9 db redis status myredis" in message
+    assert "beam redis" not in message