分享下我的code。欢迎给意见。- import heapq
- import pytest
- from typing import Dict, List
- class GPUCreditCreditID:
- def __init__(self, credit_id: str) -> None:
- self.credit_id = credit_id
- # all ranges are inclusive
- self.start_heap: List[List] = [] # [start, amount]
- self.end_heap: List[List] = [] # [end, amount]
- self.timestamp = 0
- self.credit = 0
- def check_timestamp(self, timestamp):
- if timestamp <= self.timestamp:
- raise Exception("Timestamp error")
- def add_credit(self, amount: int, timestamp: int, expiration: int) -> None:
- # self.check_timestamp(timestamp)
- # self.timestamp = timestamp
- heapq.heappush(self.start_heap, [timestamp, amount])
- heapq.heappush(self.end_heap, [timestamp + expiration, amount])
- def get_balance(self, timestamp: int) -> int:
- # self.check_timestamp(timestamp)
- # self.timestamp = timestamp
- while self.start_heap and timestamp >= self.start_heap[0][0]:
- self.credit += heapq.heappop(self.start_heap)[1]
- while self.end_heap and timestamp > self.end_heap[0][0]:
- self.credit -= heapq.heappop(self.end_heap)[1]
- return self.credit
- def use_credit(self, timestamp: int, amount: int) -> None:
- self.check_timestamp(timestamp)
- self.timestamp = timestamp
- self.get_balance(timestamp)
- if self.credit < amount:
- raise Exception("No enough credit")
- self.credit -= amount
- while self.end_heap:
- if self.end_heap[0][-1] >= amount:
- self.end_heap[0][-1] -= amount
- break
- else:
- amount -= heapq.heappop(self.end_heap)[1]
- class GPUCredit():
- def __init__(self) -> None:
- self.credit_ids: Dict[str, GPUCreditCreditID] = {}
-
- def add_credit(self, credit_id: str, amount: int, timestamp: int, expiration: int) -> None:
- if credit_id not in self.credit_ids:
- self.credit_ids[credit_id] = GPUCreditCreditID(credit_id)
- return self.credit_ids[credit_id].add_credit(amount, timestamp, expiration)
-
- def get_balance(self, credit_id: str, timestamp: int) -> int:
- if credit_id not in self.credit_ids:
- return 0
- return self.credit_ids[credit_id].get_balance(timestamp)
- def use_credit(self, credit_id: str, timestamp: int, amount: int) -> None:
- if credit_id not in self.credit_ids:
- raise Exception("No enough credit")
- return self.credit_ids[credit_id].use_credit(timestamp, amount)
- def test_gpu_credit():
- gc = GPUCredit()
- assert gc.get_balance("openai", 0) == 0
- assert gc.add_credit("apple", 10, 1, 4) == None # [1, 5]
- assert gc.get_balance("apple", 2) == 10
- assert gc.use_credit("apple", 3, 3) == None
- assert gc.get_balance("apple", 4) == 7
- assert gc.get_balance("apple", 6) == 0
- assert gc.add_credit("apple", 10, 7, 3) == None # [7, 10]
- assert gc.add_credit("apple", 10, 8, 12) == None # [8, 20]
- assert gc.use_credit("apple", 9, 5) == None
- assert gc.use_credit("apple", 10, 6) == None
- assert gc.get_balance("apple", 12) == 9
- with pytest.raises(Exception, match="No enough credit"):
- assert gc.use_credit("apple", 13, 10) == None
- assert gc.use_credit("apple", 14, 9) == None
- pytest.main()
复制代码 |