enhance(client): add env file for restored runtime (#3088)

client/scripts/sw-docker-entrypoint

@@ -14,6 +14,7 @@ STEP=${SW_TASK_STEP:-""}
 TASK_INDEX=${SW_TASK_INDEX:-0}
 TASK_NUM=${SW_TASK_NUM:-0}
 RUNTIME_RESTORED=${SW_USER_RUNTIME_RESTORED:-0}
+RUNTIME_WORKDIR=${SW_USER_RUNTIME_WORKDIR:-/opt/starwhale.user/runtime}
 
 
 welcome() {
@@ -39,7 +40,8 @@ welcome() {
 prepare(){
     if [ "${RUNTIME_RESTORED}" == "1" ]; then
         echo '-->[Preparing] Runtime has been restored, activate it...'
-        $(bash "${SW_USER_RUNTIME_WORKDIR:-/opt/starwhale.user/runtime}"/activate.sw)
+        source "${RUNTIME_WORKDIR}"/env.sw
+        $(bash "${RUNTIME_WORKDIR}"/activate.sw)
     fi
 
     if [ "${SW_INSTANCE_URI}" != "local" ]; then

client/starwhale/core/runtime/model.py

@@ -727,8 +727,16 @@ def info(self) -> LocalRuntimeVersionInfo | RuntimeInfoVo | None:
         raise NotImplementedError
 
     @classmethod
-    def restore(cls, workdir: Path, isolated_env_dir: t.Optional[Path] = None) -> None:
-        StandaloneRuntime.restore(workdir, isolated_env_dir)
+    def restore(
+        cls,
+        workdir: Path,
+        isolated_env_dir: Path | None = None,
+        verbose: bool = True,
+        runtime_uri: Resource | None = None,
+    ) -> None:
+        StandaloneRuntime.restore(
+            workdir, isolated_env_dir, verbose=verbose, runtime_uri=runtime_uri
+        )
 
     @classmethod
     def get_runtime(cls, uri: Resource) -> Runtime:
@@ -1656,7 +1664,7 @@ def activate(cls, uri: Resource, force_restore: bool = False) -> None:
 
         if force_restore or not prefix_path.exists() or is_invalid_status:
             console.print(f":safety_vest: restore runtime into {workdir}")
-            cls.restore(workdir)
+            cls.restore(workdir, runtime_uri=uri)
 
         console.print(f":carrot: activate the current shell for the runtime uri: {uri}")
         activate_python_env(
@@ -1961,8 +1969,9 @@ def _build(_manifest: t.Dict[str, t.Any]) -> None:
     def restore(
         cls,
         workdir: Path,
-        isolated_env_dir: t.Optional[Path] = None,
+        isolated_env_dir: Path | None = None,
         verbose: bool = True,
+        runtime_uri: Resource | None = None,
     ) -> None:
         if not (workdir.exists() and (workdir / DEFAULT_MANIFEST_NAME).exists()):
             raise NoSupportError("only support swrt extract workdir")
@@ -2064,6 +2073,7 @@ def restore(
                             "local_packaged_env", False
                         ),
                         verbose=verbose,
+                        runtime_uri=runtime_uri.full_uri if runtime_uri else "",
                     ),
                 ),
             ]

client/starwhale/core/runtime/process.py

@@ -161,7 +161,7 @@ def _restore_runtime(
         if force_restore or not is_valid_prefix or is_invalid_status:
             console.print(f":snail: start to restore runtime: {_uri}")
             extract_tar(tar_path=bundle_path, dest_dir=snapshot_workdir, force=True)
-            StandaloneRuntime.restore(snapshot_workdir, verbose=False)
+            StandaloneRuntime.restore(snapshot_workdir, verbose=False, runtime_uri=_uri)
 
         if venv_prefix.exists():
             prefix = venv_prefix

client/starwhale/core/runtime/view.py

@@ -366,7 +366,7 @@ def restore(cls, target: str) -> None:
             workdir = _runtime.store.snapshot_workdir
             if not workdir.exists():
                 _runtime.extract(force=True, target=workdir)
-            Runtime.restore(workdir)
+            Runtime.restore(workdir, runtime_uri=_uri)
             console.print(
                 f":ramen: runtime(from uri:{_uri}) has been restored, activate it in the current shell: \n"
                 f"\t :stars: run command: [bold green]swcli runtime activate {target}[/]"

client/starwhale/utils/venv.py

@@ -207,13 +207,23 @@ def render_python_env_activate(
     workdir: Path,
     local_packaged_env: bool = False,
     verbose: bool = True,
+    runtime_uri: str = "",
 ) -> None:
     if mode not in (PythonRunEnv.CONDA, PythonRunEnv.VENV):
         raise NoSupportError(f"mode({mode}) render python env activate scripts")
 
+    envs = {
+        SWShellActivatedRuntimeEnv.MODE: mode,
+        SWShellActivatedRuntimeEnv.PREFIX: str(prefix_path),
+        SWShellActivatedRuntimeEnv.URI: runtime_uri,
+    }
+    ensure_file(workdir / "env.sw", "\n".join({f"{k}={v}" for k, v in envs.items()}))
+
     if local_packaged_env:
         # conda local mode(conda-pack) should be activated by the source command.
-        venv_activate_render(prefix_path, workdir, relocate=mode == PythonRunEnv.VENV)
+        venv_activate_render(
+            prefix_path, workdir, relocate=mode == PythonRunEnv.VENV, verbose=verbose
+        )
     else:
         if mode == PythonRunEnv.CONDA:
             conda_activate_render(prefix_path, workdir, verbose=verbose)

client/tests/core/test_model.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import json
 import typing as t
@@ -461,7 +463,9 @@ def _run_without_package_runtime(obj: t.Any) -> None:
 
         m_env_mode.return_value = "conda"
 
-        def _restore(workdir: Path, verbose: bool = False) -> None:
+        def _restore(
+            workdir: Path, verbose: bool = False, runtime_uri: Resource | None = None
+        ) -> None:
             ensure_dir(workdir / "export/conda")
 
         m_restore.side_effect = _restore

client/tests/core/test_runtime.py

@@ -1884,6 +1884,13 @@ def test_restore_venv(
         os.environ["SW_CONTAINER"] = "1"
         Runtime.restore(Path(workdir))
 
+        env_fpath = Path(workdir) / "env.sw"
+        assert env_fpath.exists()
+        envs = env_fpath.read_text().split("\n")
+        assert "SW_ACTIVATED_RUNTIME_URI_IN_SHELL=" in envs
+        assert "SW_ACTIVATED_RUNTIME_MODE_IN_SHELL=venv" in envs
+        assert f"SW_ACTIVATED_RUNTIME_PREFIX_IN_SHELL={workdir}/export/venv" in envs
+
         assert m_command_call.call_count == 2
         assert m_command_call.call_args_list[0][0][0] == "apt-get install xxx"
         assert m_command_call.call_args_list[1][0][0] == "echo 'helloworld'"
@@ -2273,7 +2280,15 @@ def test_restore_conda(
 
         m_machine.return_value = "arm64"
         os.environ[ENV_LOG_LEVEL] = "TRACE"
-        Runtime.restore(Path(workdir))
+        runtime_uri = Resource("rttest", typ=ResourceType.runtime, refine=False)
+        Runtime.restore(Path(workdir), runtime_uri=runtime_uri)
+
+        env_fpath = Path(workdir) / "env.sw"
+        assert env_fpath.exists()
+        envs = env_fpath.read_text().split("\n")
+        assert f"SW_ACTIVATED_RUNTIME_URI_IN_SHELL={runtime_uri}" in envs
+        assert "SW_ACTIVATED_RUNTIME_MODE_IN_SHELL=conda" in envs
+        assert f"SW_ACTIVATED_RUNTIME_PREFIX_IN_SHELL={workdir}/export/conda" in envs
 
         assert "CONDARC" not in os.environ
         assert m_call.call_count == 8