From 7bfc864fa312ff7a5ed516853c9669dc2c3260ce Mon Sep 17 00:00:00 2001 From: Alyssa Coghlan Date: Fri, 1 Nov 2024 19:51:55 +1000 Subject: [PATCH] Fix automated layer versioning Closes #24 --- src/venvstacks/pack_venv.py | 20 ++- src/venvstacks/stacks.py | 176 +++++++++++++++++---------- tests/sample_project/venvstacks.toml | 8 +- 3 files changed, 130 insertions(+), 74 deletions(-) diff --git a/src/venvstacks/pack_venv.py b/src/venvstacks/pack_venv.py index d60c282..b95a091 100755 --- a/src/venvstacks/pack_venv.py +++ b/src/venvstacks/pack_venv.py @@ -196,12 +196,15 @@ def _supports_symlinks(target_path: Path) -> bool: def export_venv( source_dir: StrPath, target_dir: StrPath, + sitecustomize_source: StrPath | None = None, run_postinstall: Callable[[Path, Path], None] | None = None, ) -> Path: """Export the given build environment, skipping archive creation and unpacking * injects a suitable `postinstall.py` script for the environment being exported - * excludes __pycache__ folders and package metadata RECORD files + * excludes `__pycache__` folders and package metadata `RECORD` files + * allows build environment `sitecustomize.py` to be replaced with a deployed variant + * excludes `*sitecustomize.py` files from the tree copy if a specific source is given * replaces symlinks with copies on Windows or if the target doesn't support symlinks If supplied, *run_postinstall* is called with the path to the environment's Python @@ -213,7 +216,15 @@ def export_venv( """ source_path = as_normalized_path(source_dir) target_path = as_normalized_path(target_dir) - excluded = shutil.ignore_patterns("__pycache__", "RECORD") + patterns_to_ignore = ["__pycache__", "RECORD"] + if sitecustomize_source is not None: + # The deployed `sitecustomize.py` should be saved alongside the + # build version with a name like `_deployed_sitecustomize.py` + sc_source_path = as_normalized_path(sitecustomize_source) + sc_relative_path = sc_source_path.relative_to(source_path) + sc_target_path = target_path / sc_relative_path.with_name("sitecustomize.py") + patterns_to_ignore.append("*sitecustomize.py") + excluded = shutil.ignore_patterns(*patterns_to_ignore) # Avoid symlinks on Windows, as they need elevated privileges to create # Also avoid them if the target folder doesn't support symlink creation # (that way exports to FAT/FAT32/VFAT file systems should work, even if @@ -228,6 +239,8 @@ def export_venv( symlinks=publish_symlinks, dirs_exist_ok=True, ) + if sitecustomize_source is not None: + shutil.copy2(sc_source_path, sc_target_path) postinstall_path = _inject_postinstall_script(target_path) if run_postinstall is not None: run_postinstall(target_path, postinstall_path) @@ -237,6 +250,7 @@ def export_venv( def create_archive( source_dir: StrPath, archive_base_name: StrPath, + sitecustomize_source: StrPath | None = None, *, install_target: str | None = None, clamp_mtime: datetime | None = None, @@ -261,7 +275,7 @@ def create_archive( install_target = source_path.name with tempfile.TemporaryDirectory(dir=work_dir) as tmp_dir: target_path = Path(tmp_dir) / install_target - env_path = export_venv(source_path, target_path) + env_path = export_venv(source_path, target_path, sitecustomize_source) if not show_progress: def report_progress(_: Any) -> None: diff --git a/src/venvstacks/stacks.py b/src/venvstacks/stacks.py index 5ba244b..fe4c939 100755 --- a/src/venvstacks/stacks.py +++ b/src/venvstacks/stacks.py @@ -27,6 +27,7 @@ Any, ClassVar, Iterable, + Iterator, Literal, Mapping, MutableMapping, @@ -645,6 +646,7 @@ def create_archive( self, env_path: Path, previous_metadata: ArchiveMetadata | None = None, + sitecustomize_source_path: Path | None = None, work_path: Path | None = None, ) -> tuple[ArchiveMetadata, Path]: if env_path.name != self.env_name: @@ -669,6 +671,7 @@ def create_archive( pack_venv.create_archive( env_path, archive_base_path, + sitecustomize_source_path, clamp_mtime=last_locked, work_dir=work_path, install_target=build_metadata["install_target"], @@ -798,6 +801,7 @@ def export_environment( self, env_path: Path, previous_metadata: ExportMetadata | None = None, + sitecustomize_source_path: Path | None = None, ) -> tuple[ExportMetadata, Path]: if env_path.name != self.env_name: err_msg = ( @@ -822,6 +826,7 @@ def _run_postinstall(export_path: Path, postinstall_path: Path) -> None: exported_path = pack_venv.export_venv( env_path, export_path, + sitecustomize_source_path, _run_postinstall, ) assert self.export_path == exported_path # pack_venv ensures this is true @@ -990,6 +995,7 @@ class _PythonEnvironment(ABC): base_python_path: Path | None = field(init=False, repr=False) tools_python_path: Path | None = field(init=False, repr=False) py_version: str = field(init=False, repr=False) + sitecustomize_source_path: Path | None = field(default=None, init=False, repr=False) # Operation flags allow for requested commands to be applied only to selected layers # Notes: @@ -1023,6 +1029,11 @@ def env_name(self) -> EnvNameBuild: def install_target(self) -> EnvNameDeploy: return self.env_lock.get_deployed_name(self.env_spec.env_name) + def get_deployed_path(self, build_path: Path) -> Path: + env_deployed_path = Path(self.install_target) + relative_path = build_path.relative_to(self.env_path) + return env_deployed_path / relative_path + def __post_init__(self) -> None: self.env_path = self.build_path / self.env_name self.pylib_path = self._get_py_scheme_path("purelib") @@ -1313,7 +1324,9 @@ def create_archive( output_path, target_platform, tag_output, previous_metadata, force ) work_path = self.build_path # /tmp is likely too small for ML environments - return build_request.create_archive(env_path, previous_metadata, work_path) + return build_request.create_archive( + env_path, previous_metadata, self.sitecustomize_source_path, work_path + ) def request_export( self, @@ -1339,7 +1352,9 @@ def export_environment( # Define the input metadata that gets published in the export manifest export_request = self.request_export(output_path, previous_metadata, force) - return export_request.export_environment(env_path, previous_metadata) + return export_request.export_environment( + env_path, previous_metadata, self.sitecustomize_source_path + ) class RuntimeEnv(_PythonEnvironment): @@ -1417,11 +1432,13 @@ def create_build_environment(self, *, clean: bool = False) -> None: class _VirtualEnvironment(_PythonEnvironment): _include_system_site_packages = False + base_runtime: RuntimeEnv | None = field(init=False, repr=False) linked_constraints_paths: list[Path] = field(init=False, repr=False) def __post_init__(self) -> None: self.py_version = self.env_spec.runtime.py_version super().__post_init__() + self.base_runtime = None self.linked_constraints_paths = [] @property @@ -1430,13 +1447,17 @@ def env_spec(self) -> _VirtualEnvironmentSpec: assert isinstance(self._env_spec, _VirtualEnvironmentSpec) return self._env_spec - def link_base_runtime_paths(self, runtime: RuntimeEnv) -> None: + def link_base_runtime(self, runtime: RuntimeEnv) -> None: + if self.base_runtime is not None: + raise BuildEnvError(f"Layered environment base runtime already linked {self}") + # Link the runtime environment + self.base_runtime = runtime # Link executable paths self.base_python_path = runtime.python_path self.tools_python_path = runtime.tools_python_path - if self.linked_constraints_paths: - raise BuildEnvError("Layered environment base runtime already linked") + # Link runtime layer dependency constraints self.linked_constraints_paths[:] = [runtime.requirements_path] + print(f"Linked {self}") def get_constraint_paths(self) -> list[Path]: return self.linked_constraints_paths @@ -1464,12 +1485,12 @@ def _ensure_virtual_environment(self) -> subprocess.CompletedProcess[str]: str(self.env_path), ] result = run_python_command(command) - self._link_layered_environment() + self._generate_sitecustomize() fs_sync() print(f"Virtual environment configured in {str(self.env_path)!r}") return result - def _link_layered_environment(self) -> None: + def _generate_sitecustomize(self) -> None: pass # Nothing to do by default, subclasses override if necessary def _update_existing_environment(self, *, lock_only: bool = False) -> None: @@ -1515,8 +1536,6 @@ class ApplicationEnv(_VirtualEnvironment): category = LayerCategories.APPLICATIONS launch_module_name: str = field(init=False, repr=False) - linked_pylib_paths: list[Path] = field(init=False, repr=False) - linked_dynlib_paths: list[Path] = field(init=False, repr=False) linked_frameworks: list[FrameworkEnv] = field(init=False, repr=False) @property @@ -1527,65 +1546,69 @@ def env_spec(self) -> ApplicationSpec: def __post_init__(self) -> None: super().__post_init__() + self.sitecustomize_source_path = self.pylib_path / "_deployed_sitecustomize.py" self.launch_module_name = self.env_spec.launch_module_path.stem - self.linked_pylib_paths = [] - self.linked_dynlib_paths = [] self.linked_frameworks = [] + def _linked_environments(self) -> Iterator[_PythonEnvironment]: + for fw_env in self.linked_frameworks: + yield fw_env + runtime_env = self.base_runtime + assert runtime_env is not None + yield runtime_env + + def _linked_pylib_build_paths(self) -> Iterator[Path]: + for env in self._linked_environments(): + yield env.pylib_path + + def _linked_dynlib_build_paths(self) -> Iterator[Path]: + for env in self._linked_environments(): + dynlib_path = env.dynlib_path + if dynlib_path is not None: + yield dynlib_path + + def _linked_pylib_deployed_paths(self) -> Iterator[Path]: + for env in self._linked_environments(): + yield env.get_deployed_path(env.pylib_path) + + def _linked_dynlib_deployed_paths(self) -> Iterator[Path]: + for env in self._linked_environments(): + dynlib_path = env.dynlib_path + if dynlib_path is not None: + yield env.get_deployed_path(dynlib_path) + def link_layered_environments( self, runtime: RuntimeEnv, frameworks: Mapping[LayerBaseName, FrameworkEnv] ) -> None: - self.link_base_runtime_paths(runtime) + self.link_base_runtime(runtime) constraints_paths = self.linked_constraints_paths if not constraints_paths: raise BuildEnvError("Failed to add base environment constraints path") # The runtime site-packages folder is added here rather than via pyvenv.cfg # to ensure it appears in sys.path after the framework site-packages folders - pylib_paths = self.linked_pylib_paths - dynlib_paths = self.linked_dynlib_paths fw_envs = self.linked_frameworks - if pylib_paths or dynlib_paths or fw_envs: + if fw_envs: raise BuildEnvError("Layered application environment already linked") - for env_spec in self.env_spec.frameworks: - env = frameworks[env_spec.name] - fw_envs.append(env) - constraints_paths.append(env.requirements_path) - install_target_path = Path(env.install_target) - - def _fw_env_path(build_path: Path) -> Path: - relative_path = build_path.relative_to(env.env_path) - return install_target_path / relative_path - - pylib_paths.append(_fw_env_path(env.pylib_path)) - if env.dynlib_path is not None: - dynlib_paths.append(_fw_env_path(env.pylib_path)) - runtime_target_path = Path(runtime.install_target) - - def _runtime_path(build_path: Path) -> Path: - relative_path = build_path.relative_to(runtime.env_path) - return runtime_target_path / relative_path - - pylib_paths.append(_runtime_path(runtime.pylib_path)) - if runtime.dynlib_path is not None: - dynlib_paths.append(_runtime_path(runtime.dynlib_path)) - - def _link_layered_environment(self) -> None: - # Create sitecustomize file - sc_dir_path = self.pylib_path + for fw_env_spec in self.env_spec.frameworks: + fw_env = frameworks[fw_env_spec.name] + fw_envs.append(fw_env) + constraints_paths.append(fw_env.requirements_path) + + @staticmethod + def _render_sitecustomize( + relative_prefix: Path, + pylib_paths: Iterable[Path], + dynlib_paths: Iterable[Path], + ) -> str: sc_contents = [ "# Automatically generated by venvstacks", "import site", "import os", "from os.path import abspath, dirname, join as joinpath", - "# Allow loading modules and packages from framework environments", + "# Allow loading modules and packages from linked environments", "this_dir = dirname(abspath(__file__))", ] - # Add framework and runtime folders to sys.path - parent_path = self.env_path.parent - relative_prefix = Path( - os.path.relpath(str(parent_path), start=str(sc_dir_path)) - ) - for pylib_path in self.linked_pylib_paths: + for pylib_path in pylib_paths: relative_path = relative_prefix / pylib_path sc_contents.extend( [ @@ -1593,31 +1616,54 @@ def _link_layered_environment(self) -> None: "site.addsitedir(path_entry)", ] ) - # Add DLL search folders if needed - dynlib_paths = self.linked_dynlib_paths - if _WINDOWS_BUILD and dynlib_paths: + dynlib_entries: list[str] = [] + for dynlib_path in dynlib_paths: + if not dynlib_path.exists(): + # Nothing added DLLs to this folder at build time, so skip it + continue + relative_path = relative_prefix / dynlib_path + dynlib_entries.extend( + [ + f"dll_dir = abspath(joinpath(this_dir, {str(relative_path)!r}))", + "os.add_dll_directory(dll_dir)", + ] + ) + if dynlib_entries: sc_contents.extend( [ "", "# Allow loading misplaced DLLs on Windows", + *dynlib_entries, ] ) - for dynlib_path in dynlib_paths: - if not dynlib_path.exists(): - # Nothing added DLLs to this folder at build time, so skip it - continue - relative_path = relative_prefix / dynlib_path - sc_contents.extend( - [ - f"dll_dir = abspath(joinpath(this_dir, {str(relative_path)!r}))", - "os.add_dll_directory(dll_dir)", - ] - ) sc_contents.append("") + return "\n".join(sc_contents) + + def _generate_sitecustomize(self) -> None: + # Create build & deployment sitecustomize files + sc_dir_path = self.pylib_path + # Add framework and runtime folders to sys.path + parent_path = self.env_path.parent + relative_prefix = Path( + os.path.relpath(str(parent_path), start=str(sc_dir_path)) + ) + sc_contents = self._render_sitecustomize( + relative_prefix, + self._linked_pylib_build_paths(), + self._linked_dynlib_build_paths(), + ) + deployed_sc_contents = self._render_sitecustomize( + relative_prefix, + self._linked_pylib_deployed_paths(), + self._linked_dynlib_deployed_paths(), + ) sc_path = self.pylib_path / "sitecustomize.py" print(f"Generating {sc_path!r}...") - with open(sc_path, "w", encoding="utf-8") as f: - f.write("\n".join(sc_contents)) + sc_path.write_text(sc_contents) + deployed_sc_path = self.sitecustomize_source_path + assert deployed_sc_path is not None + print(f"Generating {deployed_sc_path!r}...") + deployed_sc_path.write_text(deployed_sc_contents) def _update_existing_environment(self, *, lock_only: bool = False) -> None: super()._update_existing_environment(lock_only=lock_only) @@ -1805,7 +1851,7 @@ def define_build_environment( ) for fw_env in frameworks.values(): runtime = runtimes[fw_env.env_spec.runtime.name] - fw_env.link_base_runtime_paths(runtime) + fw_env.link_base_runtime(runtime) print("Defining application environments:") applications = self._define_envs( build_path, index_config, ApplicationEnv, self.applications diff --git a/tests/sample_project/venvstacks.toml b/tests/sample_project/venvstacks.toml index 9ceb843..17cb920 100644 --- a/tests/sample_project/venvstacks.toml +++ b/tests/sample_project/venvstacks.toml @@ -65,9 +65,7 @@ requirements = [ [[frameworks]] name = "scipy" -# Automatic versioning currently breaks venv layering -# https://github.com/lmstudio-ai/venvstacks/issues/24 -# versioned = true +versioned = true runtime = "cpython@3.11" requirements = [ "scipy", @@ -96,9 +94,7 @@ requirements = [ [[applications]] name = "scipy-import" -# Automatic versioning currently breaks venv layering -# https://github.com/lmstudio-ai/venvstacks/issues/24 -# versioned = true +versioned = true launch_module = "launch_modules/scipy_import.py" frameworks = ["scipy"] requirements = [