protected void RequireBlocks(IComputationHandler handler, params int[] indices)
{
foreach (int index in indices)
{
if (_fetchedBlocks.ContainsKey(index))
{
continue;
}
if (_pendingFetchBlockTasks.ContainsKey(index))
{
_logger.Debug($"Waiting for already running asynchronous block fetch for index {index} to complete as it is now required...");
_pendingFetchBlockTasks[index].Wait();
IDictionary <string, INDArray> block = _pendingFetchBlockTasks[index].Result;
_pendingFetchBlockTasks.Remove(index);
_fetchedBlocks.Add(index, block);
_logger.Debug($"Done waiting for asynchronous block fetch for index {index} to complete, fetch completed.");
}
else
{
_logger.Debug($"Fetching required block with index {index}...");
_fetchedBlocks.Add(index, UnderlyingDataset.FetchBlock(index, handler));
_logger.Debug($"Done fetching required block with index {index}.");
}
}
}
protected void PrepareBlocksAsync(IComputationHandler handler, params int[] indices)
{
foreach (int index in indices)
{
if (!_pendingFetchBlockTasks.ContainsKey(index) && !_fetchedBlocks.ContainsKey(index))
{
_pendingFetchBlockTasks.Add(index, Task.Run(() =>
{
_logger.Debug($"Started asynchronous background preparation of block with index {index}.");
var block = UnderlyingDataset.FetchBlock(index, handler);
if (block == null)
{
_logger.Debug($"Unable to asynchronously prepare block with index {index}, block appears to be empty (null).");
}
else
{
_logger.Debug($"Done with asynchronous background preparation of block with index {index}.");
}
return(block);
}));
}
}
}
public override IEnumerable <IDictionary <string, INDArray> > Yield(IComputationHandler handler, SigmaEnvironment environment)
{
CheckNotNull(handler, environment);
int currentIndex = 0;
// TODO populate registry with relevant parameters
while (true)
{
RequireBlocks(handler, currentIndex);
if (_fetchedBlocks[currentIndex] == null)
{
break;
}
PrepareBlocksAsync(handler, currentIndex + 1);
var currentBlock = _fetchedBlocks[currentIndex];
_logger.Debug($"Yielding undivided block at index {currentIndex}.");
yield return(currentBlock);
UnderlyingDataset.FreeBlock(currentIndex, handler);
currentIndex++;
}
}
protected void FreeBlocks(IComputationHandler handler, params int[] indices)
{
foreach (int index in indices)
{
if (!_fetchedBlocks.ContainsKey(index))
{
continue;
}
var block = _fetchedBlocks[index];
UnderlyingDataset.FreeBlock(index, handler);
_fetchedBlocks.Remove(index);
}
}
private IDictionary <string, INDArray> FetchAndMergeFromDataset(IComputationHandler handler)
{
Dictionary <string, INDArray> unifiedBlock = new Dictionary <string, INDArray>();
Dictionary <string, IList <INDArray> > allFetchedBlocks = new Dictionary <string, IList <INDArray> >();
int currentBlockIndex = 0;
while (true)
{
IDictionary <string, INDArray> currentBlock = UnderlyingDataset.FetchBlock(currentBlockIndex, handler);
if (currentBlock != null)
{
foreach (string section in currentBlock.Keys)
{
if (!allFetchedBlocks.ContainsKey(section))
{
allFetchedBlocks.Add(section, new List <INDArray>());
}
allFetchedBlocks[section].Add(currentBlock[section]);
}
UnderlyingDataset.FreeBlock(currentBlockIndex, handler);
}
if (!UnderlyingDataset.CanFetchBlocksAfter(currentBlockIndex))
{
break;
}
currentBlockIndex++;
}
if (allFetchedBlocks.Count == 0)
{
throw new InvalidOperationException($"Cannot fetch and merge an empty block list, no blocks could be fetched from the dataset {UnderlyingDataset} for handler {handler}.");
}
foreach (string section in allFetchedBlocks.Keys)
{
unifiedBlock.Add(section, handler.MergeBatch(allFetchedBlocks[section].ToArray()));
}
return(unifiedBlock);
}
public void Dispose()
{
UnderlyingDataset.Dispose();
}
public void FreeBlock(int blockIndex, IComputationHandler handler)
{
UnderlyingDataset.FreeBlock(blockIndex, handler);
}
public IDictionary <string, INDArray> FetchBlock(int blockIndex, IComputationHandler handler, bool shouldWaitUntilAvailable = true)
{
var block = UnderlyingDataset.FetchBlock(blockIndex, handler, shouldWaitUntilAvailable);
return(block != null?GetOwnSlice(block) : null);
}
public bool CanFetchBlocksAfter(int blockIndex)
{
return(UnderlyingDataset.CanFetchBlocksAfter(blockIndex));
}
public long GetBlockSizeBytes(int blockIndex, IComputationHandler handler)
{
return(UnderlyingDataset.GetBlockSizeBytes(blockIndex, handler));
}
public bool IsBlockActive(int blockIndex, IComputationHandler handler)
{
return(UnderlyingDataset.IsBlockActive(blockIndex, handler));
}
public bool IsBlockActive(int blockIndex)
{
return(UnderlyingDataset.IsBlockActive(blockIndex));
}
public bool TrySetBlockSize(int blockSizeRecords)
{
return(UnderlyingDataset.TrySetBlockSize(blockSizeRecords));
}
public Task <IDictionary <string, INDArray> > FetchBlockAsync(int blockIndex, IComputationHandler handler, bool shouldWaitUntilAvailable = true)
{
return(UnderlyingDataset.FetchBlockAsync(MapToUnderlyingIndex(blockIndex), handler, shouldWaitUntilAvailable));
}
