Merge pull request #44 from lsst-sqre/tickets/DM-44731

tickets/DM-44731: adapt for two-python model

changelog.d/20240613_081431_athornton_DM_44731.md

@@ -0,0 +1,6 @@
+<!-- Delete the sections that don't apply -->
+
+### New features
+
+- `RESET_USER_ENV` now processed inside lsst.rsp.startup
+

changelog.d/20240613_114352_athornton_DM_44731.md

@@ -0,0 +1,5 @@
+<!-- Delete the sections that don't apply -->
+
+### New features
+
+- Add utility functions to get values that were constants in the single-Python model but now are not.

src/lsst/rsp/client.py

@@ -1,9 +1,9 @@
 """Client for other services running in the same RSP instance."""
 
-from pathlib import Path
-
 import httpx
 
+from .utils import get_access_token, get_runtime_mounts_dir
+
 
 class RSPClient(httpx.AsyncClient):
     """Configured client for other services in the RSP.
@@ -15,10 +15,9 @@ class RSPClient(httpx.AsyncClient):
     def __init__(
         self,
         service_endpoint: str,
-        *,
-        jupyterlab_dir: Path = Path("/opt/lsst/software/jupyterlab"),
     ) -> None:
-        token = (jupyterlab_dir / "secrets" / "token").read_text().strip()
+        token = get_access_token()
+        jupyterlab_dir = get_runtime_mounts_dir()
         instance_url = (
             (jupyterlab_dir / "environment" / "EXTERNAL_INSTANCE_URL")
             .read_text()

src/lsst/rsp/startup/constants.py

@@ -7,8 +7,6 @@
     "ETC_PATH",
     "PREVIOUS_LOGGING_CHECKSUMS",
     "MAX_NUMBER_OUTPUTS",
-    "NONINTERACTIVE_CONFIG_PATH",
-    "TOP_DIR_PATH",
 ]
 
 APP_NAME = "nublado"
@@ -32,17 +30,3 @@
 
 Used to prevent OOM-killing if some cell generates a lot of output.
 """
-
-TOP_DIR_PATH = Path("/opt/lsst/software")
-"""
-Location where the DM stack and our Lab machinery are rooted.
-
-Overrideable for testing.
-"""
-
-NONINTERACTIVE_CONFIG_PATH = Path(
-    TOP_DIR_PATH / "jupyterlab" / "noninteractive" / "command" / "command.json"
-)
-"""
-Location where a noninteractive pod will mount its command configuration.
-"""

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

@@ -19,13 +19,12 @@
 import symbolicmode
 
 from ... import get_access_token, get_digest
+from ...utils import get_jupyterlab_config_dir, get_runtime_mounts_dir
 from ..constants import (
     APP_NAME,
     ETC_PATH,
     MAX_NUMBER_OUTPUTS,
-    NONINTERACTIVE_CONFIG_PATH,
     PREVIOUS_LOGGING_CHECKSUMS,
-    TOP_DIR_PATH,
 )
 from ..models.noninteractive import NonInteractiveExecutor
 from ..storage.command import Command
@@ -60,28 +59,13 @@ def __init__(self) -> None:
 
     def go(self) -> None:
         """Start the user lab."""
-        # The LabRunner is not actually the first thing that launches when
-        # we start a user lab.
-        #
-        # At the moment, the only Python 3 in the Lab container is the one that
-        # is part of the DM stack conda environment.
-        #
-        # So to get far enough to even start the LabRunner, we need to have
-        # already sourced the shell magic that sets up that conda env.
-        #
-        # However, even before we do that, we check the environment to see
-        # whether it needs resetting.  That's actually kind of handy to do
-        # before we start any Python process, because it's not impossible that
-        # the user could have messed up their own Python environment so badly
-        # that Python couldn't even get this far.
-        #
-        # When we split the stack Python from the Python-that-runs-JuptyerLab,
-        # there will be no need to set up the stack before launching
-        # Jupyterlab, and the likelihood that the user environment is so
-        # corrupt that the user cannot get a terminal session open to purge
-        # it themselves will lessen greatly.  At that point moving the
-        # user-environment purge into lsst-rsp may be feasible.
-        #
+        # If the user somehow manages to screw up their local environment
+        # so badly that Jupyterlab won't even start, we will have to
+        # bail them out on the fileserver end.  Since Jupyter Lab is in
+        # its own venv, which is not writeable by the user, this should
+        # require quite a bit of creativity.
+
+        self._relocate_user_environment_if_requested()
 
         # Set up environment variables that we'll need either to launch the
         # Lab or for the user's terminal environment
@@ -94,9 +78,6 @@ def go(self) -> None:
         # Check out notebooks, and set up git-lfs
         self._setup_git()
 
