#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :resource_manager.py
# @Time      :2025/4/29 11:06
# @Author    :David
# @Company: shenyang JY

from multiprocessing import cpu_count, Manager


class ResourceController:
    def __init__(self, max_workers=None, gpu_list=None):
        self.cpu_cores = max_workers or cpu_count() // 2  # 预留一半资源
        self.gpu_pool = gpu_list or []
        self.lock = Manager().Lock()

    def get_gpu(self):
        """动态分配GPU资源"""
        with self.lock:
            if self.gpu_pool:
                return self.gpu_pool.pop()
        return None

    def release_gpu(self, gpu_id):
        """释放GPU资源"""
        if gpu_id is not None:
            with self.lock:
                self.gpu_pool.append(gpu_id)