feat(butler): Implement butler pool factory

tcjennings · tcjennings · commit 77973c7ae912 · 2025-04-10T12:10:31.000-05:00
diff --git a/docs/BUTLER.md b/docs/BUTLER.md
@@ -0,0 +1,52 @@
+# CM Service and Butler
+The CM Service uses a Butler at two levels: for its own use and for the use of batch jobs it submits.
+
+## First-Party Butlers
+CM Service uses Butlers to solve data queries, such as those used to determine group composition and splitting characteristics. CM Service also uses Butlers to assemble input collections and to chain output collections.
+
+To provide Butler services for these first-party operations, CM Service uses a Butler Factory, which provides the following functionality:
+
+- Instantiated and pre-cached Butler instances created on application startup.
+- A caching mechanism to provide Butler clone instances to factory callers.
+- Cloned Butler instances may be collection-constrained.
+
+From this factory, CM Service will obtain a short-lived Butler clone for performing first-party butler operations.
+Where these operations are synchronous and the CM event loop should not be blocked, Butler methods are executed on a threadpool or delegated to a FastAPI/starlette BackgroundTask.
+
+The CM Service establishes a Butler factory at startup so that any blocking IO involved with the instantiation of Butlers is front-loaded.
+CM Service uses some configuration details to determine the set of Butlers that the factory should support:
+
+- The environment variable `DAF_BUTLER_REPOSITORIES` is used to construct an instance of `lsst.daf.butler.ButlerRepoIndex` for the service instance.
+- The Butler registry authentication information must be available as a JSON string serialized to the `LSST_DB_AUTH_CREDENTIALS` environment variable, which is consumed by `lsst.util.db_auth.DbAuth` as its `authList`; CM Service does not support reading a DbAuth credential set from a file, as this requires file permissions incompatible with Kubernetes secrets mounted as volumes while also using a nonroot user in the container.
+
+The contents of the `DAF_BUTLER_REPOSITORIES` environment variable is a JSON string that represents a mapping of Butler repo names to their associated configuration files.
+
+<details>
+<summary>Example Repository Index JSON object</summary>
+
+```json
+{
+  "/repo/main": "/sdf/group/rubin/repo/main/butler.yaml",
+  "/repo/main+sasquatch_dev": "/sdf/group/rubin/repo/main/butler+sasquatch_dev.yaml"
+}
+```
+
+</details>
+
+## Second-Party Butlers
+When CM Service submits a batch job to run on a WMS, chances are high that the operations within that batch job will have to access a Butler for file IO.
+
+The configuration parameter `config.butler.repository_index` is used to identify a Butler repository index file *relative to the WMS batch environment* that will be implicitly used by batch operations. This repository index file is expected to provide and resolve the correct Butler repository detail for a batch job that has been configured to use a Butler repository by its label.
+
+The value of this configuration parameter is assigned to the environment variable `DAF_BUTLER_REPOSITORY_INDEX` in an environment specific to the batch submission operation, from which it is expected to be consumed by the appropriate `lsst.daf_butler` mechanisms when needed.
+
+Authentication secrets for second-party Butlers are provided to batch jobs via `PGPASSFILE` and `PGUSER` environment variables, which are variables used by the `libpq` PostgreSQL client library as standard sources when no superceding sources are available.
+
+The `PGUSER` variable is populated by the CM Service configuration parameter `config.butler.default_username`.
+The `PGPASSFILE` value is constructed using the `config.htcondor.remote_user_home` configuration parameter and the fixed value `.lsst/postgres-credentials.txt`. This file is expected to be in a standard PostgreSQL password file format and have appropriate file mode permissions (i.e., not group- or world-accessible) as required by `libpq`.
+
+> [!NOTE]
+> This approach is suboptimal and delegates too many assumptions to the eventual batch submission context, especially the presence of specific files in specific locations. The batch submission context should be more or less completely controlled by the submitting service, including the creation of a specific short-lived `PGPASSFILE` with contents provided by the service's own understanding of the Butler repository. This assumes that the service will not submit a job that depends on a Butler which the service does not itself understand.
+
+> [!NOTE]
+> Secrets for second-party Butlers may also be provided via an environment variable. By setting `LSST_DB_AUTH_CREDENTIALS` with the JSON string representation of a `db-auth.yaml` file, all dependencies on presumed filesystem objects in the submission environment are resolved.
diff --git a/src/lsst/cmservice/common/butler.py b/src/lsst/cmservice/common/butler.py
@@ -1,15 +1,34 @@
-"""Utility functions for working with butler commands"""
+"""Module in support of working with Butlers.
+
+A CM Service that has a configured ``BUTLER_REPO_INDEX`` provides a
+factory function that manages a pool of Butler instances for each of these
+repos. This factory function is inspired by the
+`lsst.daf.butler.LabeledButlerFactory` and provides ``clone()`` instances of
+available Butlers when asked to provide one.
+
+Notes
+-----
+The butler "factory" follows a "global" pattern where it is assigned to
+a module-level variable at import-time. Other modules can import this factory
+and use it to produce on-demand butler clones. It is not necessary to use a
+butler factory as an injected dependency when using this pattern, but this
+module should be imported as early as possible in the application startup;
+it does not depend on a running event loop.
+"""
 