-        # Clear EUPS cache
-        self._cmd.run("eups", "admin", "clearCache")
-
         # Decide between interactive and noninteractive start, do
         # things that change between those two, and launch the Lab
         self._launch()
@@ -110,6 +91,30 @@ def _externalize(self, setting: str) -> str:
         ext_url = self._env["EXTERNAL_INSTANCE_URL"]
         return f"{ext_url.strip('/')}/{setting.lstrip('/')}"
 
+    #
+    #
+    #
+    def _relocate_user_environment_if_requested(self) -> None:
+        if not self._env.get("RESET_USER_ENV", ""):
+            return
+        self._logger.debug("User environment relocation requested")
+        now = datetime.datetime.now(datetime.UTC).strftime("%Y%m%d%H%M%S")
+        reloc = self._home / f".user_env.{now}"
+        for candidate in ("cache", "conda", "eups", "local", "jupyter"):
+            c_path = self._home / f".{candidate}"
+            if c_path.is_dir():
+                if not reloc.is_dir():
+                    reloc.mkdir()
+                tgt = reloc / candidate
+                self._logger.debug(f"Moving {c_path.name} to {tgt.name}")
+                shutil.move(c_path, tgt)
+        u_setups = self._home / "notebooks" / ".user_setups"
+        if u_setups.is_file():
+            tgt = reloc / "notebooks" / "user_setups"
+            tgt.parent.mkdir()
+            self._logger.debug(f"Moving {u_setups.name} to {tgt}")
+            shutil.move(u_setups, tgt)
+
     #
     # Next up, a big block of setting up our subprocess environment.
     #
@@ -333,8 +338,9 @@ def _copy_logging_profile(self) -> None:
             pdir = user_profile.parent
             if not pdir.is_dir():
                 pdir.mkdir(parents=True)
+            jl_path = get_jupyterlab_config_dir()
             user_profile.write_bytes(
-                (TOP_DIR_PATH / "jupyterlab" / "20-logging.py").read_bytes()
+                (jl_path / "etc" / "20-logging.py").read_bytes()
             )
 
     def _copy_dircolors(self) -> None:
@@ -351,7 +357,8 @@ def _copy_etc_skel(self) -> None:
         self._logger.debug("Copying files from /etc/skel if they don't exist")
         etc_skel = ETC_PATH / "skel"
         # alas, Path.walk() requires Python 3.12, which isn't in the
-        # stack containers yet.
+        # stack containers yet.  Once the Lab/stack split is finalized,
+        # we can make this simpler.
         contents = os.walk(etc_skel)
         #
         # We assume that if the file exists at all, we should leave it alone.
@@ -538,6 +545,9 @@ def _check_for_git_lfs(self) -> bool:
         return False
 
     def _launch(self) -> None:
+        # We're about to start the lab: set the flag saying we're running
+        # inside the lab.  It's used by shell startup.
+        self._env["RUNNING_INSIDE_JUPYTERLAB"] = "TRUE"
         if bool(self._env.get("NONINTERACTIVE", "")):
             self._start_noninteractive()
             # We exec a lab; control never returns here
@@ -581,9 +591,7 @@ def _manage_access_token(self) -> None:
         self._logger.debug("Updating access token")
         tokfile = self._home / ".access_token"
         tokfile.unlink(missing_ok=True)
-        ctr_token = (
-            TOP_DIR_PATH / "software" / "jupyterlab" / "secrets" / "token"
-        )
+        ctr_token = get_runtime_mounts_dir() / "secrets" / "token"
         if ctr_token.exists():
             self._logger.debug(f"Symlinking {tokfile!s}->{ctr_token!s}")
             tokfile.symlink_to(ctr_token)
@@ -600,9 +608,13 @@ def _manage_access_token(self) -> None:
             self._logger.debug("Could not determine access token")
 
     def _start_noninteractive(self) -> None:
-        launcher = NonInteractiveExecutor.from_config(
-            NONINTERACTIVE_CONFIG_PATH
+        config_path = (
+            get_runtime_mounts_dir()
+            / "noninteractive"
+            / "command"
+            / "command.json"
         )
+        launcher = NonInteractiveExecutor.from_config(config_path)
         launcher.execute(env=self._env)
 
     def _set_timeout_variables(self) -> list[str]:
@@ -651,6 +663,8 @@ def _start(self) -> None:
         ]
         cmd.extend(self._set_timeout_variables())
         self._logger.debug("Command to run:", command=cmd)
+        # Set environment variable to indicate we are inside JupyterLab
+        # (we want the shell to source loadLSST.bash once we are)
         if self._debug:
             # Maybe we want to parameterize these?
             retries = 10

src/lsst/rsp/utils.py

@@ -136,6 +136,44 @@ def get_digest() -> str:
     return spec[hash_pos + len(hash_marker) :]
 
 
