Skip to content

Commit

Permalink
perf: improve wait for results performances by reducing the size of t…
Browse files Browse the repository at this point in the history 
…he request and making multiple at the same time
  • Loading branch information
@aneojgurhem
aneojgurhem committed Aug 20, 2024
1 parent 3f2c13f commit db2a4c4
Showing 1 changed file with 100 additions and 65 deletions.
165 changes: 100 additions & 65 deletions packages/csharp/ArmoniK.Api.Client/EventsClientExt.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
using ArmoniK.Api.gRPC.V1;
using ArmoniK.Api.gRPC.V1.Events;
using ArmoniK.Api.gRPC.V1.Results;
using ArmoniK.Utils;

using Grpc.Core;

Expand Down Expand Up @@ -68,87 +69,121 @@ private static FiltersAnd ResultsFilter(string resultId)
},
};


/// <summary>
/// Wait until the given results are completed
/// </summary>
/// <param name="client">gRPC result client</param>
/// <param name="sessionId">The session ID in which the results are located</param>
/// <param name="resultIds">A collection of results to wait for</param>
/// <param name="cancellationToken">Token used to cancel the execution of the method</param>
/// <exception cref="Exception">if a result is aborted</exception>
[PublicAPI]
[Obsolete("Use the overload with the bucket size and the parallelism")]
public static Task WaitForResultsAsync(this Events.EventsClient client,
string sessionId,
ICollection<string> resultIds,
CancellationToken cancellationToken = default)
=> client.WaitForResultsAsync(sessionId,
resultIds,
100,
1,
cancellationToken);


/// <summary>
/// Wait until the given results are completed
/// </summary>
/// <param name="client">gRPC result client</param>
/// <param name="sessionId">The session ID in which the results are located</param>
/// <param name="resultIds">A collection of results to wait for</param>
/// <param name="parallelism">Number of parallel threads to use. One bucket per thread.</param>
/// <param name="bucket_size">Number of results Id to use to create the request to the event API</param>
/// <param name="cancellationToken">Token used to cancel the execution of the method</param>
/// <exception cref="Exception">if a result is aborted</exception>
[PublicAPI]
public static async Task WaitForResultsAsync(this Events.EventsClient client,
string sessionId,
ICollection<string> resultIds,
CancellationToken cancellationToken)
{
var resultsNotFound = new HashSet<string>(resultIds);
while (resultsNotFound.Any())
{
using var streamingCall = client.GetEvents(new EventSubscriptionRequest
int bucket_size = 100,
int parallelism = 1,
CancellationToken cancellationToken = default)
=> await resultIds.ToChunks(bucket_size)
.ParallelForEach(new ParallelTaskOptions
{
ParallelismLimit = parallelism,
},
async results =>
{
var resultsNotFound = new HashSet<string>(results);
while (resultsNotFound.Any())
{
using var streamingCall = client.GetEvents(new EventSubscriptionRequest
{
SessionId = sessionId,
ReturnedEvents =
{
EventsEnum.ResultStatusUpdate,
EventsEnum.NewResult,
},
ResultsFilters = new Filters
{
Or =
{
resultsNotFound.Select(ResultsFilter),
},
},
},
cancellationToken: cancellationToken);
try
{
while (await streamingCall.ResponseStream.MoveNext(cancellationToken))
{
var resp = streamingCall.ResponseStream.Current;
if (resp.UpdateCase == EventSubscriptionResponse.UpdateOneofCase.ResultStatusUpdate &&
resultsNotFound.Contains(resp.ResultStatusUpdate.ResultId))
{
if (resp.ResultStatusUpdate.Status == ResultStatus.Completed)
{
SessionId = sessionId,
ReturnedEvents =
resultsNotFound.Remove(resp.ResultStatusUpdate.ResultId);
if (!resultsNotFound.Any())
{
EventsEnum.ResultStatusUpdate,
EventsEnum.NewResult,
},
ResultsFilters = new Filters
{
Or =
{
resultsNotFound.Select(ResultsFilter),
},
},
},
cancellationToken: cancellationToken);
try
{
while (await streamingCall.ResponseStream.MoveNext(cancellationToken))
{
var resp = streamingCall.ResponseStream.Current;
if (resp.UpdateCase == EventSubscriptionResponse.UpdateOneofCase.ResultStatusUpdate && resultsNotFound.Contains(resp.ResultStatusUpdate.ResultId))
{
if (resp.ResultStatusUpdate.Status == ResultStatus.Completed)
{
resultsNotFound.Remove(resp.ResultStatusUpdate.ResultId);
if (!resultsNotFound.Any())
{
break;
}
}
break;
}
}

if (resp.ResultStatusUpdate.Status == ResultStatus.Aborted)
{
throw new ResultAbortedException($"Result {resp.ResultStatusUpdate.ResultId} has been aborted");
}
}
if (resp.ResultStatusUpdate.Status == ResultStatus.Aborted)
{
throw new ResultAbortedException($"Result {resp.ResultStatusUpdate.ResultId} has been aborted");
}
}

if (resp.UpdateCase == EventSubscriptionResponse.UpdateOneofCase.NewResult && resultsNotFound.Contains(resp.NewResult.ResultId))
{
if (resp.NewResult.Status == ResultStatus.Completed)
{
resultsNotFound.Remove(resp.NewResult.ResultId);
if (!resultsNotFound.Any())
{
break;
}
}
if (resp.UpdateCase == EventSubscriptionResponse.UpdateOneofCase.NewResult &&
resultsNotFound.Contains(resp.NewResult.ResultId))
{
if (resp.NewResult.Status == ResultStatus.Completed)
{
resultsNotFound.Remove(resp.NewResult.ResultId);
if (!resultsNotFound.Any())
{
break;
}
}

if (resp.NewResult.Status == ResultStatus.Aborted)
{
throw new ResultAbortedException($"Result {resp.NewResult.ResultId} has been aborted");
}
}
}
}
catch (OperationCanceledException)
{
}
catch (RpcException)
{
}
}
}
if (resp.NewResult.Status == ResultStatus.Aborted)
{
throw new ResultAbortedException($"Result {resp.NewResult.ResultId} has been aborted");
}
}
}
}
catch (OperationCanceledException)
{
}
catch (RpcException)
{
}
}
});
}
}

0 comments on commit db2a4c4

Please sign in to comment.