File: Storage\DiskBasedResultStore.cs
Web Access
Project: src\src\Libraries\Microsoft.Extensions.AI.Evaluation.Reporting\CSharp\Microsoft.Extensions.AI.Evaluation.Reporting.csproj (Microsoft.Extensions.AI.Evaluation.Reporting)
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI.Evaluation.Reporting.JsonSerialization;
using Microsoft.Extensions.AI.Evaluation.Reporting.Utilities;
using Microsoft.Shared.Diagnostics;
 
namespace Microsoft.Extensions.AI.Evaluation.Reporting.Storage;
 
/// <summary>
/// An <see cref="IResultStore"/> implementation that stores <see cref="ScenarioRunResult"/>s on disk.
/// </summary>
public sealed class DiskBasedResultStore : IResultStore
{
    private const string DeserializationFailedMessage = "Unable to deserialize the scenario run result file at {0}.";
 
#if NET
    private static EnumerationOptions InTopDirectoryOnly { get; } =
        new EnumerationOptions
        {
            IgnoreInaccessible = true,
            MatchType = MatchType.Simple,
            RecurseSubdirectories = false,
            ReturnSpecialDirectories = false,
        };
#else
    private const SearchOption InTopDirectoryOnly = SearchOption.TopDirectoryOnly;
#endif
 
    private readonly string _resultsRootPath;
 
    /// <summary>
    /// Initializes a new instance of the <see cref="DiskBasedResultStore"/> class.
    /// </summary>
    /// <param name="storageRootPath">
    /// The path to a directory on disk under which the <see cref="ScenarioRunResult"/>s should be stored.
    /// </param>
    public DiskBasedResultStore(string storageRootPath)
    {
        storageRootPath = Path.GetFullPath(storageRootPath);
        _resultsRootPath = Path.Combine(storageRootPath, "results");
    }
 
    /// <inheritdoc/>
    public async IAsyncEnumerable<ScenarioRunResult> ReadResultsAsync(
        string? executionName = null,
        string? scenarioName = null,
        string? iterationName = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        IEnumerable<FileInfo> resultFiles =
            EnumerateResultFiles(executionName, scenarioName, iterationName, cancellationToken);
 
        foreach (FileInfo resultFile in resultFiles)
        {
            cancellationToken.ThrowIfCancellationRequested();
 
            using FileStream stream = resultFile.OpenRead();
 
            ScenarioRunResult? result =
                await JsonSerializer.DeserializeAsync<ScenarioRunResult>(
                    stream,
                    SerializerContext.Default.ScenarioRunResult,
                    cancellationToken).ConfigureAwait(false);
 
            yield return result is null
                ? throw new JsonException(
                    string.Format(CultureInfo.CurrentCulture, DeserializationFailedMessage, resultFile.FullName))
                : result;
        }
    }
 
    /// <inheritdoc/>
    public async ValueTask WriteResultsAsync(
        IEnumerable<ScenarioRunResult> results,
        CancellationToken cancellationToken = default)
    {
        _ = Throw.IfNull(results, nameof(results));
 
        foreach (ScenarioRunResult result in results)
        {
            cancellationToken.ThrowIfCancellationRequested();
 
            var resultDir =
                new DirectoryInfo(Path.Combine(_resultsRootPath, result.ExecutionName, result.ScenarioName));
 
            resultDir.Create();
 
            var resultFile = new FileInfo(Path.Combine(resultDir.FullName, $"{result.IterationName}.json"));
 
            using FileStream stream = resultFile.Create();
 
            await JsonSerializer.SerializeAsync(
                stream,
                result,
                SerializerContext.Default.ScenarioRunResult,
                cancellationToken).ConfigureAwait(false);
        }
    }
 
    /// <inheritdoc/>
    public ValueTask DeleteResultsAsync(
        string? executionName = null,
        string? scenarioName = null,
        string? iterationName = null,
        CancellationToken cancellationToken = default)
    {
        if (executionName is null && scenarioName is null && iterationName is null)
        {
            Directory.Delete(_resultsRootPath, recursive: true);
            _ = Directory.CreateDirectory(_resultsRootPath);
        }
        else if (executionName is not null && scenarioName is null && iterationName is null)
        {
            var executionDir = new DirectoryInfo(Path.Combine(_resultsRootPath, executionName));
 
            if (executionDir.Exists)
            {
                executionDir.Delete(recursive: true);
            }
        }
        else if (executionName is not null && scenarioName is not null && iterationName is null)
        {
            var scenarioDir =
                new DirectoryInfo(Path.Combine(_resultsRootPath, executionName, scenarioName));
 
            if (scenarioDir.Exists)
            {
                scenarioDir.Delete(recursive: true);
            }
        }
        else if (executionName is not null && scenarioName is not null && iterationName is not null)
        {
            var resultFile =
                new FileInfo(Path.Combine(_resultsRootPath, executionName, scenarioName, $"{iterationName}.json"));
 
            if (resultFile.Exists)
            {
                resultFile.Delete();
            }
        }
        else
        {
            IEnumerable<FileInfo> resultFiles =
                EnumerateResultFiles(executionName, scenarioName, iterationName, cancellationToken);
 
            foreach (FileInfo resultFile in resultFiles)
            {
                cancellationToken.ThrowIfCancellationRequested();
 
                DirectoryInfo scenarioDir = resultFile.Directory!;
                DirectoryInfo executionDir = scenarioDir.Parent!;
 
                resultFile.Delete();
 
                if (!scenarioDir.EnumerateFileSystemInfos().Any())
                {
                    scenarioDir.Delete(recursive: true);
 
                    if (!executionDir.EnumerateFileSystemInfos().Any())
                    {
                        executionDir.Delete(recursive: true);
                    }
                }
            }
        }
 
        return default;
    }
 
