Implement IDisposable

Author: wschin

src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs

@@ -179,19 +179,25 @@ private OnnxTransformer(IHostEnvironment env, Options options, byte[] modelBytes
             foreach (var col in options.OutputColumns)
                 Host.CheckNonWhiteSpace(col, nameof(options.OutputColumns));
 
-            // Use ONNXRuntime to figure out the right input ane output configuration.
+            // Use ONNXRuntime to figure out the right input and output configuration.
             // However, ONNXRuntime doesn't provide strongly-typed method to access the produced
             // variables, we will inspect the ONNX model file to get information regarding types.
             try
             {
                 if (modelBytes == null)
                 {
+                    // Entering this region means that the model file is passed in by the user.
                     Host.CheckNonWhiteSpace(options.ModelFile, nameof(options.ModelFile));
                     Host.CheckIO(File.Exists(options.ModelFile), "Model file {0} does not exists.", options.ModelFile);
-                    Model = new OnnxModel(options.ModelFile, options.GpuDeviceId, options.FallbackToCpu);
+                    // Because we cannot delete the user file, ownModelFile should be false.
+                    Model = new OnnxModel(options.ModelFile, options.GpuDeviceId, options.FallbackToCpu, ownModelFile: false);
                 }
                 else
+                {
+                    // Entering this region means that the byte[] is passed as the model. To feed that byte[] to ONNXRuntime, we need
+                    // to create a temporal file to store it and then call ONNXRuntime's API to load that file.
                     Model = OnnxModel.CreateFromBytes(modelBytes, options.GpuDeviceId, options.FallbackToCpu);
+                }
             }
             catch (OnnxRuntimeException e)
             {
@@ -283,7 +289,7 @@ private protected override void SaveModel(ModelSaveContext ctx)
             ctx.CheckAtModel();
             ctx.SetVersionInfo(GetVersionInfo());
 
-            ctx.SaveBinaryStream("OnnxModel", w => { w.WriteByteArray(Model.ToByteArray()); });
+            ctx.SaveBinaryStream("OnnxModel", w => { w.WriteByteArray(File.ReadAllBytes(Model.ModelFile)); });
 
             Host.CheckNonEmpty(Inputs, nameof(Inputs));
             ctx.Writer.Write(Inputs.Length);

src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs

@@ -21,7 +21,7 @@ namespace Microsoft.ML.Transforms.Onnx
     /// It provides API to open a session, score tensors (NamedOnnxValues) and return
     /// the results.
     /// </summary>
-    internal sealed class OnnxModel
+    internal sealed class OnnxModel : IDisposable
     {
         /// <summary>
         /// OnnxModelInfo contains the data that we should get from
@@ -112,18 +112,39 @@ public OnnxVariableInfo(string name, OnnxShape shape, System.Type ortType, DataV
             }
         }
 
-        public readonly OnnxModelInfo ModelInfo;
+        /// <summary>
+        /// The ONNXRuntime facility to execute the loaded ONNX model.
+        /// </summary>
         private readonly InferenceSession _session;
-        private readonly string _modelFile;
+        /// <summary>
+        /// Indicates if <see cref="ModelFile"/> is a temporal file created by <see cref="CreateFromBytes(byte[], int?, bool)"/>
+        /// or <see cref="CreateFromBytes(byte[])"/>. If <see langword="true"/>, <see cref="Dispose(bool)"/> should delete <see cref="ModelFile"/>.
+        /// </summary>
+        private bool _ownModelFile;
+        /// <summary>
+        /// The ONNX model file that <see cref="OnnxModel"/> built upon.
+        /// </summary>
+        internal OnnxModelInfo ModelInfo { get; }
+        /// <summary>
+        /// The location where the used ONNX model loaded from.
+        /// </summary>
+        internal string ModelFile { get; }
 
         /// <summary>
         /// Constructs OnnxModel object from file.
         /// </summary>
         /// <param name="modelFile">Model file path.</param>
         /// <param name="gpuDeviceId">GPU device ID to execute on. Null for CPU.</param>
         /// <param name="fallbackToCpu">If true, resumes CPU execution quitely upon GPU error.</param>
-        public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false)
+        /// <param name="ownModelFile">If true, the <paramref name="modelFile"/> will be deleted when <see cref="OnnxModel"/> is
+        /// no longer needed.</param>
+        public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false, bool ownModelFile=false)
         {
+            ModelFile = modelFile;
+            // If we don't own the model file, _disposed should be false to prevent deleting user's file.
+            _ownModelFile = ownModelFile;
+            _disposed = false;
+
             if (gpuDeviceId != null)
             {
                 try
@@ -147,7 +168,7 @@ public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu =
 
             // Load ONNX model file and parse its input and output schema. The reason of doing so is that ONNXRuntime
             // doesn't expose full type information via its C# APIs.
-            _modelFile = modelFile;
+            ModelFile = modelFile;
             var model = new OnnxCSharpToProtoWrapper.ModelProto();
             using (var modelStream = File.OpenRead(modelFile))
                 model = OnnxCSharpToProtoWrapper.ModelProto.Parser.ParseFrom(modelStream);
@@ -191,7 +212,9 @@ public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu =
         }
 
         /// <summary>
-        /// Create an OnnxModel from a byte[]
+        /// Create an OnnxModel from a byte[]. Usually, a ONNX model is consumed by <see cref="OnnxModel"/> as a file.
+        /// With <see cref="CreateFromBytes(byte[])"/> and <see cref="CreateFromBytes(byte[], int?, bool)"/>, it's possible
+        /// to use in-memory model (type: byte[]) to create <see cref="OnnxModel"/>.
         /// </summary>
         /// <param name="modelBytes">Bytes of the serialized model</param>
         /// <returns>OnnxModel</returns>
@@ -202,6 +225,9 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes)
 
         /// <summary>
         /// Create an OnnxModel from a byte[]. Set execution to GPU if required.