-from functools import partial
+from collections.abc import Callable
+from functools import cache, partial
+from typing import TYPE_CHECKING
 
-import yaml
-from anyio import Path, to_thread
-from sqlalchemy.engine import url
+from anyio import to_thread
+from botocore.exceptions import ClientError
+from sqlalchemy.exc import OperationalError
 
 from lsst.daf.butler import Butler, ButlerConfig, ButlerRepoIndex
 from lsst.daf.butler._exceptions import MissingCollectionError
+from lsst.daf.butler.direct_butler import DirectButler
+from lsst.daf.butler.registry import CollectionArgType, RegistryConfig
 from lsst.resources import ResourcePathExpression
-from lsst.utils.db_auth import DbAuth
 
 from ..config import config
 from . import errors
@@ -20,55 +39,128 @@
 
 BUTLER_REPO_INDEX = ButlerRepoIndex()
 """An index of all known butler repositories, as populated by the
-DAF_BUTLER_REPOSITORIES environment variable.
+``DAF_BUTLER_REPOSITORIES`` environment variable.
 """
 
 
-async def get_butler_config(repo: str, *, without_datastore: bool = False) -> ButlerConfig:
-    """Create a butler config object for a repo known to the service's
-    environment.
+class ButlerFactory:
+    """The ButlerFactory will create an instance of each Butler known to the
+    application during initialization. This occurs synchronously so it is best
+    performed at application startup. After initializing, the factory can hand
+    out ``clone()`` copies of available Butlers.
     """
 
-    try:
-        repo_uri: ResourcePathExpression = BUTLER_REPO_INDEX.get_repo_uri(label=repo)
-    except KeyError:
-        # No such repo known to the service
-        logger.warning("Butler repo %s not known to environment.", repo)
-        repo_uri = repo
-
-    try:
-        bc_f = partial(
-            ButlerConfig,
-            other=repo_uri,
-            without_datastore=without_datastore,
-        )
-        bc = await to_thread.run_sync(bc_f)
-    except FileNotFoundError:
-        # No such repo known to Butler
-        logger.error("Butler repo %s not known.", repo)
-        raise RuntimeError("Unknown Butler Repo %s", repo)
+    def __init__(self) -> None:
+        """Initialize a ButlerFactory by creating butler pool instances for
+        each known repository.
+        """
+        # create and cache any butler factories known to the service
+        for label in BUTLER_REPO_INDEX.get_known_repos():
+            if config.butler.eager:
+                _ = self.get_butler_factory(label)
+
+    def get_butler(self, label: str, collections: list[str] | None = None) -> Butler | None:
+        """Get a butler clone from the factory with the specified collections
+        constraint applied.
+
+        Notes
+        -----
+        This is the primary public interface to the factory object.
+        """
+        factory = self.get_butler_factory(label)
+        if factory is None:
+            return None
+        return factory(collections=collections)
+
+    @cache
+    def get_butler_factory(
+        self, label: str, *, without_datastore: bool = True
+    ) -> Callable[..., Butler] | None:
+        """Return a factory function that creates a butler clone.
+
+        Notes
+        -----
+        This method is backed by a `functools.cache`, a threadsafe cache.
+
+        If the return value is None, the service could not create a Butler for
+        the desired label or no such Butler is configured. In the former case,
+        the service log should include exception information related to the
+        failed Butler creation.
+
+        If the application's Butler ``eager`` parameter is set, the Factory
+        instance instantiates all known Butlers at initialization. If this
+        parameter is false, then calling this method the first time will block
+        until the requested Butler is ready.
+
+        Returns
+        -------
+        `lsst.daf.butler` or `None`
+            A cloned instance of a ``Butler`` or None if the labelled Butler
+            could not be created from the configuration inputs.
+        """
+        try:
+            _butler_config = self.get_butler_config(label=label)
+            _butler = Butler.from_config(_butler_config, without_datastore=without_datastore)
+            if TYPE_CHECKING:
+                assert isinstance(_butler, DirectButler)
+            _butler._preload_cache(load_dimension_record_cache=False)
+        except (ClientError, OperationalError):
+            # Case that configured butler could not be created because of an
+            # S3 or database error
+            logger.exception()
+            return None
+        except KeyError:
+            # Case that no such butler was configured
+            logger.warning("No such butler configured: %s", label)
+            return None
+
+        def factory(collections: CollectionArgType) -> Butler:
+            return _butler.clone(collections=collections)
+
+        return factory
+
+    @cache
+    def get_butler_config(self, label: str) -> ButlerConfig:
+        """Create a butler config object for a repo known to the service's
+        environment.
+
+        Returns
+        -------
+        ``lsst.daf.butler.ButlerConfig``
+        """
+
+        try:
+            repo_uri: ResourcePathExpression = BUTLER_REPO_INDEX.get_repo_uri(label=label)
+        except KeyError:
+            logger.warning("Butler repo %s not known to environment.", label)
+            repo_uri = label
+
+        try:
+            bc = ButlerConfig(other=repo_uri)
+        except FileNotFoundError:
+            logger.error("Butler repo %s not known.", label)
+            raise RuntimeError("Unknown Butler Repo %s", label)
 
-    try:
-        db_auth_info = yaml.safe_load(await Path(config.butler.authentication_file).read_text())
-    except FileNotFoundError:
-        logger.error("No Butler Registry authentication secrets found.")
-        # delegate db auth info discovery to normal toolchain
         return bc
 
-    db_url = url.make_url(bc["registry"]["db"])
-    db_auth = DbAuth(authList=db_auth_info).getAuth(
-        dialectname=db_url.drivername,
-        username=db_url.username,
-        host=db_url.host,
-        port=db_url.port,
-        database=db_url.database,
-    )
+    @cache
+    def get_butler_registry_config(self, label: str) -> RegistryConfig:
+        """Fetch the Registry Config for a Butler by label.
+
+        Registry
+        --------
+        ``lsst.daf.butler.registry.RegistryConfig``
+        """
+        return RegistryConfig(self.get_butler_config(label=label))
 
-    bc[".registry.username"] = db_auth[0]
-    bc[".registry.password"] = db_auth[1]
-    return bc
+
+BUTLER_FACTORY = ButlerFactory()
+"""A module level butler factory created at module import, available for use
+in other modules.
+"""
 
 
+# TODO: deprecate these functions that attempt to "remove" data from Butlers.
 async def remove_run_collections(
     butler_repo: str,
     collection_name: str,
diff --git a/src/lsst/cmservice/config.py b/src/lsst/cmservice/config.py
@@ -64,6 +64,9 @@ class ButlerConfiguration(BaseModel):
         default="butler",
     )
 
+    # FIXME this index is used to hydrate a WMS submission environment, so it
+    #       is not used by CM Service to construct Butlers, and may be variable
+    #       between sites, so should be relocated to a facility-specific config
     repository_index: str = Field(
         description="Fully qualified path to a butler repository index.",
         default="/sdf/group/rubin/shared/data-repos.yaml",
@@ -79,6 +82,11 @@ class ButlerConfiguration(BaseModel):
         default="rubin",
     )
 
+    eager: bool = Field(
+        description="Whether to eagerly instantiate known Butlers at Factory startup",
+        default=True,
+    )
+
     mock: bool = Field(
         description="Whether to mock out Butler calls.",
         default=False,
diff --git a/src/lsst/cmservice/daemon.py b/src/lsst/cmservice/daemon.py
@@ -10,6 +10,7 @@
 from safir.logging import configure_uvicorn_logging
 
 from . import __version__
+from .common.butler import BUTLER_FACTORY  # noqa: F401
 from .common.daemon import daemon_iteration
 from .common.logging import LOGGER
 from .common.panda import get_panda_token
@@ -30,19 +31,20 @@ async def lifespan(app: FastAPI) -> AsyncGenerator:
     os.environ |= config.panda.model_dump(by_alias=True, exclude_none=True)
     os.environ |= config.htcondor.model_dump(by_alias=True, exclude_none=True)
     app.state.tasks = set()
-    daemon = create_task(main_loop(), name="daemon")
+    daemon = create_task(main_loop(app=app), name="daemon")
     app.state.tasks.add(daemon)
     yield
     # stop
 
 
-async def main_loop() -> None:
+async def main_loop(app: FastAPI) -> None:
     """Daemon execution loop.
 
     With a database session, perform a single daemon interation and then sleep
     until the next daemon appointment.
     """
     engine = create_database_engine(config.db.url, config.db.password)
+
     sleep_time = config.daemon.processing_interval
 
     async with engine.begin():
diff --git a/src/lsst/cmservice/handlers/elements.py b/src/lsst/cmservice/handlers/elements.py
@@ -8,11 +8,10 @@
 from anyio import to_thread
 from sqlalchemy.ext.asyncio import async_scoped_session
 
-from lsst.daf.butler import Butler
-
-from ..common.butler import get_butler_config
+from ..common.butler import BUTLER_FACTORY
 from ..common.enums import StatusEnum
 from ..common.errors import CMMissingScriptInputError, test_type_and_raise
+from ..common.logging import LOGGER
 from ..config import config
 from ..db.campaign import Campaign
 from ..db.element import ElementMixin
@@ -21,6 +20,8 @@
 from ..db.script import Script
 from .script_handler import FunctionHandler
 
+logger = LOGGER.bind(module=__name__)
+
 
 class RunElementScriptHandler(FunctionHandler):
     """Shared base class to handling running and
@@ -221,14 +222,10 @@ async def split(
         if mock_butler:
             sorted_field_values = np.arange(10)
         else:
-            butler_config = await get_butler_config(butler_repo, without_datastore=True)
-            butler_f = partial(
-                Butler.from_config,
-                butler_config,
-                collections=[input_coll, campaign_input_coll],
-                without_datastore=True,
-            )
-            butler = await to_thread.run_sync(butler_f)
+            butler = BUTLER_FACTORY.get_butler(butler_repo, collections=[input_coll, campaign_input_coll])
+            if butler is None:
+                logger.error(f"butler repo {butler_repo} is not known to the application.")
+                raise RuntimeError("No such butler")
             itr_q_f = partial(
                 butler.registry.queryDataIds,
                 [split_field],
diff --git a/tests/common/test_butler.py b/tests/common/test_butler.py