+def get_jupyterlab_config_dir() -> Path:
+    """Return the directory where Jupyterlab configuration is stored.
+    For single-python images, this will be `/opt/lsst/software/jupyterlab`.
+
+    For images with split stack and Jupyterlab Pythons, it will be the
+    value of `JUPYTERLAB_CONFIG_DIR`.
+
+    Returns
+    -------
+    pathlib.Path
+        Location where Jupyterlab configuration is stored.
+    """
+    return Path(
+        os.environ.get(
+            "JUPYTERLAB_CONFIG_DIR", "/opt/lsst/software/jupyterlab"
+        )
+    )
+
+
+def get_runtime_mounts_dir() -> Path:
+    """Return the directory where Nublado runtime info is mounted.  For
+    single-python images, this will be `/opt/lsst/software/jupyterlab`.
+
+    For images with split stack and Jupyterlab Pythons, it will be the
+    value of `NUBLADO_RUNTIME_MOUNTS_DIR`.
+
+    Returns
+    -------
+    pathlib.Path
+        Location where the Nublado runtime information is mounted.
+    """
+    return Path(
+        os.environ.get(
+            "NUBLADO_RUNTIME_MOUNTS_DIR", "/opt/lsst/software/jupyterlab"
+        )
+    )
+
+
 def get_access_token(
     tokenfile: str | Path | None = None, log: Any | None = None
 ) -> str:
@@ -147,14 +185,14 @@ def get_access_token(
     was started.  Return the empty string if the token cannot be determined.
     """
     if tokenfile:
-        return Path(tokenfile).read_text()
-    base_dir = Path("/opt/lsst/software/jupyterlab")
+        return Path(tokenfile).read_text().strip()
+    base_dir = get_runtime_mounts_dir()
     for candidate in (
         base_dir / "secrets" / "token",
         base_dir / "environment" / "ACCESS_TOKEN",
     ):
         with suppress(FileNotFoundError):
-            return candidate.read_text()
+            return candidate.read_text().strip()
 
     # If we got here, we couldn't find a file. Return the environment variable
     # if set, otherwise the empty string.

tests/client_test.py

@@ -1,7 +1,5 @@
 """Test the RSPClient."""
 
-from pathlib import Path
-
 import pytest
 from pytest_httpx import HTTPXMock
 
@@ -21,10 +19,7 @@ async def test_client(httpx_mock: HTTPXMock) -> None:
             "Content-Type": "application/json",
         },
     )
-    jupyterlab_dir = (
-        Path(__file__).parent / "support" / "files" / "client" / "jupyterlab"
-    )
-    client = RSPClient("/test-service", jupyterlab_dir=jupyterlab_dir)
+    client = RSPClient("/test-service")
     await client.get("/foo")
     # The httpx mock will throw an error at teardown if we did not exercise
     # the mock, so we know the request matched both the URL and the headers.

tests/conftest.py

@@ -17,29 +17,29 @@ def _rsp_paths(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
     # For each of these, we want to cover both the "from ..constants import"
     # and the "import lsst.rsp.constants" case.
     with patch(
-        "lsst.rsp.startup.services.labrunner.TOP_DIR_PATH",
-        (Path(__file__).parent / "support" / "files" / "stack_top"),
+        "lsst.rsp.startup.services.labrunner.ETC_PATH",
+        (Path(__file__).parent / "support" / "files" / "etc"),
     ):
         with patch(
-            "lsst.rsp.startup.constants.TOP_DIR_PATH",
-            (Path(__file__).parent / "support" / "files" / "stack_top"),
+            "lsst.rsp.startup.constants.ETC_PATH",
+            (Path(__file__).parent / "support" / "files" / "etc"),
         ):
-            with patch(
-                "lsst.rsp.startup.services.labrunner.ETC_PATH",
-                (Path(__file__).parent / "support" / "files" / "etc"),
-            ):
-                with patch(
-                    "lsst.rsp.startup.constants.ETC_PATH",
-                    (Path(__file__).parent / "support" / "files" / "etc"),
-                ):
-                    yield
+            yield
 
 
 @pytest.fixture
 def _rsp_env(
     _rsp_paths: None, monkeypatch: pytest.MonkeyPatch
 ) -> Iterator[None]:
-    template = Path(__file__).parent / "support" / "files" / "homedir"
+    file_dir = Path(__file__).parent / "support" / "files"
+    template = file_dir / "homedir"
+    monkeypatch.setenv(
+        "NUBLADO_RUNTIME_MOUNTS_DIR", str(file_dir / "etc" / "nublado")
+    )
+    monkeypatch.setenv(
+        "JUPYTERLAB_CONFIG_DIR",
+        str(file_dir / "jupyterlab"),
+    )
     with TemporaryDirectory() as homedir:
         monkeypatch.setenv("HOME", homedir)
         monkeypatch.setenv("USER", "hambone")