Coverage for src/local_deep_research/embeddings/model_loading.py: 95%

1"""Keep model loading exclusive while allowing resident-model inference.""" 

2 

3from contextlib import contextmanager 

4from threading import Condition, RLock, get_ident 

5from typing import Any, Callable, TypeVar 

6 

7_Model = TypeVar("_Model") 

8 

9 

10class ModelAccessLock: 

11 """Reentrant writer lock with shared inference access. 

12 

13 Shared holders create handles for the warm model and run admitted 

14 embedding calls; the model cache lets those calls use the model one at a 

15 time. Writers take priority over new readers, so configuration changes and 

16 cleanup cannot starve behind a stream of inference requests. The cache's 

17 model constructors also use this gate: Transformers temporarily changes 

18 process-wide dtype and initialization state, which must not overlap 

19 resident inference. 

20 """ 

21 

22 def __init__(self) -> None: 

23 self._condition = Condition(RLock()) 

24 self._writer: int | None = None 

25 self._depth = 0 

26 self._readers: dict[int, int] = {} 

27 self._waiting_writers = 0 

28 

29 def acquire(self, blocking: bool = True) -> bool: 

30 if not self._condition.acquire(blocking=blocking): 30 ↛ 31line 30 didn't jump to line 31 because the condition on line 30 was never true

31 return False 

32 try: 

33 owner = get_ident() 

34 if self._writer == owner: 

35 self._depth += 1 

36 return True 

37 if not blocking and (self._writer is not None or self._readers): 

38 return False 

39 if owner in self._readers: 

40 raise RuntimeError( 

41 "Cannot load or clear a model during inference" 

42 ) 

43 self._waiting_writers += 1 

44 try: 

45 while self._writer is not None or self._readers: 

46 self._condition.wait() 

47 self._writer = owner 

48 self._depth = 1 

49 return True 

50 finally: 

51 self._waiting_writers -= 1 

52 self._condition.notify_all() 

53 finally: 

54 self._condition.release() 

55 

56 def release(self) -> None: 

57 with self._condition: 

58 if self._writer != get_ident(): 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true

59 raise RuntimeError( 

60 "Model loading lock is not owned by this thread" 

61 ) 

62 self._depth -= 1 

63 if not self._depth: 

64 self._writer = None 

65 self._condition.notify_all() 

66 

67 def __enter__(self): 

68 self.acquire() 

69 return self 

70 

71 def __exit__(self, *_exc): 

72 self.release() 

73 

74 @contextmanager 

75 def inference(self): 

76 owner = get_ident() 

77 with self._condition: 

78 while self._writer != owner and ( 

79 self._writer is not None 

80 or (self._waiting_writers and owner not in self._readers) 

81 ): 

82 self._condition.wait() 

83 self._readers[owner] = self._readers.get(owner, 0) + 1 

84 try: 

85 yield 

86 finally: 

87 with self._condition: 

88 self._readers[owner] -= 1 

89 if not self._readers[owner]: 

90 del self._readers[owner] 

91 self._condition.notify_all() 

92 

93 

94MODEL_LOAD_LOCK = ModelAccessLock() 

95 

96 

97def load_sentence_transformer( 

98 factory: Callable[..., _Model], **kwargs: Any 

99) -> _Model: 

100 """Keep Transformers' dtype and initialization helpers out of races. 

101 

102 The gate is held for the whole constructor call and blocks every local 

103 embedding call, so a constructor that downloads holds it for the whole 

104 download. 

105 """ 

106 with MODEL_LOAD_LOCK: 

107 return factory(**kwargs)