SupplyAgent/supply_infra/config.py

from __future__ import annotations

from functools import lru_cache
from pathlib import Path
from urllib.parse import quote_plus

from pydantic import Field, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict

# supply_infra/config.py -> project root; avoid depending on process cwd
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
_ENV_FILE = _PROJECT_ROOT / ".env"


class InfraSettings(BaseSettings):
    """Shared infrastructure settings."""

    model_config = SettingsConfigDict(
        env_file=str(_ENV_FILE),
        env_file_encoding="utf-8",
        extra="ignore",
    )

    # MySQL
    mysql_host: str = Field(default="127.0.0.1", alias="MYSQL_HOST")
    mysql_port: int = Field(default=3306, alias="MYSQL_PORT")
    mysql_user: str = Field(default="root", alias="MYSQL_USER")
    mysql_password: str = Field(default="", alias="MYSQL_PASSWORD")
    mysql_database: str = Field(default="supply_agent", alias="MYSQL_DATABASE")
    mysql_pool_size: int = Field(default=1, ge=1, alias="MYSQL_POOL_SIZE")
    mysql_connection_budget: int = Field(
        default=40,
        ge=30,
        le=50,
        alias="MYSQL_CONNECTION_BUDGET",
    )
    mysql_pool_size_api: int = Field(default=6, ge=1, alias="MYSQL_POOL_SIZE_API")
    mysql_pool_size_control: int = Field(
        default=2,
        ge=1,
        alias="MYSQL_POOL_SIZE_CONTROL",
    )
    mysql_max_overflow: int = Field(default=0, ge=0, alias="MYSQL_MAX_OVERFLOW")
    mysql_pool_timeout_seconds: int = Field(
        default=10,
        ge=1,
        alias="MYSQL_POOL_TIMEOUT_SECONDS",
    )
    mysql_pool_recycle_seconds: int = Field(
        default=1800,
        ge=60,
        alias="MYSQL_POOL_RECYCLE_SECONDS",
    )
    mysql_connect_timeout_seconds: int = Field(
        default=10,
        ge=1,
        alias="MYSQL_CONNECT_TIMEOUT_SECONDS",
    )
    mysql_read_timeout_seconds: int = Field(
        default=30,
        ge=1,
        alias="MYSQL_READ_TIMEOUT_SECONDS",
    )
    mysql_write_timeout_seconds: int = Field(
        default=30,
        ge=1,
        alias="MYSQL_WRITE_TIMEOUT_SECONDS",
    )
    mysql_operational_reserve: int = Field(
        default=4,
        ge=2,
        alias="MYSQL_OPERATIONAL_RESERVE",
    )
    mysql_echo: bool = Field(default=False, alias="MYSQL_ECHO")

    # ODPS (MaxCompute)
    odps_access_id: str = Field(default="", alias="ODPS_ACCESS_ID")
    odps_access_key: str = Field(default="", alias="ODPS_ACCESS_KEY")
    odps_project: str = Field(default="", alias="ODPS_PROJECT")
    odps_endpoint: str = Field(
        default="https://service.cn.maxcompute.aliyun.com/api",
        alias="ODPS_ENDPOINT",
    )

    # Scheduler
    scheduler_timezone: str = Field(default="Asia/Shanghai", alias="SCHEDULER_TIMEZONE")
    scheduler_enabled: bool = Field(default=False, alias="SCHEDULER_ENABLED")
    scheduler_cron_hour: int = Field(default=15, ge=0, le=23, alias="SCHEDULER_CRON_HOUR")
    scheduler_cron_minute: int = Field(
        default=0,
        ge=0,
        le=59,
        alias="SCHEDULER_CRON_MINUTE",
    )

    # Pipeline process/control plane
    process_role: str = Field(default="api", alias="PROCESS_ROLE")
    database_auto_create: bool = Field(default=False, alias="DATABASE_AUTO_CREATE")
    pipeline_worker_processes: int = Field(
        default=4,
        ge=1,
        le=20,
        alias="PIPELINE_WORKER_PROCESSES",
    )
    pipeline_max_active_steps: int = Field(
        default=4,
        ge=1,
        le=20,
        alias="PIPELINE_MAX_ACTIVE_STEPS",
    )
    pipeline_worker_poll_seconds: float = Field(
        default=2.0,
        gt=0,
        alias="PIPELINE_WORKER_POLL_SECONDS",
    )
    pipeline_lease_seconds: int = Field(
        default=120,
        ge=30,
        alias="PIPELINE_LEASE_SECONDS",
    )
    pipeline_heartbeat_seconds: int = Field(
        default=30,
        ge=5,
        alias="PIPELINE_HEARTBEAT_SECONDS",
    )
    pipeline_reconcile_seconds: int = Field(
        default=60,
        ge=10,
        alias="PIPELINE_RECONCILE_SECONDS",
    )
    pipeline_missed_run_grace_seconds: int = Field(
        default=600,
        ge=60,
        alias="PIPELINE_MISSED_RUN_GRACE_SECONDS",
    )
    pipeline_shutdown_grace_seconds: int = Field(
        default=300,
        ge=10,
        alias="PIPELINE_SHUTDOWN_GRACE_SECONDS",
    )
    pipeline_warn_after_hours: int = Field(
        default=12,
        ge=1,
        alias="PIPELINE_WARN_AFTER_HOURS",
    )
    pipeline_critical_before_next_minutes: int = Field(
        default=120,
        ge=1,
        alias="PIPELINE_CRITICAL_BEFORE_NEXT_MINUTES",
    )
    pipeline_final_warn_before_next_minutes: int = Field(
        default=30,
        ge=1,
        alias="PIPELINE_FINAL_WARN_BEFORE_NEXT_MINUTES",
    )
    pipeline_log_dir: str = Field(default="logs/pipeline", alias="PIPELINE_LOG_DIR")

    # Aliyun OSS (agent run log publishing)
    aliyun_oss_access_key_id: str = Field(default="", alias="ALIYUN_OSS_ACCESS_KEY_ID")
    aliyun_oss_access_key_secret: str = Field(default="", alias="ALIYUN_OSS_ACCESS_KEY_SECRET")
    aliyun_oss_region: str = Field(default="cn-hangzhou", alias="ALIYUN_OSS_REGION")
    aliyun_oss_bucket: str = Field(default="art-pubbucket", alias="ALIYUN_OSS_BUCKET")
    aliyun_oss_root_prefix: str = Field(default="supply_agent", alias="ALIYUN_OSS_ROOT_PREFIX")
    aliyun_oss_manual_log_prefix: str = Field(
        default="supply_agent/manual_logs",
        alias="ALIYUN_OSS_MANUAL_LOG_PREFIX",
    )
    aliyun_oss_public_base_url: str = Field(
        default="http://rescdn.yishihui.com",
        alias="ALIYUN_OSS_PUBLIC_BASE_URL",
    )
    aliyun_oss_connect_timeout_seconds: int = Field(
        default=30,
        ge=5,
        le=300,
        alias="ALIYUN_OSS_CONNECT_TIMEOUT_SECONDS",
    )
    log_oss_upload_enabled: bool = Field(default=True, alias="LOG_OSS_UPLOAD_ENABLED")

    # AIGC platform
    aigc_api_token: str = Field(default="", alias="AIGC_API_TOKEN")

    @model_validator(mode="after")
    def validate_pipeline_runtime(self) -> "InfraSettings":
        roles = {"api", "scheduler", "worker", "reconciler", "step", "cli"}
        if self.process_role not in roles:
            raise ValueError(
                f"PROCESS_ROLE must be one of {sorted(roles)}, got {self.process_role!r}"
            )
        if self.pipeline_heartbeat_seconds >= self.pipeline_lease_seconds:
            raise ValueError(
                "PIPELINE_HEARTBEAT_SECONDS must be smaller than PIPELINE_LEASE_SECONDS"
            )
        if self.pipeline_max_active_steps > self.pipeline_worker_processes:
            raise ValueError(
                "PIPELINE_MAX_ACTIVE_STEPS cannot exceed PIPELINE_WORKER_PROCESSES"
            )
        if self.mysql_max_overflow != 0:
            raise ValueError("MYSQL_MAX_OVERFLOW must remain 0 for connection budgeting")
        required = self.pipeline_required_connections
        if required > self.mysql_connection_budget:
            raise ValueError(
                "Pipeline connection budget exceeded: "
                f"required={required} budget={self.mysql_connection_budget}"
            )
        return self

    @property
    def mysql_url(self) -> str:
        user = quote_plus(self.mysql_user)
        password = quote_plus(self.mysql_password)
        return (
            f"mysql+pymysql://{user}:{password}"
            f"@{self.mysql_host}:{self.mysql_port}/{self.mysql_database}"
            f"?charset=utf8mb4"
        )

    @property
    def selected_mysql_pool_size(self) -> int:
        if self.process_role == "api":
            return self.mysql_pool_size_api
        if self.process_role in {"scheduler", "worker", "reconciler", "step"}:
            return self.mysql_pool_size_control
        return self.mysql_pool_size

    @property
    def pipeline_required_connections(self) -> int:
        control_processes = 2  # scheduler + reconciler
        return (
            self.mysql_pool_size_api
            + control_processes * self.mysql_pool_size_control
            + self.pipeline_worker_processes * self.mysql_pool_size_control
            + self.pipeline_max_active_steps * self.mysql_pool_size_control
            + self.mysql_operational_reserve
        )


@lru_cache
def get_infra_settings() -> InfraSettings:
    return InfraSettings()