- Modified the `load_config` function to allow launch preferences in `gpu-rent.vars` to override `.env` settings for non-secret variables, enhancing user control. - Updated the list of Python candidates in `stack_env_probe.py` to include additional paths for ComfyUI, improving the detection of Python environments. - Added new pip candidates in `tune_swarm_perf.py` to support various ComfyUI installations, ensuring better compatibility with different setups.
261 lines
9.4 KiB
Python
261 lines
9.4 KiB
Python
"""Load <project>/.env. No secrets in git."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
|
||
from dotenv import load_dotenv
|
||
|
||
from gpu_rent.errors import ConfigError
|
||
from gpu_rent.paths import (
|
||
app_root,
|
||
default_ssh_key_path,
|
||
env_path,
|
||
extensions_manifest_path,
|
||
migrate_legacy_if_needed,
|
||
models_manifest_path,
|
||
ollama_models_manifest_path,
|
||
runtime_dir,
|
||
vars_path,
|
||
)
|
||
from gpu_rent.varsfile import apply_vars_file
|
||
from gpu_rent.llm_runtime import normalize_runtime
|
||
|
||
|
||
def _as_bool(value: str | None, default: bool) -> bool:
|
||
if value is None or value.strip() == "":
|
||
return default
|
||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||
|
||
|
||
def parse_enable_swarmui(raw: str | None, *, default: bool = True) -> bool:
|
||
"""ENABLE_SWARMUI / WORKLOAD=llm|swarm|both → whether to install SwarmUI."""
|
||
wl = (os.environ.get("WORKLOAD") or "").strip().lower()
|
||
if wl in {"llm", "llm-only", "llm_only", "llama", "ollama-only"}:
|
||
return False
|
||
if wl in {"swarm", "swarmui", "ui"}:
|
||
return True
|
||
if wl in {"both", "all", "full"}:
|
||
return True
|
||
if raw is None or str(raw).strip() == "":
|
||
return default
|
||
return _as_bool(str(raw), default)
|
||
|
||
|
||
def _as_int(value: str | None, default: int) -> int:
|
||
if value is None or value.strip() == "":
|
||
return default
|
||
return int(value)
|
||
|
||
|
||
def _as_float(value: str | None, default: float) -> float:
|
||
if value is None or str(value).strip() == "":
|
||
return default
|
||
return float(str(value).strip().replace(",", "."))
|
||
|
||
|
||
def _csv(value: str | None, default: tuple[str, ...]) -> tuple[str, ...]:
|
||
if value is None or value.strip() == "":
|
||
return default
|
||
return tuple(part.strip() for part in value.split(",") if part.strip())
|
||
|
||
|
||
@dataclass
|
||
class Config:
|
||
os_auth_url: str
|
||
os_user_domain_name: str
|
||
os_username: str
|
||
os_password: str
|
||
os_project_id: str
|
||
os_region_name: str
|
||
gpu_rent_az: str
|
||
|
||
ssh_private_key_path: Path
|
||
ssh_user: str
|
||
|
||
boot_volume_id: str
|
||
data_volume_id: str
|
||
data_volume_size_gb: int
|
||
boot_snapshot_name: str
|
||
|
||
civitai_api_token: str
|
||
civitai_api_host: str
|
||
hf_token: str
|
||
models_manifest: Path
|
||
extensions_manifest: Path
|
||
git_token: str
|
||
|
||
local_models_dir: Path
|
||
local_wildcards_dir: Path
|
||
local_workflows_dir: Path
|
||
local_output_dir: Path
|
||
app_root: Path
|
||
|
||
autocomplete_enabled: bool
|
||
autocomplete_github_repo: str
|
||
autocomplete_github_path: str
|
||
autocomplete_github_ref: str
|
||
autocomplete_filename: str
|
||
|
||
swarmui_local_port: int
|
||
swarmui_image: str
|
||
update_git: bool
|
||
|
||
llm_runtime: str
|
||
enable_swarmui: bool
|
||
ollama_models_manifest: Path
|
||
ollama_local_port: int
|
||
|
||
default_flavor_id: str
|
||
flavor_preference: tuple[str, ...]
|
||
flavor_fallback: bool
|
||
scan_pools: str
|
||
default_spot: bool
|
||
keep_floating_ip: bool
|
||
idle_minutes: int
|
||
idle_grace_minutes: int
|
||
pull_output: bool
|
||
notify_ready: bool
|
||
up_stop_on_fail: bool
|
||
selectel_api_token: str
|
||
balance_notify_step_rub: float
|
||
balance_notify_low_rub: float
|
||
|
||
missing: list[str] = field(default_factory=list)
|
||
|
||
@property
|
||
def auth_ok(self) -> bool:
|
||
return not self.missing
|
||
|
||
|
||
def _required(name: str) -> str:
|
||
return (os.environ.get(name) or "").strip()
|
||
|
||
|
||
def load_config(*, require_auth: bool = True) -> Config:
|
||
root = app_root()
|
||
runtime_dir().mkdir(parents=True, exist_ok=True)
|
||
migrate_legacy_if_needed()
|
||
|
||
env_file = env_path()
|
||
if env_file.is_file():
|
||
load_dotenv(env_file, override=False)
|
||
|
||
# Launch prefs (gpu-rent.vars) win over .env for non-secrets — so
|
||
# LLM_RUNTIME=ollama in vars is not blocked by LLM_RUNTIME=none from env.example.
|
||
apply_vars_file(vars_path(), override=True)
|
||
|
||
missing: list[str] = []
|
||
required = (
|
||
"OS_AUTH_URL",
|
||
"OS_USER_DOMAIN_NAME",
|
||
"OS_USERNAME",
|
||
"OS_PASSWORD",
|
||
"OS_PROJECT_ID",
|
||
"OS_REGION_NAME",
|
||
"GPU_RENT_AZ",
|
||
)
|
||
values = {name: _required(name) for name in required}
|
||
for name, value in values.items():
|
||
if not value:
|
||
missing.append(name)
|
||
|
||
if require_auth and missing:
|
||
raise ConfigError(
|
||
f"В {env_file} не хватает: "
|
||
+ ", ".join(missing)
|
||
+ ". Как заполнить: docs/setup.md (сервисный пользователь, не X-Token)."
|
||
)
|
||
|
||
key_override = (os.environ.get("SSH_PRIVATE_KEY_PATH") or "").strip()
|
||
ssh_key = Path(key_override).expanduser() if key_override else default_ssh_key_path()
|
||
|
||
models_manifest = Path(
|
||
(os.environ.get("MODELS_MANIFEST") or "").strip() or str(models_manifest_path())
|
||
).expanduser()
|
||
extensions_manifest = Path(
|
||
(os.environ.get("EXTENSIONS_MANIFEST") or "").strip()
|
||
or str(extensions_manifest_path())
|
||
).expanduser()
|
||
ollama_manifest = Path(
|
||
(os.environ.get("OLLAMA_MODELS_MANIFEST") or "").strip()
|
||
or str(ollama_models_manifest_path())
|
||
).expanduser()
|
||
|
||
try:
|
||
llm_runtime = normalize_runtime(os.environ.get("LLM_RUNTIME"))
|
||
except ValueError:
|
||
llm_runtime = "none"
|
||
|
||
enable_swarmui = parse_enable_swarmui(os.environ.get("ENABLE_SWARMUI"), default=True)
|
||
|
||
def _dir(env_name: str, folder: str) -> Path:
|
||
raw = (os.environ.get(env_name) or "").strip()
|
||
return Path(raw).expanduser() if raw else (root / folder)
|
||
|
||
return Config(
|
||
os_auth_url=values["OS_AUTH_URL"] or "https://cloud.api.selcloud.ru/identity/v3",
|
||
os_user_domain_name=values["OS_USER_DOMAIN_NAME"],
|
||
os_username=values["OS_USERNAME"],
|
||
os_password=values["OS_PASSWORD"],
|
||
os_project_id=values["OS_PROJECT_ID"],
|
||
os_region_name=values["OS_REGION_NAME"],
|
||
gpu_rent_az=values["GPU_RENT_AZ"],
|
||
ssh_private_key_path=ssh_key,
|
||
ssh_user=(os.environ.get("SSH_USER") or "ubuntu").strip(),
|
||
boot_volume_id=(os.environ.get("BOOT_VOLUME_ID") or "").strip(),
|
||
data_volume_id=(os.environ.get("DATA_VOLUME_ID") or "").strip(),
|
||
data_volume_size_gb=_as_int(os.environ.get("DATA_VOLUME_SIZE_GB"), 100),
|
||
boot_snapshot_name=(os.environ.get("BOOT_SNAPSHOT_NAME") or "gpu-rent-boot-ok").strip(),
|
||
civitai_api_token=(os.environ.get("CIVITAI_API_TOKEN") or "").strip(),
|
||
civitai_api_host=(os.environ.get("CIVITAI_API_HOST") or "civitai.red").strip().lower(),
|
||
hf_token=(
|
||
os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or ""
|
||
).strip(),
|
||
models_manifest=models_manifest,
|
||
extensions_manifest=extensions_manifest,
|
||
git_token=(os.environ.get("GIT_TOKEN") or "").strip(),
|
||
local_models_dir=_dir("LOCAL_MODELS_DIR", "Models"),
|
||
local_wildcards_dir=_dir("LOCAL_WILDCARDS_DIR", "Wildcards"),
|
||
local_workflows_dir=_dir("LOCAL_WORKFLOWS_DIR", "CustomWorkflows"),
|
||
local_output_dir=_dir("LOCAL_OUTPUT_DIR", "Output"),
|
||
app_root=root,
|
||
autocomplete_enabled=_as_bool(os.environ.get("AUTOCOMPLETE_ENABLED"), True),
|
||
autocomplete_github_repo=(
|
||
os.environ.get("AUTOCOMPLETE_GITHUB_REPO") or "DominikDoom/a1111-sd-webui-tagcomplete"
|
||
).strip(),
|
||
autocomplete_github_path=(
|
||
os.environ.get("AUTOCOMPLETE_GITHUB_PATH") or "tags/danbooru.csv"
|
||
).strip(),
|
||
autocomplete_github_ref=(os.environ.get("AUTOCOMPLETE_GITHUB_REF") or "main").strip(),
|
||
autocomplete_filename=(os.environ.get("AUTOCOMPLETE_FILENAME") or "danbooru.csv").strip(),
|
||
swarmui_local_port=_as_int(os.environ.get("SWARMUI_LOCAL_PORT"), 17801),
|
||
swarmui_image=(os.environ.get("SWARMUI_IMAGE") or "").strip(),
|
||
update_git=_as_bool(os.environ.get("UPDATE_GIT"), True),
|
||
llm_runtime=llm_runtime,
|
||
enable_swarmui=enable_swarmui,
|
||
ollama_models_manifest=ollama_manifest,
|
||
ollama_local_port=_as_int(os.environ.get("OLLAMA_LOCAL_PORT"), 17811),
|
||
default_flavor_id=(os.environ.get("DEFAULT_FLAVOR_ID") or "").strip(),
|
||
flavor_preference=_csv(
|
||
os.environ.get("FLAVOR_PREFERENCE"),
|
||
("4090-24", "4090-48", "a5000", "a100-40"),
|
||
),
|
||
flavor_fallback=_as_bool(os.environ.get("FLAVOR_FALLBACK"), True),
|
||
scan_pools=(os.environ.get("SCAN_POOLS") or "ru-6,ru-7").strip(),
|
||
default_spot=_as_bool(os.environ.get("DEFAULT_SPOT"), True),
|
||
keep_floating_ip=_as_bool(os.environ.get("KEEP_FLOATING_IP"), False),
|
||
idle_minutes=_as_int(os.environ.get("IDLE_MINUTES"), 30),
|
||
idle_grace_minutes=_as_int(os.environ.get("IDLE_GRACE_MINUTES"), 45),
|
||
pull_output=_as_bool(os.environ.get("PULL_OUTPUT"), False),
|
||
notify_ready=_as_bool(os.environ.get("NOTIFY_READY"), True),
|
||
# Default on: failed up must not leave a billing GPU running.
|
||
up_stop_on_fail=_as_bool(os.environ.get("UP_STOP_ON_FAIL"), True),
|
||
selectel_api_token=(os.environ.get("SELECTEL_API_TOKEN") or "").strip(),
|
||
balance_notify_step_rub=_as_float(os.environ.get("BALANCE_NOTIFY_STEP_RUB"), 200.0),
|
||
balance_notify_low_rub=_as_float(os.environ.get("BALANCE_NOTIFY_LOW_RUB"), 0.0),
|
||
missing=missing,
|
||
)
|