Merge pull request #53 from lsst-sqre/tickets/DM-45035

tickets/DM-45035: Conditionally add DAF_BUTLER_CACHE_DIRECTORY to env

Author: athornton

changelog.d/20240814_154234_athornton_DM_45035.md

@@ -0,0 +1,5 @@
+<!-- Delete the sections that don't apply -->
+
+### New features
+
+- Conditionally set environment variable `DAF_BUTLER_CACHE_DIRECTORY`

src/lsst/rsp/startup/services/labrunner.py

@@ -120,6 +120,7 @@ def _configure_env(self) -> None:
         self._logger.debug("Configuring environment for JupyterLab process")
         self._set_user()
         self._set_tmpdir_if_scratch_available()
+        self._set_butler_cache()
         self._set_cpu_variables()
         self._set_image_digest()
         self._expand_panda_tilde()
@@ -169,6 +170,38 @@ def _set_tmpdir_if_scratch_available(self) -> None:
             return
         self._env["TMPDIR"] = str(user_scratch)
 
+    def _set_butler_cache(self) -> None:
+        # This should be called *after* _set_tmpdir_if_scratch_available()
+        # at least for now.  We may eventually want to force it to local
+        # ephemeral storage and demand enough ephemeral storage to cover it
+        # (currently about 500MB).
+        #
+        # For now, though, let's set it to `butler_cache` inside `TMPDIR`
+        dbcd = self._env.get("DAF_BUTLER_CACHE_DIRECTORY", "")
+        if dbcd:
+            self._logger.debug(
+                f"Not setting DAF_BUTLER_CACHE_DIRECTORY: already set to"
+                f" {dbcd}"
+            )
+            return
+        # Yes, we know that ruff doesn't like `/tmp`
+        # In any sane RSP environment, either we will have set TMPDIR, or
+        # /tmp will be on ephemeral storage.
+        tmpdir = Path(self._env.get("TMPDIR", "/tmp"))  # noqa: S108
+        dbc = tmpdir / "butler_cache"
+        try:
+            dbc.mkdir(parents=True, exist_ok=True)
+        except (OSError, FileExistsError) as exc:
+            self._logger.warning(
+                f"Could not create DAF_BUTLER_CACHE_DIRECTORY at"
+                f" {dbc!s}: {exc}"
+            )
+            return
+        if not os.access(dbc, os.W_OK):
+            self._logger.warning(f"Unable to write to {dbc}")
+            return
+        self._env["DAF_BUTLER_CACHE_DIRECTORY"] = str(dbc)
+
     def _set_cpu_variables(self) -> None:
         self._logger.debug("Setting CPU threading variables")
         try:

tests/startup_test.py

@@ -70,6 +70,48 @@ def test_set_tmpdir(monkeypatch: pytest.MonkeyPatch) -> None:
             assert "TMPDIR" not in lr._env
 
 
+@pytest.mark.usefixtures("_rsp_env")
+def test_set_butler_cache(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Happy path.
+    lr = LabRunner()
+    lr._set_tmpdir_if_scratch_available()
+    lr._set_butler_cache()
+    assert lr._env["DAF_BUTLER_CACHE_DIRECTORY"].endswith(
+        "/scratch/hambone/tmp/butler_cache"
+    )
+    # Exists, but it's not a directory
+    dbc = Path(lr._env["DAF_BUTLER_CACHE_DIRECTORY"])
+    dbc.rmdir()
+    dbc.touch()
+    lr = LabRunner()
+    lr._set_tmpdir_if_scratch_available()
+    lr._set_butler_cache()
+    assert "DAF_BUTLER_CACHE_DIRECTORY" not in lr._env
+    # Put it back the way it was
+    dbc.unlink()
+    # Pre-set DAF_BUTLER_CACHE_DIR.
+    monkeypatch.setenv("DAF_BUTLER_CACHE_DIRECTORY", "/preset")
+    lr = LabRunner()
+    lr._set_butler_cache()
+    assert lr._env["DAF_BUTLER_CACHE_DIRECTORY"] == "/preset"
+    monkeypatch.delenv("DAF_BUTLER_CACHE_DIRECTORY")
+    # Pre-set TMPDIR so cache is not writeable
+    monkeypatch.setenv("TMPDIR", "/nonexistent")
+    lr = LabRunner()
+    lr._set_butler_cache()
+    assert "DAF_BUTLER_CACHE_DIRECTORY" not in lr._env
+    monkeypatch.delenv("TMPDIR")
+    # No TMPDIR set: should go under `/tmp`
+    lr = LabRunner()
+    lr._set_butler_cache()
+    e_dbcd = lr._env["DAF_BUTLER_CACHE_DIRECTORY"]
+    # Yes, we know that ruff doesn't like `/tmp`
+    # In any sane RSP environment, either we will have set TMPDIR, or
+    # /tmp will be on ephemeral storage.
+    assert e_dbcd == "/tmp/butler_cache"  # noqa: S108
+    Path(e_dbcd).rmdir()
+
+
 @pytest.mark.usefixtures("_rsp_env")
 def test_cpu_vars(monkeypatch: pytest.MonkeyPatch) -> None:
     lr = LabRunner()