diff --git a/changes/9738.feature.md b/changes/9738.feature.md new file mode 100644 index 00000000000..1388dd21638 --- /dev/null +++ b/changes/9738.feature.md @@ -0,0 +1 @@ +Add `BACKENDAI_PERSISTENT_PATHS` environment variable in containers to display persistent vfolder mount paths on shell startup. diff --git a/src/ai/backend/agent/agent.py b/src/ai/backend/agent/agent.py index 0b7fad476c3..bb292bbe52a 100644 --- a/src/ai/backend/agent/agent.py +++ b/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) diff --git a/src/ai/backend/agent/stage/kernel_lifecycle/docker/environ.py b/src/ai/backend/agent/stage/kernel_lifecycle/docker/environ.py index cdc0a6e729a..023a0b382b1 100644 --- a/src/ai/backend/agent/stage/kernel_lifecycle/docker/environ.py +++ b/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: diff --git a/src/ai/backend/runner/.bashrc b/src/ai/backend/runner/.bashrc index ca736a42d71..c5b76e8940d 100644 --- a/src/ai/backend/runner/.bashrc +++ b/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 diff --git a/src/ai/backend/runner/.zshrc b/src/ai/backend/runner/.zshrc index 56206ae5ef9..3c63d86b779 100644 --- a/src/ai/backend/runner/.zshrc +++ b/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