From 12a180f41ce06145cdf1ddb89e975edea7fc07f7 Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Thu, 27 Mar 2025 11:04:16 +0100
Subject: [PATCH] finish codegen venv manager
---
cmake/ManageCodegenVenv.py | 88 ++++++++++++++++++++++++++-----
cmake/WalberlaCodegen.cmake | 54 +++++++++++++------
cmake/codegen-requirements.txt.in | 4 ++
cmake/test/reqs.txt | 2 +
cmake/test/state.json | 1 +
tests/CMakeLists.txt | 3 ++
6 files changed, 123 insertions(+), 29 deletions(-)
create mode 100644 cmake/test/reqs.txt
create mode 100644 cmake/test/state.json
diff --git a/cmake/ManageCodegenVenv.py b/cmake/ManageCodegenVenv.py
index 77f515d..33301ac 100644
--- a/cmake/ManageCodegenVenv.py
+++ b/cmake/ManageCodegenVenv.py
@@ -5,8 +5,9 @@ import sys
import subprocess
import json
import shutil
+import hashlib
from contextlib import contextmanager
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass, asdict, field
from argparse import ArgumentParser
from pathlib import Path
@@ -17,6 +18,14 @@ class VenvState:
main_requirements_file: str | None = None
initialized: bool = False
+ user_requirements: list[str] = field(default_factory=list)
+ user_requirements_hash: str | None = None
+
+ @property
+ def venv_python(self) -> Path:
+ assert self.venv_dir is not None
+ return Path(self.venv_dir) / "bin" / "python"
+
@staticmethod
@contextmanager
def lock(statefile: Path) -> Generator[VenvState, None, None]:
@@ -41,31 +50,79 @@ class VenvState:
statefile_bak.replace(statefile)
+def reinitialize(state: VenvState):
+ assert state.venv_dir is not None
+
+ venv_dir = Path(state.venv_dir)
+ if venv_dir.exists():
+ shutil.rmtree(venv_dir)
+
+ base_py = Path(sys.executable)
+
+ # Create the virtual environment
+ venv_args = [base_py, "-m", "venv", state.venv_dir]
+ subprocess.run(venv_args).check_returncode()
+
+ # Install base requirements
+ venv_py = str(state.venv_python)
+ install_args = [venv_py, "-m", "pip", "install", "-r", state.main_requirements_file]
+ subprocess.run(install_args).check_returncode()
+
+ state.initialized = True
+
+
def action_initialize(args):
statefile = Path(args.statefile)
with VenvState.lock(statefile) as state:
if not state.initialized:
p_venv_dir = Path(args.venv_dir).resolve()
- if p_venv_dir.exists():
- shutil.rmtree(p_venv_dir)
reqs_file = Path(args.requirements_file).resolve()
state.venv_dir = str(p_venv_dir)
state.main_requirements_file = str(reqs_file)
- base_py = Path(sys.executable)
+ reinitialize(state)
+
+ # Reset user requirements
+ state.user_requirements = []
+
+
+def action_require(args):
+ statefile = Path(args.statefile)
+
+ with VenvState.lock(statefile) as state:
+ if not state.initialized:
+ raise RuntimeError("Virtual environment is not initialized.")
+
+ for req in args.requirements:
+ state.user_requirements.append(req)
+
- # Create the virtual environment
- venv_args = [base_py, "-m", "venv", state.venv_dir]
- subprocess.run(venv_args).check_returncode()
+def action_populate(args):
+ statefile = Path(args.statefile)
+
+ with VenvState.lock(statefile) as state:
+ h = hashlib.sha256()
+ for req in state.user_requirements:
+ h.update(bytes(req, encoding="utf8"))
+ digest = h.hexdigest()
- # Install base requirements
- venv_py = Path(state.venv_dir).absolute() / "bin" / "python"
- install_args = [venv_py, "-m", "pip", "install", "-r", state.main_requirements_file]
- subprocess.run(install_args).check_returncode()
+ if digest != state.user_requirements_hash:
+ if state.user_requirements_hash is not None:
+ # Populate was run before -> rebuild entire environment
+ print(
+ "User requirements have changed - rebuilding virtual environment.",
+ flush=True
+ )
+ reinitialize(state)
- state.initialized = True
+ pip_args = [str(state.venv_python), "-m", "pip", "install"]
+ for req in state.user_requirements:
+ install_args = pip_args.copy() + [req]
+ subprocess.run(install_args).check_returncode()
+
+ state.user_requirements_hash = digest
def main():
@@ -93,6 +150,13 @@ def main():
)
parser_initialize.set_defaults(func=action_initialize)
+ parser_require = subparsers.add_parser("require")
+ parser_require.add_argument("requirements", nargs="+", type=str)
+ parser_require.set_defaults(func=action_require)
+
+ parser_populate = subparsers.add_parser("populate")
+ parser_populate.set_defaults(func=action_populate)
+
args = parser.parse_args()
args.func(args)
diff --git a/cmake/WalberlaCodegen.cmake b/cmake/WalberlaCodegen.cmake
index 435af24..ceafb3e 100644
--- a/cmake/WalberlaCodegen.cmake
+++ b/cmake/WalberlaCodegen.cmake
@@ -7,20 +7,25 @@ set(
)
if( WALBERLA_CODEGEN_PRIVATE_VENV )
+ find_package( Python COMPONENTS Interpreter REQUIRED )
+
set(WALBERLA_CODEGEN_VENV_PATH ${CMAKE_CURRENT_BINARY_DIR}/codegen-venv CACHE PATH "Location of the virtual environment used for code generation")
set(_venv_python_exe ${WALBERLA_CODEGEN_VENV_PATH}/bin/python)
- set(
- _WALBERLA_CODEGEN_VENV_MANAGER
- ${sfg_walberla_SOURCE_DIR}/cmake/ManageCodegenVenv.py
- CACHE INTERNAL
- "venv manager filepath - for internal use only"
- )
set(
_WALBERLA_CODEGEN_VENV_STATEFILE
${CMAKE_CURRENT_BINARY_DIR}/walberla-venv-state.json
CACHE INTERNAL
"venv statefile - for internal use only"
)
+ set(
+ _WALBERLA_CODEGEN_VENV_INVOKE_MANAGER
+ ${Python_EXECUTABLE}
+ ${sfg_walberla_SOURCE_DIR}/cmake/ManageCodegenVenv.py
+ -s
+ ${_WALBERLA_CODEGEN_VENV_STATEFILE}
+ CACHE INTERNAL
+ "venv manager filepath - for internal use only"
+ )
set(
WALBERLA_CODEGEN_VENV_REQUIREMENTS
@@ -30,7 +35,6 @@ if( WALBERLA_CODEGEN_PRIVATE_VENV )
)
mark_as_advanced(WALBERLA_CODEGEN_VENV_REQUIREMENTS)
- find_package( Python COMPONENTS Interpreter REQUIRED )
set(
_requirements_file
@@ -44,9 +48,7 @@ if( WALBERLA_CODEGEN_PRIVATE_VENV )
execute_process(
COMMAND
- ${Python_EXECUTABLE}
- ${_WALBERLA_CODEGEN_VENV_MANAGER}
- -s ${_WALBERLA_CODEGEN_VENV_STATEFILE}
+ $CACHE{_WALBERLA_CODEGEN_VENV_INVOKE_MANAGER}
init
${WALBERLA_CODEGEN_VENV_PATH}
${_requirements_file}
@@ -109,21 +111,39 @@ configure_file(
message( STATUS "Wrote project-wide code generator configuration to ${WALBERLA_CODEGEN_CONFIG_MODULE}" )
-#[[
-Run `pip install` in the code generation environment with the given arguments.
-#]]
-function(walberla_codegen_venv_install)
+function(walberla_codegen_venv_require)
if(NOT WALBERLA_CODEGEN_PRIVATE_VENV)
return()
endif()
- if(NOT _sfg_private_venv_done)
- message( FATAL_ERROR "The private virtual environment for code generation was not initialized yet" )
+ execute_process(
+ COMMAND
+ $CACHE{_WALBERLA_CODEGEN_VENV_INVOKE_MANAGER}
+ require
+ ${ARGV}
+ RESULT_VARIABLE _lastResult
+ )
+
+ if( ${_lastResult} )
+ message( FATAL_ERROR "venv-require: Operation failed" )
+ endif()
+endfunction()
+
+function(walberla_codegen_venv_populate)
+ if(NOT WALBERLA_CODEGEN_PRIVATE_VENV)
+ return()
endif()
execute_process(
- COMMAND ${WALBERLA_CODEGEN_PYTHON} -m pip install ${ARGV}
+ COMMAND
+ $CACHE{_WALBERLA_CODEGEN_VENV_INVOKE_MANAGER}
+ populate
+ RESULT_VARIABLE _lastResult
)
+
+ if( ${_lastResult} )
+ message( FATAL_ERROR "venv-populate: Operation failed" )
+ endif()
endfunction()
# Code Generation Functions
diff --git a/cmake/codegen-requirements.txt.in b/cmake/codegen-requirements.txt.in
index 07abf2e..3d2115e 100644
--- a/cmake/codegen-requirements.txt.in
+++ b/cmake/codegen-requirements.txt.in
@@ -1,8 +1,12 @@
+# pystencils
git+https://i10git.cs.fau.de/pycodegen/pystencils.git@v2.0-dev
+# lbmpy
git+https://i10git.cs.fau.de/pycodegen/lbmpy.git@fhennig/pystencils2.0-compat
+# pystencils-sfg
git+https://i10git.cs.fau.de/pycodegen/pystencils-sfg.git
+# walberla-codegen
-e ${sfg_walberla_SOURCE_DIR}
diff --git a/cmake/test/reqs.txt b/cmake/test/reqs.txt
new file mode 100644
index 0000000..9c593dc
--- /dev/null
+++ b/cmake/test/reqs.txt
@@ -0,0 +1,2 @@
+pystencils
+lbmpy
\ No newline at end of file
diff --git a/cmake/test/state.json b/cmake/test/state.json
new file mode 100644
index 0000000..0de44ea
--- /dev/null
+++ b/cmake/test/state.json
@@ -0,0 +1 @@
+{"venv_dir": "/media/data/fhennig/research-hpc/projects/2024_pystencils_nbackend/sfg-walberla/cmake/test/venv", "main_requirements_file": "/media/data/fhennig/research-hpc/projects/2024_pystencils_nbackend/sfg-walberla/cmake/test/reqs.txt", "initialized": true, "user_requirements": ["py-cpuinfo", "pyyaml"], "user_requirements_hash": "bebb0df474f163f05d50504ee0efc07ba064e08aabeb7b2c745d8f57200be4ac"}
\ No newline at end of file
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 34ada83..312a02f 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -44,6 +44,9 @@ endif()
add_subdirectory(${CMAKE_SOURCE_DIR}/.. ${CMAKE_BINARY_DIR}/sfg-walberla)
+walberla_codegen_venv_require( py-cpuinfo )
+walberla_codegen_venv_populate()
+
# Test Directories
include(CTest)
--
GitLab