234 lines
9.1 KiB
Python
234 lines
9.1 KiB
Python
import shutil
|
||
from pathlib import Path
|
||
from typing import Optional, List
|
||
from dataclasses import dataclass
|
||
from .Environments.venv import Venv
|
||
|
||
|
||
@dataclass
|
||
class InstanceConfig:
|
||
"""Конфигурация экземпляра приложения"""
|
||
path: Path
|
||
requirements_dir: Path
|
||
manual_requirements: Path
|
||
freezed_packages: Path
|
||
created: bool = False
|
||
|
||
|
||
class Instance:
|
||
"""Базовый класс для управления изолированными окружениями"""
|
||
|
||
def __init__(self, path: str, env: str = "venv"):
|
||
self.path = Path(path)
|
||
self.config = InstanceConfig(
|
||
path=self.path,
|
||
requirements_dir=self.path / ".vaiola" / "requirements.d",
|
||
manual_requirements=self.path / ".vaiola" / "requirements.txt",
|
||
freezed_packages=self.path / ".vaiola" / "freezed_packages.txt",
|
||
created=(self.path / ".vaiola").exists()
|
||
)
|
||
|
||
self.freezed_packages_list: List[str] = []
|
||
|
||
if env == "venv":
|
||
self.env = Venv(self.path / "venv")
|
||
|
||
if not self.config.created:
|
||
self.env.create()
|
||
self.config.requirements_dir.mkdir(parents=True, exist_ok=True)
|
||
self.config.freezed_packages.touch()
|
||
self.config.manual_requirements.touch()
|
||
else:
|
||
self._load_freezed_packages()
|
||
|
||
def _load_freezed_packages(self) -> None:
|
||
"""Загружает список замороженных пакетов из файла"""
|
||
if self.config.freezed_packages.exists():
|
||
with open(self.config.freezed_packages, 'r') as f:
|
||
self.freezed_packages_list = [
|
||
line.strip() for line in f.readlines() if line.strip()
|
||
]
|
||
|
||
def _parse_package_name(self, package_spec: str) -> str:
|
||
"""Извлекает имя пакета из спецификации (убирает версию и доп. параметры)"""
|
||
# Удаляем версии и дополнительные параметры
|
||
name = re.split(r'[<>=!~]', package_spec)[0].strip()
|
||
# Удаляем экстра-параметры (в квадратных скобках)
|
||
if '[' in name:
|
||
name = name.split('[')[0].strip()
|
||
return name
|
||
|
||
def _is_package_freezed(self, package_spec: str) -> bool:
|
||
"""Проверяет, заморожен ли пакет"""
|
||
package_name = self._parse_package_name(package_spec)
|
||
return package_name in self.freezed_packages_list
|
||
|
||
def _remove_package_from_file(self, package_spec: str, file_path: Path) -> None:
|
||
"""Удаляет пакет из указанного файла"""
|
||
package_name = self._parse_package_name(package_spec)
|
||
|
||
if file_path.exists():
|
||
with open(file_path, 'r') as f:
|
||
lines = f.readlines()
|
||
|
||
# Фильтруем строки, убирая те, что содержат этот пакет
|
||
filtered_lines = [
|
||
line for line in lines
|
||
if self._parse_package_name(line.strip()) != package_name
|
||
]
|
||
|
||
with open(file_path, 'w') as f:
|
||
f.writelines(filtered_lines)
|
||
|
||
|
||
def check_consistency(self) -> bool:
|
||
"""Проверяет целостность окружения"""
|
||
# TODO: Реализовать проверку целостности
|
||
return True
|
||
|
||
def install_requirements(
|
||
self,
|
||
requirements: List[str],
|
||
name: str,
|
||
force: bool = False,
|
||
extra_index_url: Optional[str] = None
|
||
) -> None:
|
||
"""Устанавливает требования из списка"""
|
||
requirements_file = self.config.requirements_dir / f"{name}.req"
|
||
requirements_file.write_text("\n".join(requirements))
|
||
|
||
if force:
|
||
self.env.install_requirements(requirements, extra_index_url)
|
||
|
||
def install_requirements_file(
|
||
self,
|
||
requirements_file: str,
|
||
name: str,
|
||
force: bool = False,
|
||
requirement_level: str = 'req',
|
||
extra_index_url: Optional[str] = None
|
||
) -> None:
|
||
"""Устанавливает требования из файла"""
|
||
target_file = self.config.requirements_dir / f"{name}.{requirement_level}"
|
||
shutil.copy2(requirements_file, target_file)
|
||
|
||
if force:
|
||
self.env.install_requirements_file(requirements_file, extra_index_url)
|
||
|
||
def install_package(
|
||
self,
|
||
package: str,
|
||
extra_index_url: Optional[str] = None,
|
||
force: bool = False,
|
||
freeze: bool = False
|
||
) -> None:
|
||
"""
|
||
Устанавливает пакет с учетом всех требований
|
||
|
||
Args:
|
||
package: Спецификация пакета в формате pip
|
||
extra_index_url: Дополнительный URL репозитория
|
||
force: Принудительная установка, даже если пакет заморожен
|
||
freeze: Заморозить пакет после установки
|
||
"""
|
||
# Шаг 1: Проверка флага force
|
||
if not force:
|
||
# Шаг 2: Проверка на заморозку
|
||
if self._is_package_freezed(package):
|
||
raise ValueError(f"Пакет {self._parse_package_name(package)} заморожен и не может быть установлен")
|
||
|
||
# Шаг 3: Проверка целостности
|
||
if not self.check_integrity():
|
||
raise RuntimeError("Невозможно установить пакет: проверка целостности не пройдена")
|
||
|
||
# Шаг 4: Установка пакета
|
||
try:
|
||
self.env.install_package(package, extra_index_url)
|
||
except Exception as e:
|
||
raise RuntimeError(f"Ошибка установки пакета {package}: {e}")
|
||
|
||
# Шаг 5: Удаление существующих упоминаний пакета
|
||
self._remove_package_from_file(package, self.config.manual_requirements)
|
||
|
||
# Шаг 6: Добавление пакета в manual_requirements
|
||
with open(self.config.manual_requirements, 'a') as f:
|
||
f.write(f"{package}\n")
|
||
|
||
# Шаг 7: Проверка флага freeze
|
||
if not freeze:
|
||
return
|
||
|
||
# Шаг 8: Удаление из freezed_packages
|
||
self._remove_package_from_file(package, self.config.freezed_packages)
|
||
self._load_freezed_packages() # Обновляем список
|
||
|
||
# Шаг 9: Добавление в freezed_packages (только имя пакета)
|
||
package_name = self._parse_package_name(package)
|
||
with open(self.config.freezed_packages, 'a') as f:
|
||
f.write(f"{package_name}\n")
|
||
self.freezed_packages_list.append(package_name)
|
||
|
||
def install_optional_requirements(
|
||
self,
|
||
requirements_file: str,
|
||
name: str,
|
||
force: bool = False
|
||
) -> None:
|
||
"""Устанавливает опциональные требования"""
|
||
self.install_requirements_file(
|
||
requirements_file, name, force, requirement_level='opt'
|
||
)
|
||
|
||
|
||
class GenericPyTorchInstance(Instance):
|
||
"""Базовый класс для экземпляров с PyTorch"""
|
||
|
||
BASE_URL = "https://download.pytorch.org/whl/"
|
||
|
||
def __init__(
|
||
self,
|
||
path: str,
|
||
api: str,
|
||
torch_version: str,
|
||
torchvision_version: str,
|
||
torchaudio_version: str,
|
||
env: str = "venv",
|
||
freeze_torch: bool = True
|
||
):
|
||
super().__init__(path, env)
|
||
|
||
self.url = f"{self.BASE_URL.rstrip('/')}/{api}"
|
||
self.torch_version = torch_version
|
||
self.torchvision_version = torchvision_version
|
||
self.torchaudio_version = torchaudio_version
|
||
|
||
# Установка PyTorch компонентов
|
||
self.install_package(f"torch=={torch_version}", self.url, force=True)
|
||
self.install_package(f"torchvision=={torchvision_version}", self.url, force=True)
|
||
self.install_package(f"torchaudio=={torchaudio_version}", self.url, force=True)
|
||
|
||
if freeze_torch:
|
||
self.freezed_packages.extend([
|
||
f"torch=={torch_version}",
|
||
f"torchvision=={torchvision_version}",
|
||
f"torchaudio=={torchaudio_version}"
|
||
])
|
||
|
||
|
||
class ComfyUIInstance(GenericPyTorchInstance):
|
||
"""Специализированный класс для ComfyUI"""
|
||
|
||
def __init__(
|
||
self,
|
||
path: str,
|
||
api: str,
|
||
torch_version: str,
|
||
torchvision_version: str,
|
||
torchaudio_version: str,
|
||
env: str = "venv"
|
||
):
|
||
super().__init__(
|
||
path, api, torch_version, torchvision_version, torchaudio_version, env
|
||
)
|
||
self.app_dir = self.path / 'app'
|