From c0d0756b1febb272d9f73f8206403fcaa1fa9a10 Mon Sep 17 00:00:00 2001 From: Leo Mengesha Date: Mon, 13 Jul 2026 03:59:55 +0100 Subject: [PATCH 1/3] first set of changes for next version --- .../Background/BackgroundSaver.cs | 2 +- .../Factories/StoreModelFactory.cs | 22 +++ .../Factories/StoreSessionFactory.cs | 22 +++ .../Helpers/FileSystemHelper.cs | 2 +- .../Manager/CheckpointManager.cs | 110 ++++++----- .../Manager/SessionManager.cs | 69 ++----- .../Models/CheckpointQuery.cs | 15 ++ .../Models/CheckpointSummary.cs | 10 + .../Models/Enums/CheckpointSortField.cs | 8 + .../Models/Enums/SessionSortField.cs | 10 + .../Models/Enums/SqlDbType.cs | 7 + .../Models/Enums/StoreType.cs | 7 + .../Models/SessionManifest.cs | 2 +- .../Models/SessionQuery.cs | 12 ++ .../Models/SessionSummary.cs | 10 + .../Models/TokenizerData.cs | 13 -- .../Queries/PostgresLargeObjectQueries.cs | 2 +- .../Queries/PostgresSessionQueries.cs | 2 +- .../Queries/PostgresTrainingQueries.cs | 2 +- .../Queries/SqlServerSessionQueries.cs | 4 +- .../Queries/SqlServerTrainingQueries.cs | 2 +- .../Settings/BackgroundSaveOptions.cs | 3 + .../Settings/DbStorageOptions.cs | 14 ++ .../Settings/FileSystemStoreOptions.cs | 5 + .../Settings/RetentionPolicy.cs | 22 +++ .../Settings/StorageOptions.cs | 40 ++++ .../Stores/FileSystem/FileSystemModelStore.cs | 186 ++++++++++++++++-- .../FileSystem/FileSystemSessionStore.cs | 103 ++++++++-- .../Stores/FileSystemIndex/CheckpointIndex.cs | 181 +++++++++++++++++ .../Stores/FileSystemIndex/IndexStorage.cs | 36 ++++ .../Stores/FileSystemIndex/SessionIndex.cs | 163 +++++++++++++++ .../Stores/IDbModelStore.cs | 6 + .../Stores/IDbSessionStore.cs | 6 + .../Stores/IFileSystemModelStore.cs | 5 + .../Stores/IFileSystemSessionStore.cs | 5 + .../StateCheckpoint.NET/Stores/IModelStore.cs | 18 +- .../Stores/ISessionStore.cs | 17 +- .../Stores/Postgres/PostgresModelStore.cs | 150 +++++++++++++- .../Stores/Postgres/PostgresSessionStore.cs | 116 ++++++++++- .../Stores/Postgres/PostgresStoreBase.cs | 4 +- .../Stores/SqlServer/SqlServerModelStore.cs | 123 +++++++++++- .../Stores/SqlServer/SqlServerSessionStore.cs | 113 ++++++++++- .../Stores/SqlServer/SqlServerStoreBase.cs | 4 +- 43 files changed, 1469 insertions(+), 184 deletions(-) create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreModelFactory.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreSessionFactory.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointQuery.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/CheckpointSortField.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SessionSortField.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SqlDbType.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/StoreType.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionQuery.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Settings/RetentionPolicy.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Settings/StorageOptions.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbModelStore.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbSessionStore.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemSessionStore.cs diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Background/BackgroundSaver.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Background/BackgroundSaver.cs index 41db686..68cb3fb 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Background/BackgroundSaver.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Background/BackgroundSaver.cs @@ -1,6 +1,6 @@ using System.Threading.Channels; -namespace StateCheckpoint.NET.Manager; +namespace StateCheckpoint.NET; /// /// Internal background queue for fire-and-forget save operations. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreModelFactory.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreModelFactory.cs new file mode 100644 index 0000000..659dcc3 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreModelFactory.cs @@ -0,0 +1,22 @@ +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; + +namespace StateCheckpoint.NET; + +internal static class StoreModelFactory +{ + public static IModelStore Create(StorageOptions options) + { + return options.StoreType switch + { + StoreType.Local => new FileSystemModelStore(options?.FileSystemStoreOptions?.RootPath ?? "./checkpoints"), + StoreType.SqlDb when options.SqlDbType == SqlDbType.Postgres => + new PostgresModelStore(options.DbStoreOptions?.ConnectionString ?? + throw new InvalidOperationException("ConnectionString required for Postgres.")), + StoreType.SqlDb when options.SqlDbType == SqlDbType.SqlServer => + new SqlServerModelStore(options.DbStoreOptions?.ConnectionString ?? + throw new InvalidOperationException("ConnectionString required for SQL Server.")), + _ => throw new NotSupportedException($"StoreType {options.StoreType} with SqlDbType {options.SqlDbType} is not supported.") + }; + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreSessionFactory.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreSessionFactory.cs new file mode 100644 index 0000000..b912b58 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Factories/StoreSessionFactory.cs @@ -0,0 +1,22 @@ +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; + +namespace StateCheckpoint.NET; + +internal static class StoreSessionFactory +{ + public static ISessionStore Create(StorageOptions options) + { + return options.StoreType switch + { + StoreType.Local => new FileSystemSessionStore(options?.FileSystemStoreOptions?.RootPath ?? "./sessions"), + StoreType.SqlDb when options.SqlDbType == SqlDbType.Postgres => + new PostgresSessionStore(options?.DbStoreOptions?.ConnectionString ?? + throw new InvalidOperationException("ConnectionString required for Postgres.")), + StoreType.SqlDb when options.SqlDbType == SqlDbType.SqlServer => + new SqlServerSessionStore(options?.DbStoreOptions?.ConnectionString ?? + throw new InvalidOperationException("ConnectionString required for SQL Server.")), + _ => throw new NotSupportedException($"StoreType {options.StoreType} with SqlDbType {options.SqlDbType} is not supported.") + }; + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Helpers/FileSystemHelper.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Helpers/FileSystemHelper.cs index f7d9de1..ef08ce1 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Helpers/FileSystemHelper.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Helpers/FileSystemHelper.cs @@ -1,7 +1,7 @@ using System.Text.Json; using StateCheckpoint.NET.Settings; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; internal static class FileSystemHelper { diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs index b4978de..f5763c7 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs @@ -1,64 +1,40 @@ using StateCheckpoint.NET.Models; using StateCheckpoint.NET.Settings; -using StateCheckpoint.NET.Stores; -namespace StateCheckpoint.NET.Manager; +namespace StateCheckpoint.NET; public class CheckpointManager : IAsyncDisposable { private readonly IModelStore _store; private readonly BackgroundSaver? _backgroundSaver; - private static string _defaultCheckpointPath = "./checkpoints"; /// /// Initializes the manager with a custom storage provider. /// - /// Any implementation of IModelStore (FileSystem, PostgreSQL, etc.) - public CheckpointManager(IModelStore store) + /// + /// + public CheckpointManager(StorageOptions storageOptions) { - _store = store ?? throw new ArgumentNullException(nameof(store)); - } + // Validate options + if (storageOptions.StoreType == StoreType.SqlDb && string.IsNullOrWhiteSpace(storageOptions?.DbStoreOptions?.ConnectionString)) + throw new InvalidOperationException("ConnectionString is required when StoreType is SqlDb."); - /// - /// Initializes the manager with the default FileSystem store. - /// Models are saved to ./checkpoints by default. - /// - public CheckpointManager() : this(new FileSystemModelStore(_defaultCheckpointPath)) - { - } + if (storageOptions.StoreType == StoreType.Local && string.IsNullOrWhiteSpace(storageOptions?.FileSystemStoreOptions?.RootPath)) + throw new InvalidOperationException("RootPath is required when StoreType is Local."); - /// - /// Initializes the manager with the default FileSystem store at a custom root path. - /// - /// Root directory where checkpoints will be stored. - public CheckpointManager(string rootPath) : this(new FileSystemModelStore(rootPath)) - { - } + _store = StoreModelFactory.Create(storageOptions); - /// - /// Initializes the manager with the default FileSystem store and enables background saves. - /// - /// Background save configuration (Enabled, QueueCapacity, OnError). - public CheckpointManager(BackgroundSaveOptions options) - : this(new FileSystemModelStore(_defaultCheckpointPath), options) - { - } - - /// - /// Initializes the manager with a custom storage provider and enables background saves. - /// - /// Any implementation of IModelStore (FileSystem, PostgreSQL, etc.). - /// Background save configuration (Enabled, QueueCapacity, OnError). - public CheckpointManager(IModelStore store, BackgroundSaveOptions options) - { - _store = store ?? throw new ArgumentNullException(nameof(store)); + // Ensure schema only if the store is a database store + if (storageOptions?.DbStoreOptions?.EnsureSchemaOnStartup == true && _store is IDbModelStore dbStore) + { + dbStore.EnsureSchemaAsync().GetAwaiter().GetResult(); + } - if (options != null && options.Enabled) + if (storageOptions?.BackgroundSaveOptions?.Enabled == true) { _backgroundSaver = new BackgroundSaver( - capacity: options.QueueCapacity, - onError: options.OnError - ); + capacity: storageOptions.BackgroundSaveOptions.QueueCapacity, + onError: storageOptions?.BackgroundSaveOptions?.OnError); } } @@ -101,10 +77,8 @@ public async Task SaveAsync( Tags = tags ?? new Dictionary() }; - // --- BACKGROUND MODE --- if (_backgroundSaver != null) { - // DEEP COPY the byte arrays to prevent mutation during background write var capturedCheckpoint = new ModelCheckpoint { ModelId = checkpoint.ModelId, @@ -118,16 +92,58 @@ public async Task SaveAsync( Tags = checkpoint.Tags }; - await _backgroundSaver.EnqueueAsync(async (ct) => await _store.SaveAsync(capturedCheckpoint, ct)); + await _backgroundSaver.EnqueueAsync(async (ct) => await _store.SaveAsync(capturedCheckpoint, ct), cancellationToken); - return checkpoint.ModelId; // Returns immediately! + return checkpoint.ModelId; } - // --- SYNCHRONOUS MODE (Default) --- await _store.SaveAsync(checkpoint, cancellationToken); + return checkpoint.ModelId; } + /// Finds the checkpoint with the highest score according to a user-supplied selector. + /// Queries only summaries (no binary data) and loads the winning checkpoint. + /// + /// Function that maps a CheckpointSummary to a numeric score (higher = better). + /// Optional query filter to narrow the search. + /// Cancellation token. + /// The full ModelCheckpoint with the highest score, or null if none exist. + public async Task FindBestAsync( + Func scoreSelector, + CheckpointQuery? filter = null, + CancellationToken cancellationToken = default) + { + filter ??= new CheckpointQuery(); + var summaries = await _store.QueryAsync(filter, cancellationToken); + + if (summaries.Count == 0) return null; + + var bestSummary = summaries.MaxBy(scoreSelector); + + if (bestSummary == null) return null; + + return await _store.LoadAsync(bestSummary.ModelId, cancellationToken); + } + + /// + /// Finds the checkpoint with the lowest training loss (optionally filtered). + /// + public Task FindBestByLossAsync(CheckpointQuery? filter = null, CancellationToken cancellationToken = default) + => FindBestAsync(s => -s.LastTrainingLoss, filter, cancellationToken); + + /// + /// Finds the checkpoint with the most recent epoch (optionally filtered). + /// + public Task FindLatestEpochAsync(CheckpointQuery? filter = null, CancellationToken cancellationToken = default) + => FindBestAsync(s => s.CurrentEpoch, filter, cancellationToken); + + public Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default) + => _store.QueryAsync(query, cancellationToken); + + public IAsyncEnumerable QueryStreamAsync(CheckpointQuery query, CancellationToken cancellationToken = default) + => _store.QueryStreamAsync(query, cancellationToken); + /// /// Load a checkpoint. /// diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs index 88d0087..e05cb88 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs @@ -1,64 +1,32 @@ using StateCheckpoint.NET.Models; using StateCheckpoint.NET.Settings; -using StateCheckpoint.NET.Stores; -namespace StateCheckpoint.NET.Manager; +namespace StateCheckpoint.NET; public class SessionManager : IAsyncDisposable { private readonly ISessionStore _store; private readonly BackgroundSaver? _backgroundSaver; - private static string _defaultSessionPath = "./sessions"; /// /// Initializes the manager with a custom storage provider. /// /// Any implementation of ISessionStore (FileSystem, PostgreSQL, etc.) - public SessionManager(ISessionStore store) + public SessionManager(StorageOptions storageOptions) { - _store = store ?? throw new ArgumentNullException(nameof(store)); - } - - /// - /// Initializes the manager with the default FileSystem store. - /// Sessions are saved to ./sessions by default. - /// - public SessionManager() : this(new FileSystemSessionStore(_defaultSessionPath)) - { - } - - /// - /// Initializes the manager with the default FileSystem store at a custom root path. - /// - /// Root directory where sessions will be stored. - public SessionManager(string rootPath) : this(new FileSystemSessionStore(rootPath)) - { - } + // ... same validation as CheckpointManager + _store = StoreSessionFactory.Create(storageOptions); - /// - /// Initializes the manager with the default FileSystem store and enables background saves. - /// - /// Background save configuration (Enabled, QueueCapacity, OnError). - public SessionManager(BackgroundSaveOptions options) - : this(new FileSystemSessionStore(_defaultSessionPath), options) - { - } - - /// - /// Initializes the manager with a custom storage provider and enables background saves. - /// - /// Any implementation of ISessionStore (FileSystem, PostgreSQL, etc.). - /// Background save configuration (Enabled, QueueCapacity, OnError). - public SessionManager(ISessionStore store, BackgroundSaveOptions options) - { - _store = store ?? throw new ArgumentNullException(nameof(store)); + if (storageOptions?.DbStoreOptions?.EnsureSchemaOnStartup == true && _store is IDbSessionStore dbStore) + { + dbStore.EnsureSchemaAsync().GetAwaiter().GetResult(); + } - if (options != null && options.Enabled) + if (storageOptions?.BackgroundSaveOptions?.Enabled == true) { _backgroundSaver = new BackgroundSaver( - capacity: options.QueueCapacity, - onError: options.OnError - ); + capacity: storageOptions.BackgroundSaveOptions.QueueCapacity, + onError: storageOptions?.BackgroundSaveOptions?.OnError); } } @@ -95,10 +63,8 @@ public async Task SaveAsync( Tags = tags ?? new Dictionary() }; - // --- BACKGROUND MODE --- if (_backgroundSaver != null) { - // DEEP COPY the byte array to prevent mutation during background write var capturedSession = new SessionCheckpoint { SessionId = session.SessionId, @@ -115,8 +81,8 @@ public async Task SaveAsync( return session.SessionId; } - // --- SYNCHRONOUS MODE (Default) --- await _store.SaveAsync(session, cancellationToken); + return session.SessionId; } @@ -147,6 +113,12 @@ public async Task DeleteAsync(Guid sessionId, CancellationToken cancellationToke public async Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default) => await _store.ListAsync(tagKey, tagValue, cancellationToken); + public Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default) + => _store.QueryAsync(query, cancellationToken); + + public IAsyncEnumerable QueryStreamAsync(SessionQuery query, CancellationToken cancellationToken = default) + => _store.QueryStreamAsync(query, cancellationToken); + /// /// Disposes the manager and ensures the background saver finishes all pending operations. /// @@ -160,8 +132,9 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue public async ValueTask DisposeAsync() { if (_backgroundSaver != null) - { await _backgroundSaver.DisposeAsync(); - } + + if (_store is IAsyncDisposable asyncDisposable) + await asyncDisposable.DisposeAsync(); } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointQuery.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointQuery.cs new file mode 100644 index 0000000..56c6556 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointQuery.cs @@ -0,0 +1,15 @@ +namespace StateCheckpoint.NET.Models; + +public class CheckpointQuery +{ + public int? MinEpoch { get; set; } + public int? MaxEpoch { get; set; } + public float? MinLoss { get; set; } + public float? MaxLoss { get; set; } + public DateTime? CreatedAfter { get; set; } + public DateTime? CreatedBefore { get; set; } + public Dictionary? Tags { get; set; } // AND‑match on all pairs + public int? Limit { get; set; } + public CheckpointSortField OrderBy { get; set; } = CheckpointSortField.CreatedAt; + public bool Descending { get; set; } = true; +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs new file mode 100644 index 0000000..537146e --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs @@ -0,0 +1,10 @@ +namespace StateCheckpoint.NET.Models; + +public class CheckpointSummary +{ + public Guid ModelId { get; set; } + public int CurrentEpoch { get; set; } + public float LastTrainingLoss { get; set; } + public DateTime CreatedAt { get; set; } + public Dictionary Tags { get; set; } = new(); +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/CheckpointSortField.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/CheckpointSortField.cs new file mode 100644 index 0000000..e322a24 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/CheckpointSortField.cs @@ -0,0 +1,8 @@ +namespace StateCheckpoint.NET.Models; + +public enum CheckpointSortField +{ + CreatedAt, + Epoch, + Loss +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SessionSortField.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SessionSortField.cs new file mode 100644 index 0000000..f7fedf4 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SessionSortField.cs @@ -0,0 +1,10 @@ +using System; +using System.Collections.Generic; +using System.Text; + +namespace StateCheckpoint.NET.Models; + +public enum SessionSortField +{ + LastUpdated +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SqlDbType.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SqlDbType.cs new file mode 100644 index 0000000..dc72b55 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/SqlDbType.cs @@ -0,0 +1,7 @@ +namespace StateCheckpoint.NET.Models; + +public enum SqlDbType +{ + Postgres, + SqlServer +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/StoreType.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/StoreType.cs new file mode 100644 index 0000000..a99fca0 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/Enums/StoreType.cs @@ -0,0 +1,7 @@ +namespace StateCheckpoint.NET.Models; + +public enum StoreType +{ + Local, // FileSystem + SqlDb // PostgreSQL or SQL Server +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionManifest.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionManifest.cs index 6c50273..23791ab 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionManifest.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionManifest.cs @@ -1,6 +1,6 @@ using StateCheckpoint.NET.Models; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; internal class SessionManifest { diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionQuery.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionQuery.cs new file mode 100644 index 0000000..a7c89c1 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionQuery.cs @@ -0,0 +1,12 @@ +namespace StateCheckpoint.NET.Models; + +public class SessionQuery +{ + public string? ModelFingerprint { get; set; } + public DateTime? UpdatedAfter { get; set; } + public DateTime? UpdatedBefore { get; set; } + public Dictionary? Tags { get; set; } + public int? Limit { get; set; } + public SessionSortField OrderBy { get; set; } = SessionSortField.LastUpdated; + public bool Descending { get; set; } = true; +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs new file mode 100644 index 0000000..cdbacc2 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs @@ -0,0 +1,10 @@ +namespace StateCheckpoint.NET.Models; + +public class SessionSummary +{ + public Guid SessionId { get; set; } + public string ModelFingerprint { get; set; } = string.Empty; + public DateTime LastUpdated { get; set; } + public Dictionary Tags { get; set; } = new(); +} + diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/TokenizerData.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/TokenizerData.cs index 057605b..6ed2c8a 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Models/TokenizerData.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/TokenizerData.cs @@ -2,36 +2,23 @@ public class TokenizerData { - // 1. The actual type (e.g., "BPE", "Unigram", "WordPiece") public string Type { get; set; } = "BPE"; - // 2. The core vocabulary: Token string -> Integer ID public Dictionary TokenToId { get; set; } = new(); - // 3. Reverse lookup: Integer ID -> Token string - // (Saves user from having to reverse the dictionary manually) public Dictionary IdToToken { get; set; } = new(); - // 4. BPE Merge Rules (only relevant for BPE/ByteLevelBPE). - // Order matters! The order of merges defines the tokenization priority. public List? MergeRules { get; set; } - // 5. Optional: Unigram log-probabilities (only used for Unigram tokenizers). - // Make it nullable to keep the JSON small for BPE users. public Dictionary? TokenLogProbabilities { get; set; } - // 6. Explicit Special Tokens (solves your "where is BOS?" problem). - // e.g., { "bos": 1, "eos": 2, "pad": 0, "unk": 3 } public Dictionary SpecialTokens { get; set; } = new() { { "bos", 1 }, { "eos", 2 }, { "pad", 0 }, { "unk", 3 } // sane defaults }; - // 7. Optional byte-fallback mapping (for byte-level BPE like Llama/GPT-4) - // If the user uses byte-fallback, store the byte -> token mapping here. public Dictionary? ByteToToken { get; set; } - // Helper method for the user to easily get the EOS id public int GetSpecialTokenId(string name, int defaultValue = -1) { return SpecialTokens.TryGetValue(name, out var id) ? id : defaultValue; diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresLargeObjectQueries.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresLargeObjectQueries.cs index 1da759f..13fc41d 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresLargeObjectQueries.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresLargeObjectQueries.cs @@ -1,4 +1,4 @@ -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; /// /// Low-level PostgreSQL Large Object SQL functions. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresSessionQueries.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresSessionQueries.cs index 73c184c..d7eeb15 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresSessionQueries.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresSessionQueries.cs @@ -1,4 +1,4 @@ -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; /// /// PostgreSQL queries specific to the Inference (Session) domain. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresTrainingQueries.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresTrainingQueries.cs index 969e6d0..5989bb8 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresTrainingQueries.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/PostgresTrainingQueries.cs @@ -1,4 +1,4 @@ -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; /// /// PostgreSQL queries specific to the Training (Model) domain. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerSessionQueries.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerSessionQueries.cs index 3946e83..86edb53 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerSessionQueries.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerSessionQueries.cs @@ -1,6 +1,4 @@ -using System.Security.Cryptography; - -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; /// /// SQL Server queries specific to the Inference (Session) domain. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs index 4bbfcfb..c86c344 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs @@ -1,4 +1,4 @@ -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; /// /// SQL Server queries specific to the Training (Model) domain. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/BackgroundSaveOptions.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/BackgroundSaveOptions.cs index 33f5563..c4662be 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/BackgroundSaveOptions.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/BackgroundSaveOptions.cs @@ -1,5 +1,8 @@ namespace StateCheckpoint.NET.Settings; +/// +/// +/// public class BackgroundSaveOptions { /// diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs new file mode 100644 index 0000000..7771340 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs @@ -0,0 +1,14 @@ +namespace StateCheckpoint.NET.Settings; + +/// +/// +/// +public class DbStorageOptions +{ + /// + /// + /// + public string? ConnectionString { get; set; } + + public bool EnsureSchemaOnStartup { get; set; } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/FileSystemStoreOptions.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/FileSystemStoreOptions.cs index afd472a..38b6166 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/FileSystemStoreOptions.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/FileSystemStoreOptions.cs @@ -2,6 +2,11 @@ public class FileSystemStoreOptions { + /// + /// + /// + public string? RootPath { get; set; } = "./checkpoints"; + /// /// If true, the library will create the root directory if it does not exist. /// If false, the library will throw an exception if the directory is missing. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/RetentionPolicy.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/RetentionPolicy.cs new file mode 100644 index 0000000..55d47d0 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/RetentionPolicy.cs @@ -0,0 +1,22 @@ +namespace StateCheckpoint.NET; + +public class RetentionPolicy +{ + /// Keep only the N most recent checkpoints (by CreatedAt). + public int? MaxCheckpoints { get; set; } + + /// Delete checkpoints older than this age. + public TimeSpan? MaxAge { get; set; } + + /// Delete checkpoints with loss above this threshold. + public float? MaxLossThreshold { get; set; } + + /// Delete checkpoints more than N epochs behind the latest saved epoch. + public int? MaxEpochAge { get; set; } + + /// Never delete the checkpoint with the lowest loss, even if other rules would remove it. Default: true. + public bool AlwaysKeepBest { get; set; } = true; + + /// Tags that exempt a checkpoint from all retention rules (e.g., { "status": "production" }). + public Dictionary? PinnedTags { get; set; } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/StorageOptions.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/StorageOptions.cs new file mode 100644 index 0000000..b73238d --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/StorageOptions.cs @@ -0,0 +1,40 @@ +using StateCheckpoint.NET.Models; + +namespace StateCheckpoint.NET.Settings; + +/// +/// +/// +public class StorageOptions +{ + /// + /// + /// + public StoreType StoreType { get; set; } = StoreType.Local; + + /// + /// + /// + public SqlDbType SqlDbType { get; set; } = SqlDbType.Postgres; // ignored if StorageType is not SqlDb + + /// + /// + /// + public FileSystemStoreOptions? FileSystemStoreOptions { get; set; } + + /// + /// + /// + public DbStorageOptions? DbStoreOptions { get; set; } + + /// + /// + /// + public BackgroundSaveOptions? BackgroundSaveOptions { get; set; } + + /// + /// + /// + public RetentionPolicy? RetentionPolicy { get; set; } +} + diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs index de1a58b..6694518 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs @@ -1,12 +1,17 @@ using StateCheckpoint.NET.Models; using StateCheckpoint.NET.Settings; +using StateCheckpoint.NET.Stores; +using System.Runtime.CompilerServices; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; -public class FileSystemModelStore : IModelStore +internal class FileSystemModelStore : IFileSystemModelStore { private readonly string _rootPath; private readonly FileSystemStoreOptions _options; + private readonly CheckpointIndex _index; + private readonly SemaphoreSlim _buildLock = new(1, 1); + private bool _indexLoaded; public FileSystemModelStore(string rootPath, FileSystemStoreOptions? options = null) { @@ -38,6 +43,10 @@ public FileSystemModelStore(string rootPath, FileSystemStoreOptions? options = n Directory.CreateDirectory(_rootPath); } } + + var indexFilePath = Path.Combine(_rootPath, "_index.json"); + + _index = new CheckpointIndex(indexFilePath); } /// @@ -71,6 +80,15 @@ await FileSystemHelper.SaveMultipleAsync( manifest, _options, cancellationToken: cancellationToken); + + await _index.AddOrUpdateAsync(new CheckpointSummary + { + CreatedAt = checkpoint.CreatedAt, + CurrentEpoch = checkpoint.CurrentEpoch, + LastTrainingLoss = checkpoint.LastTrainingLoss, + ModelId = checkpoint.ModelId, + Tags = checkpoint.Tags + }, cancellationToken); } /// @@ -86,7 +104,7 @@ await FileSystemHelper.SaveMultipleAsync( var (binaries, manifest) = await FileSystemHelper.LoadMultipleAsync( _rootPath, modelId, - new[] { "weights.bin", "optimizer.bin" }, // ✅ Hardcoded constants + new[] { "weights.bin", "optimizer.bin" }, "manifest.json", cancellationToken); @@ -109,28 +127,168 @@ await FileSystemHelper.SaveMultipleAsync( } } - public Task DeleteAsync(Guid modelId, CancellationToken cancellationToken = default) - => FileSystemHelper.DeleteAsync(_rootPath, modelId, cancellationToken); + public async Task DeleteAsync(Guid modelId, CancellationToken cancellationToken = default) + { + await FileSystemHelper.DeleteAsync(_rootPath, modelId, cancellationToken); - public async Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default) + await _index.RemoveAsync(modelId, cancellationToken); + } + + public async Task DeleteManyAsync(IEnumerable modelIds, CancellationToken cancellationToken = default) { - var allIds = await FileSystemHelper.ListAsync(_rootPath, cancellationToken); + foreach (var id in modelIds) + await DeleteAsync(id, cancellationToken); + } - // If no filter, return all + public async Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default) + { if (string.IsNullOrWhiteSpace(tagKey) || string.IsNullOrWhiteSpace(tagValue)) - return allIds; + return await FileSystemHelper.ListAsync(_rootPath, cancellationToken); + + await EnsureIndexLoadedAsync(cancellationToken); + + return (await _index.GetAllAsync(cancellationToken)) + .Where(s => s.Tags != null && s.Tags.TryGetValue(tagKey, out var val) && val == tagValue) + .Select(s => s.ModelId) + .ToList(); + } + + public async Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default) + { + var allIds = await FileSystemHelper.ListAsync(_rootPath, cancellationToken); + var summaries = new List(); - // Filter in-memory by loading each manifest - var filtered = new List(); foreach (var id in allIds) { cancellationToken.ThrowIfCancellationRequested(); var manifest = await FileSystemHelper.LoadManifestOnlyAsync(_rootPath, id, "manifest.json", cancellationToken); - if (manifest != null && manifest.Tags != null && manifest.Tags.TryGetValue(tagKey, out var val) && val == tagValue) - filtered.Add(id); + if (manifest == null) continue; + + // Apply filters + if (query.MinEpoch.HasValue && manifest.CurrentEpoch < query.MinEpoch) continue; + if (query.MaxEpoch.HasValue && manifest.CurrentEpoch > query.MaxEpoch) continue; + if (query.MinLoss.HasValue && manifest.LastTrainingLoss < query.MinLoss) continue; + if (query.MaxLoss.HasValue && manifest.LastTrainingLoss > query.MaxLoss) continue; + if (query.CreatedAfter.HasValue && manifest.CreatedAt < query.CreatedAfter) continue; + if (query.CreatedBefore.HasValue && manifest.CreatedAt > query.CreatedBefore) continue; + if (query.Tags != null && !TagsMatch(manifest.Tags, query.Tags)) continue; + + summaries.Add(new CheckpointSummary + { + ModelId = id, + CurrentEpoch = manifest.CurrentEpoch, + LastTrainingLoss = manifest.LastTrainingLoss, + CreatedAt = manifest.CreatedAt, + Tags = manifest.Tags ?? new() + }); } - return filtered; + // Sorting + summaries = query.OrderBy switch + { + CheckpointSortField.CreatedAt => query.Descending + ? summaries.OrderByDescending(s => s.CreatedAt).ToList() + : summaries.OrderBy(s => s.CreatedAt).ToList(), + CheckpointSortField.Epoch => query.Descending + ? summaries.OrderByDescending(s => s.CurrentEpoch).ToList() + : summaries.OrderBy(s => s.CurrentEpoch).ToList(), + CheckpointSortField.Loss => query.Descending + ? summaries.OrderByDescending(s => s.LastTrainingLoss).ToList() + : summaries.OrderBy(s => s.LastTrainingLoss).ToList(), + _ => summaries + }; + + if (query.Limit.HasValue && query.Limit.Value > 0) + summaries = summaries.Take(query.Limit.Value).ToList(); + + return summaries; + } + + public async IAsyncEnumerable QueryStreamAsync( + CheckpointQuery query, + [EnumeratorCancellation] CancellationToken ct = default) + { + await EnsureIndexLoadedAsync(ct); + var all = await _index.GetAllAsync(ct); + + // Filter in memory (no sorting/limiting – streaming preserves order of storage) + foreach (var summary in all) + { + ct.ThrowIfCancellationRequested(); + + if (query.MinEpoch.HasValue && summary.CurrentEpoch < query.MinEpoch) continue; + if (query.MaxEpoch.HasValue && summary.CurrentEpoch > query.MaxEpoch) continue; + if (query.MinLoss.HasValue && summary.LastTrainingLoss < query.MinLoss) continue; + if (query.MaxLoss.HasValue && summary.LastTrainingLoss > query.MaxLoss) continue; + if (query.CreatedAfter.HasValue && summary.CreatedAt < query.CreatedAfter) continue; + if (query.CreatedBefore.HasValue && summary.CreatedAt > query.CreatedBefore) continue; + if (query.Tags != null && query.Tags.Count > 0 && !TagsMatch(summary.Tags, query.Tags)) continue; + + yield return summary; + } + } + + private static bool TagsMatch(Dictionary? actual, Dictionary? filter) + { + if (filter == null) + return true; + + if (actual == null) + return false; + + foreach (var pair in filter) + if (!actual.TryGetValue(pair.Key, out var val) || val != pair.Value) + return false; + + return true; + } + + private async Task EnsureIndexLoadedAsync(CancellationToken ct) + { + // We can check if _items is empty or use a flag; simplest: always load if not loaded. + // We'll have a separate field to track loaded state. + // Here we just call LoadAsync if needed. Since it's cheap, we can call it every time. + // But better to check a flag. + if (!_indexLoaded) + { + await _index.LoadAsync(ct); + _indexLoaded = true; + } + } + + private async Task RebuildIndexAsync(CancellationToken ct) + { + await _buildLock.WaitAsync(ct); + try + { + // double-check after lock + await EnsureIndexLoadedAsync(ct); + var existing = await _index.GetAllAsync(ct); + if (existing.Any()) return; + + async Task LoadManifestSummaryAsync(string rootPath, Guid id, CancellationToken ct) + { + var manifest = await FileSystemHelper.LoadManifestOnlyAsync(rootPath, id, "manifest.json", ct); + + if (manifest == null) return null; + + return new CheckpointSummary + { + ModelId = id, + CurrentEpoch = manifest.CurrentEpoch, + LastTrainingLoss = manifest.LastTrainingLoss, + CreatedAt = manifest.CreatedAt, + Tags = manifest.Tags ?? new() + }; + } + + // Rebuild + await _index.RebuildAsync(_rootPath, LoadManifestSummaryAsync, ct); + } + finally + { + _buildLock.Release(); + } } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs index 9060cee..1194778 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs @@ -1,12 +1,15 @@ using StateCheckpoint.NET.Models; using StateCheckpoint.NET.Settings; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; -public class FileSystemSessionStore : ISessionStore +internal class FileSystemSessionStore : IFileSystemSessionStore { private readonly string _rootPath; private readonly FileSystemStoreOptions _options; + private readonly SessionIndex _index; + private bool _indexLoaded; + private readonly SemaphoreSlim _buildLock = new(1, 1); public FileSystemSessionStore(string rootPath, FileSystemStoreOptions? options = null) { @@ -18,7 +21,7 @@ public FileSystemSessionStore(string rootPath, FileSystemStoreOptions? options = if (!FileSystemHelper.TryValidateWriteAccess(_rootPath, out var error)) { // If fallback is provided, update the root path - if (!string.IsNullOrEmpty(_options.FallbackPath)) + if (!string.IsNullOrWhiteSpace(_options.FallbackPath)) { _rootPath = Path.Combine(_options.FallbackPath, "sessions"); @@ -38,9 +41,14 @@ public FileSystemSessionStore(string rootPath, FileSystemStoreOptions? options = Directory.CreateDirectory(_rootPath); } } + + // Index file path + var indexFilePath = Path.Combine(_rootPath, "_sessions_index.json"); + + _index = new SessionIndex(indexFilePath); } - public async Task SaveAsync(SessionCheckpoint session, CancellationToken ct = default) + public async Task SaveAsync(SessionCheckpoint session, CancellationToken cancellationToken = default) { var manifest = new SessionManifest { @@ -59,7 +67,19 @@ await FileSystemHelper.SaveAsync( _options, binaryFileName: "kv.bin", metaFileName: "meta.json", - cancellationToken: ct); + cancellationToken: cancellationToken); + + var summary = new SessionSummary + { + SessionId = session.SessionId, + ModelFingerprint = session.ModelFingerprint, + LastUpdated = session.LastUpdated, + Tags = session.Tags ?? new Dictionary() + }; + + await _index.AddOrUpdateAsync(summary, cancellationToken); + + if (!_indexLoaded) _indexLoaded = true; // index is now loaded } public async Task LoadAsync(Guid sessionId, CancellationToken cancellationToken = default) @@ -86,25 +106,76 @@ await FileSystemHelper.SaveAsync( } } - public Task DeleteAsync(Guid sessionId, CancellationToken cancellationToken = default) - => FileSystemHelper.DeleteAsync(_rootPath, sessionId, cancellationToken); + public async Task DeleteAsync(Guid sessionId, CancellationToken cancellationToken = default) + { + await FileSystemHelper.DeleteAsync(_rootPath, sessionId, cancellationToken); + + await _index.RemoveAsync(sessionId, cancellationToken); + } public async Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default) { - var allIds = await FileSystemHelper.ListAsync(_rootPath, cancellationToken); if (string.IsNullOrWhiteSpace(tagKey) || string.IsNullOrWhiteSpace(tagValue)) - return allIds; + return await FileSystemHelper.ListAsync(_rootPath, cancellationToken); - var filtered = new List(); - foreach (var id in allIds) + await EnsureIndexLoadedAsync(cancellationToken); + + return (await _index.GetAllAsync(cancellationToken)) + .Where(s => s.Tags != null && s.Tags.TryGetValue(tagKey, out var val) && val == tagValue) + .Select(s => s.SessionId) + .ToList(); + } + + public async Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default) + { + await EnsureIndexLoadedAsync(cancellationToken); + + return await _index.QueryAsync(query, cancellationToken); + } + + // Rebuild index from all session manifest files (fallback) + public async Task RebuildIndexAsync(CancellationToken cancellationToken) + { + await _buildLock.WaitAsync(cancellationToken); + try { - cancellationToken.ThrowIfCancellationRequested(); + // Double-check after acquiring the lock + await EnsureIndexLoadedAsync(cancellationToken); + + var existing = await _index.GetAllAsync(cancellationToken); + + if (existing.Any()) return; // already rebuilt by another thread - var manifest = await FileSystemHelper.LoadManifestOnlyAsync(_rootPath, id, "meta.json", cancellationToken); - if (manifest != null && manifest.Tags != null && manifest.Tags.TryGetValue(tagKey, out var val) && val == tagValue) - filtered.Add(id); + await _index.RebuildAsync(_rootPath, LoadManifestSummaryAsync, cancellationToken); } + finally + { + _buildLock.Release(); + } + } - return filtered; + // Ensure the index is loaded into memory (lazy) + private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) + { + if (!_indexLoaded) + { + await _index.LoadAsync(cancellationToken); + + _indexLoaded = true; + } + } + + private async Task LoadManifestSummaryAsync(string rootPath, Guid id, CancellationToken ct) + { + var manifest = await FileSystemHelper.LoadManifestOnlyAsync( + rootPath, id, "meta.json", ct); + if (manifest == null) return null; + return new SessionSummary + { + SessionId = id, + ModelFingerprint = manifest.ModelFingerprint, + LastUpdated = manifest.LastUpdated, + Tags = manifest.Tags ?? new Dictionary() + }; } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs new file mode 100644 index 0000000..09812db --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs @@ -0,0 +1,181 @@ +using StateCheckpoint.NET.Models; + +namespace StateCheckpoint.NET; + +internal sealed class CheckpointIndex +{ + private readonly IndexStorage _storage; + private readonly SemaphoreSlim _lock = new(1, 1); + private List _items; + + public CheckpointIndex(string indexFilePath) + : this(new IndexStorage(indexFilePath)) + { + } + + internal CheckpointIndex(IndexStorage storage) // for testing + { + _storage = storage; + _items = new List(); + } + + /// + /// Loads the index from storage; if no file exists, starts with an empty list. + /// Must be called before any other operations. + /// + public async Task LoadAsync(CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + var data = await _storage.ReadAsync(ct); + _items = data.ToList(); + } + finally + { + _lock.Release(); + } + } + + public async Task AddOrUpdateAsync(CheckpointSummary summary, CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + var existing = _items.FirstOrDefault(e => e.ModelId == summary.ModelId); + if (existing != null) + { + existing.CurrentEpoch = summary.CurrentEpoch; + existing.LastTrainingLoss = summary.LastTrainingLoss; + existing.CreatedAt = summary.CreatedAt; + existing.Tags = summary.Tags; + } + else + { + _items.Add(summary); + } + await _storage.WriteAsync(_items, ct); + } + finally + { + _lock.Release(); + } + } + + public async Task RemoveAsync(Guid modelId, CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + var removed = _items.RemoveAll(e => e.ModelId == modelId); + if (removed > 0) + await _storage.WriteAsync(_items, ct); + } + finally + { + _lock.Release(); + } + } + + /// + /// Rebuilds the index by enumerating all manifest files. + /// Acquires the lock and overwrites the file. + /// + public async Task RebuildAsync( + string rootPath, + Func> loadManifestFn, + CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + var allIds = await FileSystemHelper.ListAsync(rootPath, ct); + var newItems = new List(); + foreach (var id in allIds) + { + ct.ThrowIfCancellationRequested(); + var summary = await loadManifestFn(rootPath, id, ct); + if (summary != null) + newItems.Add(summary); + } + _items = newItems; + await _storage.WriteAsync(_items, ct); + } + finally + { + _lock.Release(); + } + } + + /// + /// Returns a copy of all items (for querying). + /// + public async Task> GetAllAsync(CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + return _items.ToList(); // defensive copy + } + finally + { + _lock.Release(); + } + } + + /// + /// Queries the index with a predicate and optional sorting/limit. + /// + public async Task> QueryAsync( + CheckpointQuery query, + CancellationToken ct = default) + { + var all = await GetAllAsync(ct); + var filtered = all.AsEnumerable(); + + // Apply filters (same as before) + if (query.MinEpoch.HasValue) filtered = filtered.Where(s => s.CurrentEpoch >= query.MinEpoch.Value); + if (query.MaxEpoch.HasValue) filtered = filtered.Where(s => s.CurrentEpoch <= query.MaxEpoch.Value); + if (query.MinLoss.HasValue) filtered = filtered.Where(s => s.LastTrainingLoss >= query.MinLoss.Value); + if (query.MaxLoss.HasValue) filtered = filtered.Where(s => s.LastTrainingLoss <= query.MaxLoss.Value); + if (query.CreatedAfter.HasValue) filtered = filtered.Where(s => s.CreatedAt >= query.CreatedAfter.Value); + if (query.CreatedBefore.HasValue) filtered = filtered.Where(s => s.CreatedAt <= query.CreatedBefore.Value); + if (query.Tags != null && query.Tags.Count > 0) + filtered = filtered.Where(s => TagsMatch(s.Tags, query.Tags)); + + // Sort + var sorted = query.OrderBy switch + { + CheckpointSortField.CreatedAt => query.Descending + ? filtered.OrderByDescending(s => s.CreatedAt) + : filtered.OrderBy(s => s.CreatedAt), + CheckpointSortField.Epoch => query.Descending + ? filtered.OrderByDescending(s => s.CurrentEpoch) + : filtered.OrderBy(s => s.CurrentEpoch), + CheckpointSortField.Loss => query.Descending + ? filtered.OrderByDescending(s => s.LastTrainingLoss) + : filtered.OrderBy(s => s.LastTrainingLoss), + _ => filtered + }; + + if (query.Limit.HasValue) + sorted = sorted.Take(query.Limit.Value); + + return sorted.ToList(); + } + + private static bool TagsMatch(Dictionary? actual, Dictionary? filter) + { + if (filter == null) + return true; + + if (actual == null) + return false; + + foreach (var pair in filter) + if (!actual.TryGetValue(pair.Key, out var val) || val != pair.Value) + return false; + + return true; + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs new file mode 100644 index 0000000..f270815 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs @@ -0,0 +1,36 @@ +using System.Text.Json; + +namespace StateCheckpoint.NET; + + +internal sealed class IndexStorage +{ + private readonly string _filePath; + private readonly JsonSerializerOptions _jsonOpts; + + public IndexStorage(string filePath, JsonSerializerOptions? jsonOpts = null) + { + _filePath = filePath; + _jsonOpts = jsonOpts ?? new JsonSerializerOptions { WriteIndented = true }; + } + + public async Task> ReadAsync(CancellationToken ct = default) + { + if (!File.Exists(_filePath)) + return Array.Empty(); + + var json = await File.ReadAllTextAsync(_filePath, ct); + + return JsonSerializer.Deserialize>(json, _jsonOpts) ?? new List(); + } + + public async Task WriteAsync(IEnumerable data, CancellationToken ct = default) + { + var json = JsonSerializer.Serialize(data, _jsonOpts); + var tempPath = _filePath + ".tmp"; + + await File.WriteAllTextAsync(tempPath, json, ct); + + File.Move(tempPath, _filePath, overwrite: true); + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs new file mode 100644 index 0000000..c29d5a2 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs @@ -0,0 +1,163 @@ +using StateCheckpoint.NET.Models; + +namespace StateCheckpoint.NET; + +internal sealed class SessionIndex +{ + private readonly IndexStorage _storage; + private readonly SemaphoreSlim _lock = new(1, 1); + private List _items; + + public SessionIndex(string indexFilePath) + : this(new IndexStorage(indexFilePath)) + { + } + + internal SessionIndex(IndexStorage storage) + { + _storage = storage; + _items = new List(); + } + + public async Task LoadAsync(CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + var data = await _storage.ReadAsync(ct); + _items = data.ToList(); + } + finally + { + _lock.Release(); + } + } + + public async Task AddOrUpdateAsync(SessionSummary summary, CancellationToken cancellationToken = default) + { + await _lock.WaitAsync(cancellationToken); + try + { + var existing = _items.FirstOrDefault(e => e.SessionId == summary.SessionId); + if (existing != null) + { + existing.ModelFingerprint = summary.ModelFingerprint; + existing.LastUpdated = summary.LastUpdated; + existing.Tags = summary.Tags; + } + else + { + _items.Add(summary); + } + + await _storage.WriteAsync(_items, cancellationToken); + } + finally + { + _lock.Release(); + } + } + + public async Task RemoveAsync(Guid sessionId, CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + var removed = _items.RemoveAll(e => e.SessionId == sessionId); + if (removed > 0) + await _storage.WriteAsync(_items, ct); + } + finally + { + _lock.Release(); + } + } + + public async Task RebuildAsync( + string rootPath, + Func> loadManifestFn, + CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + try + { + var allIds = await FileSystemHelper.ListAsync(rootPath, ct); + + var newItems = new List(); + foreach (var id in allIds) + { + ct.ThrowIfCancellationRequested(); + var summary = await loadManifestFn(rootPath, id, ct); + + if (summary != null) + newItems.Add(summary); + } + + _items = newItems; + + await _storage.WriteAsync(_items, ct); + } + finally + { + _lock.Release(); + } + } + + public async Task> GetAllAsync(CancellationToken ct = default) + { + await _lock.WaitAsync(ct); + + try + { + return _items.ToList(); + } + finally + { + _lock.Release(); + } + } + + public async Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default) + { + var all = await GetAllAsync(cancellationToken); + var filtered = all.AsEnumerable(); + + if (!string.IsNullOrWhiteSpace(query.ModelFingerprint)) + filtered = filtered.Where(s => s.ModelFingerprint == query.ModelFingerprint); + + if (query.UpdatedAfter.HasValue) + filtered = filtered.Where(s => s.LastUpdated >= query.UpdatedAfter.Value); + + if (query.UpdatedBefore.HasValue) + filtered = filtered.Where(s => s.LastUpdated <= query.UpdatedBefore.Value); + + if (query.Tags != null && query.Tags.Count > 0) + filtered = filtered.Where(s => TagsMatch(s.Tags, query.Tags)); + + // Sort (only LastUpdated for sessions) + var sorted = query.Descending + ? filtered.OrderByDescending(s => s.LastUpdated) + : filtered.OrderBy(s => s.LastUpdated); + + if (query.Limit.HasValue) + sorted = query.Descending ? sorted.Take(query.Limit.Value).OrderByDescending(s => s.LastUpdated) + : sorted.Take(query.Limit.Value).OrderBy(s => s.LastUpdated); + + return sorted.ToList(); + } + + private static bool TagsMatch(Dictionary? actual, Dictionary? filter) + { + if (filter == null) + return true; + + if (actual == null) + return false; + + foreach (var pair in filter) + if (!actual.TryGetValue(pair.Key, out var val) || val != pair.Value) + return false; + + return true; + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbModelStore.cs new file mode 100644 index 0000000..7db5f6c --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbModelStore.cs @@ -0,0 +1,6 @@ +namespace StateCheckpoint.NET; + +internal interface IDbModelStore: IModelStore +{ + Task EnsureSchemaAsync(CancellationToken cancellationToken = default); +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbSessionStore.cs new file mode 100644 index 0000000..5970071 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IDbSessionStore.cs @@ -0,0 +1,6 @@ +namespace StateCheckpoint.NET; + +internal interface IDbSessionStore: ISessionStore +{ + Task EnsureSchemaAsync(CancellationToken cancellationToken = default); +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs new file mode 100644 index 0000000..41fa60e --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs @@ -0,0 +1,5 @@ +namespace StateCheckpoint.NET; + +internal interface IFileSystenModelStore: IModelStore +{ +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemSessionStore.cs new file mode 100644 index 0000000..f23dfa1 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemSessionStore.cs @@ -0,0 +1,5 @@ +namespace StateCheckpoint.NET; + +internal interface IFileSystemSessionStore: ISessionStore +{ +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs index 766d451..91a4c57 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs @@ -1,12 +1,14 @@ using StateCheckpoint.NET.Models; -namespace StateCheckpoint.NET.Stores +namespace StateCheckpoint.NET; + +internal interface IModelStore { - public interface IModelStore - { - Task SaveAsync(ModelCheckpoint checkpoint, CancellationToken cancellationToken = default); - Task LoadAsync(Guid modelId, CancellationToken cancellationToken = default); - Task DeleteAsync(Guid modelId, CancellationToken cancellationToken = default); - Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default); - } + Task SaveAsync(ModelCheckpoint checkpoint, CancellationToken cancellationToken = default); + Task LoadAsync(Guid modelId, CancellationToken cancellationToken = default); + Task DeleteAsync(Guid modelId, CancellationToken cancellationToken = default); + Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default); + Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default); + IAsyncEnumerable QueryStreamAsync(CheckpointQuery query, CancellationToken cancellationToken = default); + Task DeleteManyAsync(IEnumerable modelIds, CancellationToken cancellationToken = default); } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs index a6e0533..0487c30 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs @@ -1,12 +1,13 @@ using StateCheckpoint.NET.Models; -namespace StateCheckpoint.NET.Stores +namespace StateCheckpoint.NET; + +internal interface ISessionStore { - public interface ISessionStore - { - Task SaveAsync(SessionCheckpoint session, CancellationToken ct = default); - Task LoadAsync(Guid sessionId, CancellationToken ct = default); - Task DeleteAsync(Guid sessionId, CancellationToken ct = default); - Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken ct = default); - } + Task SaveAsync(SessionCheckpoint session, CancellationToken cancellationToken = default); + Task LoadAsync(Guid sessionId, CancellationToken cancellationToken = default); + Task DeleteAsync(Guid sessionId, CancellationToken cancellationToken = default); + Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default); + Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default); + IAsyncEnumerable QueryStreamAsync(SessionQuery query, CancellationToken cancellationToken = default); } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs index 846a86c..858caf6 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs @@ -1,11 +1,13 @@ -using StateCheckpoint.NET.Models; -using Npgsql; +using Npgsql; using NpgsqlTypes; +using StateCheckpoint.NET.Models; +using System.Runtime.CompilerServices; +using System.Text; using System.Text.Json; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; -public class PostgresModelStore : PostgresStoreBase, IModelStore +internal class PostgresModelStore : PostgresStoreBase, IDbModelStore { private static readonly JsonSerializerOptions _jsonOpts = new() { WriteIndented = false }; @@ -253,6 +255,68 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue return modelIds; } + public async Task> QueryAsync(CheckpointQuery query, CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new NpgsqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + + var results = new List(); + while (await reader.ReadAsync(ct)) + { + results.Add(new CheckpointSummary + { + ModelId = reader.GetGuid(0), + CurrentEpoch = reader.GetInt32(1), + LastTrainingLoss = (float)reader.GetDouble(2), + CreatedAt = reader.GetDateTime(3), + Tags = JsonSerializer.Deserialize>(reader.GetString(4)) ?? new() + }); + } + + return results; + } + + public async IAsyncEnumerable QueryStreamAsync( + CheckpointQuery query, + [EnumeratorCancellation] CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new NpgsqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + while (await reader.ReadAsync(ct)) + { + ct.ThrowIfCancellationRequested(); + yield return new CheckpointSummary + { + ModelId = reader.GetGuid(0), + CurrentEpoch = reader.GetInt32(1), + LastTrainingLoss = (float)reader.GetDouble(2), + CreatedAt = reader.GetDateTime(3), + Tags = JsonSerializer.Deserialize>(reader.GetString(4)) ?? new() + }; + } + } + + public async Task DeleteManyAsync(IEnumerable modelIds, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + await using var transaction = await connection.BeginTransactionAsync(cancellationToken); + + foreach (var id in modelIds) + await DeleteAsync(id, cancellationToken); + + await transaction.CommitAsync(cancellationToken); + } + //--------- private methods ----------------------------------- // --- Large Object Helpers (Use PostgresLargeObjectQueries) --- @@ -365,4 +429,82 @@ private static async Task ReadLargeObjectAsync( return memoryStream.ToArray(); } + + // Helper to build SQL and parameters for queries + private (string Sql, List Parameters) BuildQuerySql(CheckpointQuery query, bool includeOrderBy = true, bool includeLimit = true) + { + var sql = new StringBuilder(@" + SELECT m.model_id, m.epoch, m.loss, m.created_at, m.tags + FROM model_manifests m + JOIN model_blobs b ON m.model_id = b.model_id + WHERE 1=1 + "); + + var parameters = new List(); + int paramIndex = 0; + + if (query.MinEpoch.HasValue) + { + sql.Append($" AND m.epoch >= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.MinEpoch.Value)); + paramIndex++; + } + if (query.MaxEpoch.HasValue) + { + sql.Append($" AND m.epoch <= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.MaxEpoch.Value)); + paramIndex++; + } + if (query.MinLoss.HasValue) + { + sql.Append($" AND m.loss >= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.MinLoss.Value)); + paramIndex++; + } + if (query.MaxLoss.HasValue) + { + sql.Append($" AND m.loss <= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.MaxLoss.Value)); + paramIndex++; + } + if (query.CreatedAfter.HasValue) + { + sql.Append($" AND m.created_at >= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.CreatedAfter.Value)); + paramIndex++; + } + if (query.CreatedBefore.HasValue) + { + sql.Append($" AND m.created_at <= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.CreatedBefore.Value)); + paramIndex++; + } + if (query.Tags != null && query.Tags.Count > 0) + { + var json = JsonSerializer.Serialize(query.Tags); + sql.Append($" AND m.tags @> @p{paramIndex}::jsonb"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", json)); + paramIndex++; + } + + if (includeOrderBy) + { + var orderColumn = query.OrderBy switch + { + CheckpointSortField.CreatedAt => "m.created_at", + CheckpointSortField.Epoch => "m.epoch", + CheckpointSortField.Loss => "m.loss", + _ => "m.created_at" + }; + sql.Append($" ORDER BY {orderColumn} {(query.Descending ? "DESC" : "ASC")}"); + } + + if (includeLimit && query.Limit.HasValue && query.Limit.Value > 0) + { + sql.Append($" LIMIT @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.Limit.Value)); + } + + return (sql.ToString(), parameters); + } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs index d12455b..8f4bf1c 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs @@ -1,10 +1,12 @@ -using System.Text.Json; -using Npgsql; +using Npgsql; using StateCheckpoint.NET.Models; +using System.Runtime.CompilerServices; +using System.Text; +using System.Text.Json; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; -public class PostgresSessionStore : PostgresStoreBase, ISessionStore +internal class PostgresSessionStore : PostgresStoreBase, IDbSessionStore { private static readonly JsonSerializerOptions _jsonOpts = new() { WriteIndented = false }; @@ -130,4 +132,110 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue return sessionIds; } + + public async Task> QueryAsync(SessionQuery query, CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new NpgsqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + var results = new List(); + while (await reader.ReadAsync(ct)) + { + results.Add(new SessionSummary + { + SessionId = reader.GetGuid(0), + ModelFingerprint = reader.GetString(1), + LastUpdated = reader.GetDateTime(2), + Tags = JsonSerializer.Deserialize>(reader.GetString(3)) ?? new() + }); + } + return results; + } + + // New streaming method + public async IAsyncEnumerable QueryStreamAsync( + SessionQuery query, + [EnumeratorCancellation] CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new NpgsqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + while (await reader.ReadAsync(ct)) + { + ct.ThrowIfCancellationRequested(); + yield return new SessionSummary + { + SessionId = reader.GetGuid(0), + ModelFingerprint = reader.GetString(1), + LastUpdated = reader.GetDateTime(2), + Tags = JsonSerializer.Deserialize>(reader.GetString(3)) ?? new() + }; + } + } + + //---------------- Private Methods ---------------// + + private (string Sql, List Parameters) BuildQuerySql(SessionQuery query, bool includeOrderBy = true, bool includeLimit = true) + { + var sql = new StringBuilder(@" + SELECT session_id, model_fingerprint, last_updated, tags + FROM inference_sessions + WHERE 1=1 + "); + + var parameters = new List(); + int paramIndex = 0; + + if (!string.IsNullOrEmpty(query.ModelFingerprint)) + { + sql.Append($" AND model_fingerprint = @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.ModelFingerprint)); + paramIndex++; + } + if (query.UpdatedAfter.HasValue) + { + sql.Append($" AND last_updated >= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.UpdatedAfter.Value)); + paramIndex++; + } + if (query.UpdatedBefore.HasValue) + { + sql.Append($" AND last_updated <= @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.UpdatedBefore.Value)); + paramIndex++; + } + if (query.Tags != null && query.Tags.Count > 0) + { + var json = JsonSerializer.Serialize(query.Tags); + sql.Append($" AND tags @> @p{paramIndex}::jsonb"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", json)); + paramIndex++; + } + + if (includeOrderBy) + { + var orderColumn = query.OrderBy switch + { + SessionSortField.LastUpdated => "last_updated", + _ => "last_updated" + }; + sql.Append($" ORDER BY {orderColumn} {(query.Descending ? "DESC" : "ASC")}"); + } + + if (includeLimit && query.Limit.HasValue && query.Limit.Value > 0) + { + sql.Append($" LIMIT @p{paramIndex}"); + parameters.Add(new NpgsqlParameter($"@p{paramIndex}", query.Limit.Value)); + } + + return (sql.ToString(), parameters); + } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs index 8003855..4364baa 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs @@ -1,12 +1,12 @@ using Npgsql; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; /// /// Abstract base class for PostgreSQL stores. /// Supports connection string OR an externally managed NpgsqlDataSource. /// -public abstract class PostgresStoreBase : IAsyncDisposable +internal abstract class PostgresStoreBase : IAsyncDisposable { private readonly NpgsqlDataSource _dataSource; private readonly bool _ownsDataSource; diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs index 65a0ebb..4fed938 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs @@ -1,10 +1,12 @@ -using StateCheckpoint.NET.Models; -using Microsoft.Data.SqlClient; +using Microsoft.Data.SqlClient; +using StateCheckpoint.NET.Models; +using System.Runtime.CompilerServices; +using System.Text; using System.Text.Json; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; -public class SqlServerModelStore : SqlServerStoreBase, IModelStore +internal class SqlServerModelStore : SqlServerStoreBase, IDbModelStore { private static readonly JsonSerializerOptions _jsonOpts = new() { WriteIndented = false }; @@ -155,4 +157,117 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue return list; } } + + public async Task> QueryAsync(CheckpointQuery query, CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new SqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + var results = new List(); + while (await reader.ReadAsync(ct)) + { + results.Add(new CheckpointSummary + { + ModelId = reader.GetGuid(0), + CurrentEpoch = reader.GetInt32(1), + LastTrainingLoss = (float)reader.GetDouble(2), + CreatedAt = reader.GetDateTime(3), + Tags = JsonSerializer.Deserialize>(reader.GetString(4)) ?? new() + }); + } + + return results; + } + + public async IAsyncEnumerable QueryStreamAsync( + CheckpointQuery query, + [EnumeratorCancellation] CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new SqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + while (await reader.ReadAsync(ct)) + { + ct.ThrowIfCancellationRequested(); + yield return new CheckpointSummary + { + ModelId = reader.GetGuid(0), + CurrentEpoch = reader.GetInt32(1), + LastTrainingLoss = (float)reader.GetDouble(2), + CreatedAt = reader.GetDateTime(3), + Tags = JsonSerializer.Deserialize>(reader.GetString(4)) ?? new() + }; + } + } + + public async Task DeleteManyAsync(IEnumerable modelIds, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + await using var tx = await connection.BeginTransactionAsync(cancellationToken); + + foreach (var id in modelIds) + await DeleteAsync(id, cancellationToken); + + await tx.CommitAsync(cancellationToken); + } + + //------------- Private Methods -----------------// + + private (string Sql, List Parameters) BuildQuerySql(CheckpointQuery query, bool includeOrderBy = true, bool includeLimit = true) + { + var sql = new StringBuilder(@" + SELECT model_id, epoch, loss, created_at, tags + FROM ModelManifests + WHERE 1=1 + "); + + var parameters = new List(); + + if (query.MinEpoch.HasValue) { sql.Append(" AND epoch >= @MinEpoch"); parameters.Add(new SqlParameter("@MinEpoch", query.MinEpoch.Value)); } + if (query.MaxEpoch.HasValue) { sql.Append(" AND epoch <= @MaxEpoch"); parameters.Add(new SqlParameter("@MaxEpoch", query.MaxEpoch.Value)); } + if (query.MinLoss.HasValue) { sql.Append(" AND loss >= @MinLoss"); parameters.Add(new SqlParameter("@MinLoss", query.MinLoss.Value)); } + if (query.MaxLoss.HasValue) { sql.Append(" AND loss <= @MaxLoss"); parameters.Add(new SqlParameter("@MaxLoss", query.MaxLoss.Value)); } + if (query.CreatedAfter.HasValue) { sql.Append(" AND created_at >= @CreatedAfter"); parameters.Add(new SqlParameter("@CreatedAfter", query.CreatedAfter.Value)); } + if (query.CreatedBefore.HasValue) { sql.Append(" AND created_at <= @CreatedBefore"); parameters.Add(new SqlParameter("@CreatedBefore", query.CreatedBefore.Value)); } + + if (query.Tags != null && query.Tags.Count > 0) + { + foreach (var pair in query.Tags) + { + var keyParam = $"@tagKey_{pair.Key}"; + var valParam = $"@tagVal_{pair.Key}"; + sql.Append($" AND JSON_VALUE(tags, '$.{pair.Key}') = {valParam}"); + parameters.Add(new SqlParameter(valParam, pair.Value)); + } + } + + if (includeOrderBy) + { + var orderColumn = query.OrderBy switch + { + CheckpointSortField.CreatedAt => "created_at", + CheckpointSortField.Epoch => "epoch", + CheckpointSortField.Loss => "loss", + _ => "created_at" + }; + + sql.Append($" ORDER BY {orderColumn} {(query.Descending ? "DESC" : "ASC")}"); + } + + if (includeLimit && query.Limit.HasValue && query.Limit.Value > 0) + { + sql.Append(" OFFSET 0 ROWS FETCH NEXT @Limit ROWS ONLY"); + parameters.Add(new SqlParameter("@Limit", query.Limit.Value)); + } + + return (sql.ToString(), parameters); + } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs index d7bbb5f..2341075 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs @@ -1,10 +1,12 @@ -using StateCheckpoint.NET.Models; -using Microsoft.Data.SqlClient; +using Microsoft.Data.SqlClient; +using StateCheckpoint.NET.Models; +using System.Runtime.CompilerServices; +using System.Text; using System.Text.Json; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; -public class SqlServerSessionStore : SqlServerStoreBase, ISessionStore +internal class SqlServerSessionStore : SqlServerStoreBase, IDbSessionStore { private static readonly JsonSerializerOptions _jsonOpts = new() { WriteIndented = false }; @@ -140,4 +142,107 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue return list; } } + + public async Task> QueryAsync(SessionQuery query, CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new SqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + var results = new List(); + while (await reader.ReadAsync(ct)) + { + results.Add(new SessionSummary + { + SessionId = reader.GetGuid(0), + ModelFingerprint = reader.GetString(1), + LastUpdated = reader.GetDateTime(2), + Tags = JsonSerializer.Deserialize>(reader.GetString(3)) ?? new() + }); + } + return results; + } + + public async IAsyncEnumerable QueryStreamAsync( + SessionQuery query, + [EnumeratorCancellation] CancellationToken ct = default) + { + await using var connection = await GetConnectionAsync(ct); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); + + await using var command = new SqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); + + await using var reader = await command.ExecuteReaderAsync(ct); + while (await reader.ReadAsync(ct)) + { + ct.ThrowIfCancellationRequested(); + yield return new SessionSummary + { + SessionId = reader.GetGuid(0), + ModelFingerprint = reader.GetString(1), + LastUpdated = reader.GetDateTime(2), + Tags = JsonSerializer.Deserialize>(reader.GetString(3)) ?? new() + }; + } + } + + //------------ Private Methods --------------// + + private (string Sql, List Parameters) BuildQuerySql(SessionQuery query, bool includeOrderBy = true, bool includeLimit = true) + { + var sql = new StringBuilder(@" + SELECT session_id, model_fingerprint, last_updated, tags + FROM InferenceSessions + WHERE 1=1 + "); + + var parameters = new List(); + + if (!string.IsNullOrEmpty(query.ModelFingerprint)) + { + sql.Append(" AND model_fingerprint = @ModelFingerprint"); + parameters.Add(new SqlParameter("@ModelFingerprint", query.ModelFingerprint)); + } + if (query.UpdatedAfter.HasValue) + { + sql.Append(" AND last_updated >= @UpdatedAfter"); + parameters.Add(new SqlParameter("@UpdatedAfter", query.UpdatedAfter.Value)); + } + if (query.UpdatedBefore.HasValue) + { + sql.Append(" AND last_updated <= @UpdatedBefore"); + parameters.Add(new SqlParameter("@UpdatedBefore", query.UpdatedBefore.Value)); + } + if (query.Tags != null && query.Tags.Count > 0) + { + foreach (var pair in query.Tags) + { + var valParam = $"@tagVal_{pair.Key}"; + sql.Append($" AND JSON_VALUE(tags, '$.{pair.Key}') = {valParam}"); + parameters.Add(new SqlParameter(valParam, pair.Value)); + } + } + + if (includeOrderBy) + { + var orderColumn = query.OrderBy switch + { + SessionSortField.LastUpdated => "last_updated", + _ => "last_updated" + }; + sql.Append($" ORDER BY {orderColumn} {(query.Descending ? "DESC" : "ASC")}"); + } + + if (includeLimit && query.Limit.HasValue && query.Limit.Value > 0) + { + sql.Append(" OFFSET 0 ROWS FETCH NEXT @Limit ROWS ONLY"); + parameters.Add(new SqlParameter("@Limit", query.Limit.Value)); + } + + return (sql.ToString(), parameters); + } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs index 1700146..9b56a08 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs @@ -1,13 +1,13 @@ using Microsoft.Data.SqlClient; using System.Data; -namespace StateCheckpoint.NET.Stores; +namespace StateCheckpoint.NET; /// /// Abstract base class for SQL Server stores. /// Manages connection lifecycle and lazy opening. /// -public abstract class SqlServerStoreBase : IAsyncDisposable +internal abstract class SqlServerStoreBase : IAsyncDisposable { private readonly string _connectionString = string.Empty; private SqlConnection? _connection; From 86a27c4838d25c87ecab5e8bd7825927f705b7e8 Mon Sep 17 00:00:00 2001 From: Leo Mengesha Date: Mon, 13 Jul 2026 22:34:06 +0100 Subject: [PATCH 2/3] many fixes to tests and minor bugs --- .../Background/BackgroundSaverTests.cs | 3 +- .../Integration/DbHelper.cs | 16 + .../Integration/FileSystemTests.cs | 64 +++ .../Integration/IntegrationTestsBase.cs | 217 ++++++++ .../Integration/PostgresTests.cs | 141 +++-- .../Integration/SqlServerTests.cs | 137 ++++- .../Integration/TestBase.cs | 108 ---- .../Managers/CheckpointManagerTests.cs | 508 +++++++++--------- .../Managers/SessionManagerTests.cs | 465 ++++++---------- .../Stores/Postgres/PostgresTestBase.cs | 26 +- .../Stores/SqlServer/SqlServerTestBase.cs | 38 +- .../StateCheckpoint.NET/Helpers/TagsHelper.cs | 19 + .../Manager/CheckpointManager.cs | 56 +- .../Manager/SessionManager.cs | 69 ++- .../Queries/SqlServerTrainingQueries.cs | 1 - .../Settings/DbStorageOptions.cs | 3 + .../Stores/FileSystem/FileSystemBase.cs | 87 +++ .../Stores/FileSystem/FileSystemModelStore.cs | 156 ++---- .../FileSystem/FileSystemSessionStore.cs | 78 +-- .../Stores/FileSystemIndex/CheckpointIndex.cs | 40 +- .../Stores/FileSystemIndex/IndexStorage.cs | 1 - .../Stores/FileSystemIndex/SessionIndex.cs | 28 +- .../Stores/IFileSystemModelStore.cs | 2 +- .../Stores/SqlServer/SqlServerModelStore.cs | 6 +- 24 files changed, 1238 insertions(+), 1031 deletions(-) create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/DbHelper.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/FileSystemTests.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/IntegrationTestsBase.cs delete mode 100644 StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/TestBase.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Helpers/TagsHelper.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Background/BackgroundSaverTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Background/BackgroundSaverTests.cs index c58251c..20c3fac 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Background/BackgroundSaverTests.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Background/BackgroundSaverTests.cs @@ -1,5 +1,4 @@ -using StateCheckpoint.NET.Manager; -using FluentAssertions; +using FluentAssertions; namespace StateCheckpoint.NET.Tests.Manager; diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/DbHelper.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/DbHelper.cs new file mode 100644 index 0000000..0cf4c49 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/DbHelper.cs @@ -0,0 +1,16 @@ +namespace StateCheckpoint.NET.Tests.Integration +{ + internal static class DbHelper + { + public static string BuildConnectionString() + { + var host = Environment.GetEnvironmentVariable("PG_HOST") ?? "localhost"; + var port = Environment.GetEnvironmentVariable("PG_PORT") ?? "5432"; + var db = Environment.GetEnvironmentVariable("PG_DATABASE") ?? "postgres"; + var user = Environment.GetEnvironmentVariable("PG_USER") ?? "postgres"; + var password = Environment.GetEnvironmentVariable("PG_PASSWORD") ?? "postgres"; + + return $"Host={host};Port={port};Database={db};Username={user};Password={password}"; + } + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/FileSystemTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/FileSystemTests.cs new file mode 100644 index 0000000..de5831c --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/FileSystemTests.cs @@ -0,0 +1,64 @@ +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; +using System.IO; + +namespace StateCheckpoint.NET.Tests.Integration.FileSystem; + +[Collection("NonParallel")] +public class FileSystemTests : IntegrationTestsBase +{ + private readonly string _testRoot; + + public FileSystemTests() + { + _testRoot = Path.Combine(Path.GetTempPath(), "StateCheckpointIntegrationTests", Guid.NewGuid().ToString()); + } + + protected override StorageOptions GetCheckpointStorageOptions() + { + return new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions + { + RootPath = _testRoot, + EnsureDirectoryExists = true, + ValidatePermissionsOnStartup = true + } + }; + } + + protected override StorageOptions GetSessionStorageOptions() + { + return new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions + { + RootPath = _testRoot, + EnsureDirectoryExists = true, + ValidatePermissionsOnStartup = true + } + }; + } + + protected override async Task CleanupAsync() + { + if (Directory.Exists(_testRoot)) + Directory.Delete(_testRoot, recursive: true); + + await Task.CompletedTask; + } + + [Fact] + public async Task Checkpoint_RoundTrip_ShouldWork() => await RoundTrip_Checkpoint_ShouldSaveLoadDeleteList(); + + [Fact] + public async Task Session_RoundTrip_ShouldWork() => await RoundTrip_Session_ShouldSaveLoadDeleteList(); + + [Fact] + public async Task Query_ShouldReturnFilteredSummaries() => await base.Query_ShouldReturnFilteredSummaries(); + + [Fact] + public async Task FindBest_ShouldReturnBestCheckpoint() => await base.FindBest_ShouldReturnBestCheckpoint(); +} \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/IntegrationTestsBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/IntegrationTestsBase.cs new file mode 100644 index 0000000..39296e3 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/IntegrationTestsBase.cs @@ -0,0 +1,217 @@ +using FluentAssertions; +using Microsoft.Data.SqlClient; +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; + +namespace StateCheckpoint.NET.Tests.Integration; + +[Collection("NonParallel")] +public abstract class IntegrationTestsBase : IAsyncLifetime +{ + protected CheckpointManager CheckpointManager { get; set; } = null!; + protected SessionManager SessionManager { get; set; } = null!; + + private bool _disposed; + + // Derived classes provide storage configuration. + protected abstract StorageOptions GetCheckpointStorageOptions(); + protected abstract StorageOptions GetSessionStorageOptions(); + + public virtual async Task InitializeAsync() + { + var checkpointOptions = GetCheckpointStorageOptions(); + var sessionOptions = GetSessionStorageOptions(); + + CheckpointManager = new CheckpointManager(checkpointOptions); + SessionManager = new SessionManager(sessionOptions); + await Task.CompletedTask; + } + + public async Task DisposeAsync() + { + if (_disposed) return; + _disposed = true; + + try + { + await CheckpointManager.DisposeAsync(); + await SessionManager.DisposeAsync(); + await CleanupAsync(); + } + catch (Exception) + { + // Log or ignore – but do not rethrow to avoid test runner crashes. + } + } + + protected virtual async Task CleanupAsync() + { + try + { + var connectionString = DbHelper.BuildConnectionString(); + + await using var connection = new SqlConnection(connectionString); + + await connection.OpenAsync(); + + await using var command = new SqlCommand( + "TRUNCATE TABLE ModelManifests; TRUNCATE TABLE ModelBlobs; TRUNCATE TABLE InferenceSessions;", + connection); + + await command.ExecuteNonQueryAsync(); + } + catch (Exception) + { + // Ignore if tables don't exist. + } + } + + // ----- Shared integration test helpers ----- + + protected async Task RoundTrip_Checkpoint_ShouldSaveLoadDeleteList() + { + var checkpoint = new ModelCheckpoint + { + ModelId = Guid.NewGuid(), + WeightsBytes = new byte[] { 1, 2, 3, 4 }, + OptimizerBytes = new byte[] { 5, 6, 7 }, + HyperParams = new HyperParameters { HiddenSize = 768 }, + Tokenizer = new TokenizerData { Type = "BPE" }, + CurrentEpoch = 1, + LastTrainingLoss = 0.5f, + CreatedAt = DateTime.UtcNow, + Tags = new Dictionary { { "env", "integration" } } + }; + + var modelId = await CheckpointManager.SaveAsync(checkpoint); + + var loaded = await CheckpointManager.LoadAsync(modelId); + loaded.Should().NotBeNull(); + loaded.ModelId.Should().Be(modelId); + loaded.WeightsBytes.Should().BeEquivalentTo(checkpoint.WeightsBytes); + loaded.OptimizerBytes.Should().BeEquivalentTo(checkpoint.OptimizerBytes); + loaded.HyperParams.Should().BeEquivalentTo(checkpoint.HyperParams); + loaded.Tokenizer.Should().BeEquivalentTo(checkpoint.Tokenizer); + loaded.CurrentEpoch.Should().Be(1); + loaded.LastTrainingLoss.Should().Be(0.5f); + loaded.Tags.Should().BeEquivalentTo(checkpoint.Tags); + + var ids = await CheckpointManager.ListAsync(); + ids.Should().Contain(modelId); + + await CheckpointManager.DeleteAsync(modelId); + + var deleted = await CheckpointManager.LoadAsync(modelId); + deleted.Should().BeNull(); + + var idsAfterDelete = await CheckpointManager.ListAsync(); + idsAfterDelete.Should().NotContain(modelId); + } + + protected async Task RoundTrip_Session_ShouldSaveLoadDeleteList() + { + var session = new SessionCheckpoint + { + SessionId = Guid.NewGuid(), + KvCacheBytes = new byte[] { 100, 200, 230 }, + TokenHistory = new int[] { 1, 2, 3 }, + ModelFingerprint = "test-model", + SamplingConfig = new SamplingData { Temperature = 0.8f }, + LastUpdated = DateTime.UtcNow, + Tags = new Dictionary { { "user", "alice" } } + }; + + var returnedId = await SessionManager.SaveAsync(session); + returnedId.Should().Be(session.SessionId); + + var loaded = await SessionManager.LoadAsync(session.SessionId); + loaded.Should().NotBeNull(); + loaded.SessionId.Should().Be(session.SessionId); + loaded.KvCacheBytes.Should().BeEquivalentTo(session.KvCacheBytes); + loaded.TokenHistory.Should().BeEquivalentTo(session.TokenHistory); + loaded.ModelFingerprint.Should().Be(session.ModelFingerprint); + loaded.SamplingConfig.Should().BeEquivalentTo(session.SamplingConfig); + loaded.Tags.Should().BeEquivalentTo(session.Tags); + + var ids = await SessionManager.ListAsync(); + ids.Should().Contain(session.SessionId); + + await SessionManager.DeleteAsync(session.SessionId); + var deleted = await SessionManager.LoadAsync(session.SessionId); + deleted.Should().BeNull(); + + var idsAfterDelete = await SessionManager.ListAsync(); + idsAfterDelete.Should().NotContain(session.SessionId); + } + + protected async Task Query_ShouldReturnFilteredSummaries() + { + var cp1 = new ModelCheckpoint + { + ModelId = Guid.NewGuid(), + WeightsBytes = new byte[] { 1 }, + OptimizerBytes = new byte[] { 2 }, + HyperParams = new HyperParameters(), + Tokenizer = new TokenizerData(), + CurrentEpoch = 1, + LastTrainingLoss = 0.5f, + CreatedAt = DateTime.UtcNow + }; + + var cp2 = new ModelCheckpoint + { + ModelId = Guid.NewGuid(), + WeightsBytes = new byte[] { 3 }, + OptimizerBytes = new byte[] { 4 }, + HyperParams = new HyperParameters(), + Tokenizer = new TokenizerData(), + CurrentEpoch = 2, + LastTrainingLoss = 0.2f, + CreatedAt = DateTime.UtcNow + }; + + await CheckpointManager.SaveAsync(cp1); + await CheckpointManager.SaveAsync(cp2); + + var query = new CheckpointQuery { MinEpoch = 2 }; + + var summaries = await CheckpointManager.QueryAsync(query); + + summaries.Count.Should().Be(1); + summaries.Should().Contain(s => s.CurrentEpoch == 2); + } + + protected async Task FindBest_ShouldReturnBestCheckpoint() + { + var cp1 = new ModelCheckpoint + { + ModelId = Guid.NewGuid(), + WeightsBytes = new byte[] { 1 }, + OptimizerBytes = new byte[] { 2 }, + HyperParams = new HyperParameters(), + Tokenizer = new TokenizerData(), + CurrentEpoch = 1, + LastTrainingLoss = 0.5f, + CreatedAt = DateTime.UtcNow + }; + var cp2 = new ModelCheckpoint + { + ModelId = Guid.NewGuid(), + WeightsBytes = new byte[] { 3 }, + OptimizerBytes = new byte[] { 4 }, + HyperParams = new HyperParameters(), + Tokenizer = new TokenizerData(), + CurrentEpoch = 2, + LastTrainingLoss = 0.2f, + CreatedAt = DateTime.UtcNow + }; + + await CheckpointManager.SaveAsync(cp1); + await CheckpointManager.SaveAsync(cp2); + + var best = await CheckpointManager.FindBestByLossAsync(); + best.Should().NotBeNull(); + best.ModelId.Should().Be(cp2.ModelId); + best.LastTrainingLoss.Should().Be(0.2f); + } +} \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/PostgresTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/PostgresTests.cs index 6bc58c0..f4ddd1e 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/PostgresTests.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/PostgresTests.cs @@ -1,6 +1,6 @@ -using StateCheckpoint.NET.Manager; -using StateCheckpoint.NET.Stores; -using Npgsql; +using Npgsql; +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; using Testcontainers.PostgreSql; namespace StateCheckpoint.NET.Tests.Integration.Postgres; @@ -9,7 +9,7 @@ namespace StateCheckpoint.NET.Tests.Integration.Postgres; public class PostgresTests : IntegrationTestsBase { private readonly PostgreSqlContainer? _container; - private readonly string _connectionString = string.Empty; + private string _connectionString = string.Empty; private readonly bool _useTestcontainers; public PostgresTests() @@ -27,39 +27,55 @@ public PostgresTests() } else { - var host = Environment.GetEnvironmentVariable("PG_HOST") ?? "localhost"; - var port = Environment.GetEnvironmentVariable("PG_PORT") ?? "5432"; - var db = Environment.GetEnvironmentVariable("PG_DATABASE") ?? "postgres"; - var user = Environment.GetEnvironmentVariable("PG_USER") ?? "postgres"; - var password = Environment.GetEnvironmentVariable("PG_PASSWORD") ?? "postgres"; - - _connectionString = $"Host={host};Port={port};Database={db};Username={user};Password={password}"; + _connectionString = DbHelper.BuildConnectionString(); } } - protected override async Task InitializeStoresAsync() + protected override StorageOptions GetCheckpointStorageOptions() { - string finalConnectionString; + return new StorageOptions + { + StoreType = StoreType.SqlDb, + SqlDbType = SqlDbType.Postgres, + DbStoreOptions = new DbStorageOptions + { + ConnectionString = GetConnectionString(), + EnsureSchemaOnStartup = true + } + }; + } - if (_useTestcontainers) + protected override StorageOptions GetSessionStorageOptions() + { + return new StorageOptions { + StoreType = StoreType.SqlDb, + SqlDbType = SqlDbType.Postgres, + DbStoreOptions = new DbStorageOptions + { + ConnectionString = GetConnectionString(), + EnsureSchemaOnStartup = true + } + }; + } + + private string GetConnectionString() + { + if (_useTestcontainers) + return _container!.GetConnectionString(); + else + return _connectionString; + } + + public override async Task InitializeAsync() + { + if (_useTestcontainers) await _container!.StartAsync(); - finalConnectionString = _container.GetConnectionString(); - } else - { await EnsureLocalDatabaseExistsAsync(); - finalConnectionString = _connectionString; - } - var modelStore = new PostgresModelStore(finalConnectionString); - var sessionStore = new PostgresSessionStore(finalConnectionString); - - await modelStore.EnsureSchemaAsync(); - await sessionStore.EnsureSchemaAsync(); - - CheckpointManager = new CheckpointManager(modelStore); - SessionManager = new SessionManager(sessionStore); + await base.InitializeAsync(); // ensures schema + await TruncateTablesAsync(); // now tables exist } private async Task EnsureLocalDatabaseExistsAsync() @@ -69,11 +85,13 @@ private async Task EnsureLocalDatabaseExistsAsync() var builder = new NpgsqlConnectionStringBuilder(_connectionString); var testDbName = "checkpoint_integration"; builder.Database = "postgres"; + await using var connection = new NpgsqlConnection(builder.ConnectionString); + await connection.OpenAsync(); - var createDbQuery = $"CREATE DATABASE {testDbName};"; - await using var command = new NpgsqlCommand(createDbQuery, connection); + await using var command = new NpgsqlCommand($"CREATE DATABASE {testDbName};", connection); + try { await command.ExecuteNonQueryAsync(); @@ -83,26 +101,79 @@ private async Task EnsureLocalDatabaseExistsAsync() // Database already exists } - // Update connection string to point to the test database. - // We'll just use the test DB name in the final connection string. + builder.Database = testDbName; + _connectionString = builder.ConnectionString; } catch (Exception ex) { - throw new InvalidOperationException("Could not connectionect to local PostgreSQL. Ensure it is running.", ex); + throw new InvalidOperationException("Could not connect to local PostgreSQL. Ensure it is running.", ex); } } - protected override async Task CleanupStoresAsync() + private async Task TruncateTablesAsync() { - if (_useTestcontainers && _container != null) + try { - await _container.DisposeAsync(); + var connString = GetConnectionString(); + + await using var connection = new NpgsqlConnection(connString); + + await connection.OpenAsync(); + + // Check if tables exist + var checkTables = @" + SELECT COUNT(*) FROM information_schema.tables + WHERE table_name IN ('model_manifests', 'model_blobs', 'inference_sessions')"; + + await using var checkCmd = new NpgsqlCommand(checkTables, connection); + + var tableCount = Convert.ToInt32(await checkCmd.ExecuteScalarAsync()); + if (tableCount < 3) + return; + + + try + { + await using var truncateCmd = new NpgsqlCommand( + "TRUNCATE TABLE model_blobs, model_manifests, inference_sessions;", + connection); + await truncateCmd.ExecuteNonQueryAsync(); + } + catch (PostgresException ex) when (ex.SqlState == "23503") + { + await using var deleteCmd = new NpgsqlCommand( + "DELETE FROM model_blobs; DELETE FROM model_manifests; DELETE FROM inference_sessions;", + connection); + + await deleteCmd.ExecuteNonQueryAsync(); + } + catch (Exception ex) + { + throw; + } + } + catch (Exception ex) + { + throw; } } + protected override async Task CleanupAsync() + { + if (_useTestcontainers && _container != null) + await _container.DisposeAsync(); + } + + // ----- Tests ----- [Fact] public async Task Checkpoint_RoundTrip_ShouldWork() => await RoundTrip_Checkpoint_ShouldSaveLoadDeleteList(); [Fact] public async Task Session_RoundTrip_ShouldWork() => await RoundTrip_Session_ShouldSaveLoadDeleteList(); + + [Fact] + public async Task Query_ShouldReturnFilteredSummaries() => await base.Query_ShouldReturnFilteredSummaries(); + + [Fact] + public async Task FindBest_ShouldReturnBestCheckpoint() => await base.FindBest_ShouldReturnBestCheckpoint(); } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/SqlServerTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/SqlServerTests.cs index ea560d5..23e77bd 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/SqlServerTests.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/SqlServerTests.cs @@ -1,6 +1,6 @@ -using StateCheckpoint.NET.Manager; -using StateCheckpoint.NET.Stores; -using Microsoft.Data.SqlClient; +using Microsoft.Data.SqlClient; +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; using Testcontainers.MsSql; namespace StateCheckpoint.NET.Tests.Integration.SqlServer; @@ -8,8 +8,8 @@ namespace StateCheckpoint.NET.Tests.Integration.SqlServer; [Collection("NonParallel")] public class SqlServerTests : IntegrationTestsBase { - private readonly MsSqlContainer? _container = null; - private readonly string _connectionString = string.Empty; + private readonly MsSqlContainer? _container; + private string _connectionString = string.Empty; private readonly bool _useTestcontainers; public SqlServerTests() @@ -25,39 +25,62 @@ public SqlServerTests() } else { - _connectionString = "Data Source=(localdb)\\MSSQLLocalDB;Initial Catalog=CheckpointIntegration;Integrated Security=True;"; + _connectionString = "Data Source=(localdb)\\MSSQLLocalDB;Initial Catalog=master;Integrated Security=True;"; } } - protected override async Task InitializeStoresAsync() + protected override StorageOptions GetCheckpointStorageOptions() { - string finalConnectionString; + return new StorageOptions + { + StoreType = StoreType.SqlDb, + SqlDbType = SqlDbType.SqlServer, + DbStoreOptions = new DbStorageOptions + { + ConnectionString = GetConnectionString(), + EnsureSchemaOnStartup = true + } + }; + } - if (_useTestcontainers) + protected override StorageOptions GetSessionStorageOptions() + { + return new StorageOptions { + StoreType = StoreType.SqlDb, + SqlDbType = SqlDbType.SqlServer, + DbStoreOptions = new DbStorageOptions + { + ConnectionString = GetConnectionString(), + EnsureSchemaOnStartup = true + } + }; + } + + private string GetConnectionString() + { + if (_useTestcontainers) + return _container!.GetConnectionString(); + + else return _connectionString; + } + + public override async Task InitializeAsync() + { + if (_useTestcontainers) await _container!.StartAsync(); - finalConnectionString = _container.GetConnectionString(); - } - else - { - await EnsureLocalDbDatabaseExistsAsync(); - finalConnectionString = _connectionString; - } - var modelStore = new SqlServerModelStore(finalConnectionString); - var sessionStore = new SqlServerSessionStore(finalConnectionString); + else await EnsureLocalDbDatabaseExistsAsync(); - await modelStore.EnsureSchemaAsync(); - await sessionStore.EnsureSchemaAsync(); + await base.InitializeAsync(); - CheckpointManager = new CheckpointManager(modelStore); - SessionManager = new SessionManager(sessionStore); + await TruncateTablesAsync(); } - private static async Task EnsureLocalDbDatabaseExistsAsync() + private async Task EnsureLocalDbDatabaseExistsAsync() { const string masterConnectionString = "Data Source=(localdb)\\MSSQLLocalDB;Initial Catalog=master;Integrated Security=True;"; - using var connection = new SqlConnection(masterConnectionString); + await using var connection = new SqlConnection(masterConnectionString); await connection.OpenAsync(); var createDbQuery = @" @@ -65,21 +88,79 @@ IF NOT EXISTS (SELECT name FROM sys.databases WHERE name = 'CheckpointIntegratio BEGIN CREATE DATABASE [CheckpointIntegration]; END"; - using var command = new SqlCommand(createDbQuery, connection); + await using var command = new SqlCommand(createDbQuery, connection); await command.ExecuteNonQueryAsync(); + + _connectionString = "Data Source=(localdb)\\MSSQLLocalDB;Initial Catalog=CheckpointIntegration;Integrated Security=True;"; } - protected override async Task CleanupStoresAsync() + private async Task TruncateTablesAsync() { - if (_useTestcontainers && _container != null) + try { - await _container.DisposeAsync(); + var connString = GetConnectionString(); + + await using var connection = new SqlConnection(connString); + + await connection.OpenAsync(); + + // Check if tables exist + var checkTables = @" + SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES + WHERE TABLE_NAME IN ('ModelManifests', 'ModelBlobs', 'InferenceSessions')"; + + await using var checkCmd = new SqlCommand(checkTables, connection); + + var tableCount = (int)await checkCmd.ExecuteScalarAsync(); + + if (tableCount < 3) + return; + + try + { + await using var truncateCmd = new SqlCommand( + "TRUNCATE TABLE ModelBlobs; TRUNCATE TABLE ModelManifests; TRUNCATE TABLE InferenceSessions;", + connection); + + await truncateCmd.ExecuteNonQueryAsync(); + } + catch (SqlException ex) when (ex.Number == 4712) + { + await using var deleteCmd = new SqlCommand( + "DELETE FROM ModelBlobs; DELETE FROM ModelManifests; DELETE FROM InferenceSessions;", + connection); + + await deleteCmd.ExecuteNonQueryAsync(); + } + catch (Exception ex) + { + throw; // rethrow to make test fail so we can see the error + } + } + catch (Exception ex) + { + throw; // rethrow so test fails and we can diagnose } } + protected override async Task CleanupAsync() + { + // Truncate again after the test (for safety) + await TruncateTablesAsync(); + + if (_useTestcontainers && _container != null) + await _container.DisposeAsync(); + } + [Fact] public async Task Checkpoint_RoundTrip_ShouldWork() => await RoundTrip_Checkpoint_ShouldSaveLoadDeleteList(); [Fact] public async Task Session_RoundTrip_ShouldWork() => await RoundTrip_Session_ShouldSaveLoadDeleteList(); + + [Fact] + public async Task Query_ShouldReturnFilteredSummaries() => await base.Query_ShouldReturnFilteredSummaries(); + + [Fact] + public async Task FindBest_ShouldReturnBestCheckpoint() => await base.FindBest_ShouldReturnBestCheckpoint(); } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/TestBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/TestBase.cs deleted file mode 100644 index a6715a2..0000000 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Integration/TestBase.cs +++ /dev/null @@ -1,108 +0,0 @@ -using StateCheckpoint.NET.Manager; -using StateCheckpoint.NET.Models; -using FluentAssertions; - -namespace StateCheckpoint.NET.Tests.Integration; - -[Collection("NonParallel")] -public abstract class IntegrationTestsBase : IAsyncLifetime -{ - protected CheckpointManager CheckpointManager { get; set; } = null!; - protected SessionManager SessionManager { get; set; } = null!; - - // Derived classes must implement these to provide the specific stores. - protected abstract Task InitializeStoresAsync(); - protected abstract Task CleanupStoresAsync(); - - public async Task InitializeAsync() - { - await InitializeStoresAsync(); - } - - public async Task DisposeAsync() - { - await CleanupStoresAsync(); - } - - // --- Shared integration tests (will be called by derived classes) --- - - protected async Task RoundTrip_Checkpoint_ShouldSaveLoadDeleteList() - { - // Arrange - var weights = new byte[] { 1, 2, 3, 4 }; - var optimizer = new byte[] { 5, 6, 7 }; - var hyperParams = new HyperParameters { HiddenSize = 768 }; - var tokenizer = new TokenizerData { Type = "BPE" }; - var tags = new Dictionary { { "env", "integration" } }; - - // Act - Save - var modelId = await CheckpointManager.SaveAsync(weights, optimizer, hyperParams, tokenizer, epoch: 1, loss: 0.5f, tags: tags); - - // Assert - Load - var loaded = await CheckpointManager.LoadAsync(modelId); - - loaded.Should().NotBeNull(); - loaded.ModelId.Should().Be(modelId); - loaded.WeightsBytes.Should().BeEquivalentTo(weights); - loaded.OptimizerBytes.Should().BeEquivalentTo(optimizer); - loaded.HyperParams.Should().BeEquivalentTo(hyperParams); - loaded.Tokenizer.Should().BeEquivalentTo(tokenizer); - loaded.CurrentEpoch.Should().Be(1); - loaded.LastTrainingLoss.Should().Be(0.5f); - loaded.Tags.Should().BeEquivalentTo(tags); - - // Assert - List - var ids = await CheckpointManager.ListAsync(); - ids.Should().Contain(modelId); - - // Act - Delete - await CheckpointManager.DeleteAsync(modelId); - - // Assert - Deleted - var deleted = await CheckpointManager.LoadAsync(modelId); - deleted.Should().BeNull(); - - var idsAfterDelete = await CheckpointManager.ListAsync(); - idsAfterDelete.Should().NotContain(modelId); - } - - protected async Task RoundTrip_Session_ShouldSaveLoadDeleteList() - { - // Arrange - var sessionId = Guid.NewGuid(); - var kvCache = new byte[] { 100, 200, 230 }; - var tokenHistory = new int[] { 1, 2, 3 }; - var fingerprint = "test-model"; - var sampling = new SamplingData { Temperature = 0.8f }; - var tags = new Dictionary { { "user", "alice" } }; - - // Act - Save - var returnedId = await SessionManager.SaveAsync(sessionId, kvCache, tokenHistory, fingerprint, sampling, tags); - - // Assert - returnedId.Should().Be(sessionId); - - // Load - var loaded = await SessionManager.LoadAsync(sessionId); - - loaded.Should().NotBeNull(); - loaded.SessionId.Should().Be(sessionId); - loaded.KvCacheBytes.Should().BeEquivalentTo(kvCache); - loaded.TokenHistory.Should().BeEquivalentTo(tokenHistory); - loaded.ModelFingerprint.Should().Be(fingerprint); - loaded.SamplingConfig.Should().BeEquivalentTo(sampling); - loaded.Tags.Should().BeEquivalentTo(tags); - - // List - var ids = await SessionManager.ListAsync(); - ids.Should().Contain(sessionId); - - // Delete - await SessionManager.DeleteAsync(sessionId); - var deleted = await SessionManager.LoadAsync(sessionId); - deleted.Should().BeNull(); - - var idsAfterDelete = await SessionManager.ListAsync(); - idsAfterDelete.Should().NotContain(sessionId); - } -} \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/CheckpointManagerTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/CheckpointManagerTests.cs index 7e5ca0e..94f163d 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/CheckpointManagerTests.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/CheckpointManagerTests.cs @@ -1,359 +1,349 @@ -using StateCheckpoint.NET.Manager; +using FluentAssertions; using StateCheckpoint.NET.Models; using StateCheckpoint.NET.Settings; -using StateCheckpoint.NET.Stores; -using FluentAssertions; -using Moq; using System; using System.Collections.Generic; -using System.Threading; +using System.IO; +using System.Linq; using System.Threading.Tasks; using Xunit; namespace StateCheckpoint.NET.Tests.Manager; -public class CheckpointManagerTests +public class CheckpointManagerTests : IAsyncLifetime { - private readonly Mock _mockStore; - private readonly HyperParameters _hyperParams; - private readonly TokenizerData _tokenizer; - private readonly byte[] _weights; - private readonly byte[] _optimizer; - private readonly Dictionary _savedCheckpoints; + private readonly string _testRoot; + private StorageOptions _storageOptions; + private CheckpointManager _manager; public CheckpointManagerTests() { - _mockStore = new Mock(MockBehavior.Strict); - _savedCheckpoints = new Dictionary(); - - // Setup default behavior for SaveAsync - _mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback((checkpoint, ct) => - { - _savedCheckpoints[checkpoint.ModelId] = checkpoint; - }) - .Returns(Task.CompletedTask); + _testRoot = Path.Combine(Path.GetTempPath(), "StateCheckpointTests", Guid.NewGuid().ToString()); + } - // Setup default behavior for LoadAsync - _mockStore.Setup(x => x.LoadAsync(It.IsAny(), It.IsAny())) - .Returns((id, ct) => - { - _savedCheckpoints.TryGetValue(id, out var checkpoint); - return Task.FromResult(checkpoint); - }); - - // Setup default behavior for DeleteAsync - _mockStore.Setup(x => x.DeleteAsync(It.IsAny(), It.IsAny())) - .Callback((id, ct) => _savedCheckpoints.Remove(id)) - .Returns(Task.CompletedTask); - - // Setup default behavior for ListAsync - _mockStore.Setup(x => x.ListAsync(It.IsAny(), It.IsAny(), It.IsAny())) - .Returns((tagKey, tagValue, ct) => + public async Task InitializeAsync() + { + Directory.CreateDirectory(_testRoot); + _storageOptions = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { - var ids = new List(_savedCheckpoints.Keys); - return Task.FromResult(ids); - }); - - _hyperParams = new HyperParameters { HiddenSize = 768 }; - _tokenizer = new TokenizerData { Type = "BPE" }; - _weights = new byte[] { 1, 2, 3 }; - _optimizer = new byte[] { 4, 5, 6 }; + RootPath = _testRoot, + EnsureDirectoryExists = true, + ValidatePermissionsOnStartup = true + } + }; + _manager = new CheckpointManager(_storageOptions); + await Task.CompletedTask; } - // -------------------- Constructors -------------------- - - [Fact] - public void Constructor_WithStore_ShouldSetStore() + public async Task DisposeAsync() { - // Act - var manager = new CheckpointManager(_mockStore.Object); - - // Assert - var id = manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f).GetAwaiter().GetResult(); - _mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Once); - _savedCheckpoints.Keys.Should().Contain(id); + await _manager.DisposeAsync(); + if (Directory.Exists(_testRoot)) + Directory.Delete(_testRoot, recursive: true); } - [Fact] - public async Task Constructor_WithBackgroundOptions_ShouldCreateBackgroundSaver() + private ModelCheckpoint CreateTestCheckpoint( + Guid? id = null, + int epoch = 1, + float loss = 0.5f, + Dictionary? tags = null) { - // Arrange - var saveCalled = new TaskCompletionSource(); - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback(() => saveCalled.TrySetResult(true)) - .Returns(Task.CompletedTask); - - var options = new BackgroundSaveOptions { Enabled = true, QueueCapacity = 3 }; - var manager = new CheckpointManager(mockStore.Object, options); - - // Act – SaveAsync returns immediately (fire‑and‑forget) - var saveTask = manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f); - - // Assert – SaveAsync completed synchronously - saveTask.IsCompletedSuccessfully.Should().BeTrue(); - - // Store should NOT have been called yet (background hasn't run) - saveCalled.Task.IsCompleted.Should().BeFalse(); - - // Flush the background queue - await ((IAsyncDisposable)manager).DisposeAsync(); - - // Now the store should have been called - saveCalled.Task.IsCompleted.Should().BeTrue(); + return new ModelCheckpoint + { + ModelId = id ?? Guid.NewGuid(), + WeightsBytes = new byte[] { 1, 2, 3 }, + OptimizerBytes = new byte[] { 4, 5, 6 }, + HyperParams = new HyperParameters { HiddenSize = 768 }, + Tokenizer = new TokenizerData { Type = "BPE" }, + CurrentEpoch = epoch, + LastTrainingLoss = loss, + CreatedAt = DateTime.UtcNow, + Tags = tags ?? new Dictionary() + }; } - // -------------------- SaveAsync (Synchronous Mode) -------------------- + #region Constructor Validation [Fact] - public async Task SaveAsync_Synchronous_ShouldSaveCheckpoint() + public void Constructor_WhenSqlDbWithoutConnectionString_Throws() { - // Arrange - var manager = new CheckpointManager(_mockStore.Object); - - // Act - var id = await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 5, 2.345f); - - // Assert - _mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Once); - var saved = _savedCheckpoints[id]; - saved.WeightsBytes.Should().BeEquivalentTo(_weights); - saved.OptimizerBytes.Should().BeEquivalentTo(_optimizer); - saved.HyperParams.Should().BeEquivalentTo(_hyperParams); - saved.Tokenizer.Should().BeEquivalentTo(_tokenizer); - saved.CurrentEpoch.Should().Be(5); - saved.LastTrainingLoss.Should().Be(2.345f); + var options = new StorageOptions + { + StoreType = StoreType.SqlDb, + DbStoreOptions = new DbStorageOptions { ConnectionString = null } + }; + Action act = () => new CheckpointManager(options); + act.Should().Throw() + .WithMessage("ConnectionString is required when StoreType is SqlDb."); } [Fact] - public async Task SaveAsync_Synchronous_WithExistingId_ShouldUseProvidedId() + public void Constructor_WhenLocalWithoutRootPath_Throws() { - // Arrange - var existingId = Guid.NewGuid(); - var manager = new CheckpointManager(_mockStore.Object); - - // Act - var id = await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f, existingId); + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = null } + }; + Action act = () => new CheckpointManager(options); + act.Should().Throw() + .WithMessage("RootPath is required when StoreType is Local."); + } - // Assert - id.Should().Be(existingId); - _mockStore.Verify(x => x.SaveAsync(It.Is(c => c.ModelId == existingId), It.IsAny()), Times.Once); + [Fact] + public void Constructor_WithValidLocalOptions_CreatesStore() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = _testRoot } + }; + var manager = new CheckpointManager(options); + manager.Should().NotBeNull(); + Directory.Exists(_testRoot).Should().BeTrue(); } [Fact] - public async Task SaveAsync_Synchronous_WithTags_ShouldStoreTags() + public void Constructor_WithBackgroundSaveEnabled_CreatesBackgroundSaver() { - // Arrange - var tags = new Dictionary { { "key", "value" } }; - var manager = new CheckpointManager(_mockStore.Object); + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = _testRoot }, + BackgroundSaveOptions = new BackgroundSaveOptions { Enabled = true, QueueCapacity = 5 } + }; + var manager = new CheckpointManager(options); + var field = typeof(CheckpointManager).GetField("_backgroundSaver", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); + var saver = field?.GetValue(manager); + saver.Should().NotBeNull(); + } - // Act - await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f, tags: tags); + #endregion - // Assert - _mockStore.Verify(x => x.SaveAsync(It.Is(c => c.Tags == tags), It.IsAny()), Times.Once); - } + #region Save and Load [Fact] - public async Task SaveAsync_Synchronous_ShouldRespectCancellation() + public async Task SaveAsync_WithNewCheckpoint_SavesAndReturnsId() { - // Arrange - var cts = new CancellationTokenSource(); - cts.Cancel(); - var manager = new CheckpointManager(_mockStore.Object); + var checkpoint = CreateTestCheckpoint(); + var id = await _manager.SaveAsync(checkpoint); - // Act - var act = async () => await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f, cancellationToken: cts.Token); + id.Should().Be(checkpoint.ModelId); - // Assert - await act.Should().ThrowAsync(); - _mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Never); + var loaded = await _manager.LoadAsync(id); + loaded.Should().NotBeNull(); + loaded.ModelId.Should().Be(id); + loaded.WeightsBytes.Should().BeEquivalentTo(checkpoint.WeightsBytes); + loaded.OptimizerBytes.Should().BeEquivalentTo(checkpoint.OptimizerBytes); + loaded.HyperParams.Should().BeEquivalentTo(checkpoint.HyperParams); + loaded.Tokenizer.Should().BeEquivalentTo(checkpoint.Tokenizer); + loaded.CurrentEpoch.Should().Be(checkpoint.CurrentEpoch); + loaded.LastTrainingLoss.Should().Be(checkpoint.LastTrainingLoss); + loaded.Tags.Should().BeEquivalentTo(checkpoint.Tags); } - // -------------------- SaveAsync (Background Mode) -------------------- - [Fact] - public async Task SaveAsync_Background_ShouldReturnImmediatelyAndSaveInBackground() + public async Task SaveAsync_WithExistingId_UpdatesCheckpoint() { - // Arrange - Use a ManualResetEvent or TaskCompletionSource to block the save - var saveCompleted = new TaskCompletionSource(); - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback((c, ct) => - { - // Simulate work - saveCompleted.SetResult(true); - }) - .Returns(Task.CompletedTask); + var checkpoint = CreateTestCheckpoint(); + var id = await _manager.SaveAsync(checkpoint); - var options = new BackgroundSaveOptions { Enabled = true }; - var manager = new CheckpointManager(mockStore.Object, options); + var updated = CreateTestCheckpoint(id, epoch: 5, loss: 0.1f); + var returnedId = await _manager.SaveAsync(updated); - // Act - var id = await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f); - id.Should().NotBe(Guid.Empty); + returnedId.Should().Be(id); - // Wait for background save to complete - await saveCompleted.Task.WaitAsync(TimeSpan.FromSeconds(5)); - - // Assert - store was called - mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Once); - - await manager.DisposeAsync(); + var loaded = await _manager.LoadAsync(id); + loaded.CurrentEpoch.Should().Be(5); + loaded.LastTrainingLoss.Should().Be(0.1f); } [Fact] - public async Task SaveAsync_Background_ShouldPerformDeepCopy() + public async Task LoadAsync_WhenNotExists_ReturnsNull() { - // Arrange - var capturedCheckpoint = (ModelCheckpoint?)null; - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback((c, ct) => capturedCheckpoint = c) - .Returns(Task.CompletedTask); - - var options = new BackgroundSaveOptions { Enabled = true }; - var manager = new CheckpointManager(mockStore.Object, options); + var result = await _manager.LoadAsync(Guid.NewGuid()); + result.Should().BeNull(); + } - var mutableWeights = new byte[] { 10, 20, 30 }; - var mutableOptimizer = new byte[] { 40, 50, 60 }; + #endregion - // Act - await manager.SaveAsync(mutableWeights, mutableOptimizer, _hyperParams, _tokenizer, 1, 0.5f); - mutableWeights[0] = 99; // Mutate after save call - mutableOptimizer[0] = 99; + #region Delete - // Wait for background to process - await Task.Delay(200); + [Fact] + public async Task DeleteAsync_RemovesCheckpoint() + { + var checkpoint = CreateTestCheckpoint(); + var id = await _manager.SaveAsync(checkpoint); - // Assert - saved bytes should be the original, not mutated - capturedCheckpoint.Should().NotBeNull(); - capturedCheckpoint!.WeightsBytes.Should().BeEquivalentTo(new byte[] { 10, 20, 30 }); - capturedCheckpoint.OptimizerBytes.Should().BeEquivalentTo(new byte[] { 40, 50, 60 }); + await _manager.DeleteAsync(id); - await manager.DisposeAsync(); + var loaded = await _manager.LoadAsync(id); + loaded.Should().BeNull(); } + #endregion + + #region List + [Fact] - public async Task SaveAsync_Background_WithErrorCallback_ShouldInvokeOnError() + public async Task ListAsync_WithoutFilters_ReturnsAllIds() { - // Arrange - Exception? captured = null; - var options = new BackgroundSaveOptions + var ids = new List(); + for (int i = 0; i < 3; i++) { - Enabled = true, - OnError = ex => captured = ex - }; - - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Throws(new InvalidOperationException("Save failed")); + var cp = CreateTestCheckpoint(); + ids.Add(await _manager.SaveAsync(cp)); + } - var manager = new CheckpointManager(mockStore.Object, options); + var all = await _manager.ListAsync(); + all.Should().BeEquivalentTo(ids); + } - // Act - await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f); + [Fact] + public async Task ListAsync_WithTagFilter_ReturnsMatchingIds() + { + var cp1 = CreateTestCheckpoint(tags: new Dictionary { { "env", "prod" } }); + var cp2 = CreateTestCheckpoint(tags: new Dictionary { { "env", "dev" } }); - // Wait for background to process - await Task.Delay(200); + var id1 = await _manager.SaveAsync(cp1); + var id2 = await _manager.SaveAsync(cp2); - // Assert - captured.Should().NotBeNull(); - captured.Should().BeOfType(); - captured!.Message.Should().Be("Save failed"); + var result = await _manager.ListAsync("env", "prod"); - await manager.DisposeAsync(); + result.Count.Should().Be(1); + result.Should().Contain(id1); + result.Should().NotContain(id2); } - // -------------------- LoadAsync -------------------- + #endregion + + #region Query [Fact] - public async Task LoadAsync_ShouldReturnCheckpointFromStore() + public async Task QueryAsync_ReturnsSummaries() { - // Arrange - var manager = new CheckpointManager(_mockStore.Object); - var id = await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f); - - // Reset the mock call counter for LoadAsync - _mockStore.Invocations.Clear(); + var cp1 = CreateTestCheckpoint(epoch: 1, loss: 0.5f); + var cp2 = CreateTestCheckpoint(epoch: 2, loss: 0.3f); + await _manager.SaveAsync(cp1); + await _manager.SaveAsync(cp2); - // Act - var loaded = await manager.LoadAsync(id); + var query = new CheckpointQuery { MinEpoch = 2 }; + var summaries = await _manager.QueryAsync(query); - // Assert - _mockStore.Verify(x => x.LoadAsync(id, It.IsAny()), Times.Once); - loaded.Should().NotBeNull(); - loaded!.ModelId.Should().Be(id); + summaries.Should().ContainSingle(s => s.CurrentEpoch == 2); + summaries.Should().NotContain(s => s.CurrentEpoch == 1); } [Fact] - public async Task LoadAsync_WhenNotFound_ShouldReturnNull() + public async Task QueryStreamAsync_StreamsSummaries() { - // Arrange - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.LoadAsync(It.IsAny(), It.IsAny())) - .Returns((id, ct) => Task.FromResult(null)); - - var manager = new CheckpointManager(mockStore.Object); + var cp1 = CreateTestCheckpoint(epoch: 1); + var cp2 = CreateTestCheckpoint(epoch: 2); + await _manager.SaveAsync(cp1); + await _manager.SaveAsync(cp2); - // Act - var loaded = await manager.LoadAsync(Guid.NewGuid()); + var stream = _manager.QueryStreamAsync(new CheckpointQuery()); + var list = await stream.ToListAsync(); - // Assert - loaded.Should().BeNull(); + list.Should().HaveCount(2); + list.Should().Contain(s => s.CurrentEpoch == 1); + list.Should().Contain(s => s.CurrentEpoch == 2); } - // -------------------- DeleteAsync -------------------- + #endregion + + #region FindBest [Fact] - public async Task DeleteAsync_ShouldCallStoreDelete() + public async Task FindBestAsync_ReturnsBestByScore() { - // Arrange - var manager = new CheckpointManager(_mockStore.Object); - var id = await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f); - _savedCheckpoints.Keys.Should().Contain(id); - - _mockStore.Invocations.Clear(); - - // Act - await manager.DeleteAsync(id); - - // Assert - _mockStore.Verify(x => x.DeleteAsync(id, It.IsAny()), Times.Once); - _savedCheckpoints.Keys.Should().NotContain(id); + var cp1 = CreateTestCheckpoint(epoch: 1, loss: 0.5f); + var cp2 = CreateTestCheckpoint(epoch: 2, loss: 0.2f); + var cp3 = CreateTestCheckpoint(epoch: 3, loss: 0.4f); + await _manager.SaveAsync(cp1); + await _manager.SaveAsync(cp2); + await _manager.SaveAsync(cp3); + + // By loss (lowest = highest score when using negative loss) + var best = await _manager.FindBestAsync(s => -s.LastTrainingLoss); + best.Should().NotBeNull(); + best.ModelId.Should().Be(cp2.ModelId); + best.LastTrainingLoss.Should().Be(0.2f); } - // -------------------- ListAsync -------------------- - [Fact] - public async Task ListAsync_ShouldReturnAllIds() + public async Task FindBestByLossAsync_ReturnsLowestLoss() { - // Arrange - var manager = new CheckpointManager(_mockStore.Object); - var id1 = await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 1, 0.5f); - var id2 = await manager.SaveAsync(_weights, _optimizer, _hyperParams, _tokenizer, 2, 0.6f); - - _mockStore.Invocations.Clear(); + var cp1 = CreateTestCheckpoint(loss: 0.5f); + var cp2 = CreateTestCheckpoint(loss: 0.1f); + await _manager.SaveAsync(cp1); + await _manager.SaveAsync(cp2); + + var best = await _manager.FindBestByLossAsync(); + best.Should().NotBeNull(); + best.ModelId.Should().Be(cp2.ModelId); + } - // Act - var ids = await manager.ListAsync(); + [Fact] + public async Task FindLatestEpochAsync_ReturnsHighestEpoch() + { + var cp1 = CreateTestCheckpoint(epoch: 1); + var cp2 = CreateTestCheckpoint(epoch: 10); + var cp3 = CreateTestCheckpoint(epoch: 5); + await _manager.SaveAsync(cp1); + await _manager.SaveAsync(cp2); + await _manager.SaveAsync(cp3); + + var best = await _manager.FindLatestEpochAsync(); + best.Should().NotBeNull(); + best.ModelId.Should().Be(cp2.ModelId); + best.CurrentEpoch.Should().Be(10); + } - // Assert - _mockStore.Verify(x => x.ListAsync(null, null, It.IsAny()), Times.Once); - ids.Should().BeEquivalentTo(new[] { id1, id2 }); + [Fact] + public async Task FindBestAsync_WhenNoCheckpoints_ReturnsNull() + { + var result = await _manager.FindBestAsync(s => s.CurrentEpoch); + result.Should().BeNull(); } [Fact] - public async Task ListAsync_WithTagFilter_ShouldPassTagsToStore() + public async Task FindBestAsync_WithFilter_AppliesFilter() { - // Arrange - var manager = new CheckpointManager(_mockStore.Object); + var cp1 = CreateTestCheckpoint(epoch: 1, loss: 0.5f); + var cp2 = CreateTestCheckpoint(epoch: 2, loss: 0.2f); + var cp3 = CreateTestCheckpoint(epoch: 3, loss: 0.4f); + await _manager.SaveAsync(cp1); + await _manager.SaveAsync(cp2); + await _manager.SaveAsync(cp3); + + var filter = new CheckpointQuery { MinEpoch = 2 }; + var best = await _manager.FindBestAsync(s => -s.LastTrainingLoss, filter); + + best.Should().NotBeNull(); + best.ModelId.Should().Be(cp2.ModelId); + best.CurrentEpoch.Should().Be(2); + best.LastTrainingLoss.Should().Be(0.2f); + } - // Act - await manager.ListAsync("key", "value"); + #endregion - // Assert - _mockStore.Verify(x => x.ListAsync("key", "value", It.IsAny()), Times.Once); + #region Dispose + + [Fact] + public async Task DisposeAsync_DisposesStore() + { + // Already called in DisposeAsync of the test class + // We'll just verify that the store is disposed by trying to use it? Not needed. + // We can check that the directory is gone? The test already does. + // This test just ensures the method exists and doesn't throw. + await _manager.DisposeAsync(); + // If the store was already disposed, calling again should be safe. + await _manager.DisposeAsync(); } + + #endregion } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/SessionManagerTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/SessionManagerTests.cs index b1db14b..75147ed 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/SessionManagerTests.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Managers/SessionManagerTests.cs @@ -1,372 +1,261 @@ -using StateCheckpoint.NET.Manager; +using FluentAssertions; using StateCheckpoint.NET.Models; using StateCheckpoint.NET.Settings; -using StateCheckpoint.NET.Stores; -using FluentAssertions; -using Moq; namespace StateCheckpoint.NET.Tests.Manager; -public class SessionManagerTests +public class SessionManagerTests : IAsyncLifetime { - private readonly Mock _mockStore; - private readonly byte[] _kvCache; - private readonly int[] _tokenHistory; - private readonly string _modelFingerprint; - private readonly SamplingData _samplingConfig; - private readonly Guid _sessionId; - private readonly Dictionary _savedSessions; + private readonly string _testRoot; + private StorageOptions _storageOptions; + private SessionManager _manager; public SessionManagerTests() { - _mockStore = new Mock(MockBehavior.Strict); - _savedSessions = new Dictionary(); - - // Default setup for SaveAsync - _mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback((session, ct) => - { - _savedSessions[session.SessionId] = session; - }) - .Returns(Task.CompletedTask); + _testRoot = Path.Combine(Path.GetTempPath(), "StateCheckpointSessionTests", Guid.NewGuid().ToString()); + } - // Default setup for LoadAsync - _mockStore.Setup(x => x.LoadAsync(It.IsAny(), It.IsAny())) - .Returns((id, ct) => - { - _savedSessions.TryGetValue(id, out var session); - return Task.FromResult(session); - }); - - // Default setup for DeleteAsync - _mockStore.Setup(x => x.DeleteAsync(It.IsAny(), It.IsAny())) - .Callback((id, ct) => _savedSessions.Remove(id)) - .Returns(Task.CompletedTask); - - // Default setup for ListAsync - _mockStore.Setup(x => x.ListAsync(It.IsAny(), It.IsAny(), It.IsAny())) - .Returns((tagKey, tagValue, ct) => + public async Task InitializeAsync() + { + Directory.CreateDirectory(_testRoot); + _storageOptions = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { - var ids = new List(_savedSessions.Keys); - return Task.FromResult(ids); - }); - - _kvCache = new byte[] { 10, 20, 30, 40 }; - _tokenHistory = new int[] { 1, 2, 3, 100, 200 }; - _modelFingerprint = "llama-2-7b-v1"; - _samplingConfig = new SamplingData { Temperature = 0.8f, TopP = 0.95f }; - _sessionId = Guid.NewGuid(); + RootPath = _testRoot, + EnsureDirectoryExists = true, + ValidatePermissionsOnStartup = true + } + }; + _manager = new SessionManager(_storageOptions); + await Task.CompletedTask; } - // -------------------- Constructors -------------------- - - [Fact] - public void Constructor_WithStore_ShouldSetStore() + public async Task DisposeAsync() { - // Act - var manager = new SessionManager(_mockStore.Object); - - // Assert - var id = manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint).GetAwaiter().GetResult(); - _mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Once); - _savedSessions.Keys.Should().Contain(id); + await _manager.DisposeAsync(); + if (Directory.Exists(_testRoot)) + Directory.Delete(_testRoot, recursive: true); } - [Fact] - public async Task Constructor_WithBackgroundOptions_ShouldCreateBackgroundSaver() + private SessionCheckpoint CreateTestSession( + Guid? id = null, + string modelFingerprint = "test-model", + byte[]? kvCache = null, + int[]? tokenHistory = null, + Dictionary? tags = null) { - // Arrange - var saveCalled = new TaskCompletionSource(); - var mockStore = new Mock(MockBehavior.Loose); - - // Use a ManualResetEventSlim to block the worker thread. - using var blockEvent = new ManualResetEventSlim(false); - - // Setup the mock's SaveAsync to block until we signal the event. - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback((s, ct) => - { - // Block the worker thread until we release it. - blockEvent.Wait(ct); - saveCalled.TrySetResult(true); - }) - .Returns(Task.CompletedTask); - - var options = new BackgroundSaveOptions { Enabled = true, QueueCapacity = 1 }; - var manager = new SessionManager(mockStore.Object, options); - - // Act – SaveAsync enqueues the save, but the worker is blocked on the event. - var saveTask = manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint); - - // Assert – SaveAsync completed synchronously (fire-and-forget). - saveTask.IsCompletedSuccessfully.Should().BeTrue(); - - // The worker is blocked, so the save should NOT have been called yet. - saveCalled.Task.IsCompleted.Should().BeFalse(); - - // Now release the worker by signaling the event. - blockEvent.Set(); - - // Wait for the background save to complete. - await saveCalled.Task.TimeoutAfter(2000); - - // The save should have been called. - saveCalled.Task.IsCompleted.Should().BeTrue(); - - // Clean up - await manager.DisposeAsync(); + return new SessionCheckpoint + { + SessionId = id ?? Guid.NewGuid(), + KvCacheBytes = kvCache ?? new byte[] { 10, 20, 30 }, + TokenHistory = tokenHistory ?? new int[] { 1, 2, 3 }, + ModelFingerprint = modelFingerprint, + SamplingConfig = new SamplingData { Temperature = 0.8f }, + LastUpdated = DateTime.UtcNow, + Tags = tags ?? new Dictionary() + }; } - // -------------------- SaveAsync (Synchronous) -------------------- + #region Constructor Validation [Fact] - public async Task SaveAsync_Synchronous_ShouldSaveSession() + public void Constructor_WhenSqlDbWithoutConnectionString_Throws() { - // Arrange - var manager = new SessionManager(_mockStore.Object); - - // Act - var id = await manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint, _samplingConfig); - - // Assert - _mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Once); - var saved = _savedSessions[id]; - saved.KvCacheBytes.Should().BeEquivalentTo(_kvCache); - saved.TokenHistory.Should().BeEquivalentTo(_tokenHistory); - saved.ModelFingerprint.Should().Be(_modelFingerprint); - saved.SamplingConfig.Temperature.Should().Be(_samplingConfig.Temperature); + var options = new StorageOptions + { + StoreType = StoreType.SqlDb, + DbStoreOptions = new DbStorageOptions { ConnectionString = null } + }; + Action act = () => new SessionManager(options); + act.Should().Throw() + .WithMessage("ConnectionString is required when StoreType is SqlDb."); } [Fact] - public async Task SaveAsync_Synchronous_WithTags_ShouldStoreTags() + public void Constructor_WhenLocalWithoutRootPath_Throws() { - // Arrange - var tags = new Dictionary { { "User", "Alice" } }; - var manager = new SessionManager(_mockStore.Object); - - // Act - await manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint, tags: tags); + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = null } + }; + Action act = () => new SessionManager(options); + act.Should().Throw() + .WithMessage("RootPath is required when StoreType is Local."); + } - // Assert - _mockStore.Verify(x => x.SaveAsync( - It.Is(s => s.Tags == tags), - It.IsAny()), Times.Once); + [Fact] + public void Constructor_WithValidLocalOptions_CreatesStore() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = _testRoot } + }; + var manager = new SessionManager(options); + manager.Should().NotBeNull(); + Directory.Exists(_testRoot).Should().BeTrue(); } [Fact] - public async Task SaveAsync_Synchronous_ShouldRespectCancellation() + public void Constructor_WithBackgroundSaveEnabled_CreatesBackgroundSaver() { - // Arrange - var cts = new CancellationTokenSource(); - cts.Cancel(); - var manager = new SessionManager(_mockStore.Object); - - // Act - var act = async () => await manager.SaveAsync( - _sessionId, _kvCache, _tokenHistory, _modelFingerprint, - cancellationToken: cts.Token); - - // Assert - await act.Should().ThrowAsync(); - _mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Never); + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = _testRoot }, + BackgroundSaveOptions = new BackgroundSaveOptions { Enabled = true, QueueCapacity = 5 } + }; + var manager = new SessionManager(options); + var field = typeof(SessionManager).GetField("_backgroundSaver", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); + var saver = field?.GetValue(manager); + saver.Should().NotBeNull(); } + #endregion + + #region Save and Load + [Fact] - public async Task SaveAsync_Synchronous_WhenSamplingConfigNull_UsesDefault() + public async Task SaveAsync_WithNewSession_SavesAndReturnsId() { - // Arrange - var manager = new SessionManager(_mockStore.Object); + var session = CreateTestSession(); + var id = await _manager.SaveAsync(session); - // Act - var id = await manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint, samplingConfig: null); + id.Should().Be(session.SessionId); - // Assert - _mockStore.Verify(x => x.SaveAsync( - It.Is(s => s.SamplingConfig.Temperature == 0.7f && s.SamplingConfig.TopP == 0.9f), - It.IsAny()), Times.Once); + var loaded = await _manager.LoadAsync(id); + loaded.Should().NotBeNull(); + loaded.SessionId.Should().Be(id); + loaded.KvCacheBytes.Should().BeEquivalentTo(session.KvCacheBytes); + loaded.TokenHistory.Should().BeEquivalentTo(session.TokenHistory); + loaded.ModelFingerprint.Should().Be(session.ModelFingerprint); + loaded.SamplingConfig.Should().BeEquivalentTo(session.SamplingConfig); + loaded.Tags.Should().BeEquivalentTo(session.Tags); } - // -------------------- SaveAsync (Background) -------------------- - [Fact] - public async Task SaveAsync_Background_ShouldReturnImmediatelyAndSaveInBackground() + public async Task SaveAsync_WithExistingId_UpdatesSession() { - // Arrange - var saveCompleted = new TaskCompletionSource(); - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback((s, ct) => saveCompleted.SetResult(true)) - .Returns(Task.CompletedTask); - - var options = new BackgroundSaveOptions { Enabled = true }; - var manager = new SessionManager(mockStore.Object, options); - - // Act - var id = await manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint); - id.Should().NotBe(Guid.Empty); - - // Wait for background - await saveCompleted.Task.WaitAsync(TimeSpan.FromSeconds(5)); - - // Assert - mockStore.Verify(x => x.SaveAsync(It.IsAny(), It.IsAny()), Times.Once); - await manager.DisposeAsync(); + var session = CreateTestSession(); + var id = await _manager.SaveAsync(session); + + var updated = CreateTestSession(id, modelFingerprint: "updated-model", kvCache: new byte[] { 99, 88 }); + var returnedId = await _manager.SaveAsync(updated); + + returnedId.Should().Be(id); + + var loaded = await _manager.LoadAsync(id); + loaded.ModelFingerprint.Should().Be("updated-model"); + loaded.KvCacheBytes.Should().BeEquivalentTo(new byte[] { 99, 88 }); } [Fact] - public async Task SaveAsync_Background_ShouldPerformDeepCopy() + public async Task LoadAsync_WhenNotExists_ReturnsNull() { - // Arrange - SessionCheckpoint? captured = null; - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Callback((s, ct) => captured = s) - .Returns(Task.CompletedTask); - - var options = new BackgroundSaveOptions { Enabled = true }; - var manager = new SessionManager(mockStore.Object, options); - var mutableKv = new byte[] { 50, 60, 70, 80 }; - - // Act - await manager.SaveAsync(_sessionId, mutableKv, _tokenHistory, _modelFingerprint); - mutableKv[0] = 99; // Mutate - - // Wait for background - await Task.Delay(200); - - // Assert - captured.Should().NotBeNull(); - captured!.KvCacheBytes.Should().BeEquivalentTo(new byte[] { 50, 60, 70, 80 }); - await manager.DisposeAsync(); + var result = await _manager.LoadAsync(Guid.NewGuid()); + result.Should().BeNull(); } + #endregion + + #region Delete + [Fact] - public async Task SaveAsync_Background_WithErrorCallback_ShouldInvokeOnError() + public async Task DeleteAsync_RemovesSession() { - // Arrange - Exception? captured = null; - var options = new BackgroundSaveOptions - { - Enabled = true, - OnError = ex => captured = ex - }; - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.SaveAsync(It.IsAny(), It.IsAny())) - .Throws(new InvalidOperationException("Save failed")); + var session = CreateTestSession(); + var id = await _manager.SaveAsync(session); - var manager = new SessionManager(mockStore.Object, options); + await _manager.DeleteAsync(id); - // Act - await manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint); - await Task.Delay(200); - - // Assert - captured.Should().NotBeNull(); - captured.Should().BeOfType(); - captured!.Message.Should().Be("Save failed"); - await manager.DisposeAsync(); + var loaded = await _manager.LoadAsync(id); + loaded.Should().BeNull(); } - // -------------------- LoadAsync -------------------- + #endregion + + #region List [Fact] - public async Task LoadAsync_ShouldReturnSessionFromStore() + public async Task ListAsync_WithoutFilters_ReturnsAllIds() { - // Arrange - var manager = new SessionManager(_mockStore.Object); - var id = await manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint); - _mockStore.Invocations.Clear(); - - // Act - var loaded = await manager.LoadAsync(id); + var ids = new List(); + for (int i = 0; i < 3; i++) + { + var s = CreateTestSession(); + ids.Add(await _manager.SaveAsync(s)); + } - // Assert - _mockStore.Verify(x => x.LoadAsync(id, It.IsAny()), Times.Once); - loaded.Should().NotBeNull(); - loaded!.SessionId.Should().Be(id); + var all = await _manager.ListAsync(); + all.Should().BeEquivalentTo(ids); } [Fact] - public async Task LoadAsync_WhenNotFound_ShouldReturnNull() + public async Task ListAsync_WithTagFilter_ReturnsMatchingIds() { - // Arrange - var mockStore = new Mock(MockBehavior.Loose); - mockStore.Setup(x => x.LoadAsync(It.IsAny(), It.IsAny())) - .Returns((id, ct) => Task.FromResult(null)); + var s1 = CreateTestSession(tags: new Dictionary { { "env", "prod" } }); + var s2 = CreateTestSession(tags: new Dictionary { { "env", "dev" } }); - var manager = new SessionManager(mockStore.Object); + var id1 = await _manager.SaveAsync(s1); + var id2 = await _manager.SaveAsync(s2); - // Act - var loaded = await manager.LoadAsync(Guid.NewGuid()); + var result = await _manager.ListAsync("env", "prod"); - // Assert - loaded.Should().BeNull(); + result.Count.Should().Be(1); + result.Should().Contain(id1); + result.Should().NotContain(id2); } - // -------------------- DeleteAsync -------------------- + #endregion + + #region Query [Fact] - public async Task DeleteAsync_ShouldCallStoreDelete() + public async Task QueryAsync_ReturnsSummaries() { - // Arrange - var manager = new SessionManager(_mockStore.Object); - var id = await manager.SaveAsync(_sessionId, _kvCache, _tokenHistory, _modelFingerprint); - _savedSessions.Keys.Should().Contain(id); - _mockStore.Invocations.Clear(); - - // Act - await manager.DeleteAsync(id); - - // Assert - _mockStore.Verify(x => x.DeleteAsync(id, It.IsAny()), Times.Once); - _savedSessions.Keys.Should().NotContain(id); - } + var s1 = CreateTestSession(modelFingerprint: "model-A"); + var s2 = CreateTestSession(modelFingerprint: "model-B"); + await _manager.SaveAsync(s1); + await _manager.SaveAsync(s2); - // -------------------- ListAsync -------------------- + var query = new SessionQuery { ModelFingerprint = "model-B" }; + var summaries = await _manager.QueryAsync(query); - [Fact] - public async Task ListAsync_ShouldReturnAllIds() - { - // Arrange - var manager = new SessionManager(_mockStore.Object); - var id1 = await manager.SaveAsync(Guid.NewGuid(), _kvCache, _tokenHistory, _modelFingerprint); - var id2 = await manager.SaveAsync(Guid.NewGuid(), _kvCache, _tokenHistory, _modelFingerprint); - _mockStore.Invocations.Clear(); - - // Act - var ids = await manager.ListAsync(); - - // Assert - _mockStore.Verify(x => x.ListAsync(null, null, It.IsAny()), Times.Once); - ids.Should().BeEquivalentTo(new[] { id1, id2 }); + summaries.Should().ContainSingle(s => s.ModelFingerprint == "model-B"); + summaries.Should().NotContain(s => s.ModelFingerprint == "model-A"); } [Fact] - public async Task ListAsync_WithTagFilter_ShouldPassTagsToStore() + public async Task QueryStreamAsync_StreamsSummaries() { - // Arrange - var manager = new SessionManager(_mockStore.Object); + var s1 = CreateTestSession(modelFingerprint: "model-A"); + var s2 = CreateTestSession(modelFingerprint: "model-B"); + await _manager.SaveAsync(s1); + await _manager.SaveAsync(s2); - // Act - await manager.ListAsync("User", "Alice"); + var stream = _manager.QueryStreamAsync(new SessionQuery()); + var list = await stream.ToListAsync(); - // Assert - _mockStore.Verify(x => x.ListAsync("User", "Alice", It.IsAny()), Times.Once); + list.Should().HaveCount(2); + list.Should().Contain(s => s.ModelFingerprint == "model-A"); + list.Should().Contain(s => s.ModelFingerprint == "model-B"); } - // -------------------- DisposeAsync -------------------- + #endregion + + #region Dispose [Fact] - public async Task DisposeAsync_ShouldDisposeBackgroundSaver() + public async Task DisposeAsync_DisposesStore() { - // Arrange - var options = new BackgroundSaveOptions { Enabled = true }; - var manager = new SessionManager(_mockStore.Object, options); - - // Act - await manager.DisposeAsync(); - - // Assert - verify the store wasn't harmed, and disposing twice is safe - await manager.DisposeAsync(); // Should not throw + // Already called in DisposeAsync of the test class. + // We'll just call it again to ensure no exception. + await _manager.DisposeAsync(); + await _manager.DisposeAsync(); // should be safe } + + #endregion } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresTestBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresTestBase.cs index 75fbbc2..d32f334 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresTestBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresTestBase.cs @@ -1,12 +1,12 @@ -using StateCheckpoint.NET.Stores; + using Npgsql; namespace StateCheckpoint.NET.Tests.Stores.Postgres; public abstract class PostgresTestBase : IAsyncLifetime { - protected PostgresModelStore ModelStore { get; private set; } = null!; - protected PostgresSessionStore SessionStore { get; private set; } = null!; + internal PostgresModelStore ModelStore { get; private set; } = null!; + internal PostgresSessionStore SessionStore { get; private set; } = null!; private string _connectionString = string.Empty; public async Task InitializeAsync() @@ -15,7 +15,7 @@ public async Task InitializeAsync() _connectionString = connectionString; - await ClearTablesAsync(); + await DisposeAsync(); ModelStore = new PostgresModelStore(connectionString); await ModelStore.EnsureSchemaAsync(); @@ -24,24 +24,6 @@ public async Task InitializeAsync() await SessionStore.EnsureSchemaAsync(); } - private async Task ClearTablesAsync() - { - try - { - await using var connection = new NpgsqlConnection(_connectionString); - await connection.OpenAsync(); - // TRUNCATE CASCADE handles foreign keys automatically in PostgreSQL - await using var command = new NpgsqlCommand( - "TRUNCATE TABLE model_manifests, model_blobs, inference_sessions CASCADE;", - connection); - await command.ExecuteNonQueryAsync(); - } - catch (Exception) - { - // Ignore if tables don't exist yet (they will be created later in InitializeAsync) - } - } - public async Task DisposeAsync() { // Truncate all tables to clean up after each test. diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerTestBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerTestBase.cs index a596633..6e516c8 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerTestBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerTestBase.cs @@ -1,13 +1,13 @@ -using StateCheckpoint.NET.Stores; -using Microsoft.Data.SqlClient; +using Microsoft.Data.SqlClient; namespace StateCheckpoint.NET.Tests.Stores.SqlServer; public abstract class SqlServerTestBase : IAsyncLifetime { - protected SqlServerModelStore ModelStore { get; private set; } = null!; - protected SqlServerSessionStore SessionStore { get; private set; } = null!; + internal SqlServerModelStore ModelStore { get; private set; } = null!; + internal SqlServerSessionStore SessionStore { get; private set; } = null!; private string _connectionString = string.Empty; + private bool _disposed; public async Task InitializeAsync() { @@ -15,7 +15,7 @@ public async Task InitializeAsync() _connectionString = connectionString; - await ClearTablesAsync(); + await CleanupAsync(); ModelStore = new SqlServerModelStore(connectionString); await ModelStore.EnsureSchemaAsync(); @@ -24,7 +24,7 @@ public async Task InitializeAsync() await SessionStore.EnsureSchemaAsync(); } - private async Task ClearTablesAsync() + protected virtual async Task CleanupAsync() { try { @@ -32,38 +32,32 @@ private async Task ClearTablesAsync() await connection.OpenAsync(); - // SQL Server truncate order: child tables first to avoid FK violations. - // ModelBlobs references ModelManifests, so truncate it first. await using var command = new SqlCommand( - "TRUNCATE TABLE ModelBlobs; TRUNCATE TABLE ModelManifests; TRUNCATE TABLE InferenceSessions;", + "TRUNCATE TABLE ModelManifests; TRUNCATE TABLE ModelBlobs; TRUNCATE TABLE InferenceSessions;", connection); await command.ExecuteNonQueryAsync(); } - catch (Exception) + catch { - // Ignore if tables don't exist yet. + // Ignore if tables don't exist. } } - public async Task DisposeAsync() { + if (_disposed) return; + try { - await using var connection = new SqlConnection(_connectionString); - - await connection.OpenAsync(); + await CleanupAsync(); - await using var command = new SqlCommand( - "TRUNCATE TABLE ModelManifests, ModelBlobs, InferenceSessions;", - connection); - - await command.ExecuteNonQueryAsync(); + _disposed = true; } - catch (Exception) + catch { - // Ignore if tables don't exist. + Console.WriteLine("err"); + // Ignore exceptions during cleanup to prevent test runner crash. } } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Helpers/TagsHelper.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Helpers/TagsHelper.cs new file mode 100644 index 0000000..e6f92f3 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Helpers/TagsHelper.cs @@ -0,0 +1,19 @@ +namespace StateCheckpoint.NET; + +internal static class TagsHelper +{ + public static bool TagsMatch(Dictionary? actual, Dictionary? filter) + { + if (filter == null) + return true; + + if (actual == null) + return false; + + foreach (var pair in filter) + if (!actual.TryGetValue(pair.Key, out var val) || val != pair.Value) + return false; + + return true; + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs index f5763c7..866fb26 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs @@ -52,54 +52,21 @@ public CheckpointManager(StorageOptions storageOptions) /// Cancellation token. /// Returns the ModelId (GUID) of the saved checkpoint. public async Task SaveAsync( - byte[] weights, - byte[] optimizer, - HyperParameters hyperParams, - TokenizerData tokenizer, - int epoch, - float loss, - Guid? existingId = null, - Dictionary? tags = null, + ModelCheckpoint modelCheckpoint, CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); - var checkpoint = new ModelCheckpoint - { - ModelId = existingId ?? Guid.NewGuid(), - WeightsBytes = weights, - OptimizerBytes = optimizer, - HyperParams = hyperParams, - Tokenizer = tokenizer, - CurrentEpoch = epoch, - LastTrainingLoss = loss, - CreatedAt = DateTime.UtcNow, - Tags = tags ?? new Dictionary() - }; - if (_backgroundSaver != null) { - var capturedCheckpoint = new ModelCheckpoint - { - ModelId = checkpoint.ModelId, - WeightsBytes = checkpoint.WeightsBytes.ToArray(), - OptimizerBytes = checkpoint.OptimizerBytes.ToArray(), - HyperParams = checkpoint.HyperParams, - Tokenizer = checkpoint.Tokenizer, - CurrentEpoch = checkpoint.CurrentEpoch, - LastTrainingLoss = checkpoint.LastTrainingLoss, - CreatedAt = checkpoint.CreatedAt, - Tags = checkpoint.Tags - }; - - await _backgroundSaver.EnqueueAsync(async (ct) => await _store.SaveAsync(capturedCheckpoint, ct), cancellationToken); - - return checkpoint.ModelId; + await _backgroundSaver.EnqueueAsync(async (ct) => await _store.SaveAsync(modelCheckpoint, ct), cancellationToken); + + return modelCheckpoint.ModelId; } - await _store.SaveAsync(checkpoint, cancellationToken); + await _store.SaveAsync(modelCheckpoint, cancellationToken); - return checkpoint.ModelId; + return modelCheckpoint.ModelId; } /// Finds the checkpoint with the highest score according to a user-supplied selector. @@ -141,6 +108,12 @@ public async Task SaveAsync( public Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default) => _store.QueryAsync(query, cancellationToken); + /// + /// + /// + /// + /// + /// public IAsyncEnumerable QueryStreamAsync(CheckpointQuery query, CancellationToken cancellationToken = default) => _store.QueryStreamAsync(query, cancellationToken); @@ -178,8 +151,9 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue public async ValueTask DisposeAsync() { if (_backgroundSaver != null) - { await _backgroundSaver.DisposeAsync(); - } + + if (_store is IAsyncDisposable asyncDisposable) + await asyncDisposable.DisposeAsync(); } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs index e05cb88..4d3f083 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs @@ -3,6 +3,9 @@ namespace StateCheckpoint.NET; +/// +/// +/// public class SessionManager : IAsyncDisposable { private readonly ISessionStore _store; @@ -14,19 +17,24 @@ public class SessionManager : IAsyncDisposable /// Any implementation of ISessionStore (FileSystem, PostgreSQL, etc.) public SessionManager(StorageOptions storageOptions) { - // ... same validation as CheckpointManager + if (storageOptions.StoreType == StoreType.SqlDb && string.IsNullOrWhiteSpace(storageOptions.DbStoreOptions?.ConnectionString)) + throw new InvalidOperationException("ConnectionString is required when StoreType is SqlDb."); + + if (storageOptions.StoreType == StoreType.Local && string.IsNullOrWhiteSpace(storageOptions.FileSystemStoreOptions?.RootPath)) + throw new InvalidOperationException("RootPath is required when StoreType is Local."); + _store = StoreSessionFactory.Create(storageOptions); - if (storageOptions?.DbStoreOptions?.EnsureSchemaOnStartup == true && _store is IDbSessionStore dbStore) + if (storageOptions.DbStoreOptions?.EnsureSchemaOnStartup == true && _store is IDbSessionStore dbStore) { dbStore.EnsureSchemaAsync().GetAwaiter().GetResult(); } - if (storageOptions?.BackgroundSaveOptions?.Enabled == true) + if (storageOptions.BackgroundSaveOptions?.Enabled == true) { _backgroundSaver = new BackgroundSaver( capacity: storageOptions.BackgroundSaveOptions.QueueCapacity, - onError: storageOptions?.BackgroundSaveOptions?.OnError); + onError: storageOptions.BackgroundSaveOptions?.OnError); } } @@ -42,48 +50,21 @@ public SessionManager(StorageOptions storageOptions) /// Cancellation token. /// Returns the SessionId (GUID) of the saved session. public async Task SaveAsync( - Guid sessionId, - byte[] kvCacheBytes, - int[] tokenHistory, - string modelFingerprint, - SamplingData? samplingConfig = null, - Dictionary? tags = null, + SessionCheckpoint sessionCheckpoint, CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); - var session = new SessionCheckpoint - { - SessionId = sessionId, - KvCacheBytes = kvCacheBytes, - TokenHistory = tokenHistory, - ModelFingerprint = modelFingerprint, - SamplingConfig = samplingConfig ?? new SamplingData(), - LastUpdated = DateTime.UtcNow, - Tags = tags ?? new Dictionary() - }; - if (_backgroundSaver != null) - { - var capturedSession = new SessionCheckpoint - { - SessionId = session.SessionId, - KvCacheBytes = session.KvCacheBytes.ToArray(), - TokenHistory = session.TokenHistory, - ModelFingerprint = session.ModelFingerprint, - SamplingConfig = session.SamplingConfig, - LastUpdated = session.LastUpdated, - Tags = session.Tags - }; - - await _backgroundSaver.EnqueueAsync(async (cToken) => await _store.SaveAsync(capturedSession, cToken)); - - return session.SessionId; + { + await _backgroundSaver.EnqueueAsync(async (cToken) => await _store.SaveAsync(sessionCheckpoint, cToken)); + + return sessionCheckpoint.SessionId; } - await _store.SaveAsync(session, cancellationToken); + await _store.SaveAsync(sessionCheckpoint, cancellationToken); - return session.SessionId; + return sessionCheckpoint.SessionId; } /// @@ -113,9 +94,21 @@ public async Task DeleteAsync(Guid sessionId, CancellationToken cancellationToke public async Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default) => await _store.ListAsync(tagKey, tagValue, cancellationToken); + /// + /// + /// + /// + /// + /// public Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default) => _store.QueryAsync(query, cancellationToken); + /// + /// + /// + /// + /// + /// public IAsyncEnumerable QueryStreamAsync(SessionQuery query, CancellationToken cancellationToken = default) => _store.QueryStreamAsync(query, cancellationToken); diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs index c86c344..023a042 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Queries/SqlServerTrainingQueries.cs @@ -84,6 +84,5 @@ FROM ModelManifests m // --- Listing --- public const string ListAllModelIds = "SELECT ModelId FROM ModelManifests;"; - //public const string ListModelIdsByTag = "SELECT ModelId FROM ModelManifests WHERE Tags LIKE @TagPattern;"; public const string ListModelIdsByTag = "SELECT ModelId FROM ModelManifests WHERE JSON_VALUE(Tags, '$.{tagKey}') = @TagValue"; } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs index 7771340..07fdcd8 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Settings/DbStorageOptions.cs @@ -10,5 +10,8 @@ public class DbStorageOptions /// public string? ConnectionString { get; set; } + /// + /// + /// public bool EnsureSchemaOnStartup { get; set; } } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs new file mode 100644 index 0000000..f4aa5f5 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs @@ -0,0 +1,87 @@ +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; + +namespace StateCheckpoint.NET; + +internal abstract class FileSystemBase +{ + private readonly string _rootPath; + private readonly FileSystemStoreOptions _options; + private readonly TIndex _index; + private bool _indexLoaded; + private readonly SemaphoreSlim _buildLock = new(1, 1); + + public FileSystemBase(string rootPath, FileSystemStoreOptions options, Func indexSelector) + { + _options = options ?? new FileSystemStoreOptions(); + _rootPath = Path.Combine(rootPath, "models"); + + if (_options.ValidatePermissionsOnStartup) + { + if (!FileSystemHelper.TryValidateWriteAccess(_rootPath, out var error)) + { + // If fallback is provided, update the root path + if (!string.IsNullOrWhiteSpace(_options.FallbackPath)) + { + _rootPath = Path.Combine(_options.FallbackPath, "models"); + + Directory.CreateDirectory(_rootPath); + } + else + { + throw error!; + } + } + } + + // Still ensure the directory exists if required + else if (_options.EnsureDirectoryExists) + Directory.CreateDirectory(_rootPath); + + var indexFilePath = Path.Combine(_rootPath, "_index.json"); + + _index = indexSelector(indexFilePath); + } + + private async Task EnsureIndexLoadedAsync(Action getAll, Func>> getAllEntries, Func rebuild, CancellationToken cancellationToken) + { + if (!_indexLoaded) + { + await _buildLock.WaitAsync(cancellationToken); + try + { + if (_indexLoaded) return; + + + await _index.LoadAsync(cancellationToken); + + var entries = await getAllEntries(); // await _index.GetAllAsync(cancellationToken); + if (!entries.Any()) + { + await _index.RebuildAsync(_rootPath, async (rootPath, id, ct) => + { + var manifest = await FileSystemHelper + .LoadManifestOnlyAsync(rootPath, id, "meta.json", ct); + + if (manifest == null) return null; + + return new SessionSummary + { + SessionId = id, + ModelFingerprint = manifest.ModelFingerprint, + LastUpdated = manifest.LastUpdated, + Tags = manifest.Tags ?? new Dictionary() + }; + }, cancellationToken); + + rebuild(_rootPath); + } + _indexLoaded = true; + } + finally + { + _buildLock.Release(); + } + } + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs index 6694518..c056eaa 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs @@ -23,7 +23,7 @@ public FileSystemModelStore(string rootPath, FileSystemStoreOptions? options = n if (!FileSystemHelper.TryValidateWriteAccess(_rootPath, out var error)) { // If fallback is provided, update the root path - if (!string.IsNullOrEmpty(_options.FallbackPath)) + if (!string.IsNullOrWhiteSpace(_options.FallbackPath)) { _rootPath = Path.Combine(_options.FallbackPath, "models"); @@ -35,14 +35,10 @@ public FileSystemModelStore(string rootPath, FileSystemStoreOptions? options = n } } } - else - { - // Still ensure the directory exists if required - if (_options.EnsureDirectoryExists) - { + + // Still ensure the directory exists if required + else if (_options.EnsureDirectoryExists) Directory.CreateDirectory(_rootPath); - } - } var indexFilePath = Path.Combine(_rootPath, "_index.json"); @@ -155,67 +151,23 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue public async Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default) { - var allIds = await FileSystemHelper.ListAsync(_rootPath, cancellationToken); - var summaries = new List(); - - foreach (var id in allIds) - { - cancellationToken.ThrowIfCancellationRequested(); - - var manifest = await FileSystemHelper.LoadManifestOnlyAsync(_rootPath, id, "manifest.json", cancellationToken); - if (manifest == null) continue; - - // Apply filters - if (query.MinEpoch.HasValue && manifest.CurrentEpoch < query.MinEpoch) continue; - if (query.MaxEpoch.HasValue && manifest.CurrentEpoch > query.MaxEpoch) continue; - if (query.MinLoss.HasValue && manifest.LastTrainingLoss < query.MinLoss) continue; - if (query.MaxLoss.HasValue && manifest.LastTrainingLoss > query.MaxLoss) continue; - if (query.CreatedAfter.HasValue && manifest.CreatedAt < query.CreatedAfter) continue; - if (query.CreatedBefore.HasValue && manifest.CreatedAt > query.CreatedBefore) continue; - if (query.Tags != null && !TagsMatch(manifest.Tags, query.Tags)) continue; - - summaries.Add(new CheckpointSummary - { - ModelId = id, - CurrentEpoch = manifest.CurrentEpoch, - LastTrainingLoss = manifest.LastTrainingLoss, - CreatedAt = manifest.CreatedAt, - Tags = manifest.Tags ?? new() - }); - } - - // Sorting - summaries = query.OrderBy switch - { - CheckpointSortField.CreatedAt => query.Descending - ? summaries.OrderByDescending(s => s.CreatedAt).ToList() - : summaries.OrderBy(s => s.CreatedAt).ToList(), - CheckpointSortField.Epoch => query.Descending - ? summaries.OrderByDescending(s => s.CurrentEpoch).ToList() - : summaries.OrderBy(s => s.CurrentEpoch).ToList(), - CheckpointSortField.Loss => query.Descending - ? summaries.OrderByDescending(s => s.LastTrainingLoss).ToList() - : summaries.OrderBy(s => s.LastTrainingLoss).ToList(), - _ => summaries - }; - - if (query.Limit.HasValue && query.Limit.Value > 0) - summaries = summaries.Take(query.Limit.Value).ToList(); + await EnsureIndexLoadedAsync(cancellationToken); - return summaries; + return await _index.QueryAsync(query, cancellationToken); } public async IAsyncEnumerable QueryStreamAsync( CheckpointQuery query, - [EnumeratorCancellation] CancellationToken ct = default) + [EnumeratorCancellation] CancellationToken cancellationToken = default) { - await EnsureIndexLoadedAsync(ct); - var all = await _index.GetAllAsync(ct); + await EnsureIndexLoadedAsync(cancellationToken); + + var all = await _index.GetAllAsync(cancellationToken); // Filter in memory (no sorting/limiting – streaming preserves order of storage) foreach (var summary in all) { - ct.ThrowIfCancellationRequested(); + cancellationToken.ThrowIfCancellationRequested(); if (query.MinEpoch.HasValue && summary.CurrentEpoch < query.MinEpoch) continue; if (query.MaxEpoch.HasValue && summary.CurrentEpoch > query.MaxEpoch) continue; @@ -223,72 +175,52 @@ public async IAsyncEnumerable QueryStreamAsync( if (query.MaxLoss.HasValue && summary.LastTrainingLoss > query.MaxLoss) continue; if (query.CreatedAfter.HasValue && summary.CreatedAt < query.CreatedAfter) continue; if (query.CreatedBefore.HasValue && summary.CreatedAt > query.CreatedBefore) continue; - if (query.Tags != null && query.Tags.Count > 0 && !TagsMatch(summary.Tags, query.Tags)) continue; + if (query.Tags != null && query.Tags.Count > 0 && !TagsHelper.TagsMatch(summary.Tags, query.Tags)) continue; yield return summary; } } - private static bool TagsMatch(Dictionary? actual, Dictionary? filter) + private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) { - if (filter == null) - return true; - - if (actual == null) - return false; - - foreach (var pair in filter) - if (!actual.TryGetValue(pair.Key, out var val) || val != pair.Value) - return false; - - return true; - } - - private async Task EnsureIndexLoadedAsync(CancellationToken ct) - { - // We can check if _items is empty or use a flag; simplest: always load if not loaded. - // We'll have a separate field to track loaded state. - // Here we just call LoadAsync if needed. Since it's cheap, we can call it every time. - // But better to check a flag. if (!_indexLoaded) { - await _index.LoadAsync(ct); - _indexLoaded = true; - } - } - - private async Task RebuildIndexAsync(CancellationToken ct) - { - await _buildLock.WaitAsync(ct); - try - { - // double-check after lock - await EnsureIndexLoadedAsync(ct); - var existing = await _index.GetAllAsync(ct); - if (existing.Any()) return; - - async Task LoadManifestSummaryAsync(string rootPath, Guid id, CancellationToken ct) + await _buildLock.WaitAsync(cancellationToken); + try { - var manifest = await FileSystemHelper.LoadManifestOnlyAsync(rootPath, id, "manifest.json", ct); + if (_indexLoaded) return; // double-check - if (manifest == null) return null; + // Load the index from disk + await _index.LoadAsync(cancellationToken); - return new CheckpointSummary + var entries = await _index.GetAllAsync(cancellationToken); + + // If index is empty, rebuild it from manifests + if (!entries.Any()) { - ModelId = id, - CurrentEpoch = manifest.CurrentEpoch, - LastTrainingLoss = manifest.LastTrainingLoss, - CreatedAt = manifest.CreatedAt, - Tags = manifest.Tags ?? new() - }; - } + await _index.RebuildAsync(_rootPath, async (string rootPath, Guid id, CancellationToken ct) => + { + var manifest = await FileSystemHelper.LoadManifestOnlyAsync(rootPath, id, "manifest.json", ct); + + if (manifest == null) return null; + + return new CheckpointSummary + { + ModelId = id, + CurrentEpoch = manifest.CurrentEpoch, + LastTrainingLoss = manifest.LastTrainingLoss, + CreatedAt = manifest.CreatedAt, + Tags = manifest.Tags ?? new() + }; + }, cancellationToken); + } - // Rebuild - await _index.RebuildAsync(_rootPath, LoadManifestSummaryAsync, ct); - } - finally - { - _buildLock.Release(); + _indexLoaded = true; + } + finally + { + _buildLock.Release(); + } } } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs index 1194778..a40de77 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs @@ -1,5 +1,6 @@ using StateCheckpoint.NET.Models; using StateCheckpoint.NET.Settings; +using System.Runtime.CompilerServices; namespace StateCheckpoint.NET; @@ -133,49 +134,66 @@ public async Task> QueryAsync(SessionQuery query, Cancellat return await _index.QueryAsync(query, cancellationToken); } - // Rebuild index from all session manifest files (fallback) - public async Task RebuildIndexAsync(CancellationToken cancellationToken) + public async IAsyncEnumerable QueryStreamAsync( + SessionQuery query, + [EnumeratorCancellation] CancellationToken cancellationToken = default) { - await _buildLock.WaitAsync(cancellationToken); - try - { - // Double-check after acquiring the lock - await EnsureIndexLoadedAsync(cancellationToken); + await EnsureIndexLoadedAsync(cancellationToken); + var all = await _index.GetAllAsync(cancellationToken); - var existing = await _index.GetAllAsync(cancellationToken); + foreach (var summary in all) + { + cancellationToken.ThrowIfCancellationRequested(); - if (existing.Any()) return; // already rebuilt by another thread + // Apply filters + if (!string.IsNullOrEmpty(query.ModelFingerprint) && summary.ModelFingerprint != query.ModelFingerprint) continue; + if (query.UpdatedAfter.HasValue && summary.LastUpdated < query.UpdatedAfter.Value) continue; + if (query.UpdatedBefore.HasValue && summary.LastUpdated > query.UpdatedBefore.Value) continue; + if (query.Tags != null && query.Tags.Count > 0 && !TagsHelper.TagsMatch(summary.Tags, query.Tags)) continue; - await _index.RebuildAsync(_rootPath, LoadManifestSummaryAsync, cancellationToken); - } - finally - { - _buildLock.Release(); + yield return summary; } } + //---------------- Private Methods ----------------// + // Ensure the index is loaded into memory (lazy) private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) { if (!_indexLoaded) { - await _index.LoadAsync(cancellationToken); + await _buildLock.WaitAsync(cancellationToken); + try + { + if (_indexLoaded) return; - _indexLoaded = true; - } - } + await _index.LoadAsync(cancellationToken); - private async Task LoadManifestSummaryAsync(string rootPath, Guid id, CancellationToken ct) - { - var manifest = await FileSystemHelper.LoadManifestOnlyAsync( - rootPath, id, "meta.json", ct); - if (manifest == null) return null; - return new SessionSummary - { - SessionId = id, - ModelFingerprint = manifest.ModelFingerprint, - LastUpdated = manifest.LastUpdated, - Tags = manifest.Tags ?? new Dictionary() - }; + var entries = await _index.GetAllAsync(cancellationToken); + if (!entries.Any()) + { + await _index.RebuildAsync(_rootPath, async (rootPath, id, ct) => + { + var manifest = await FileSystemHelper + .LoadManifestOnlyAsync(rootPath, id, "meta.json", ct); + + if (manifest == null) return null; + + return new SessionSummary + { + SessionId = id, + ModelFingerprint = manifest.ModelFingerprint, + LastUpdated = manifest.LastUpdated, + Tags = manifest.Tags ?? new Dictionary() + }; + }, cancellationToken); + } + _indexLoaded = true; + } + finally + { + _buildLock.Release(); + } + } } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs index 09812db..9acc1bc 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs @@ -84,22 +84,25 @@ public async Task RemoveAsync(Guid modelId, CancellationToken ct = default) public async Task RebuildAsync( string rootPath, Func> loadManifestFn, - CancellationToken ct = default) + CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { - var allIds = await FileSystemHelper.ListAsync(rootPath, ct); + var allIds = await FileSystemHelper.ListAsync(rootPath, cancellationToken); + var newItems = new List(); foreach (var id in allIds) { - ct.ThrowIfCancellationRequested(); - var summary = await loadManifestFn(rootPath, id, ct); + cancellationToken.ThrowIfCancellationRequested(); + var summary = await loadManifestFn(rootPath, id, cancellationToken); if (summary != null) newItems.Add(summary); } + _items = newItems; - await _storage.WriteAsync(_items, ct); + + await _storage.WriteAsync(_items, cancellationToken); } finally { @@ -110,9 +113,9 @@ public async Task RebuildAsync( /// /// Returns a copy of all items (for querying). /// - public async Task> GetAllAsync(CancellationToken ct = default) + public async Task> GetAllAsync(CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { return _items.ToList(); // defensive copy @@ -128,9 +131,9 @@ public async Task> GetAllAsync(CancellationToke /// public async Task> QueryAsync( CheckpointQuery query, - CancellationToken ct = default) + CancellationToken cancellationToken = default) { - var all = await GetAllAsync(ct); + var all = await GetAllAsync(cancellationToken); var filtered = all.AsEnumerable(); // Apply filters (same as before) @@ -141,7 +144,7 @@ public async Task> QueryAsync( if (query.CreatedAfter.HasValue) filtered = filtered.Where(s => s.CreatedAt >= query.CreatedAfter.Value); if (query.CreatedBefore.HasValue) filtered = filtered.Where(s => s.CreatedAt <= query.CreatedBefore.Value); if (query.Tags != null && query.Tags.Count > 0) - filtered = filtered.Where(s => TagsMatch(s.Tags, query.Tags)); + filtered = filtered.Where(s => TagsHelper.TagsMatch(s.Tags, query.Tags)); // Sort var sorted = query.OrderBy switch @@ -163,19 +166,4 @@ public async Task> QueryAsync( return sorted.ToList(); } - - private static bool TagsMatch(Dictionary? actual, Dictionary? filter) - { - if (filter == null) - return true; - - if (actual == null) - return false; - - foreach (var pair in filter) - if (!actual.TryGetValue(pair.Key, out var val) || val != pair.Value) - return false; - - return true; - } } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs index f270815..3eb9b85 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs @@ -2,7 +2,6 @@ namespace StateCheckpoint.NET; - internal sealed class IndexStorage { private readonly string _filePath; diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs index c29d5a2..9a8832f 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/SessionIndex.cs @@ -19,12 +19,12 @@ internal SessionIndex(IndexStorage storage) _items = new List(); } - public async Task LoadAsync(CancellationToken ct = default) + public async Task LoadAsync(CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { - var data = await _storage.ReadAsync(ct); + var data = await _storage.ReadAsync(cancellationToken); _items = data.ToList(); } finally @@ -58,14 +58,14 @@ public async Task AddOrUpdateAsync(SessionSummary summary, CancellationToken can } } - public async Task RemoveAsync(Guid sessionId, CancellationToken ct = default) + public async Task RemoveAsync(Guid sessionId, CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { var removed = _items.RemoveAll(e => e.SessionId == sessionId); if (removed > 0) - await _storage.WriteAsync(_items, ct); + await _storage.WriteAsync(_items, cancellationToken); } finally { @@ -76,18 +76,18 @@ public async Task RemoveAsync(Guid sessionId, CancellationToken ct = default) public async Task RebuildAsync( string rootPath, Func> loadManifestFn, - CancellationToken ct = default) + CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { - var allIds = await FileSystemHelper.ListAsync(rootPath, ct); + var allIds = await FileSystemHelper.ListAsync(rootPath, cancellationToken); var newItems = new List(); foreach (var id in allIds) { - ct.ThrowIfCancellationRequested(); - var summary = await loadManifestFn(rootPath, id, ct); + cancellationToken.ThrowIfCancellationRequested(); + var summary = await loadManifestFn(rootPath, id, cancellationToken); if (summary != null) newItems.Add(summary); @@ -95,7 +95,7 @@ public async Task RebuildAsync( _items = newItems; - await _storage.WriteAsync(_items, ct); + await _storage.WriteAsync(_items, cancellationToken); } finally { @@ -103,9 +103,9 @@ public async Task RebuildAsync( } } - public async Task> GetAllAsync(CancellationToken ct = default) + public async Task> GetAllAsync(CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs index 41fa60e..77a38b0 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IFileSystemModelStore.cs @@ -1,5 +1,5 @@ namespace StateCheckpoint.NET; -internal interface IFileSystenModelStore: IModelStore +internal interface IFileSystemModelStore: IModelStore { } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs index 4fed938..08a6820 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs @@ -224,7 +224,7 @@ public async Task DeleteManyAsync(IEnumerable modelIds, CancellationToken private (string Sql, List Parameters) BuildQuerySql(CheckpointQuery query, bool includeOrderBy = true, bool includeLimit = true) { var sql = new StringBuilder(@" - SELECT model_id, epoch, loss, created_at, tags + SELECT modelId, epoch, loss, createdAt, tags FROM ModelManifests WHERE 1=1 "); @@ -253,10 +253,10 @@ FROM ModelManifests { var orderColumn = query.OrderBy switch { - CheckpointSortField.CreatedAt => "created_at", + CheckpointSortField.CreatedAt => "createdAt", CheckpointSortField.Epoch => "epoch", CheckpointSortField.Loss => "loss", - _ => "created_at" + _ => "createdAt" }; sql.Append($" ORDER BY {orderColumn} {(query.Descending ? "DESC" : "ASC")}"); From ab38119a0ba31873c360c2379939ec5b03772eb3 Mon Sep 17 00:00:00 2001 From: Leo Mengesha Date: Tue, 14 Jul 2026 01:20:31 +0100 Subject: [PATCH 3/3] more changes --- .../Stores/Postgres/PostgresStoreBaseTests.cs | 4 +- .../Stores/SqlServer/SqlServerStoreBase.cs | 4 +- .../StorageOptionsValidationExtensionTests.cs | 367 ++++++++++++++++++ .../Extensions/ManagerExtensions.cs | 237 +++++++++++ .../Manager/CheckpointManager.cs | 41 +- .../Manager/SessionManager.cs | 32 +- .../Models/CheckpointDiff.cs | 66 ++++ .../Models/CheckpointSummary.cs | 1 + .../StateCheckpoint.NET/Models/SessionDiff.cs | 71 ++++ .../Models/SessionSummary.cs | 2 + .../Stores/FileSystem/FileSystemBase.cs | 35 +- .../Stores/FileSystem/FileSystemModelStore.cs | 157 ++++---- .../FileSystem/FileSystemSessionStore.cs | 122 ++---- .../Stores/FileSystemIndex/CheckpointIndex.cs | 18 +- .../Stores/FileSystemIndex/IndexStorage.cs | 8 +- .../StateCheckpoint.NET/Stores/IModelStore.cs | 3 +- .../Stores/ISessionStore.cs | 2 + .../Stores/Postgres/PostgresModelStore.cs | 55 ++- .../Stores/Postgres/PostgresSessionStore.cs | 52 ++- .../Stores/Postgres/PostgresStoreBase.cs | 4 +- .../Stores/SqlServer/SqlServerModelStore.cs | 42 +- .../Stores/SqlServer/SqlServerSessionStore.cs | 54 ++- .../Stores/SqlServer/SqlServerStoreBase.cs | 8 +- .../StoreOptionsValidationExtension.cs | 61 +++ 24 files changed, 1197 insertions(+), 249 deletions(-) create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET.Tests/Validation/StorageOptionsValidationExtensionTests.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Extensions/ManagerExtensions.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointDiff.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionDiff.cs create mode 100644 StateCheckpoint.NET/StateCheckpoint.NET/Validation/StoreOptionsValidationExtension.cs diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresStoreBaseTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresStoreBaseTests.cs index 951257c..acaed6d 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresStoreBaseTests.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/Postgres/PostgresStoreBaseTests.cs @@ -12,8 +12,8 @@ private class TestablePostgresStore : PostgresStoreBase public TestablePostgresStore(string connectionString) : base(connectionString) { } public TestablePostgresStore(NpgsqlDataSource dataSource) : base(dataSource) { } - public new async Task GetConnectionAsync(CancellationToken ct = default) - => await base.GetConnectionAsync(ct); + public new async Task GetConnectionAsync(CancellationToken cancellationToken = default) + => await base.GetConnectionAsync(cancellationToken); } private const string DummyConnectionString = "Host=localhost;Database=dummy"; diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerStoreBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerStoreBase.cs index 4b5098e..bfd060b 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerStoreBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Stores/SqlServer/SqlServerStoreBase.cs @@ -32,8 +32,8 @@ private class TestableSqlServerStore : SqlServerStoreBase public TestableSqlServerStore(string connectionString) : base(connectionString) { } public TestableSqlServerStore(SqlConnection connection) : base(connection) { } - public new async Task GetConnectionAsync(CancellationToken ct = default) - => await base.GetConnectionAsync(ct); + public new async Task GetConnectionAsync(CancellationToken cancellationToken = default) + => await base.GetConnectionAsync(cancellationToken); } [SkippableFact] diff --git a/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Validation/StorageOptionsValidationExtensionTests.cs b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Validation/StorageOptionsValidationExtensionTests.cs new file mode 100644 index 0000000..a5b66b7 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET.Tests/Validation/StorageOptionsValidationExtensionTests.cs @@ -0,0 +1,367 @@ +using System; +using System.Collections.Generic; +using FluentAssertions; +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; + +namespace StateCheckpoint.NET; + +public class StorageOptionsValidationTests +{ + #region StoreType Local + + [Fact] + public void Validate_WhenStoreTypeLocalAndFileSystemOptionsNull_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = null + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("FileSystemStoreOptions is required when StoreType is Local."); + } + + [Fact] + public void Validate_WhenStoreTypeLocalAndRootPathNull_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = null } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RootPath is required when StoreType is Local."); + } + + [Fact] + public void Validate_WhenStoreTypeLocalAndRootPathEmpty_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "" } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RootPath is required when StoreType is Local."); + } + + [Fact] + public void Validate_WhenStoreTypeLocalWithValidOptions_Passes() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" } + }; + Action act = () => options.Validate(); + act.Should().NotThrow(); + } + + #endregion + + #region StoreType SqlDb + + [Fact] + public void Validate_WhenStoreTypeSqlDbAndDbOptionsNull_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.SqlDb, + DbStoreOptions = null + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("DbStoreOptions is required when StoreType is SqlDb."); + } + + [Fact] + public void Validate_WhenStoreTypeSqlDbAndConnectionStringNull_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.SqlDb, + DbStoreOptions = new DbStorageOptions { ConnectionString = null } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("ConnectionString is required when StoreType is SqlDb."); + } + + [Fact] + public void Validate_WhenStoreTypeSqlDbAndConnectionStringEmpty_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.SqlDb, + DbStoreOptions = new DbStorageOptions { ConnectionString = "" } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("ConnectionString is required when StoreType is SqlDb."); + } + + [Fact] + public void Validate_WhenStoreTypeSqlDbWithValidOptions_Passes() + { + var options = new StorageOptions + { + StoreType = StoreType.SqlDb, + DbStoreOptions = new DbStorageOptions { ConnectionString = "Host=localhost;Database=test" } + }; + Action act = () => options.Validate(); + act.Should().NotThrow(); + } + + #endregion + + #region Invalid StoreType + + // StoreType is an enum, so we can force an invalid value via casting. + [Fact] + public void Validate_WhenStoreTypeInvalid_Throws() + { + var options = new StorageOptions + { + StoreType = (StoreType)999, // invalid value + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("Unsupported StoreType '999'. Valid values are Local or SqlDb."); + } + + #endregion + + #region BackgroundSaveOptions + + [Fact] + public void Validate_WhenBackgroundSaveEnabledAndQueueCapacityZero_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + BackgroundSaveOptions = new BackgroundSaveOptions { Enabled = true, QueueCapacity = 0 } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("BackgroundSaveOptions.QueueCapacity must be greater than 0 when background saves are enabled."); + } + + [Fact] + public void Validate_WhenBackgroundSaveEnabledAndQueueCapacityNegative_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + BackgroundSaveOptions = new BackgroundSaveOptions { Enabled = true, QueueCapacity = -5 } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("BackgroundSaveOptions.QueueCapacity must be greater than 0 when background saves are enabled."); + } + + [Fact] + public void Validate_WhenBackgroundSaveDisabled_IgnoresQueueCapacity() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + BackgroundSaveOptions = new BackgroundSaveOptions { Enabled = false, QueueCapacity = 0 } + }; + Action act = () => options.Validate(); + act.Should().NotThrow(); + } + + [Fact] + public void Validate_WhenBackgroundSaveOptionsNull_Passes() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + BackgroundSaveOptions = null + }; + Action act = () => options.Validate(); + act.Should().NotThrow(); + } + + #endregion + + #region RetentionPolicy + + [Fact] + public void Validate_WhenMaxCheckpointsZero_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy { MaxCheckpoints = 0 } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RetentionPolicy.MaxCheckpoints must be greater than 0."); + } + + [Fact] + public void Validate_WhenMaxCheckpointsNegative_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy { MaxCheckpoints = -1 } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RetentionPolicy.MaxCheckpoints must be greater than 0."); + } + + [Fact] + public void Validate_WhenMaxAgeZero_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy { MaxAge = TimeSpan.Zero } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RetentionPolicy.MaxAge must be a positive TimeSpan."); + } + + [Fact] + public void Validate_WhenMaxAgeNegative_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy { MaxAge = TimeSpan.FromDays(-1) } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RetentionPolicy.MaxAge must be a positive TimeSpan."); + } + + [Fact] + public void Validate_WhenMaxEpochAgeZero_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy { MaxEpochAge = 0 } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RetentionPolicy.MaxEpochAge must be greater than 0."); + } + + [Fact] + public void Validate_WhenMaxEpochAgeNegative_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy { MaxEpochAge = -5 } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RetentionPolicy.MaxEpochAge must be greater than 0."); + } + + [Fact] + public void Validate_WhenMaxLossThresholdNegative_Throws() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy { MaxLossThreshold = -1.0f } + }; + Action act = () => options.Validate(); + act.Should().Throw() + .WithMessage("RetentionPolicy.MaxLossThreshold must be a positive number."); + } + + [Fact] + public void Validate_WhenRetentionPolicyNull_Passes() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = null + }; + Action act = () => options.Validate(); + act.Should().NotThrow(); + } + + [Fact] + public void Validate_WhenRetentionPolicyValid_Passes() + { + var options = new StorageOptions + { + StoreType = StoreType.Local, + FileSystemStoreOptions = new FileSystemStoreOptions { RootPath = "./data" }, + RetentionPolicy = new RetentionPolicy + { + MaxCheckpoints = 10, + MaxAge = TimeSpan.FromDays(30), + MaxEpochAge = 5, + MaxLossThreshold = 0.5f, + AlwaysKeepBest = true, + PinnedTags = new Dictionary { { "status", "production" } } + } + }; + Action act = () => options.Validate(); + act.Should().NotThrow(); + } + + #endregion + + #region Combined Complex Validations + + [Fact] + public void Validate_WithAllOptionsValid_Passes() + { + var options = new StorageOptions + { + StoreType = StoreType.SqlDb, + SqlDbType = SqlDbType.Postgres, + DbStoreOptions = new DbStorageOptions + { + ConnectionString = "Host=localhost;Database=test", + EnsureSchemaOnStartup = true + }, + BackgroundSaveOptions = new BackgroundSaveOptions + { + Enabled = true, + QueueCapacity = 20, + OnError = ex => Console.WriteLine(ex.Message) + }, + RetentionPolicy = new RetentionPolicy + { + MaxCheckpoints = 5, + MaxAge = TimeSpan.FromDays(7), + MaxEpochAge = 3, + MaxLossThreshold = 0.8f, + AlwaysKeepBest = true + } + }; + Action act = () => options.Validate(); + act.Should().NotThrow(); + } + + #endregion +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Extensions/ManagerExtensions.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Extensions/ManagerExtensions.cs new file mode 100644 index 0000000..e7e119e --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Extensions/ManagerExtensions.cs @@ -0,0 +1,237 @@ +using StateCheckpoint.NET.Models; +using System.Reflection; + +namespace StateCheckpoint.NET; + +internal static class ManagerExtensions +{ + public static async Task InternalApplyRetentionPolicyAsync(this CheckpointManager service, CancellationToken cancellationToken = default) + { + if (service.RetentionPolicy == null) return 0; + + // Get all summaries (metadata only) + var summaries = await service.Store.QueryAsync(new CheckpointQuery(), cancellationToken); + if (summaries.Count == 0) return 0; + + // Identify protected checkpoints (pinned, best) + var protectedIds = new HashSet(); + if (service.RetentionPolicy.PinnedTags != null) + { + var pinned = summaries.Where(s => IsPinned(s.Tags, service.RetentionPolicy.PinnedTags)); + foreach (var s in pinned) protectedIds.Add(s.ModelId); + } + if (service.RetentionPolicy.AlwaysKeepBest) + { + var best = summaries.MinBy(s => s.LastTrainingLoss); + if (best != null) protectedIds.Add(best.ModelId); + } + + // Build candidate list + var candidates = summaries.Where(s => !protectedIds.Contains(s.ModelId)).ToList(); + var toDelete = new HashSet(); + + // Apply rules + if (service.RetentionPolicy.MaxCheckpoints.HasValue && service.RetentionPolicy.MaxCheckpoints.Value > 0) + { + var sorted = candidates.OrderByDescending(s => s.CreatedAt).ToList(); + var keep = sorted.Take(service.RetentionPolicy.MaxCheckpoints.Value).Select(s => s.ModelId).ToHashSet(); + foreach (var s in candidates.Where(s => !keep.Contains(s.ModelId))) + toDelete.Add(s.ModelId); + } + if (service.RetentionPolicy.MaxAge.HasValue) + { + var cutoff = DateTime.UtcNow - service.RetentionPolicy.MaxAge.Value; + foreach (var s in candidates.Where(s => s.CreatedAt < cutoff)) + toDelete.Add(s.ModelId); + } + if (service.RetentionPolicy.MaxLossThreshold.HasValue) + { + foreach (var s in candidates.Where(s => s.LastTrainingLoss > service.RetentionPolicy.MaxLossThreshold.Value)) + toDelete.Add(s.ModelId); + } + if (service.RetentionPolicy.MaxEpochAge.HasValue && service.RetentionPolicy.MaxEpochAge.Value > 0) + { + var latestEpoch = summaries.Max(s => s.CurrentEpoch); + var minEpoch = latestEpoch - service.RetentionPolicy.MaxEpochAge.Value; + foreach (var s in candidates.Where(s => s.CurrentEpoch < minEpoch)) + toDelete.Add(s.ModelId); + } + + if (toDelete.Count == 0) return 0; + + await service.Store.DeleteManyAsync(toDelete, cancellationToken); + + return toDelete.Count; + } + + public static async Task InternalApplyRetentionPolicyAsync(this SessionManager service, CancellationToken cancellationToken = default) + { + if (service.RetentionPolicy == null) return 0; + + // Get all session summaries (metadata only) + var summaries = await service.Store.QueryAsync(new SessionQuery(), cancellationToken); + if (summaries.Count == 0) return 0; + + var protectedIds = new HashSet(); + + // Protect pinned sessions + if (service.RetentionPolicy.PinnedTags != null) + { + var pinned = summaries.Where(s => IsPinned(s.Tags, service.RetentionPolicy.PinnedTags)); + foreach (var s in pinned) protectedIds.Add(s.SessionId); + } + + // No "best" concept for sessions – skip AlwaysKeepBest + + // Build candidate list (unprotected) + var candidates = summaries.Where(s => !protectedIds.Contains(s.SessionId)).ToList(); + var toDelete = new HashSet(); + + // MaxCheckpoints – keep only N most recent sessions + if (service.RetentionPolicy.MaxCheckpoints.HasValue && service.RetentionPolicy.MaxCheckpoints.Value > 0) + { + var sorted = candidates.OrderByDescending(s => s.LastUpdated).ToList(); + var keep = sorted.Take(service.RetentionPolicy.MaxCheckpoints.Value).Select(s => s.SessionId).ToHashSet(); + foreach (var s in candidates.Where(s => !keep.Contains(s.SessionId))) + toDelete.Add(s.SessionId); + } + + // MaxAge – delete sessions older than the cutoff + if (service.RetentionPolicy.MaxAge.HasValue) + { + var cutoff = DateTime.UtcNow - service.RetentionPolicy.MaxAge.Value; + foreach (var s in candidates.Where(s => s.LastUpdated < cutoff)) + toDelete.Add(s.SessionId); + } + + // (Loss/Epoch rules do not apply to sessions) + if (toDelete.Count == 0) return 0; + + await service.Store.DeleteManyAsync(toDelete, cancellationToken); + + return toDelete.Count; + } + + public static async Task InternalCompareAsync(this CheckpointManager service, Guid baseLineId, Guid CandidateId, CancellationToken cancellationToken = default) + { + // Load both summaries in parallel + var summaryA = await service.Store.GetCheckpointSummaryAsync(baseLineId, cancellationToken); + var summaryB = await service.Store.GetCheckpointSummaryAsync(CandidateId, cancellationToken); + + if (summaryA == null) + throw new ArgumentException($"Checkpoint with ID '{baseLineId}' not found.", nameof(baseLineId)); + + if (summaryB == null) + throw new ArgumentException($"Checkpoint with ID '{CandidateId}' not found.", nameof(CandidateId)); + + var diff = new CheckpointDiff + { + BaselineId = baseLineId, + CandidateId = CandidateId, + EpochDelta = summaryB.CurrentEpoch - summaryA.CurrentEpoch, + LossDelta = summaryB.LastTrainingLoss - summaryA.LastTrainingLoss, + AgeDelta = summaryB.CreatedAt - summaryA.CreatedAt + }; + + // Compare hyperparameters (using reflection on HyperParameters) + diff.HyperParamChanges = CompareHyperParams(summaryA.HyperParams, summaryB.HyperParams); + + // Compare tags + diff.TagsAddedInCandidate = summaryB.Tags + .Where(kvp => !summaryA.Tags.ContainsKey(kvp.Key)) + .ToDictionary(kvp => kvp.Key, kvp => kvp.Value); + + diff.TagsRemovedInCandidate = summaryA.Tags + .Where(kvp => !summaryB.Tags.ContainsKey(kvp.Key)) + .ToDictionary(kvp => kvp.Key, kvp => kvp.Value); + + return diff; + } + + public static async Task InternalCompareAsync(this SessionManager service, Guid baselineId, Guid candidateId, CancellationToken cancellationToken = default) + { + var baseline = await service.Store.GetSessionSummaryAsync(baselineId, cancellationToken); + var candidate = await service.Store.GetSessionSummaryAsync(candidateId, cancellationToken); + + if (baseline == null) + throw new ArgumentException($"Session with ID '{baselineId}' not found.", nameof(baselineId)); + if (candidate == null) + throw new ArgumentException($"Session with ID '{candidateId}' not found.", nameof(candidateId)); + + var diff = new SessionDiff + { + BaselineId = baselineId, + CandidateId = candidateId, + LastUpdatedDelta = candidate.LastUpdated - baseline.LastUpdated, + TokenHistoryLengthDelta = candidate.TokenHistoryLength - baseline.TokenHistoryLength, + ModelFingerprintChange = baseline.ModelFingerprint != candidate.ModelFingerprint + ? $"{baseline.ModelFingerprint} → {candidate.ModelFingerprint}" + : null, + SamplingConfigChanges = CompareSamplingConfigs(baseline.SamplingConfig, candidate.SamplingConfig) + }; + + // Tags diff + diff.TagsAddedInCandidate = candidate.Tags + .Where(kvp => !baseline.Tags.ContainsKey(kvp.Key)) + .ToDictionary(kvp => kvp.Key, kvp => kvp.Value); + + diff.TagsRemovedInCandidate = baseline.Tags + .Where(kvp => !candidate.Tags.ContainsKey(kvp.Key)) + .ToDictionary(kvp => kvp.Key, kvp => kvp.Value); + + return diff; + } + + private static Dictionary CompareSamplingConfigs( + SamplingData? baseline, + SamplingData? candidate) + { + var changes = new Dictionary(); + if (baseline == null && candidate == null) return changes; + if (baseline == null) throw new ArgumentException("Baseline sampling config is null."); + if (candidate == null) throw new ArgumentException("Candidate sampling config is null."); + + var props = typeof(SamplingData).GetProperties(BindingFlags.Public | BindingFlags.Instance); + foreach (var prop in props) + { + var oldVal = prop.GetValue(baseline); + var newVal = prop.GetValue(candidate); + if (!Equals(oldVal, newVal)) + changes[prop.Name] = (oldVal, newVal); + } + return changes; + } + + private static Dictionary CompareHyperParams(HyperParameters? oldHp, HyperParameters? newHp) + { + var changes = new Dictionary(); + if (oldHp == null && newHp == null) return changes; + if (oldHp == null) throw new ArgumentException("Old hyperparameters are null."); + if (newHp == null) throw new ArgumentException("New hyperparameters are null."); + + // Get all public instance properties + var props = typeof(HyperParameters).GetProperties(BindingFlags.Public | BindingFlags.Instance); + foreach (var prop in props) + { + var oldVal = prop.GetValue(oldHp); + var newVal = prop.GetValue(newHp); + if (!Equals(oldVal, newVal)) + { + changes[prop.Name] = (oldVal, newVal); + } + } + return changes; + } + + private static bool IsPinned(Dictionary? tags, Dictionary? pinnedTags) + { + if (pinnedTags == null || tags == null) + return false; + + foreach (var pair in pinnedTags) + if (!tags.TryGetValue(pair.Key, out var val) || val != pair.Value) + return false; + + return true; + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs index 866fb26..f248d61 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/CheckpointManager.cs @@ -7,6 +7,7 @@ public class CheckpointManager : IAsyncDisposable { private readonly IModelStore _store; private readonly BackgroundSaver? _backgroundSaver; + private readonly RetentionPolicy _retentionPolicy; /// /// Initializes the manager with a custom storage provider. @@ -15,12 +16,9 @@ public class CheckpointManager : IAsyncDisposable /// public CheckpointManager(StorageOptions storageOptions) { - // Validate options - if (storageOptions.StoreType == StoreType.SqlDb && string.IsNullOrWhiteSpace(storageOptions?.DbStoreOptions?.ConnectionString)) - throw new InvalidOperationException("ConnectionString is required when StoreType is SqlDb."); + storageOptions.Validate(); - if (storageOptions.StoreType == StoreType.Local && string.IsNullOrWhiteSpace(storageOptions?.FileSystemStoreOptions?.RootPath)) - throw new InvalidOperationException("RootPath is required when StoreType is Local."); + _retentionPolicy = storageOptions.RetentionPolicy ?? new RetentionPolicy(); _store = StoreModelFactory.Create(storageOptions); @@ -66,6 +64,9 @@ public async Task SaveAsync( await _store.SaveAsync(modelCheckpoint, cancellationToken); + if (_retentionPolicy != null) + await this.InternalApplyRetentionPolicyAsync(cancellationToken); + return modelCheckpoint.ModelId; } @@ -144,6 +145,24 @@ public async Task DeleteAsync(Guid modelId, CancellationToken cancellationToken public async Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default) => await _store.ListAsync(tagKey, tagValue, cancellationToken); + /// + /// + /// + /// + /// + public async Task ApplyRetentionPolicyAsync(CancellationToken cancellationToken = default) + => await this.InternalApplyRetentionPolicyAsync(cancellationToken); + + /// + /// + /// + /// + /// + /// + /// + public async Task CompareAsync(Guid baseLineId, Guid CandidateId, CancellationToken cancellationToken = default) + => await this.InternalCompareAsync(baseLineId, CandidateId, cancellationToken); + /// /// Disposes the manager and ensures the background saver finishes all pending operations. /// Must be called if background saves are enabled. @@ -156,4 +175,16 @@ public async ValueTask DisposeAsync() if (_store is IAsyncDisposable asyncDisposable) await asyncDisposable.DisposeAsync(); } + + //----------- Internal -------------// + + internal RetentionPolicy RetentionPolicy + { + get { return _retentionPolicy; } + } + + internal IModelStore Store + { + get { return _store; } + } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs index 4d3f083..9cec313 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Manager/SessionManager.cs @@ -10,6 +10,7 @@ public class SessionManager : IAsyncDisposable { private readonly ISessionStore _store; private readonly BackgroundSaver? _backgroundSaver; + private readonly RetentionPolicy _retentionPolicy; /// /// Initializes the manager with a custom storage provider. @@ -17,11 +18,9 @@ public class SessionManager : IAsyncDisposable /// Any implementation of ISessionStore (FileSystem, PostgreSQL, etc.) public SessionManager(StorageOptions storageOptions) { - if (storageOptions.StoreType == StoreType.SqlDb && string.IsNullOrWhiteSpace(storageOptions.DbStoreOptions?.ConnectionString)) - throw new InvalidOperationException("ConnectionString is required when StoreType is SqlDb."); + storageOptions.Validate(); - if (storageOptions.StoreType == StoreType.Local && string.IsNullOrWhiteSpace(storageOptions.FileSystemStoreOptions?.RootPath)) - throw new InvalidOperationException("RootPath is required when StoreType is Local."); + _retentionPolicy = storageOptions.RetentionPolicy ?? new RetentionPolicy(); _store = StoreSessionFactory.Create(storageOptions); @@ -64,6 +63,9 @@ public async Task SaveAsync( await _store.SaveAsync(sessionCheckpoint, cancellationToken); + if (_retentionPolicy != null) + await this.InternalApplyRetentionPolicyAsync(cancellationToken); + return sessionCheckpoint.SessionId; } @@ -112,6 +114,16 @@ public Task> QueryAsync(SessionQuery query, CancellationTok public IAsyncEnumerable QueryStreamAsync(SessionQuery query, CancellationToken cancellationToken = default) => _store.QueryStreamAsync(query, cancellationToken); + /// + /// + /// + /// + /// + /// + /// + public async Task CompareAsync(Guid baseLineId, Guid CandidateId, CancellationToken cancellationToken = default) + => await this.InternalCompareAsync(baseLineId, CandidateId, cancellationToken); + /// /// Disposes the manager and ensures the background saver finishes all pending operations. /// @@ -130,4 +142,16 @@ public async ValueTask DisposeAsync() if (_store is IAsyncDisposable asyncDisposable) await asyncDisposable.DisposeAsync(); } + + //----------- Internal -------------// + + internal RetentionPolicy RetentionPolicy + { + get { return _retentionPolicy; } + } + + internal ISessionStore Store + { + get { return _store; } + } } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointDiff.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointDiff.cs new file mode 100644 index 0000000..0fdc928 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointDiff.cs @@ -0,0 +1,66 @@ +using System.Text; + +namespace StateCheckpoint.NET.Models; + +public class CheckpointDiff +{ + /// ID of the reference checkpoint (the 'left' side). + public Guid BaselineId { get; set; } + + /// ID of the compared checkpoint (the 'right' side, usually newer). + public Guid CandidateId { get; set; } + + /// Difference in epoch: Candidate.Epoch - Baseline.Epoch (positive means candidate is later). + public int EpochDelta { get; set; } + + /// Difference in loss: Candidate.Loss - Baseline.Loss (negative means candidate has lower loss). + public float LossDelta { get; set; } + + /// Time difference: Candidate.CreatedAt - Baseline.CreatedAt. + public TimeSpan AgeDelta { get; set; } + + /// Hyperparameters that changed, with old (baseline) and new (candidate) values. + public Dictionary HyperParamChanges { get; set; } = new(); + + /// Tags present in the candidate but not in the baseline. + public Dictionary TagsAddedInCandidate { get; set; } = new(); + + /// Tags present in the baseline but not in the candidate. + public Dictionary TagsRemovedInCandidate { get; set; } = new(); + + /// Returns a human-readable summary of the differences. + public string GetSummary() + { + var sb = new StringBuilder(); + sb.AppendLine($"Comparing Baseline '{BaselineId}' vs Candidate '{CandidateId}':"); + sb.AppendLine($" Epoch delta: {EpochDelta} ({(EpochDelta > 0 ? "Candidate is newer" : "Baseline is newer")})"); + sb.AppendLine($" Loss delta: {LossDelta:F4} ({(LossDelta < 0 ? "Candidate has lower loss (better)" : "Baseline has lower loss (better)")})"); + sb.AppendLine($" Age delta: {AgeDelta.TotalHours:F1} hours"); + + if (HyperParamChanges.Any()) + { + sb.AppendLine(" Hyperparameter changes:"); + foreach (var kvp in HyperParamChanges) + sb.AppendLine($" {kvp.Key}: {kvp.Value.BaselineValue} → {kvp.Value.CandidateValue}"); + } + + if (TagsAddedInCandidate.Any()) + { + sb.AppendLine(" Tags added in candidate:"); + foreach (var kvp in TagsAddedInCandidate) + sb.AppendLine($" + {kvp.Key}={kvp.Value}"); + } + + if (TagsRemovedInCandidate.Any()) + { + sb.AppendLine(" Tags removed in candidate:"); + foreach (var kvp in TagsRemovedInCandidate) + sb.AppendLine($" - {kvp.Key}={kvp.Value}"); + } + + if (!HyperParamChanges.Any() && !TagsAddedInCandidate.Any() && !TagsRemovedInCandidate.Any()) + sb.AppendLine(" No other differences found."); + + return sb.ToString(); + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs index 537146e..516622c 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/CheckpointSummary.cs @@ -6,5 +6,6 @@ public class CheckpointSummary public int CurrentEpoch { get; set; } public float LastTrainingLoss { get; set; } public DateTime CreatedAt { get; set; } + public HyperParameters? HyperParams { get; set; } public Dictionary Tags { get; set; } = new(); } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionDiff.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionDiff.cs new file mode 100644 index 0000000..1061dc5 --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionDiff.cs @@ -0,0 +1,71 @@ +using System.Text; + +namespace StateCheckpoint.NET.Models; + +public class SessionDiff +{ + /// ID of the reference (baseline) session. + public Guid BaselineId { get; set; } + + /// ID of the compared (candidate) session. + public Guid CandidateId { get; set; } + + /// Time difference: Candidate.LastUpdated - Baseline.LastUpdated. + public TimeSpan LastUpdatedDelta { get; set; } + + /// Model fingerprint difference (null if same). + public string? ModelFingerprintChange { get; set; } + + /// Difference in token history length: Candidate.TokenHistory.Length - Baseline.TokenHistory.Length. + public int TokenHistoryLengthDelta { get; set; } + + /// Sampling configuration changes. + public Dictionary SamplingConfigChanges { get; set; } = new(); + + /// Tags present in the candidate but not in the baseline. + public Dictionary TagsAddedInCandidate { get; set; } = new(); + + /// Tags present in the baseline but not in the candidate. + public Dictionary TagsRemovedInCandidate { get; set; } = new(); + + /// Returns a human-readable summary. + public string GetSummary() + { + var sb = new StringBuilder(); + + sb.AppendLine($"Comparing Baseline '{BaselineId}' vs Candidate '{CandidateId}':"); + sb.AppendLine($" Last updated delta: {LastUpdatedDelta.TotalMinutes:F1} minutes"); + + if (!string.IsNullOrEmpty(ModelFingerprintChange)) + sb.AppendLine($" Model fingerprint changed: {ModelFingerprintChange}"); + + sb.AppendLine($" Token history length delta: {TokenHistoryLengthDelta} ({(TokenHistoryLengthDelta > 0 ? "Candidate has more tokens" : "Candidate has fewer tokens")})"); + + + if (SamplingConfigChanges.Any()) + { + sb.AppendLine(" Sampling config changes:"); + foreach (var kvp in SamplingConfigChanges) + sb.AppendLine($" {kvp.Key}: {kvp.Value.BaselineValue} → {kvp.Value.CandidateValue}"); + } + + if (TagsAddedInCandidate.Any()) + { + sb.AppendLine(" Tags added in candidate:"); + foreach (var kvp in TagsAddedInCandidate) + sb.AppendLine($" + {kvp.Key}={kvp.Value}"); + } + + if (TagsRemovedInCandidate.Any()) + { + sb.AppendLine(" Tags removed in candidate:"); + foreach (var kvp in TagsRemovedInCandidate) + sb.AppendLine($" - {kvp.Key}={kvp.Value}"); + } + + if (!SamplingConfigChanges.Any() && !TagsAddedInCandidate.Any() && !TagsRemovedInCandidate.Any()) + sb.AppendLine(" No other differences found."); + + return sb.ToString(); + } +} diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs index cdbacc2..558ac06 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Models/SessionSummary.cs @@ -5,6 +5,8 @@ public class SessionSummary public Guid SessionId { get; set; } public string ModelFingerprint { get; set; } = string.Empty; public DateTime LastUpdated { get; set; } + public int TokenHistoryLength { get; set; } + public SamplingData? SamplingConfig { get; set; } public Dictionary Tags { get; set; } = new(); } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs index f4aa5f5..c60c437 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemBase.cs @@ -5,13 +5,13 @@ namespace StateCheckpoint.NET; internal abstract class FileSystemBase { - private readonly string _rootPath; - private readonly FileSystemStoreOptions _options; - private readonly TIndex _index; - private bool _indexLoaded; - private readonly SemaphoreSlim _buildLock = new(1, 1); + protected readonly string _rootPath; + protected readonly FileSystemStoreOptions _options; + protected readonly TIndex _index; + protected bool _indexLoaded; + protected readonly SemaphoreSlim _buildLock = new(1, 1); - public FileSystemBase(string rootPath, FileSystemStoreOptions options, Func indexSelector) + public FileSystemBase(string rootPath, FileSystemStoreOptions? options, Func indexSelector) { _options = options ?? new FileSystemStoreOptions(); _rootPath = Path.Combine(rootPath, "models"); @@ -43,7 +43,7 @@ public FileSystemBase(string rootPath, FileSystemStoreOptions options, Func(Action getAll, Func>> getAllEntries, Func rebuild, CancellationToken cancellationToken) + protected async Task EnsureIndexLoadedAsync(Func load, Func>> getAllEntries, Func rebuild, CancellationToken cancellationToken) { if (!_indexLoaded) { @@ -53,28 +53,13 @@ private async Task EnsureIndexLoadedAsync(Action getAll, Func - { - var manifest = await FileSystemHelper - .LoadManifestOnlyAsync(rootPath, id, "meta.json", ct); - - if (manifest == null) return null; - - return new SessionSummary - { - SessionId = id, - ModelFingerprint = manifest.ModelFingerprint, - LastUpdated = manifest.LastUpdated, - Tags = manifest.Tags ?? new Dictionary() - }; - }, cancellationToken); - - rebuild(_rootPath); + await rebuild(_rootPath); } _indexLoaded = true; } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs index c056eaa..7f1aef5 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemModelStore.cs @@ -5,44 +5,11 @@ namespace StateCheckpoint.NET; -internal class FileSystemModelStore : IFileSystemModelStore +internal class FileSystemModelStore : FileSystemBase, IFileSystemModelStore { - private readonly string _rootPath; - private readonly FileSystemStoreOptions _options; - private readonly CheckpointIndex _index; - private readonly SemaphoreSlim _buildLock = new(1, 1); - private bool _indexLoaded; - - public FileSystemModelStore(string rootPath, FileSystemStoreOptions? options = null) + public FileSystemModelStore(string rootPath, FileSystemStoreOptions? options = null): + base(rootPath, options, indexFilePath => new CheckpointIndex(indexFilePath)) { - _options = options ?? new FileSystemStoreOptions(); - _rootPath = Path.Combine(rootPath, "models"); - - if (_options.ValidatePermissionsOnStartup) - { - if (!FileSystemHelper.TryValidateWriteAccess(_rootPath, out var error)) - { - // If fallback is provided, update the root path - if (!string.IsNullOrWhiteSpace(_options.FallbackPath)) - { - _rootPath = Path.Combine(_options.FallbackPath, "models"); - - Directory.CreateDirectory(_rootPath); - } - else - { - throw error!; - } - } - } - - // Still ensure the directory exists if required - else if (_options.EnsureDirectoryExists) - Directory.CreateDirectory(_rootPath); - - var indexFilePath = Path.Combine(_rootPath, "_index.json"); - - _index = new CheckpointIndex(indexFilePath); } /// @@ -83,6 +50,7 @@ await _index.AddOrUpdateAsync(new CheckpointSummary CurrentEpoch = checkpoint.CurrentEpoch, LastTrainingLoss = checkpoint.LastTrainingLoss, ModelId = checkpoint.ModelId, + HyperParams = checkpoint.HyperParams, Tags = checkpoint.Tags }, cancellationToken); } @@ -181,46 +149,85 @@ public async IAsyncEnumerable QueryStreamAsync( } } - private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) + public async Task GetCheckpointSummaryAsync(Guid modelId, CancellationToken cancellationToken = default) { - if (!_indexLoaded) + var manifest = await FileSystemHelper.LoadManifestOnlyAsync( + _rootPath, modelId, "manifest.json", cancellationToken); + + if (manifest == null) return null; + + return new CheckpointSummary { - await _buildLock.WaitAsync(cancellationToken); - try - { - if (_indexLoaded) return; // double-check - - // Load the index from disk - await _index.LoadAsync(cancellationToken); - - var entries = await _index.GetAllAsync(cancellationToken); - - // If index is empty, rebuild it from manifests - if (!entries.Any()) - { - await _index.RebuildAsync(_rootPath, async (string rootPath, Guid id, CancellationToken ct) => - { - var manifest = await FileSystemHelper.LoadManifestOnlyAsync(rootPath, id, "manifest.json", ct); - - if (manifest == null) return null; - - return new CheckpointSummary - { - ModelId = id, - CurrentEpoch = manifest.CurrentEpoch, - LastTrainingLoss = manifest.LastTrainingLoss, - CreatedAt = manifest.CreatedAt, - Tags = manifest.Tags ?? new() - }; - }, cancellationToken); - } - - _indexLoaded = true; - } - finally - { - _buildLock.Release(); - } - } + ModelId = modelId, + CurrentEpoch = manifest.CurrentEpoch, + LastTrainingLoss = manifest.LastTrainingLoss, + CreatedAt = manifest.CreatedAt, + HyperParams = manifest.HyperParams, + Tags = manifest.Tags ?? new Dictionary() + }; } + + private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) => await EnsureIndexLoadedAsync( + () => _index.LoadAsync(cancellationToken), + () => _index.GetAllAsync(cancellationToken), + async path => await _index.RebuildAsync(_rootPath, async (string rootPath, Guid id, CancellationToken cancellationToken) => + { + var manifest = await FileSystemHelper.LoadManifestOnlyAsync(rootPath, id, "manifest.json", cancellationToken); + + if (manifest == null) return null; + + return new CheckpointSummary + { + ModelId = id, + CurrentEpoch = manifest.CurrentEpoch, + LastTrainingLoss = manifest.LastTrainingLoss, + CreatedAt = manifest.CreatedAt, + Tags = manifest.Tags ?? new() + }; + }, cancellationToken), + cancellationToken + ); + + //private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) + //{ + // if (!_indexLoaded) + // { + // await _buildLock.WaitAsync(cancellationToken); + // try + // { + // if (_indexLoaded) return; // double-check + + // // Load the index from disk + // await _index.LoadAsync(cancellationToken); + + // var entries = await _index.GetAllAsync(cancellationToken); + + // // If index is empty, rebuild it from manifests + // if (!entries.Any()) + // { + // await _index.RebuildAsync(_rootPath, async (string rootPath, Guid id, CancellationToken cancellationToken) => + // { + // var manifest = await FileSystemHelper.LoadManifestOnlyAsync(rootPath, id, "manifest.json", cancellationToken); + + // if (manifest == null) return null; + + // return new CheckpointSummary + // { + // ModelId = id, + // CurrentEpoch = manifest.CurrentEpoch, + // LastTrainingLoss = manifest.LastTrainingLoss, + // CreatedAt = manifest.CreatedAt, + // Tags = manifest.Tags ?? new() + // }; + // }, cancellationToken); + // } + + // _indexLoaded = true; + // } + // finally + // { + // _buildLock.Release(); + // } + // } + //} } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs index a40de77..68a184c 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystem/FileSystemSessionStore.cs @@ -4,49 +4,11 @@ namespace StateCheckpoint.NET; -internal class FileSystemSessionStore : IFileSystemSessionStore +internal class FileSystemSessionStore : FileSystemBase, IFileSystemSessionStore { - private readonly string _rootPath; - private readonly FileSystemStoreOptions _options; - private readonly SessionIndex _index; - private bool _indexLoaded; - private readonly SemaphoreSlim _buildLock = new(1, 1); - - public FileSystemSessionStore(string rootPath, FileSystemStoreOptions? options = null) + public FileSystemSessionStore(string rootPath, FileSystemStoreOptions? options = null) + : base(rootPath, options, indexFilePath => new SessionIndex(indexFilePath)) { - _options = options ?? new FileSystemStoreOptions(); - _rootPath = Path.Combine(rootPath, "sessions"); - - if (_options.ValidatePermissionsOnStartup) - { - if (!FileSystemHelper.TryValidateWriteAccess(_rootPath, out var error)) - { - // If fallback is provided, update the root path - if (!string.IsNullOrWhiteSpace(_options.FallbackPath)) - { - _rootPath = Path.Combine(_options.FallbackPath, "sessions"); - - Directory.CreateDirectory(_rootPath); - } - else - { - throw error!; - } - } - } - else - { - // Still ensure the directory exists if required - if (_options.EnsureDirectoryExists) - { - Directory.CreateDirectory(_rootPath); - } - } - - // Index file path - var indexFilePath = Path.Combine(_rootPath, "_sessions_index.json"); - - _index = new SessionIndex(indexFilePath); } public async Task SaveAsync(SessionCheckpoint session, CancellationToken cancellationToken = default) @@ -114,6 +76,12 @@ public async Task DeleteAsync(Guid sessionId, CancellationToken cancellationToke await _index.RemoveAsync(sessionId, cancellationToken); } + public async Task DeleteManyAsync(IEnumerable sessionIds, CancellationToken cancellationToken = default) + { + foreach (var id in sessionIds) + await DeleteAsync(id, cancellationToken); + } + public async Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default) { if (string.IsNullOrWhiteSpace(tagKey) || string.IsNullOrWhiteSpace(tagValue)) @@ -146,7 +114,7 @@ public async IAsyncEnumerable QueryStreamAsync( cancellationToken.ThrowIfCancellationRequested(); // Apply filters - if (!string.IsNullOrEmpty(query.ModelFingerprint) && summary.ModelFingerprint != query.ModelFingerprint) continue; + if (!string.IsNullOrWhiteSpace(query.ModelFingerprint) && summary.ModelFingerprint != query.ModelFingerprint) continue; if (query.UpdatedAfter.HasValue && summary.LastUpdated < query.UpdatedAfter.Value) continue; if (query.UpdatedBefore.HasValue && summary.LastUpdated > query.UpdatedBefore.Value) continue; if (query.Tags != null && query.Tags.Count > 0 && !TagsHelper.TagsMatch(summary.Tags, query.Tags)) continue; @@ -155,45 +123,39 @@ public async IAsyncEnumerable QueryStreamAsync( } } - //---------------- Private Methods ----------------// - - // Ensure the index is loaded into memory (lazy) - private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) + public async Task GetSessionSummaryAsync(Guid sessionId, CancellationToken cancellationToken = default) { - if (!_indexLoaded) + var manifest = await FileSystemHelper.LoadManifestOnlyAsync( + _rootPath, sessionId, "meta.json", cancellationToken); + + if (manifest == null) return null; + + return new SessionSummary { - await _buildLock.WaitAsync(cancellationToken); - try - { - if (_indexLoaded) return; - - await _index.LoadAsync(cancellationToken); - - var entries = await _index.GetAllAsync(cancellationToken); - if (!entries.Any()) - { - await _index.RebuildAsync(_rootPath, async (rootPath, id, ct) => - { - var manifest = await FileSystemHelper - .LoadManifestOnlyAsync(rootPath, id, "meta.json", ct); - - if (manifest == null) return null; - - return new SessionSummary - { - SessionId = id, - ModelFingerprint = manifest.ModelFingerprint, - LastUpdated = manifest.LastUpdated, - Tags = manifest.Tags ?? new Dictionary() - }; - }, cancellationToken); - } - _indexLoaded = true; - } - finally - { - _buildLock.Release(); - } - } + SessionId = sessionId, + ModelFingerprint = manifest.ModelFingerprint, + LastUpdated = manifest.LastUpdated, + Tags = manifest.Tags ?? new Dictionary() + }; } + + private async Task EnsureIndexLoadedAsync(CancellationToken cancellationToken) => await EnsureIndexLoadedAsync( + () => _index.LoadAsync(cancellationToken), + () => _index.GetAllAsync(cancellationToken), + async path => await _index.RebuildAsync(_rootPath, async (string rootPath, Guid id, CancellationToken _) => + { + var manifest = await FileSystemHelper.LoadManifestOnlyAsync(rootPath, id, "meta.json", cancellationToken); + + if (manifest == null) return null; + + return new SessionSummary + { + SessionId = id, + ModelFingerprint = manifest.ModelFingerprint, + LastUpdated = manifest.LastUpdated, + Tags = manifest.Tags ?? new Dictionary() + }; + }, cancellationToken), + cancellationToken + ); } \ No newline at end of file diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs index 9acc1bc..43f7277 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/CheckpointIndex.cs @@ -23,12 +23,12 @@ public CheckpointIndex(string indexFilePath) /// Loads the index from storage; if no file exists, starts with an empty list. /// Must be called before any other operations. /// - public async Task LoadAsync(CancellationToken ct = default) + public async Task LoadAsync(CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { - var data = await _storage.ReadAsync(ct); + var data = await _storage.ReadAsync(cancellationToken); _items = data.ToList(); } finally @@ -37,9 +37,9 @@ public async Task LoadAsync(CancellationToken ct = default) } } - public async Task AddOrUpdateAsync(CheckpointSummary summary, CancellationToken ct = default) + public async Task AddOrUpdateAsync(CheckpointSummary summary, CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { var existing = _items.FirstOrDefault(e => e.ModelId == summary.ModelId); @@ -54,7 +54,7 @@ public async Task AddOrUpdateAsync(CheckpointSummary summary, CancellationToken { _items.Add(summary); } - await _storage.WriteAsync(_items, ct); + await _storage.WriteAsync(_items, cancellationToken); } finally { @@ -62,14 +62,14 @@ public async Task AddOrUpdateAsync(CheckpointSummary summary, CancellationToken } } - public async Task RemoveAsync(Guid modelId, CancellationToken ct = default) + public async Task RemoveAsync(Guid modelId, CancellationToken cancellationToken = default) { - await _lock.WaitAsync(ct); + await _lock.WaitAsync(cancellationToken); try { var removed = _items.RemoveAll(e => e.ModelId == modelId); if (removed > 0) - await _storage.WriteAsync(_items, ct); + await _storage.WriteAsync(_items, cancellationToken); } finally { diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs index 3eb9b85..8f9224f 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/FileSystemIndex/IndexStorage.cs @@ -13,22 +13,22 @@ public IndexStorage(string filePath, JsonSerializerOptions? jsonOpts = null) _jsonOpts = jsonOpts ?? new JsonSerializerOptions { WriteIndented = true }; } - public async Task> ReadAsync(CancellationToken ct = default) + public async Task> ReadAsync(CancellationToken cancellationToken = default) { if (!File.Exists(_filePath)) return Array.Empty(); - var json = await File.ReadAllTextAsync(_filePath, ct); + var json = await File.ReadAllTextAsync(_filePath, cancellationToken); return JsonSerializer.Deserialize>(json, _jsonOpts) ?? new List(); } - public async Task WriteAsync(IEnumerable data, CancellationToken ct = default) + public async Task WriteAsync(IEnumerable data, CancellationToken cancellationToken = default) { var json = JsonSerializer.Serialize(data, _jsonOpts); var tempPath = _filePath + ".tmp"; - await File.WriteAllTextAsync(tempPath, json, ct); + await File.WriteAllTextAsync(tempPath, json, cancellationToken); File.Move(tempPath, _filePath, overwrite: true); } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs index 91a4c57..75e18dd 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/IModelStore.cs @@ -7,8 +7,9 @@ internal interface IModelStore Task SaveAsync(ModelCheckpoint checkpoint, CancellationToken cancellationToken = default); Task LoadAsync(Guid modelId, CancellationToken cancellationToken = default); Task DeleteAsync(Guid modelId, CancellationToken cancellationToken = default); + Task DeleteManyAsync(IEnumerable modelIds, CancellationToken cancellationToken = default); Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default); Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default); IAsyncEnumerable QueryStreamAsync(CheckpointQuery query, CancellationToken cancellationToken = default); - Task DeleteManyAsync(IEnumerable modelIds, CancellationToken cancellationToken = default); + Task GetCheckpointSummaryAsync(Guid modelId, CancellationToken cancellationToken = default); } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs index 0487c30..bd9b83a 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/ISessionStore.cs @@ -7,7 +7,9 @@ internal interface ISessionStore Task SaveAsync(SessionCheckpoint session, CancellationToken cancellationToken = default); Task LoadAsync(Guid sessionId, CancellationToken cancellationToken = default); Task DeleteAsync(Guid sessionId, CancellationToken cancellationToken = default); + Task DeleteManyAsync(IEnumerable modelIds, CancellationToken cancellationToken = default); Task> ListAsync(string? tagKey = null, string? tagValue = null, CancellationToken cancellationToken = default); Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default); IAsyncEnumerable QueryStreamAsync(SessionQuery query, CancellationToken cancellationToken = default); + Task GetSessionSummaryAsync(Guid sessionId, CancellationToken cancellationToken = default); } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs index 858caf6..3a3c7a0 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresModelStore.cs @@ -255,18 +255,19 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue return modelIds; } - public async Task> QueryAsync(CheckpointQuery query, CancellationToken ct = default) + public async Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default) { - await using var connection = await GetConnectionAsync(ct); + await using var connection = await GetConnectionAsync(cancellationToken); + var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); await using var command = new NpgsqlCommand(sql, connection); command.Parameters.AddRange(parameters.ToArray()); - await using var reader = await command.ExecuteReaderAsync(ct); + await using var reader = await command.ExecuteReaderAsync(cancellationToken); var results = new List(); - while (await reader.ReadAsync(ct)) + while (await reader.ReadAsync(cancellationToken)) { results.Add(new CheckpointSummary { @@ -283,18 +284,19 @@ public async Task> QueryAsync(CheckpointQuery query, Can public async IAsyncEnumerable QueryStreamAsync( CheckpointQuery query, - [EnumeratorCancellation] CancellationToken ct = default) + [EnumeratorCancellation] CancellationToken cancellationToken = default) { - await using var connection = await GetConnectionAsync(ct); + await using var connection = await GetConnectionAsync(cancellationToken); var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); await using var command = new NpgsqlCommand(sql, connection); command.Parameters.AddRange(parameters.ToArray()); - await using var reader = await command.ExecuteReaderAsync(ct); - while (await reader.ReadAsync(ct)) + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + while (await reader.ReadAsync(cancellationToken)) { - ct.ThrowIfCancellationRequested(); + cancellationToken.ThrowIfCancellationRequested(); + yield return new CheckpointSummary { ModelId = reader.GetGuid(0), @@ -317,6 +319,33 @@ public async Task DeleteManyAsync(IEnumerable modelIds, CancellationToken await transaction.CommitAsync(cancellationToken); } + public async Task GetCheckpointSummaryAsync(Guid modelId, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + const string sql = @" + SELECT m.epoch, m.loss, m.created_at, m.tags + FROM model_manifests m + JOIN model_blobs b ON m.model_id = b.model_id + WHERE m.model_id = @id"; + + await using var command = new NpgsqlCommand(sql, connection); + command.Parameters.AddWithValue("@id", modelId); + + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + + if (!await reader.ReadAsync(cancellationToken)) return null; + + return new CheckpointSummary + { + ModelId = modelId, + CurrentEpoch = reader.GetInt32(0), + LastTrainingLoss = (float)reader.GetDouble(1), + CreatedAt = reader.GetDateTime(2), + HyperParams = JsonSerializer.Deserialize(reader.GetString(3)), + Tags = JsonSerializer.Deserialize>(reader.GetString(3)) ?? new() + }; + } + //--------- private methods ----------------------------------- // --- Large Object Helpers (Use PostgresLargeObjectQueries) --- @@ -383,10 +412,10 @@ private static async Task UnlinkLargeObjectAsync( } private static async Task ReadLargeObjectAsync( - NpgsqlConnection connection, - uint objectId, - NpgsqlTransaction transaction, // required – caller manages it - CancellationToken cancellationToken = default) + NpgsqlConnection connection, + uint objectId, + NpgsqlTransaction transaction, // required – caller manages it + CancellationToken cancellationToken = default) { // All commands use the provided transaction – do NOT commit or dispose it here. await using var openCommand = new NpgsqlCommand(PostgresLargeObjectQueries.OpenRead, connection, transaction); diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs index 8f4bf1c..ccf4952 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresSessionStore.cs @@ -96,6 +96,17 @@ public async Task DeleteAsync(Guid sessionId, CancellationToken cancellationToke await command.ExecuteNonQueryAsync(cancellationToken); } + public async Task DeleteManyAsync(IEnumerable sessionIds, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + await using var transaction = await connection.BeginTransactionAsync(cancellationToken); + + foreach (var id in sessionIds) + await DeleteAsync(id, cancellationToken); + + await transaction.CommitAsync(cancellationToken); + } + /// /// Gets a list of all session ids /// @@ -133,17 +144,17 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue return sessionIds; } - public async Task> QueryAsync(SessionQuery query, CancellationToken ct = default) + public async Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default) { - await using var connection = await GetConnectionAsync(ct); + await using var connection = await GetConnectionAsync(cancellationToken); var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); await using var command = new NpgsqlCommand(sql, connection); command.Parameters.AddRange(parameters.ToArray()); - await using var reader = await command.ExecuteReaderAsync(ct); + await using var reader = await command.ExecuteReaderAsync(cancellationToken); var results = new List(); - while (await reader.ReadAsync(ct)) + while (await reader.ReadAsync(cancellationToken)) { results.Add(new SessionSummary { @@ -181,15 +192,40 @@ public async IAsyncEnumerable QueryStreamAsync( } } + public async Task GetSessionSummaryAsync(Guid sessionId, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + + const string sql = @" + SELECT model_fingerprint, last_updated, tags + FROM inference_sessions + WHERE session_id = @id"; + await using var command = new NpgsqlCommand(sql, connection); + + command.Parameters.AddWithValue("@id", sessionId); + + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + + if (!await reader.ReadAsync(cancellationToken)) return null; + + return new SessionSummary + { + SessionId = sessionId, + ModelFingerprint = reader.GetString(0), + LastUpdated = reader.GetDateTime(1), + Tags = JsonSerializer.Deserialize>(reader.GetString(2)) ?? new() + }; + } + //---------------- Private Methods ---------------// private (string Sql, List Parameters) BuildQuerySql(SessionQuery query, bool includeOrderBy = true, bool includeLimit = true) { var sql = new StringBuilder(@" - SELECT session_id, model_fingerprint, last_updated, tags - FROM inference_sessions - WHERE 1=1 - "); + SELECT session_id, model_fingerprint, last_updated, tags + FROM inference_sessions + WHERE 1=1 + "); var parameters = new List(); int paramIndex = 0; diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs index 4364baa..ebd085d 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/Postgres/PostgresStoreBase.cs @@ -34,9 +34,9 @@ protected PostgresStoreBase(NpgsqlDataSource dataSource) /// /// Gets an open PostgreSQL connection from the DataSource pool. /// - protected async Task GetConnectionAsync(CancellationToken ct = default) + protected async Task GetConnectionAsync(CancellationToken cancellationToken = default) { - return await _dataSource.OpenConnectionAsync(ct); + return await _dataSource.OpenConnectionAsync(cancellationToken); } /// diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs index 08a6820..35efcf7 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerModelStore.cs @@ -158,25 +158,31 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue } } - public async Task> QueryAsync(CheckpointQuery query, CancellationToken ct = default) + public async Task> QueryAsync(CheckpointQuery query, CancellationToken cancellationToken = default) { - await using var connection = await GetConnectionAsync(ct); + await using var connection = await GetConnectionAsync(cancellationToken); var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); await using var command = new SqlCommand(sql, connection); + command.Parameters.AddRange(parameters.ToArray()); - await using var reader = await command.ExecuteReaderAsync(ct); + await using var reader = await command.ExecuteReaderAsync(cancellationToken); var results = new List(); - while (await reader.ReadAsync(ct)) + + while (await reader.ReadAsync(cancellationToken)) { + var hyperParamsJson = reader.GetString(5); + var hyperParams = JsonSerializer.Deserialize(hyperParamsJson); + results.Add(new CheckpointSummary { ModelId = reader.GetGuid(0), CurrentEpoch = reader.GetInt32(1), LastTrainingLoss = (float)reader.GetDouble(2), CreatedAt = reader.GetDateTime(3), - Tags = JsonSerializer.Deserialize>(reader.GetString(4)) ?? new() + Tags = JsonSerializer.Deserialize>(reader.GetString(4)) ?? new(), + HyperParams = hyperParams }); } @@ -219,12 +225,36 @@ public async Task DeleteManyAsync(IEnumerable modelIds, CancellationToken await tx.CommitAsync(cancellationToken); } + public async Task GetCheckpointSummaryAsync(Guid modelId, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + const string sql = @" + SELECT epoch, loss, created_at, hyper_params, tags + FROM ModelManifests + WHERE model_id = @id"; + await using var command = new SqlCommand(sql, connection); + command.Parameters.AddWithValue("@id", modelId); + + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + if (!await reader.ReadAsync(cancellationToken)) return null; + + return new CheckpointSummary + { + ModelId = modelId, + CurrentEpoch = reader.GetInt32(0), + LastTrainingLoss = (float)reader.GetDouble(1), + CreatedAt = reader.GetDateTime(2), + HyperParams = JsonSerializer.Deserialize(reader.GetString(3)), + Tags = JsonSerializer.Deserialize>(reader.GetString(4)) ?? new() + }; + } + //------------- Private Methods -----------------// private (string Sql, List Parameters) BuildQuerySql(CheckpointQuery query, bool includeOrderBy = true, bool includeLimit = true) { var sql = new StringBuilder(@" - SELECT modelId, epoch, loss, createdAt, tags + SELECT modelId, epoch, loss, createdAt, tags, hyperParams FROM ModelManifests WHERE 1=1 "); diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs index 2341075..af088b5 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerSessionStore.cs @@ -103,6 +103,17 @@ public async Task DeleteAsync(Guid sessionId, CancellationToken cancellationToke await command.ExecuteNonQueryAsync(cancellationToken); } + public async Task DeleteManyAsync(IEnumerable sesionIds, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + await using var tx = await connection.BeginTransactionAsync(cancellationToken); + + foreach (var id in sesionIds) + await DeleteAsync(id, cancellationToken); + + await tx.CommitAsync(cancellationToken); + } + /// /// Gets a list of all session ids /// @@ -143,17 +154,17 @@ public async Task> ListAsync(string? tagKey = null, string? tagValue } } - public async Task> QueryAsync(SessionQuery query, CancellationToken ct = default) + public async Task> QueryAsync(SessionQuery query, CancellationToken cancellationToken = default) { - await using var connection = await GetConnectionAsync(ct); + await using var connection = await GetConnectionAsync(cancellationToken); var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); await using var command = new SqlCommand(sql, connection); command.Parameters.AddRange(parameters.ToArray()); - await using var reader = await command.ExecuteReaderAsync(ct); + await using var reader = await command.ExecuteReaderAsync(cancellationToken); var results = new List(); - while (await reader.ReadAsync(ct)) + while (await reader.ReadAsync(cancellationToken)) { results.Add(new SessionSummary { @@ -168,18 +179,19 @@ public async Task> QueryAsync(SessionQuery query, Cancellat public async IAsyncEnumerable QueryStreamAsync( SessionQuery query, - [EnumeratorCancellation] CancellationToken ct = default) + [EnumeratorCancellation] CancellationToken cancellationToken = default) { - await using var connection = await GetConnectionAsync(ct); + await using var connection = await GetConnectionAsync(cancellationToken); var (sql, parameters) = BuildQuerySql(query, includeOrderBy: true, includeLimit: true); await using var command = new SqlCommand(sql, connection); command.Parameters.AddRange(parameters.ToArray()); - await using var reader = await command.ExecuteReaderAsync(ct); - while (await reader.ReadAsync(ct)) + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + while (await reader.ReadAsync(cancellationToken)) { - ct.ThrowIfCancellationRequested(); + cancellationToken.ThrowIfCancellationRequested(); + yield return new SessionSummary { SessionId = reader.GetGuid(0), @@ -190,6 +202,30 @@ public async IAsyncEnumerable QueryStreamAsync( } } + public async Task GetSessionSummaryAsync(Guid sessionId, CancellationToken cancellationToken = default) + { + await using var connection = await GetConnectionAsync(cancellationToken); + const string sql = @" + SELECT model_fingerprint, last_updated, tags + FROM InferenceSessions + WHERE session_id = @id"; + + await using var command = new SqlCommand(sql, connection); + + command.Parameters.AddWithValue("@id", sessionId); + + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + if (!await reader.ReadAsync(cancellationToken)) return null; + + return new SessionSummary + { + SessionId = sessionId, + ModelFingerprint = reader.GetString(0), + LastUpdated = reader.GetDateTime(1), + Tags = JsonSerializer.Deserialize>(reader.GetString(2)) ?? new() + }; + } + //------------ Private Methods --------------// private (string Sql, List Parameters) BuildQuerySql(SessionQuery query, bool includeOrderBy = true, bool includeLimit = true) diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs index 9b56a08..9dd621c 100644 --- a/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Stores/SqlServer/SqlServerStoreBase.cs @@ -37,17 +37,17 @@ protected SqlServerStoreBase(SqlConnection connection) /// /// Gets an open SQL Server connection. Opens it lazily if not already open. /// - protected async Task GetConnectionAsync(CancellationToken ct = default) + protected async Task GetConnectionAsync(CancellationToken cancellationToken = default) { if (!_ownsConnection) { // External connection – ensure thread‑safe access - await _connectionLock.WaitAsync(ct); + await _connectionLock.WaitAsync(cancellationToken); try { if (_connection!.State != ConnectionState.Open) - await _connection.OpenAsync(ct); + await _connection.OpenAsync(cancellationToken); return _connection; } @@ -60,7 +60,7 @@ protected async Task GetConnectionAsync(CancellationToken ct = de // Owned connection – create a new one per call var connection = new SqlConnection(_connectionString!); - await connection.OpenAsync(ct); + await connection.OpenAsync(cancellationToken); return connection; } diff --git a/StateCheckpoint.NET/StateCheckpoint.NET/Validation/StoreOptionsValidationExtension.cs b/StateCheckpoint.NET/StateCheckpoint.NET/Validation/StoreOptionsValidationExtension.cs new file mode 100644 index 0000000..55ec81f --- /dev/null +++ b/StateCheckpoint.NET/StateCheckpoint.NET/Validation/StoreOptionsValidationExtension.cs @@ -0,0 +1,61 @@ +using StateCheckpoint.NET.Models; +using StateCheckpoint.NET.Settings; + +namespace StateCheckpoint.NET; + +internal static class StoreOptionsValidationExtension +{ + public static void Validate(this StorageOptions value) + { + // Store type specific validation + if (value?.StoreType == StoreType.Local) + { + if (value?.FileSystemStoreOptions == null) + throw new InvalidOperationException("FileSystemStoreOptions is required when StoreType is Local."); + + if (string.IsNullOrWhiteSpace(value?.FileSystemStoreOptions?.RootPath)) + throw new InvalidOperationException("RootPath is required when StoreType is Local."); + } + else if (value?.StoreType == StoreType.SqlDb) + { + if (value?.DbStoreOptions == null) + throw new InvalidOperationException("DbStoreOptions is required when StoreType is SqlDb."); + + if (string.IsNullOrWhiteSpace(value?.DbStoreOptions?.ConnectionString)) + throw new InvalidOperationException("ConnectionString is required when StoreType is SqlDb."); + + // SqlDbType enum validation (already enforced by type, but we can check) + if (!Enum.IsDefined(typeof(SqlDbType), value.SqlDbType)) + throw new InvalidOperationException($"Invalid SqlDbType '{value?.SqlDbType.ToString()}'. Valid values are Postgres or SqlServer."); + } + else + { + throw new InvalidOperationException($"Unsupported StoreType '{value?.StoreType.ToString()}'. Valid values are Local or SqlDb."); + } + + // Validate BackgroundSaveOptions if present + if (value?.BackgroundSaveOptions != null && value?.BackgroundSaveOptions?.Enabled == true) + { + if (value?.BackgroundSaveOptions?.QueueCapacity <= 0) + throw new InvalidOperationException("BackgroundSaveOptions.QueueCapacity must be greater than 0 when background saves are enabled."); + } + + // Validate RetentionPolicy if present + if (value?.RetentionPolicy != null) + { + if (value?.RetentionPolicy?.MaxCheckpoints != null && value?.RetentionPolicy?.MaxCheckpoints.Value <= 0) + throw new InvalidOperationException("RetentionPolicy.MaxCheckpoints must be greater than 0."); + + if (value?.RetentionPolicy?.MaxAge != null && value?.RetentionPolicy?.MaxAge <= TimeSpan.Zero) + throw new InvalidOperationException("RetentionPolicy.MaxAge must be a positive TimeSpan."); + + if (value?.RetentionPolicy?.MaxEpochAge != null && value?.RetentionPolicy?.MaxEpochAge <= 0) + throw new InvalidOperationException("RetentionPolicy.MaxEpochAge must be greater than 0."); + + if (value?.RetentionPolicy?.MaxLossThreshold != null && value?.RetentionPolicy?.MaxLossThreshold < 0) + throw new InvalidOperationException("RetentionPolicy.MaxLossThreshold must be a positive number."); + + // If PinnedTags is provided, it's a dictionary; no further validation needed. + } + } +}