[Microsoft.Extensions.ML] PredictionEnginePool.GetModel(string modelName) does not ensure the named model is loaded

[wjwilimowski]

Preface

We're using a named prediction engine with the name model1 in a PredictionEnginePool.

services
    .AddPredictionEnginePool<ModelInput, ModelOutput>()
    .FromFile(modelName: "model1", filePath: "MLModel.zip", watchForChanges: false);

Example usage:

private readonly PredictionEnginePool<ModelInput, ModelOutput> _pool;

public void ThisWorks()
{
    // This will initialize the "model1" prediction engine if it hasn't been used yet
    // See source code:
    // https://github.com/dotnet/machinelearning/blob/58450d4f0709c237de95f31f8f05d46983c7a5c0/src/Microsoft.Extensions.ML/PredictionEnginePool.cs#L81
    PredictionEngine<ModelInput, ModelOutput> engine = _pool.GetPredictionEngine("model1");
    
    // We can now use the engine, then return it to the pool
}

In the above snippet, the PredictionEnginePool.GetPredictionEngine(string modelName) method ensures the engine is loaded before use.

machinelearning/src/Microsoft.Extensions.ML/PredictionEnginePool.cs

Lines 101 to 105 in 58450d4

	//Here we are in the world of named models where the model hasn't been built yet.
	var options = _predictionEngineOptions.Create(modelName);
	var pool = new PoolLoader<TData, TPrediction>(_serviceProvider, options);
	pool = _namedPools.GetOrAdd(modelName, pool);
	return pool.PredictionEnginePool.Get();

The issue

If _pool.GetPredictionEngine("model1") has never been called before calling _pool.GetModel("model1"), the model1 prediction engine is not loaded and KeyNotFoundException will be thrown.

private readonly PredictionEnginePool<ModelInput, ModelOutput> _pool;

public void ThisDoesNotWork()
{
    // Throws KeyNotFoundException unless _pool.GetPredictionEngine("model1") has been called before
    // See source code:
    // https://github.com/dotnet/machinelearning/blob/58450d4f0709c237de95f31f8f05d46983c7a5c0/src/Microsoft.Extensions.ML/PredictionEnginePool.cs#L51
    ITransformer model = _pool.GetModel("model1");
}

This inconsistent behavior can easily catch someone off-guard - one expects every public method on PredictionEnginePool to have the same "ensure initialized" behavior.

machinelearning/src/Microsoft.Extensions.ML/PredictionEnginePool.cs

Lines 51 to 54 in 58450d4

	public ITransformer GetModel(string modelName)
	{
	return _namedPools[modelName].Loader.GetModel();
	}

Code example

https://github.com/wjwilimowski/UnexpectedPredictionEnginePool