Skip to content

Commit

Permalink
ft: Created environments module, which handle the creation, staging o…
Browse files Browse the repository at this point in the history 
…f conda/pip/docker environments, as well as the running a source code within.

tests: added tests for EnvironmentHandler classess and subclasses
build: added packaging as requirement. Harmonized requirements.txt and requirements_dev.txt
  • Loading branch information
@pabloitu
pabloitu committed Jul 30, 2024
1 parent 6bd1b5c commit 7fdde2d
Show file tree
Hide file tree
Showing 7 changed files with 221 additions and 62 deletions.
204 changes: 178 additions & 26 deletions floatcsep/environments.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,52 +13,120 @@


class EnvironmentManager(ABC):
"""
Abstract base class for managing different types of environments.
This class defines the interface for creating, checking existence,
running commands, and installing dependencies in various environment types.
"""

@abstractmethod
def __init__(self, base_name: str, model_directory: str):
"""
Initializes the environment manager with a base name and model directory.
Args:
base_name (str): The base name for the environment.
model_directory (str): The directory containing the model files.
"""
self.base_name = base_name
self.model_directory = model_directory

@abstractmethod
def create_environment(self, force=False):
"""
Creates the environment. If 'force' is True, it will remove any existing
environment with the same name before creating a new one.
Args:
force (bool): Whether to forcefully remove an existing environment.
"""
pass

@abstractmethod
def env_exists(self):
"""
Checks if the environment already exists.
Returns:
bool: True if the environment exists, False otherwise.
"""
pass

@abstractmethod
def run_command(self, command):
"""
Executes a command within the context of the environment.
Args:
command (str): The command to be executed.
"""
pass

@abstractmethod
def install_dependencies(self):
"""
Installs the necessary dependencies for the environment based on the
specified configuration or requirements.
"""
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
"""
Generates a unique environment name by hashing the model directory
and appending it to the base name.
Returns:
str: A unique name for the environment.
"""
dir_hash = hashlib.md5(self.model_directory.encode()).hexdigest()[:8]
return f"{self.base_name}_{dir_hash}"


class CondaEnvironmentManager(EnvironmentManager):
"""
Manages a conda (or mamba) environment, providing methods to create, check,
and manipulate conda environments specifically.
"""

def __init__(self, base_name: str, model_directory: str):
"""
Initializes the Conda environment manager with the specified base name
and model directory. It also generates the environment name and detects
the package manager (conda or mamba) to install dependencies..
Args:
base_name (str): The base name, i.e., model name, for the conda environment.
model_directory (str): The directory containing the model files.
"""
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():
"""
Detects whether 'mamba' or 'conda' is available as the package manager.
Returns:
str: The name of the detected package manager ('mamba' or 'conda').
"""
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):
"""
Creates a conda environment using either an environment.yml file or
the specified Python version in setup.py/setup.cfg or project/toml.
If 'force' is True, any existing environment with the same name will
be removed first.
Args:
force (bool): Whether to forcefully remove an existing environment.
"""
if force and self.env_exists():
log.info(f"Removing existing conda environment: {self.env_name}")
subprocess.run(
Expand All @@ -75,9 +143,7 @@ def create_environment(self, force=False):
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"
)
log.info(f"Creating sub-conda environment {self.env_name} from environment.yml")
subprocess.run(
[
self.package_manager,
Expand Down Expand Up @@ -109,14 +175,31 @@ def create_environment(self, force=False):
self.install_dependencies()

def env_exists(self) -> bool:
"""
Checks if the conda environment exists by querying the list of
existing conda environments.
Returns:
bool: True if the conda environment exists, False otherwise.
"""
result = subprocess.run(["conda", "env", "list"], stdout=subprocess.PIPE)
return self.env_name in result.stdout.decode()

def detect_python_version(self) -> str:
"""
Determines the required Python version from setup files in the model directory.
It checks 'setup.py', 'pyproject.toml', and 'setup.cfg' (in that order), for
version specifications.
Returns:
str: The detected or default Python version.
"""
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}"
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
Expand All @@ -138,12 +221,8 @@ def is_version_compatible(requirement, current_version):
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}"
)
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)

