Respect Marked exception during model loading. (#2574)

* Respect Marked exception during model loading.

Author: Ivanidzo4ka

src/Microsoft.ML.Core/ComponentModel/ComponentCatalog.cs

@@ -185,16 +185,25 @@ internal LoadableClassInfo(LoadableClassAttributeBase attr, MethodInfo getter, C
             internal object CreateInstanceCore(object[] ctorArgs)
             {
                 Contracts.Assert(Utils.Size(ctorArgs) == CtorTypes.Length + ((RequireEnvironment) ? 1 : 0));
-
-                if (InstanceGetter != null)
+                try
                 {
-                    Contracts.Assert(Utils.Size(ctorArgs) == 0);
-                    return InstanceGetter.Invoke(null, null);
+                    if (InstanceGetter != null)
+                    {
+                        Contracts.Assert(Utils.Size(ctorArgs) == 0);
+                        return InstanceGetter.Invoke(null, null);
+                    }
+                    if (Constructor != null)
+                        return Constructor.Invoke(ctorArgs);
+                    if (CreateMethod != null)
+                        return CreateMethod.Invoke(null, ctorArgs);
+                }
+                catch (TargetInvocationException ex)
+                {
+                    if (ex.InnerException != null && ex.InnerException.IsMarked())
+                        throw Contracts.Except(ex, "Error during class instantiation");
+                    else
+                        throw;
                 }
-                if (Constructor != null)
-                    return Constructor.Invoke(ctorArgs);
-                if (CreateMethod != null)
-                    return CreateMethod.Invoke(null, ctorArgs);
                 throw Contracts.Except("Can't instantiate class '{0}'", Type.Name);
             }
 

src/Microsoft.ML.Data/DataLoadSave/TransformerChain.cs

@@ -257,8 +257,10 @@ public static TransformerChain<ITransformer> LoadFrom(IHostEnvironment env, Stre
                     ModelLoadContext.LoadModel<TransformerChain<ITransformer>, SignatureLoadModel>(env, out var transformerChain, rep, LoaderSignature);
                     return transformerChain;
                 }
-                catch
+                catch (FormatException ex)
                 {
+                    if (!ex.IsMarked())
+                        throw;
                     var chain = ModelFileUtils.LoadPipeline(env, stream, new MultiFileSource(null), extractInnerPipe: false);
                     TransformerChain<ITransformer> transformChain = (chain as CompositeDataLoader).GetTransformer();
                     var predictor = ModelFileUtils.LoadPredictorOrNull(env, stream);

src/Microsoft.ML.Transforms/LambdaTransform.cs

@@ -69,6 +69,8 @@ private static ITransformer Create(IHostEnvironment env, ModelLoadContext ctx)
             var contractName = ctx.LoadString();
 
             var composition = env.GetCompositionContainer();
+            if (composition == null)
+                throw Contracts.Except("Unable to get the MEF composition container");
             ITransformer transformer = composition.GetExportedValue<ITransformer>(contractName);
             return transformer;
         }

test/Microsoft.ML.Tests/Transformers/CustomMappingTests.cs

@@ -69,9 +69,10 @@ public void TestCustomTransformer()
                 TestEstimatorCore(customEst, data);
                 Assert.True(false, "Cannot work without MEF injection");
             }
-            catch (Exception)
+            catch (InvalidOperationException ex)
             {
-                // REVIEW: we should have a common mechanism that will make sure this is 'our' exception thrown.
+                if (!ex.IsMarked())
+                    throw;
             }
             ML.CompositionContainer = new CompositionContainer(new TypeCatalog(typeof(MyLambda)));
             TestEstimatorCore(customEst, data);