ft: added environment handlers

floatcsep/environments.py

@@ -0,0 +1,317 @@
+import logging
+import sys
+import os
+import subprocess
+import hashlib
+import shutil
+import venv
+import configparser
+from abc import ABC, abstractmethod
+from packaging.specifiers import SpecifierSet
+
+log = logging.getLogger("floatLogger")
+
+
+class EnvironmentManager(ABC):
+
+    @abstractmethod
+    def __init__(self, base_name: str, model_directory: str):
+        self.base_name = base_name
+        self.model_directory = model_directory
+
+    @abstractmethod
+    def create_environment(self, force=False):
+        pass
+
+    @abstractmethod
+    def env_exists(self):
+        pass
+
+    @abstractmethod
+    def run_command(self, command):
+        pass
+
+    @abstractmethod
+    def install_dependencies(self):
+        pass
+
+    def generate_env_name(self) -> str:
+        # Generate a hash from the model directory
+        dir_hash = hashlib.md5(self.model_directory.encode()).hexdigest()[
+            :8
+        ]  # Shorten the hash
+        return f"{self.base_name}_{dir_hash}"
+
+
+class CondaEnvironmentManager(EnvironmentManager):
+    def __init__(self, base_name: str, model_directory: str):
+        self.base_name = base_name
+        self.model_directory = model_directory
+        self.env_name = self.generate_env_name()
+        self.package_manager = self.detect_package_manager()
+
+    @staticmethod
+    def detect_package_manager():
+        if shutil.which("mamba"):
+            log.info("Mamba detected, using mamba as package manager.")
+            return "mamba"
+        log.info("Mamba not detected, using conda as package manager.")
+        return "conda"
+
+    def create_environment(self, force=False):
+        if force and self.env_exists():
+            log.info(f"Removing existing conda environment: {self.env_name}")
+            subprocess.run(
+                [
+                    self.package_manager,
+                    "env",
+                    "remove",
+                    "--name",
+                    self.env_name,
+                    "--yes",
+                ]
+            )
+
+        if not self.env_exists():
+            env_file = os.path.join(self.model_directory, "environment.yml")
+            if os.path.exists(env_file):
+                log.info(
+                    f"Creating sub-conda environment {self.env_name} from environment.yml"
+                )
+                subprocess.run(
+                    [
+                        self.package_manager,
+                        "env",
+                        "create",
+                        "--name",
+                        self.env_name,
+                        "--file",
+                        env_file,
+                    ]
+                )
+            else:
+                python_version = self.detect_python_version()
+                log.info(
+                    f"Creating sub-conda environment {self.env_name} with Python {python_version}"
+                )
+                subprocess.run(
+                    [
+                        self.package_manager,
+                        "create",
+                        "--name",
+                        self.env_name,
+                        "--yes",
+                        f"python={python_version}",
+                    ]
+                )
+            log.info(f"\tSub-conda environment created: {self.env_name}")
+
+        self.install_dependencies()
+
+    def env_exists(self) -> bool:
+        result = subprocess.run(["conda", "env", "list"], stdout=subprocess.PIPE)
+        return self.env_name in result.stdout.decode()
+
+    def detect_python_version(self) -> str:
+        setup_py = os.path.join(self.model_directory, "setup.py")
+        pyproject_toml = os.path.join(self.model_directory, "pyproject.toml")
+        setup_cfg = os.path.join(self.model_directory, "setup.cfg")
+        current_python_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+
+        def parse_version(version_str):
+            # Extract the first valid version number
+            import re
+
+            match = re.search(r"\d+(.\d+)*", version_str)
+            return match.group(0) if match else current_python_version
+
+        def is_version_compatible(requirement, current_version):
+            try:
+                specifier = SpecifierSet(requirement)
+                return current_version in specifier
+            except Exception as e:
+                log.error(f"Invalid specifier: {requirement}. Error: {e}")
+                return False
+
+        if os.path.exists(setup_py):
+            with open(setup_py) as f:
+                for line in f:
+                    if "python_requires" in line:
+                        required_version = line.split("=")[1].strip()
+                        if is_version_compatible(
+                            required_version, current_python_version
+                        ):
+                            log.info(
+                                f"Using current Python version: {current_python_version}"
+                            )
+                            return current_python_version
+                        return parse_version(required_version)
+
+        if os.path.exists(pyproject_toml):
+            with open(pyproject_toml) as f:
+                for line in f:
+                    if "python" in line and "=" in line:
+                        required_version = line.split("=")[1].strip()
+                        if is_version_compatible(
+                            required_version, current_python_version
+                        ):
+                            log.info(
+                                f"Using current Python version: {current_python_version}"
+                            )
+                            return current_python_version
+                        return parse_version(required_version)
+
+        if os.path.exists(setup_cfg):
+            config = configparser.ConfigParser()
+            config.read(setup_cfg)
+            if "options" in config and "python_requires" in config["options"]:
+                required_version = config["options"]["python_requires"].strip()
+                if is_version_compatible(required_version, current_python_version):
+                    log.info(f"Using current Python version: {current_python_version}")
+                    return current_python_version
+                return parse_version(required_version)
+
+        return current_python_version
+
+    def install_dependencies(self):
+        log.info(f"Installing dependencies in conda environment: {self.env_name}")
+        cmd = [
+            self.package_manager,
+            "run",
+            "-n",
+            self.env_name,
+            "pip",
+            "install",
+            "-e",
+            self.model_directory,
+        ]
+        subprocess.run(cmd, check=True)
+
+    def run_command(self, command):
+        cmd = ["bash", "-c", f"{self.package_manager} run -n {self.env_name} {command}"]
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+        )
+        for line in process.stdout:
+            log.info(f"\t{line[:-1]}")
+        process.wait()
+
+
+class VenvEnvironmentManager(EnvironmentManager):
+    def __init__(self, base_name: str, model_directory: str):
+        self.base_name = base_name
+        self.model_directory = model_directory
+        self.env_name = self.generate_env_name()
+        self.env_path = os.path.join(model_directory, self.env_name)
+
+    def create_environment(self, force=False):
+        if force and self.env_exists():
+            log.info(f"Removing existing virtual environment: {self.env_name}")
+            shutil.rmtree(self.env_path)
+
+        if not self.env_exists():
+            log.info(f"Creating virtual environment: {self.env_name}")
+            venv.create(self.env_path, with_pip=True)
+            log.info(f"\tVirtual environment created: {self.env_name}")
+            self.install_dependencies()
+
+    def env_exists(self) -> bool:
+        return os.path.isdir(self.env_path)
+
+    def install_dependencies(self):
+        log.info(f"Installing dependencies in virtual environment: {self.env_name}")
+        pip_executable = os.path.join(self.env_path, "bin", "pip")
+        cmd = f"{pip_executable} install -e {os.path.abspath(self.model_directory)}"
+        self.run_command(cmd)
+
+    def run_command(self, command):
+        env = os.environ.copy()
+        env.pop("PYTHONPATH", None)
+        process = subprocess.Popen(
+            command,
+            shell=True,
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+        )
+        for line in process.stdout:
+            log.info(line.strip())
+        process.wait()
+
+
+class DockerEnvironmentManager(EnvironmentManager):
+
+    def __init__(self, base_name: str, model_directory: str):
+        self.base_name = base_name
+        self.model_directory = model_directory
+
+    def create_environment(self, force=False):
+        pass
+
+    def env_exists(self):
+        pass
+
+    def run_command(self, command):
+        pass
+
+    def install_dependencies(self):
+        pass
+
+
+class EnvironmentFactory:
+
+    @staticmethod
+    def get_env(
+        build: str = None, model_name: str = "model", model_path: str = None
+    ) -> EnvironmentManager:
+
+        run_env = EnvironmentFactory.check_environment_type()
+        if run_env != build and build and build != "docker":
+            log.warning(
+                f"Selected build environment ({build}) for this model is different than that of"
+                f" the experiment run. Consider selecting the same environment."
+            )
+        if build == "conda" or (not build and run_env == "conda"):
+            return CondaEnvironmentManager(
+                base_name=f"{model_name}", model_directory=os.path.abspath(model_path)
+            )
+        elif build == "venv" or (not build and run_env == "venv"):
+            return VenvEnvironmentManager(
+                base_name=f"{model_name}", model_directory=os.path.abspath(model_path)
+            )
+        elif build == "docker":
+            return DockerEnvironmentManager(
+                base_name=f"{model_name}", model_directory=os.path.abspath(model_path)
+            )
+        else:
+            raise Exception(
+                "Wrong environment selection. Please choose between "
+                '"conda", "venv" or "docker".'
+            )
+
+    @staticmethod
+    def check_environment_type():
+        if "VIRTUAL_ENV" in os.environ:
+            return "venv"
+        try:
+            subprocess.run(
+                ["conda", "info"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            return "conda"
+        except FileNotFoundError:
+            pass
+        return None
+
+
+if __name__ == "__main__":
+
+    env = EnvironmentFactory.get_env(
+        "conda", model_path="../examples/case_h/models/pymock_poisson"
+    )
+    env.create_environment(force=True)

floatcsep/model.py

@@ -1,7 +1,13 @@
 import json
 import logging
+import sys
 import os
 import subprocess
+import hashlib
+import shutil
+import venv
+import configparser
+
 from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import List, Callable, Union, Mapping, Sequence
@@ -598,3 +604,6 @@ def create_model(model_cfg) -> Model:
 
         else:
             return TimeIndependentModel.from_dict(model_cfg)
+
+
+