modalDeploy/AutoDL/autodl_scheduling/entity/result_map.py

36 lines
1022 B
Python

from typing import Union
class ResultMap:
def __init__(self, max_size:int=1000):
self.map = dict()
self.max_size = max_size
def get(self, uid:str) -> Union[dict,None]:
if uid in self.map:
if self.map[uid]:
r = self.map[uid]
# self.remove(uid)
return r
else:
return None
else:
raise KeyError(uid)
def set(self, uid:str, result:dict=None) -> None:
if uid in self.map:
raise RuntimeError(f"Key {uid} already exists")
else:
if len(self.map) + 1 > self.max_size:
raise RuntimeError(f"ResultMap exceeds max size {self.max_size}")
self.map[uid] = result
def update(self, uid:str, result:dict) -> None:
if uid in self.map:
self.map[uid] = result
else:
raise KeyError(uid)
def remove(self, uid:str) -> None:
if uid in self.map:
del self.map[uid]