src/loaded_models.py - master - maruga/mlops

src/loaded_models.py
74 строки2 KB
s546371
modified: .gitlab-ci.yml
29 мар 2026, 21:35
29 мар 2026, 21:3547c4e99
import os
from pathlib import Path
from typing import Any

import boto3
import joblib
import pandas as pd

from src.config import TrainConfig, config_s3

_ranker = None
_item_features = None


def _loaded_file(dir: str, file_name: str) -> Path:
    temp_dir = Path("/tmp") / dir
    os.makedirs(temp_dir, exist_ok=True)
    local_path: Path = temp_dir / file_name

    if local_path.exists():
        if local_path.stat().st_size == 0:
            print(f"Предупреждение: файл {local_path} пуст, скачиваем заново")
        else:
            print(f"Файл уже загружен: {local_path}")
            return local_path

    session = boto3.session.Session()
    s3 = session.client(
        service_name="s3",
        endpoint_url=config_s3.S3_ENDPOINT_URL,
        aws_access_key_id=config_s3.S3_ACCESS_KEY,
        aws_secret_access_key=config_s3.S3_SECRET_KEY,
    )

    bucket: str = config_s3.BUCKET_NAME
    key: str = f"{config_s3.S3_COMMON_PATH}/{dir}/{file_name}"

    try:
        s3.download_file(bucket, key, str(local_path))
        print(f"Файл успешно загружен из S3: {local_path}")
        return local_path
    except Exception as e:
        if local_path.exists():
            print(f"Используем локальный файл после ошибки загрузки: {local_path}")
            return local_path
        else:
            raise FileNotFoundError(
                f"Не удалось загрузить файл {key} из S3: {e}"
            ) from e


def load_ranker() -> object:
    global _ranker
    if _ranker is None:
        local_path = _loaded_file(config_s3.S3_MODEL_DIR, TrainConfig.MODEL_PATH)
        try:
            _ranker = joblib.load(local_path)
            print("Ранкер успешно загружен через joblib")
        except Exception as e:
            raise RuntimeError(f"Ошибка загрузки ранкера из {local_path}: {e}") from e
    return _ranker


def load_item_features() -> pd.DataFrame:
    global _item_features
    if _item_features is None:
        local_path = _loaded_file(config_s3.S3_DATA_DIR, config_s3.ITEM_FEATS_PATH)
        _item_features = pd.read_parquet(local_path)
        print("Данные о товарах успешно загружены")
    return _item_features


ranker: Any = load_ranker()
item_features = load_item_features()