@@ -7,86 +7,136 @@ namespace SignsOfAI.Perplexity.Api.Engine;
77
88/// <summary>
99/// Computes causal-LM perplexity with a Qwen2.5-0.5B ONNX model via ONNX Runtime (CPU) and the
10- /// Hugging Face tokenizer. The model is a merged (with-past) export, so a from-scratch forward pass
11- /// supplies empty KV-cache tensors. Loaded once (see <see cref="InitializeAsync"/>) and reused;
10+ /// Hugging Face tokenizer. To keep the server light when idle, the model is <b>lazily loaded</b> on
11+ /// the first request and <b>unloaded from RAM after an idle period</b> (see <see cref="TryUnloadIfIdle"/>),
12+ /// reloading from disk (~1-2s) on the next request. Loads/unloads are serialized by a gate; inference
13+ /// runs outside the gate and an active-request counter prevents unloading mid-inference.
1214/// <see cref="InferenceSession.Run(RunOptions, IReadOnlyDictionary{string, OrtValue}, IReadOnlyCollection{string})"/>
1315/// is thread-safe, and each call builds its own inputs, so concurrent scoring is safe.
1416/// </summary>
15- public sealed class OnnxPerplexityEngine ( PerplexityOptions options , IHostEnvironment env , ILogger < OnnxPerplexityEngine > log )
16- : IPerplexityEngine , IDisposable
17+ public sealed class OnnxPerplexityEngine : IPerplexityEngine , IDisposable
1718{
18- private readonly PerplexityOptions _o = options ;
19+ private readonly PerplexityOptions _o ;
20+ private readonly ILogger < OnnxPerplexityEngine > _log ;
21+ private readonly string _dir , _modelPath , _tokPath ;
22+
23+ private readonly SemaphoreSlim _gate = new ( 1 , 1 ) ;
1924 private InferenceSession ? _session ;
2025 private Tokenizer ? _tokenizer ;
21- private volatile bool _ready ;
26+ private int _active ;
27+ private DateTime _lastUsedUtc = DateTime . UtcNow ;
28+ private volatile bool _filesReady ;
29+ private volatile bool _loaded ;
30+
31+ public OnnxPerplexityEngine ( PerplexityOptions options , IHostEnvironment env , ILogger < OnnxPerplexityEngine > log )
32+ {
33+ _o = options ; _log = log ;
34+ _dir = Path . IsPathRooted ( _o . ModelDir ) ? _o . ModelDir : Path . Combine ( env . ContentRootPath , _o . ModelDir ) ;
35+ _modelPath = Path . Combine ( _dir , _o . ModelFile ) ;
36+ _tokPath = Path . Combine ( _dir , _o . TokenizerFile ) ;
37+ }
2238
2339 public string ModelId => _o . ModelId ;
24- public bool IsReady => _ready ;
40+ public bool IsReady => _filesReady ;
41+ public bool IsLoaded => _loaded ;
2542
26- /// <summary>Resolves + downloads the model/tokenizer if needed, then loads them. Call once at startup.</summary>
27- public async Task InitializeAsync ( CancellationToken ct = default )
43+ /// <summary>Ensures the model + tokenizer FILES exist (downloading once if needed), without loading
44+ /// them into RAM. Call at startup. If <see cref="PerplexityOptions.PreloadModel"/> is set, also warms.</summary>
45+ public async Task EnsureFilesAsync ( CancellationToken ct = default )
2846 {
29- var dir = Path . IsPathRooted ( _o . ModelDir ) ? _o . ModelDir : Path . Combine ( env . ContentRootPath , _o . ModelDir ) ;
30- Directory . CreateDirectory ( dir ) ;
31- var modelPath = Path . Combine ( dir , _o . ModelFile ) ;
32- var tokPath = Path . Combine ( dir , _o . TokenizerFile ) ;
33-
34- await EnsureFileAsync ( modelPath , _o . ModelUrl , ct ) ;
35- await EnsureFileAsync ( tokPath , _o . TokenizerUrl , ct ) ;
47+ Directory . CreateDirectory ( _dir ) ;
48+ await EnsureFileAsync ( _modelPath , _o . ModelUrl , ct ) ;
49+ await EnsureFileAsync ( _tokPath , _o . TokenizerUrl , ct ) ;
3650
37- if ( ! File . Exists ( modelPath ) || ! File . Exists ( tokPath ) )
51+ _filesReady = File . Exists ( _modelPath ) && File . Exists ( _tokPath ) ;
52+ if ( ! _filesReady )
3853 {
39- log . LogError ( "Perplexity model/tokenizer missing at {Dir} and no download URL configured; engine stays not- ready." , dir ) ;
54+ _log . LogError ( "Perplexity model/tokenizer missing at {Dir} and no download URL configured; engine not ready." , _dir ) ;
4055 return ;
4156 }
57+ _log . LogInformation ( "Perplexity engine ready (files present): {Model}. Idle-unload={Idle}s, preload={Preload}." ,
58+ _o . ModelFile , _o . IdleUnloadSeconds , _o . PreloadModel ) ;
4259
43- var so = new Microsoft . ML . OnnxRuntime . SessionOptions ( ) ;
44- if ( _o . IntraOpThreads > 0 ) so . IntraOpNumThreads = _o . IntraOpThreads ;
60+ if ( _o . PreloadModel ) { var _ = await AcquireAsync ( ct ) ; Release ( ) ; }
61+ }
4562
46- var sw = Stopwatch . StartNew ( ) ;
47- _tokenizer = new Tokenizer ( tokPath ) ;
48- _session = new InferenceSession ( modelPath , so ) ;
49- sw . Stop ( ) ;
63+ public async Task < PerplexityRaw > ScoreAsync ( string text , CancellationToken ct = default )
64+ {
65+ if ( ! _filesReady ) throw new InvalidOperationException ( "Perplexity engine is not ready." ) ;
5066
51- _ready = true ;
52- log . LogInformation ( "Perplexity engine ready: {Model} loaded in {Ms} ms." , _o . ModelFile , sw . ElapsedMilliseconds ) ;
67+ var ( session , tokenizer ) = await AcquireAsync ( ct ) ;
68+ try
69+ {
70+ return await Task . Run ( ( ) => Forward ( session , tokenizer , text , ct ) , ct ) ;
71+ }
72+ finally { Release ( ) ; }
5373 }
5474
55- public Task < PerplexityRaw > ScoreAsync ( string text , CancellationToken ct = default )
75+ // ── Load / unload lifecycle ──────────────────────────────────────────────
76+ private async Task < ( InferenceSession , Tokenizer ) > AcquireAsync ( CancellationToken ct )
5677 {
57- if ( ! _ready || _session is null || _tokenizer is null )
58- throw new InvalidOperationException ( "Perplexity engine is not ready." ) ;
59-
60- // CPU-bound work; hand off so we don't block the request pipeline thread.
61- return Task . Run ( ( ) =>
78+ await _gate . WaitAsync ( ct ) ;
79+ try
6280 {
63- var ids = Array . ConvertAll ( _tokenizer . Encode ( text ) , x => ( long ) x ) ;
64- if ( ids . Length > _o . MaxTokens )
65- ids = ids [ .._o . MaxTokens ] ; // right-truncate for bounded latency
81+ if ( _session is null ) LoadLocked ( ) ;
82+ _active ++ ;
83+ return ( _session ! , _tokenizer ! ) ;
84+ }
85+ finally { _gate . Release ( ) ; }
86+ }
6687
67- if ( ids . Length < 2 )
68- return new PerplexityRaw ( 1.0 , 0.0 , 0 , 0 ) ;
88+ private void Release ( )
89+ {
90+ _gate . Wait ( ) ;
91+ try { if ( _active > 0 ) _active -- ; _lastUsedUtc = DateTime . UtcNow ; }
92+ finally { _gate . Release ( ) ; }
93+ }
6994
70- var sw = Stopwatch . StartNew ( ) ;
71- var ( ppl , meanLogProb , scored ) = Forward ( ids , ct ) ;
72- sw . Stop ( ) ;
73- return new PerplexityRaw ( ppl , meanLogProb , scored , sw . ElapsedMilliseconds ) ;
74- } , ct ) ;
95+ private void LoadLocked ( )
96+ {
97+ var so = new Microsoft . ML . OnnxRuntime . SessionOptions ( ) ;
98+ if ( _o . IntraOpThreads > 0 ) so . IntraOpNumThreads = _o . IntraOpThreads ;
99+ var sw = Stopwatch . StartNew ( ) ;
100+ _tokenizer = new Tokenizer ( _tokPath ) ;
101+ _session = new InferenceSession ( _modelPath , so ) ;
102+ _loaded = true ;
103+ _log . LogInformation ( "Perplexity model loaded into RAM in {Ms} ms." , sw . ElapsedMilliseconds ) ;
75104 }
76105
77- private ( double ppl , double meanLogProb , int scored ) Forward ( long [ ] ids , CancellationToken ct )
106+ /// <summary>If the model is loaded, no request is in flight, and it's been idle longer than
107+ /// <paramref name="idle"/>, dispose it to free RAM. Non-blocking; returns whether it unloaded.</summary>
108+ public bool TryUnloadIfIdle ( TimeSpan idle )
78109 {
110+ if ( ! _gate . Wait ( 0 ) ) return false ;
111+ try
112+ {
113+ if ( _session is null || _active != 0 || DateTime . UtcNow - _lastUsedUtc <= idle )
114+ return false ;
115+ _session . Dispose ( ) ; _session = null ;
116+ _tokenizer ? . Dispose ( ) ; _tokenizer = null ;
117+ _loaded = false ;
118+ GC . Collect ( ) ; GC . WaitForPendingFinalizers ( ) ; GC . Collect ( ) ;
119+ return true ;
120+ }
121+ finally { _gate . Release ( ) ; }
122+ }
123+
124+ // ── Inference ─────────────────────────────────────────────────────────────
125+ private PerplexityRaw Forward ( InferenceSession session , Tokenizer tokenizer , string text , CancellationToken ct )
126+ {
127+ var ids = Array . ConvertAll ( tokenizer . Encode ( text ) , x => ( long ) x ) ;
128+ if ( ids . Length > _o . MaxTokens ) ids = ids [ .._o . MaxTokens ] ; // right-truncate for bounded latency
129+ if ( ids . Length < 2 ) return new PerplexityRaw ( 1.0 , 0.0 , 0 , 0 ) ;
130+
79131 int n = ids . Length , vocab = _o . Vocab ;
80132 var toDispose = new List < OrtValue > ( 2 * _o . NumLayers + 3 ) ;
81133 OrtValue Mk < T > ( T [ ] data , long [ ] shape ) where T : unmanaged
82134 { var v = OrtValue . CreateTensorValueFromMemory ( data , shape ) ; toDispose . Add ( v ) ; return v ; }
83135
136+ var sw = Stopwatch . StartNew ( ) ;
84137 try
85138 {
86- var inputs = new Dictionary < string , OrtValue > ( 2 * _o . NumLayers + 3 )
87- {
88- [ "input_ids" ] = Mk ( ids , [ 1 , n ] ) ,
89- } ;
139+ var inputs = new Dictionary < string , OrtValue > ( 2 * _o . NumLayers + 3 ) { [ "input_ids" ] = Mk ( ids , [ 1 , n ] ) } ;
90140 var attn = new long [ n ] ; Array . Fill ( attn , 1L ) ;
91141 inputs [ "attention_mask" ] = Mk ( attn , [ 1 , n ] ) ;
92142 var pos = new long [ n ] ; for ( int i = 0 ; i < n ; i ++ ) pos [ i ] = i ;
@@ -100,10 +150,9 @@ OrtValue Mk<T>(T[] data, long[] shape) where T : unmanaged
100150 }
101151
102152 using var ro = new RunOptions ( ) ;
103- using var results = _session ! . Run ( ro , inputs , [ "logits" ] ) ;
153+ using var results = session . Run ( ro , inputs , [ "logits" ] ) ;
104154 var logits = results [ 0 ] . GetTensorDataAsSpan < float > ( ) ; // [1, n, vocab] row-major
105155
106- // Shift: logits at position t predict token ids[t+1]. Score t = 0..n-2.
107156 double sumNll = 0 ; int scored = 0 ;
108157 for ( int t = 0 ; t < n - 1 ; t ++ )
109158 {
@@ -117,40 +166,54 @@ OrtValue Mk<T>(T[] data, long[] shape) where T : unmanaged
117166 sumNll += - ( logits [ off + ( int ) ids [ t + 1 ] ] - logZ ) ;
118167 scored ++ ;
119168 }
120-
169+ sw . Stop ( ) ;
121170 double meanNll = sumNll / scored ;
122- return ( Math . Exp ( meanNll ) , - meanNll , scored ) ;
123- }
124- finally
125- {
126- foreach ( var d in toDispose ) d . Dispose ( ) ;
171+ return new PerplexityRaw ( Math . Exp ( meanNll ) , - meanNll , scored , sw . ElapsedMilliseconds ) ;
127172 }
173+ finally { foreach ( var d in toDispose ) d . Dispose ( ) ; }
128174 }
129175
130176 private async Task EnsureFileAsync ( string path , string ? url , CancellationToken ct )
131177 {
132178 if ( File . Exists ( path ) || string . IsNullOrWhiteSpace ( url ) ) return ;
133- log . LogInformation ( "Downloading {Url} → {Path} …" , url , path ) ;
179+ _log . LogInformation ( "Downloading {Url} → {Path} …" , url , path ) ;
134180 using var http = new HttpClient { Timeout = TimeSpan . FromMinutes ( 20 ) } ;
135181 using var resp = await http . GetAsync ( url , HttpCompletionOption . ResponseHeadersRead , ct ) ;
136182 resp . EnsureSuccessStatusCode ( ) ;
137183 var tmp = path + ".part" ;
138184 await using ( var fs = File . Create ( tmp ) )
139185 await resp . Content . CopyToAsync ( fs , ct ) ;
140186 File . Move ( tmp , path , overwrite : true ) ;
141- log . LogInformation ( "Downloaded {Path} ({Bytes:N0} bytes)." , path , new FileInfo ( path ) . Length ) ;
187+ _log . LogInformation ( "Downloaded {Path} ({Bytes:N0} bytes)." , path , new FileInfo ( path ) . Length ) ;
142188 }
143189
144- public void Dispose ( ) => _session ? . Dispose ( ) ;
190+ public void Dispose ( )
191+ {
192+ _session ? . Dispose ( ) ;
193+ _tokenizer ? . Dispose ( ) ;
194+ _gate . Dispose ( ) ;
195+ }
145196}
146197
147- /// <summary>Loads the model at startup (downloading if needed) without blocking app boot.</summary>
148- public sealed class ModelWarmupService ( OnnxPerplexityEngine engine , ILogger < ModelWarmupService > log ) : BackgroundService
198+ /// <summary>At startup: ensures the model file is on disk (download once). Then, if idle-unloading is
199+ /// enabled, periodically frees the model from RAM after it goes idle.</summary>
200+ public sealed class ModelLifecycleService ( OnnxPerplexityEngine engine , PerplexityOptions options , ILogger < ModelLifecycleService > log )
201+ : BackgroundService
149202{
150203 protected override async Task ExecuteAsync ( CancellationToken ct )
151204 {
152- try { await engine . InitializeAsync ( ct ) ; }
153- catch ( Exception ex ) when ( ex is not OperationCanceledException )
154- { log . LogError ( ex , "Perplexity engine failed to initialize." ) ; }
205+ try { await engine . EnsureFilesAsync ( ct ) ; }
206+ catch ( Exception ex ) when ( ex is not OperationCanceledException ) { log . LogError ( ex , "Perplexity engine failed to initialize." ) ; }
207+
208+ if ( options . IdleUnloadSeconds <= 0 ) return ; // keep resident once loaded
209+
210+ var idle = TimeSpan . FromSeconds ( options . IdleUnloadSeconds ) ;
211+ var interval = TimeSpan . FromSeconds ( Math . Clamp ( options . IdleUnloadSeconds / 4.0 , 15 , 120 ) ) ;
212+ while ( ! ct . IsCancellationRequested )
213+ {
214+ try { await Task . Delay ( interval , ct ) ; } catch ( OperationCanceledException ) { break ; }
215+ if ( engine . TryUnloadIfIdle ( idle ) )
216+ log . LogInformation ( "Perplexity model unloaded after {Idle}s idle — RAM freed." , options . IdleUnloadSeconds ) ;
217+ }
155218 }
156219}
0 commit comments