+        /// Usually, a ONNX model is consumed by <see cref="OnnxModel"/> as a file.
+        /// With <see cref="CreateFromBytes(byte[])"/> and <see cref="CreateFromBytes(byte[], int?, bool)"/>,
+        /// it's possible to use in-memory model (type: byte[]) to create <see cref="OnnxModel"/>.
         /// </summary>
         /// <param name="modelBytes">Bytes of the serialized model.</param>
         /// <param name="gpuDeviceId">GPU device ID to execute on. Null for CPU.</param>
@@ -214,12 +240,7 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes, int? gpuDeviceId = nu
 
             var tempModelFile = Path.Combine(tempModelDir, "model.onnx");
             File.WriteAllBytes(tempModelFile, modelBytes);
-            return new OnnxModel(tempModelFile, gpuDeviceId, fallbackToCpu);
-
-            // TODO:
-            // tempModelFile is needed in case the model needs to be saved
-            // Either have to save the modelbytes and delete the temp dir/file,
-            // or keep the dir/file and write proper cleanup when application closes
+            return new OnnxModel(tempModelFile, gpuDeviceId, fallbackToCpu, ownModelFile: true);
         }
 
         /// <summary>
@@ -233,12 +254,42 @@ public IReadOnlyCollection<NamedOnnxValue> Run(List<NamedOnnxValue> inputNamedOn
         }
 
         /// <summary>
-        /// Convert the model to a byte array.
+        /// Flag used to indicate if the unmanaged resources (aka the model file <see cref="ModelFile"/>
+        /// and <see cref="_session"/>) have been deleted.
         /// </summary>
-        /// <returns>byte[]</returns>
-        public byte[] ToByteArray()
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        /// <summary>
+        /// There are two unmanaged resources we can dispose, <see cref="_session"/> and <see cref="ModelFile"/>
+        /// if <see cref="_ownModelFile"/> is <see langword="true"/>.
+        /// </summary>
+        /// <param name="disposing"></param>
+        private void Dispose(bool disposing)
+        {
+            if (!_disposed)
+            {
+                // There are two things to be disposed.
+                if (disposing)
+                {
+                    // First, we release the resource token by ONNXRuntime.
+                    _session.Dispose();
+                    // Second, we delete the model file if that file is not created by the user.
+                    if (_ownModelFile && File.Exists(ModelFile))
+                        File.Delete(ModelFile);
+                }
+                _disposed = true;
+            }
+        }
+
+        ~OnnxModel()
         {
-            return File.ReadAllBytes(_modelFile);
+            Dispose(false);
         }
     }
 

test/Microsoft.ML.OnnxTransformerTest/OnnxTransformTests.cs

@@ -518,5 +518,44 @@ public void TestOnnxZipMapWithStringKeys()
                 Assert.Equal(dataPoints[i].Input[2], dictionary["C"]);
             }
         }
+
+        [Fact]
+        public void TestOnnxModelDisposal()
+        {
+            // Create a ONNX model as a byte[].
+            var modelFile = Path.Combine(Directory.GetCurrentDirectory(), "zipmap", "TestZipMapInt64.onnx");
+            var modelInBytes = File.ReadAllBytes(modelFile);
+
+            // Create ONNX model from the byte[].
+            var onnxModel = OnnxModel.CreateFromBytes(modelInBytes);
+
+            // Check if a temporal file is crated for storing the byte[].
+            Assert.True(File.Exists(onnxModel.ModelFile));
+
+            // Delete the temporal file.
+            onnxModel.Dispose();
+
+            // Make sure the temporal file is deleted.
+            Assert.False(File.Exists(onnxModel.ModelFile));
+        }
+
+        [Fact]
+        public void TestOnnxModelNotDisposal()
+        {
+            // Declare the path the tested ONNX model file.
+            var modelFile = Path.Combine(Directory.GetCurrentDirectory(), "zipmap", "TestZipMapInt64.onnx");
+
+            // Create ONNX model from the model file.
+            var onnxModel = new OnnxModel(modelFile);
+
+            // Check if a temporal file is crated for storing the byte[].
+            Assert.True(File.Exists(onnxModel.ModelFile));
+
+            // Don't delete the temporal file!
+            onnxModel.Dispose();
+
+            // Make sure the temporal file still exists.
+            Assert.True(File.Exists(onnxModel.ModelFile));
+        }
     }
 }