feat(BA-4766): inject BACKENDAI_PERSISTENT_PATHS env var and show persistent paths on shell startup

changes/9738.feature.md

@@ -0,0 +1 @@
+Add `BACKENDAI_PERSISTENT_PATHS` environment variable in containers to display persistent vfolder mount paths on shell startup.

src/ai/backend/agent/agent.py

@@ -2746,6 +2746,12 @@ async def create_kernel(
                     vfolder_mounts = [
                         VFolderMount.from_json(item) for item in kernel_config["mounts"]
                     ]
+                    # NOTE: Inline injection until EnvironProvisioner is wired
+                    #  into the kernel creation pipeline. See also:
+                    #  agent/stage/kernel_lifecycle/docker/environ.py
+                    if vfolder_mounts:
+                        persistent_paths = ":".join(str(m.kernel_path) for m in vfolder_mounts)
+                        environ["BACKENDAI_PERSISTENT_PATHS"] = persistent_paths
                     if not restarting:
                         await ctx.mount_vfolders(vfolder_mounts, resource_spec)
                         await ctx.mount_krunner(resource_spec, environ)

src/ai/backend/agent/stage/kernel_lifecycle/docker/environ.py

@@ -16,12 +16,14 @@
     DeviceName,
     KernelCreationConfig,
     SlotName,
+    VFolderMount,
 )
 
 LD_PRELOAD: Final[str] = "LD_PRELOAD"
 LOCAL_USER_ID: Final[str] = "LOCAL_USER_ID"
 LOCAL_GROUP_ID: Final[str] = "LOCAL_GROUP_ID"
 ADDITIONAL_GIDS: Final[str] = "ADDITIONAL_GIDS"
+BACKENDAI_PERSISTENT_PATHS: Final[str] = "BACKENDAI_PERSISTENT_PATHS"
 
 
 @dataclass
@@ -120,8 +122,15 @@ async def setup(self, spec: EnvironSpec) -> EnvironResult:
 
         hook_paths = await self._get_container_hooks(spec)
         device_environ = await self._get_device_environ(spec)
-        environ = environ.append_values(LD_PRELOAD, hook_paths, separator=":").update_always(
-            device_environ
+        # TODO: Once EnvironProvisioner is wired into the kernel creation pipeline,
+        #  remove the inline BACKENDAI_PERSISTENT_PATHS injection in agent.py.
+        environ = (
+            environ.append_values(LD_PRELOAD, hook_paths, separator=":")
+            .update_always(device_environ)
+            .set_value(
+                BACKENDAI_PERSISTENT_PATHS,
+                self._get_persistent_paths(spec),
+            )
         )
         return EnvironResult(environ=environ.to_dict())
 
@@ -159,6 +168,14 @@ def _get_core_count(self, spec: EnvironSpec) -> dict[str, str]:
         )
         return {k: str(cpu_core_count) for k in envs_corecount}
 
+    def _get_persistent_paths(self, spec: EnvironSpec) -> str | None:
+        vfolder_mounts = [
+            VFolderMount.from_json(item)
+            for item in spec.kernel_info.kernel_creation_config["mounts"]
+        ]
+        paths = [str(m.kernel_path) for m in vfolder_mounts]
+        return ":".join(paths) if paths else None
+
     async def _get_container_hooks(self, spec: EnvironSpec) -> set[str]:
         container_hook_path_set: set[str] = set()
         for device_view in spec.kernel_info.resource_spec.device_list:

src/ai/backend/runner/.bashrc

@@ -1,5 +1,15 @@
 export PS1="\[\033[01;32m\]\u@${BACKENDAI_CLUSTER_HOST:-main}\[\033[01;33m\][${BACKENDAI_SESSION_NAME}]\[\033[00m\]:\[\033[01;34m\]\w\[\033[00m\]\$ "
 
+if [[ -n "${BACKENDAI_PERSISTENT_PATHS:-}" ]]; then
+  echo -e "\e[33m⚠ Only the following vfolder paths are persistent (all other files will be lost on session termination):\e[0m"
+  IFS=':' read -ra _paths <<< "$BACKENDAI_PERSISTENT_PATHS"
+  for _p in "${_paths[@]}"; do
+    echo -e "\e[33m   - $_p\e[0m"
+  done
+else
+  echo -e "\e[33m⚠ No persistent storage mounted. All files will be lost when the session is terminated.\e[0m"
+fi
+
 if [[ `uname` == "Linux"  ]]; then
     alias ls="ls --color"
 fi

src/ai/backend/runner/.zshrc

@@ -1,5 +1,15 @@
 export PS1="%F{green}%n@${BACKENDAI_CLUSTER_HOST:-main}%F{yellow}[${BACKENDAI_SESSION_NAME}]%f:%F{blue}%~%f\$ "
 
+if [[ -n "${BACKENDAI_PERSISTENT_PATHS:-}" ]]; then
+  echo "\e[33m⚠ Only the following vfolder paths are persistent (all other files will be lost on session termination):\e[0m"
+  local IFS=':'
+  for _p in ${=BACKENDAI_PERSISTENT_PATHS}; do
+    echo "\e[33m   - $_p\e[0m"
+  done
+else
+  echo "\e[33m⚠ No persistent storage mounted. All files will be lost when the session is terminated.\e[0m"
+fi
+
 # Set up autocompletion
 autoload -Uz compinit
 compinit