Delete my_chute.py

my_chute.py

@@ -1,442 +0,0 @@
-#!/usr/bin/env python3
-
-import sys
-from importlib.machinery import PathFinder
-from importlib.util import module_from_spec, spec_from_file_location
-from inspect import signature
-from pathlib import Path
-from site import getsitepackages, getusersitepackages
-from sysconfig import get_paths
-from tempfile import NamedTemporaryFile
-from typing import Any, Generator
-from urllib.parse import quote
-
-import torch
-from aiohttp import ClientSession
-from chutes.chute import Chute, NodeSelector
-from chutes.image import Image
-from base64 import b64decode
-from cv2 import CAP_PROP_FRAME_COUNT, VideoCapture, imdecode, IMREAD_COLOR
-from huggingface_hub import snapshot_download
-from numpy import ndarray, frombuffer, uint8
-from pydantic import BaseModel
-from yaml import safe_load
-
-##============VARIABLES=======================
-HF_REPO_NAME = "TerryStewart/beverage-detection"
-HF_REPO_REVISION = "main"
-CHUTES_USERNAME = "kevinStewart"
-CHUTE_NAME = "beverage-detection"
-FILENAME = "miner.py"
-CLASSNAME = "Miner"
-CHUTE_CONFIG_FILENAME = "chute_config.yml"
-MAX_LOADED_MODEL_SIZE_GB = 5.0
-
-
-##=========DATA CLASSES================
-class TVFrame(BaseModel):
-    frame_id: int
-    url: str | None = None
-    data: str | None = None
-
-
-class TVPredictInput(BaseModel):
-    url: str | None = None
-    frames: list[TVFrame] | None = None
-    meta: dict[str, Any] = {}
-
-
-class TVPredictOutput(BaseModel):
-    success: bool
-    predictions: dict[str, list[dict]] | None = None
-    error: str | None = None
-
-
-##=========UTILITY FUNCTIONS=================
-def get_whitelisted_paths() -> set[Path]:
-    STDLIB_ROOT = Path(get_paths()["stdlib"]).resolve()
-    SITE_PACKAGES_ROOTS = {Path(p).resolve() for p in getsitepackages()} 
-    SITE_PACKAGES_ROOTS.add(Path(getusersitepackages()).resolve())
-
-    return {STDLIB_ROOT, *SITE_PACKAGES_ROOTS}
-
-
-def post_check_files_against_whitelist(*, before: set[str], miner_path: Path) -> None:
-    miner_path = miner_path.resolve()
-    WHITELIST = get_whitelisted_paths()
-    after = set(sys.modules.keys()) - before
-    offenders = []
-    for module_loaded_name in sorted(after):
-        module_loaded = sys.modules.get(module_loaded_name)
-        if module_loaded is None:
-            continue
-        module_loaded_file = getattr(module_loaded, "__file__", None)
-        if not module_loaded_file:
-            continue  
-        module_loaded_filepath = Path(module_loaded_file).resolve()
-        if module_loaded_filepath == miner_path:
-            continue
-        if any(module_loaded_filepath.is_relative_to(root) for root in WHITELIST):
-            continue
-        offenders.append(module_loaded_name)
-
-    if offenders:
-        raise ImportError(
-            f"Miner loaded modules outside the whitelist: {', '.join(offenders)}"
-        )
-
-
-def verify_cuda_version() -> str:
-    min_version = (12, 8)
-    if not torch.cuda.is_available():
-        return "❌ CUDA is not available. Make sure you have a CUDA-capable GPU and correct drivers installed"
-    device_count = torch.cuda.device_count()
-    if device_count < 1:
-        return "❌ No CUDA devices detected"
-    cuda_version_str = torch.version.cuda
-    if cuda_version_str is None:
-        return "❌ PyTorch was not compiled with CUDA support"
-    try:
-        cuda_version = tuple(map(int, cuda_version_str.split(".")))
-    except Exception as e:
-        return f"❌ Unable to parse CUDA version string: {e}"
-    if cuda_version < min_version:
-        return f"❌ CUDA version {cuda_version_str} is too old. Requires >= {min_version[0]}.{min_version[1]}"
-    return f"CUDA {cuda_version_str} detected_with {device_count} device(s)"
-
-
-def load_chute_config_from_hf_repo(config_path: Path) -> dict:
-    try:
-        if not config_path.exists():
-            raise ValueError("No config file found in repo")
-
-        with config_path.open() as configuration_file:
-            config = safe_load(configuration_file)
-            print("✅ Loaded Chutes Configuration")
-            return config or {}
-    except Exception as e:
-        print(f"⚠️ Failed to load Chutes Configuration. Using defaults: {e}")
-        return {}
-
-
-def safe_instantiate(cls, config: dict):
-    params = {k: v for k, v in config.items() if k in signature(cls).parameters}
-    print(
-        f"The following parameters will be applied when instantiating {cls.__name__}: {params}"
-    )
-    obj = cls(**params)
-    for function_name, value in config.items():
-        if not hasattr(obj, function_name) or not callable(getattr(obj, function_name)):
-            continue
-        print(f"applying .{function_name}({value})")
-        method = getattr(obj, function_name)
-        if isinstance(value, list):
-            for v in value:
-                if isinstance(v, (tuple, list)):
-                    method(*v)
-                else:
-                    method(v)
-        else:
-            method(value)
-    return obj
-
-
-def load_chute(
-    hf_repo: str,
-    hf_revision: str,
-    config_filename: str,
-    chutes_username: str,
-    chute_name: str,
-) -> Chute:
-    hf_repo_path = Path(snapshot_download(hf_repo, revision=hf_revision))
-    print("✅ Huggingface Hub repo downloaded")
-
-    config = load_chute_config_from_hf_repo(config_path=hf_repo_path / config_filename)
-    print("✅ Config file loaded")
-
-    image_config = config.get("Image", {})
-    image_config.update(
-        {
-            "username": chutes_username,
-            "name": chute_name,
-            "tag":"latest",
-            "readme": "README.md",
-        }
-    )
-    chute_image = safe_instantiate(cls=Image, config=image_config)
-    print("✅ Image instantiated")
-
-    machine_specs = safe_instantiate(
-        cls=NodeSelector, config=config.get("NodeSelector", {})
-    )
-    print("✅ NodeSelector instantiated")
-
-    chute_config = config.get("Chute", {})
-    chute_config.update(
-        {
-            "username": chutes_username,
-            "name": chute_name,
-            "image": chute_image,
-            "node_selector": machine_specs,
-            "allow_external_egress":True,
-            "readme": "README.md",
-        }
-    )
-    chute = safe_instantiate(cls=Chute, config=chute_config)
-    print("✅ Chute instantiated")
-    return chute
-
-
-def load_miner_from_hf_repo(path_hf_repo: Path, filename: str, classname: str):
-    path_hf_repo = path_hf_repo.resolve()
-    path_miner = path_hf_repo / filename
-    if not path_miner.is_file():
-        raise FileNotFoundError(
-            f"❌ Required file '{filename}' not in repo {path_hf_repo}"
-        )
-
-    WHITELIST = get_whitelisted_paths()
-
-    class HFRepoImportBlocker(PathFinder):
-
-        @classmethod
-        def find_spec(cls, fullname, path=None, target=None):
-            spec = super().find_spec(fullname, path, target)
-            if spec is None or spec.origin is None:
-                return spec
-            origin_path = Path(spec.origin).resolve()
-            if any(origin_path.is_relative_to(root) for root in WHITELIST):
-                return spec
-            raise ImportError(
-                f"❌ Import blocked: '{fullname}' (file: {origin_path}) is not in the whitelist. All code related to the miner MUST be inside {filename}"
-            )
-
-    try:
-        before_exec = set(sys.modules.keys())
-        sys.meta_path.insert(0, HFRepoImportBlocker)
-        spec = spec_from_file_location("miner", path_miner)
-        module = module_from_spec(spec)
-        sys.modules["miner"] = module
-        spec.loader.exec_module(module)
-        miner_object = getattr(module, classname)
-        miner = miner_object(path_hf_repo)
-        post_check_files_against_whitelist(before=before_exec, miner_path=path_miner)
-    except Exception as e:
-        raise e
-    finally:
-        sys.meta_path.remove(HFRepoImportBlocker)
-    return miner
-
-
-def object_size_in_bytes(current_obj: Any, seen_ids: set[int], limit: int) -> int:
-    obj_id = id(current_obj)
-    if obj_id in seen_ids or limit <= 0:
-        return 0
-    seen_ids.add(obj_id)
-    total_size = 0
-    if isinstance(current_obj, torch.Tensor):
-        return current_obj.numel() * current_obj.element_size()
-    if isinstance(current_obj, ndarray):
-        return current_obj.nbytes
-    try:
-        total_size += sys.getsizeof(current_obj)
-    except TypeError:
-        pass
-    if isinstance(current_obj, dict):
-        for key, value in current_obj.items():
-            total_size += object_size_in_bytes(key, seen_ids, limit - 1)
-            total_size += object_size_in_bytes(value, seen_ids, limit - 1)
-    elif isinstance(current_obj, (list, tuple, set)):
-        try:
-            for item in current_obj:
-                total_size += object_size_in_bytes(item, seen_ids, limit - 1)
-        except TypeError:
-            pass
-    elif hasattr(current_obj, "__dict__"):
-        for value in current_obj.__dict__.values():
-            total_size += object_size_in_bytes(value, seen_ids, limit - 1)
-    return total_size
-
-
-def get_miner_size_gb(miner) -> float:
-    total_size_bytes = object_size_in_bytes(
-        current_obj=miner, seen_ids=set(), limit=100
-    )
-    return total_size_bytes / (1024 ** 3)
-
-
-def get_video_frames_in_batches(
-    video: VideoCapture, batch_size: int
-) -> Generator[list[ndarray], None, None]:
-    batch = []
-    while True:
-        ok, frame = video.read()
-        if not ok:
-            if batch:
-                yield batch
-            break
-
-        batch.append(frame)
-
-        if len(batch) >= batch_size:
-            yield batch
-            batch = []
-
-
-##=========ENDPOINTS================
-chute = load_chute(
-    hf_repo=HF_REPO_NAME,
-    hf_revision=HF_REPO_REVISION,
-    chutes_username=CHUTES_USERNAME,
-    chute_name=CHUTE_NAME,
-    config_filename=CHUTE_CONFIG_FILENAME,
-)
-
-
-@chute.on_startup()
-async def load_model(self) -> None:
-    try:
-        self.cuda_status = verify_cuda_version()
-    except Exception as e:
-        self.cuda_status = f"❌ Failed to check cuda status: {e}"
-    print(self.cuda_status)
-
-    self.miner_size = None
-    try:
-        hf_repo_path = Path(snapshot_download(HF_REPO_NAME, revision=HF_REPO_REVISION))
-        print("✅ Huggingface Hub repo downloaded")
-
-        self.miner = load_miner_from_hf_repo(
-            path_hf_repo=hf_repo_path, filename=FILENAME, classname=CLASSNAME
-        )
-        print(f"✅ Miner loaded {self.miner}")
-        self.miner_size = get_miner_size_gb(self.miner)
-        if self.miner_size > MAX_LOADED_MODEL_SIZE_GB:
-            raise ValueError(
-                f"Loaded Miner memory footprint {self.miner_size:.2f} GB exceeds the allowed limit of {MAX_LOADED_MODEL_SIZE_GB:.2f} GB. "
-            )
-        self.status = "healthy"
-    except Exception as e:
-        self.status = f"❌ Failed to load miner: {e}"
-        self.miner = None
-        print(self.status)
-
-
-@chute.cord(public_api_path="/health")
-async def health(self, *args, **kwargs) -> dict[str, Any]:
-    return {
-        "status": self.status,
-        "model_loaded": str(self.miner),
-        "memory_gb": self.miner_size,
-        "cuda": self.cuda_status,
-    }
-
-
-@chute.cord(
-    public_api_path="/predict",
-)
-async def predict(self, data: TVPredictInput) -> dict:
-    try:
-        if self.miner is None:
-            raise ValueError("Models were not properly loaded!")
-
-        metadata = data.meta
-        batch_size = metadata.get("batch_size", 64)
-        n_keypoints = metadata.get("n_keypoints", 32)
-
-        frame_results = []
-
-        if data.frames:
-            frame_entries = [(f.frame_id, f.url, f.data) for f in data.frames]
-            async with ClientSession() as session:
-                loaded_frames: list[tuple[int, ndarray]] = []
-                for frame_id, frame_url, frame_data in frame_entries:
-                    if frame_data:
-                        try:
-                            content = b64decode(frame_data)
-                        except Exception as e:
-                            raise ValueError(f"Failed to decode frame {frame_id}: {e}")
-                    elif frame_url:
-                        url = f"https://proxy.chutes.ai/misc/proxy?url={quote(frame_url, safe='')}"
-                        async with session.get(url) as response:
-                            response.raise_for_status()
-                            content = await response.read()
-                    else:
-                        raise ValueError(f"Frame {frame_id} missing url or data")
-
-                    arr = frombuffer(content, dtype=uint8)
-                    img = imdecode(arr, IMREAD_COLOR)
-                    if img is None:
-                        raise ValueError(f"Failed to decode frame {frame_id}")
-                    loaded_frames.append((frame_id, img))
-
-            for i in range(0, len(loaded_frames), batch_size):
-                batch = loaded_frames[i: i + batch_size]
-                batch_images = [img for _, img in batch]
-                batch_frame_ids = [fid for fid, _ in batch]
-                print(f"Predicting Batch: {i // batch_size + 1}")
-                batch_frame_results = self.miner.predict_batch(
-                    batch_images=batch_images,
-                    offset=0,
-                    n_keypoints=n_keypoints,
-                )
-                for frame_id, frame_result in zip(
-                    batch_frame_ids, batch_frame_results, strict=True
-                ):
-                    fr = frame_result.model_dump()
-                    fr["frame_id"] = frame_id
-                    frame_results.append(fr)
-        else:
-            if not data.url:
-                raise ValueError("Missing both video url and frames payload")
-
-            url = f"https://proxy.chutes.ai/misc/proxy?url={quote(data.url, safe='')}"
-            async with ClientSession() as session:
-                async with session.get(url) as response:
-                    response.raise_for_status()
-                    content = await response.read()
-                    print("✅ Challenge video downloaded. Saving video to temporary file.")
-
-                    with NamedTemporaryFile(prefix="sv_video_", suffix=".mp4") as f:
-                        f.write(content)
-                        cap = VideoCapture(f.name)
-                        if not cap.isOpened():
-                            raise ValueError(
-                                f"Problem accessing downloaded video {f.name}"
-                            )
-
-                        n_frames = int(cap.get(CAP_PROP_FRAME_COUNT))
-                        print(
-                            f"Processing video with {n_frames} frames in batches of {batch_size}"
-                        )
-
-                        for batch_number, images in enumerate(
-                            get_video_frames_in_batches(
-                                video=cap, batch_size=batch_size
-                            )
-                        ):
-                            frame_number = batch_size * batch_number
-                            print(f"Predicting Batch: {batch_number + 1}")
-                            batch_frame_results = self.miner.predict_batch(
-                                batch_images=images,
-                                offset=frame_number,
-                                n_keypoints=n_keypoints,
-                            )
-                            frame_results.extend(
-                                [
-                                    frame_result.model_dump()
-                                    for frame_result in batch_frame_results
-                                ]
-                            )
-
-                        cap.release()
-                        print("✅ Frame Predictions Completed")
-
-        result = TVPredictOutput(success=True, predictions={"frames": frame_results})
-
-    except Exception as e:
-        result = TVPredictOutput(
-            success=False,
-            error=f"❌ There was a problem during prediction: {e}",
-        )
-    return result.model_dump(mode="json")