    /// <inheritdoc/>
    public async IAsyncEnumerable<string> GetLatestExecutionNamesAsync(
        int? count = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        await Task.CompletedTask.ConfigureAwait(false);
 
        if (count.HasValue && count <= 0)
        {
            yield break;
        }
 
        IEnumerable<DirectoryInfo> executionDirs = EnumerateExecutionDirs(cancellationToken: cancellationToken);
 
        if (count.HasValue)
        {
            executionDirs = executionDirs.Take(count.Value);
        }
 
        foreach (DirectoryInfo executionDir in executionDirs)
        {
            cancellationToken.ThrowIfCancellationRequested();
 
            yield return executionDir.Name;
        }
    }
 
    /// <inheritdoc/>
    public async IAsyncEnumerable<string> GetScenarioNamesAsync(
        string executionName,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        await Task.CompletedTask.ConfigureAwait(false);
 
        IEnumerable<DirectoryInfo> executionDirs = EnumerateExecutionDirs(executionName, cancellationToken);
 
        IEnumerable<DirectoryInfo> scenarioDirs =
            EnumerateScenarioDirs(executionDirs, cancellationToken: cancellationToken);
 
        foreach (DirectoryInfo scenarioDir in scenarioDirs)
        {
            cancellationToken.ThrowIfCancellationRequested();
 
            yield return scenarioDir.Name;
        }
    }
 
    /// <inheritdoc/>
    public async IAsyncEnumerable<string> GetIterationNamesAsync(
        string executionName,
        string scenarioName,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        await Task.CompletedTask.ConfigureAwait(false);
 
        IEnumerable<FileInfo> resultFiles =
            EnumerateResultFiles(executionName, scenarioName, cancellationToken: cancellationToken);
 
        foreach (FileInfo resultFile in resultFiles)
        {
            cancellationToken.ThrowIfCancellationRequested();
 
            yield return Path.GetFileNameWithoutExtension(resultFile.Name);
        }
    }
 
    private IEnumerable<DirectoryInfo> EnumerateExecutionDirs(
        string? executionName = null,
        CancellationToken cancellationToken = default)
    {
        var resultsDir = new DirectoryInfo(_resultsRootPath);
        if (!resultsDir.Exists)
        {
            yield break;
        }
 
        if (executionName is null)
        {
            IEnumerable<DirectoryInfo> executionDirs =
                resultsDir.EnumerateDirectories("*", InTopDirectoryOnly).OrderByDescending(d => d.CreationTimeUtc);
 
            foreach (DirectoryInfo executionDir in executionDirs)
            {
                cancellationToken.ThrowIfCancellationRequested();
 
                yield return executionDir;
            }
        }
        else
        {
            var executionDir = new DirectoryInfo(Path.Combine(_resultsRootPath, executionName));
            if (executionDir.Exists)
            {
                yield return executionDir;
            }
        }
    }
 
#pragma warning disable SA1204 // Static elements should appear before instance elements.
    private static IEnumerable<DirectoryInfo> EnumerateScenarioDirs(
        IEnumerable<DirectoryInfo> executionDirs,
        string? scenarioName = null,
        CancellationToken cancellationToken = default)
#pragma warning restore SA1204
    {
        foreach (DirectoryInfo executionDir in executionDirs)
        {
            cancellationToken.ThrowIfCancellationRequested();
 
            if (scenarioName is null)
            {
                IEnumerable<DirectoryInfo> scenarioDirs =
                    executionDir.EnumerateDirectories("*", InTopDirectoryOnly).OrderBy(d => d.Name);
 
                foreach (DirectoryInfo scenarioDir in scenarioDirs)
                {
                    cancellationToken.ThrowIfCancellationRequested();
 
                    yield return scenarioDir;
                }
            }
            else
            {
                var scenarioDir = new DirectoryInfo(Path.Combine(executionDir.FullName, scenarioName));
                if (scenarioDir.Exists)
                {
                    yield return scenarioDir;
                }
            }
        }
    }
 
#pragma warning disable SA1204 // Static elements should appear before instance elements.
    private static IEnumerable<FileInfo> EnumerateResultFiles(
        IEnumerable<DirectoryInfo> scenarioDirs,
        string? iterationName = null,
        CancellationToken cancellationToken = default)
#pragma warning restore SA1204
    {
        foreach (DirectoryInfo scenarioDir in scenarioDirs)
        {
            cancellationToken.ThrowIfCancellationRequested();
 
            if (iterationName is null)
            {
                IEnumerable<FileInfo> resultFiles =
                    scenarioDir
                        .EnumerateFiles("*.json", InTopDirectoryOnly)
                        .OrderBy(f => f.Name, IterationNameComparer.Default);
 
                foreach (FileInfo resultFile in resultFiles)
                {
                    cancellationToken.ThrowIfCancellationRequested();
 
                    yield return resultFile;
                }
            }
            else
            {
                var resultFile = new FileInfo(Path.Combine(scenarioDir.FullName, $"{iterationName}.json"));
                if (resultFile.Exists)
                {
                    yield return resultFile;
                }
            }
        }
    }
 
    private IEnumerable<FileInfo> EnumerateResultFiles(
        string? executionName = null,
        string? scenarioName = null,
        string? iterationName = null,
        CancellationToken cancellationToken = default)
    {
        IEnumerable<DirectoryInfo> executionDirs = EnumerateExecutionDirs(executionName, cancellationToken);
 
        IEnumerable<DirectoryInfo> scenarioDirs =
            EnumerateScenarioDirs(executionDirs, scenarioName, cancellationToken);
 
        IEnumerable<FileInfo> resultFiles = EnumerateResultFiles(scenarioDirs, iterationName, cancellationToken);
 
        return resultFiles;
    }
}