Expand All @@ -152,12 +231,8 @@ def is_version_compatible(requirement, current_version):
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}"
)
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)

Expand All @@ -174,6 +249,10 @@ def is_version_compatible(requirement, current_version):
return current_python_version

def install_dependencies(self):
"""
Installs dependencies in the conda environment using pip, based on the
model setup file
"""
log.info(f"Installing dependencies in conda environment: {self.env_name}")
cmd = [
self.package_manager,
Expand All @@ -188,7 +267,16 @@ def install_dependencies(self):
subprocess.run(cmd, check=True)

def run_command(self, command):
cmd = ["bash", "-c", f"{self.package_manager} run -n {self.env_name} {command}"]
"""
Runs a specified command within the conda environment
Args:
command (str): The command to be executed in the conda environment.
"""
cmd = [
"bash",
"-c",
f"{self.package_manager} run -n {self.env_name} {command}",
]
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
Expand All @@ -201,13 +289,34 @@ def run_command(self, command):


class VenvEnvironmentManager(EnvironmentManager):
"""
Manages a virtual environment created using Python's venv module.
Provides methods to create, check, and manipulate virtual environments.
"""

def __init__(self, base_name: str, model_directory: str):
"""
Initializes the virtual environment manager with the specified base name
and model directory.
Args:
base_name (str): The base name (i.e., model name) for the virtual environment.
model_directory (str): The directory containing the model files.
"""

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):
"""
Creates a virtual environment in the specified model directory. If 'force'
is True, any existing virtual environment will be removed before creation.
Args:
force (bool): Whether to forcefully remove an existing virtual environment.
"""
if force and self.env_exists():
log.info(f"Removing existing virtual environment: {self.env_name}")
shutil.rmtree(self.env_path)
Expand All @@ -219,15 +328,31 @@ def create_environment(self, force=False):
self.install_dependencies()

def env_exists(self) -> bool:
"""
Checks if the virtual environment exists by verifying the presence of its directory.
Returns:
bool: True if the virtual environment exists, False otherwise.
"""
return os.path.isdir(self.env_path)

def install_dependencies(self):
"""
Installs dependencies in the virtual environment using pip, based on the
model directory's configuration.
"""
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):
"""
Executes a specified command in the virtual environment and logs the output.
Args:
command (str): The command to be executed in the virtual environment.
"""
env = os.environ.copy()
env.pop("PYTHONPATH", None)
process = subprocess.Popen(
Expand All @@ -239,11 +364,17 @@ def run_command(self, command):
universal_newlines=True,
)
for line in process.stdout:
log.info(line.strip())
stripped_line = line.strip()
print(f"Logging: {stripped_line}") # Debug statement
log.info(stripped_line)
process.wait()


class DockerEnvironmentManager(EnvironmentManager):
"""
Manages a Docker environment, providing methods to create, check,
and manipulate Docker containers for the environment.
"""

def __init__(self, base_name: str, model_directory: str):
self.base_name = base_name
Expand All @@ -263,12 +394,30 @@ def install_dependencies(self):


class EnvironmentFactory:
"""
Factory class for creating instances of environment managers based on the specified type.
"""

@staticmethod
def get_env(
build: str = None, model_name: str = "model", model_path: str = None
) -> EnvironmentManager:

"""
Returns an instance of an environment manager based on the specified build type.
It checks the current environment type and can return a conda, venv, or Docker
environment manager.
Args:
build (str): The desired type of environment ('conda', 'venv', or 'docker').
model_name (str): The name of the model for which the environment is being created.
model_path (str): The path to the model directory.
Returns:
EnvironmentManager: An instance of the appropriate environment manager.
Raises:
Exception: If an invalid environment type is specified.
"""
run_env = EnvironmentFactory.check_environment_type()
if run_env != build and build and build != "docker":
log.warning(
Expand All @@ -277,15 +426,18 @@ def get_env(
)
if build == "conda" or (not build and run_env == "conda"):
return CondaEnvironmentManager(
base_name=f"{model_name}", model_directory=os.path.abspath(model_path)
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)
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)
base_name=f"{model_name}",
model_directory=os.path.abspath(model_path),
)
else:
raise Exception(
Expand Down
Loading

0 comments on commit 7fdde2d

Please sign in to comment.