Don't redirect stderr to stdout

2 files changed, 9 insertions, 7 deletions

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java b/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java
index 2ef81e3f1fa..2e2ebdeb98f 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java
@@ -48,7 +48,7 @@ public class OnnxModelProbe {
                 }
             }
 
-        } catch (IOException | InterruptedException e) {
+        } catch (IllegalArgumentException | IOException | InterruptedException e) {
             e.printStackTrace(System.err);
         }
 
@@ -127,7 +127,6 @@ public class OnnxModelProbe {
 
     private static String callVespaAnalyzeOnnxModel(String jsonInput) throws IOException, InterruptedException {
         ProcessBuilder processBuilder = new ProcessBuilder(binary, "--probe-types");
-        processBuilder.redirectErrorStream(true);
         StringBuilder output = new StringBuilder();
         Process process = processBuilder.start();
 
@@ -136,7 +135,7 @@ public class OnnxModelProbe {
         os.write(jsonInput.getBytes(StandardCharsets.UTF_8));
         os.close();
 
-        // Read output from stdout/stderr
+        // Read output from stdout
         InputStream inputStream = process.getInputStream();
         while (true) {
             int b = inputStream.read();
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/ml/OnnxModelProbeTest.java b/config-model/src/test/java/com/yahoo/vespa/model/ml/OnnxModelProbeTest.java
index 6c4c919a229..b2d4953f8e2 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/ml/OnnxModelProbeTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/ml/OnnxModelProbeTest.java
@@ -29,11 +29,14 @@ public class OnnxModelProbeTest {
                     "input2", TensorType.fromSpec("tensor<float>(d0[1],d1[2])"));
             TensorType expected = TensorType.fromSpec("tensor<float>(d0[1],d1[2],d2[2])");
 
-            // Can't test model probing directly as 'vespa-analyze-onnx-model' is not available
             TensorType outputType = OnnxModelProbe.probeModel(app, modelPath, output, inputTypes);
-            assertEquals(outputType, TensorType.empty);
-            
-            OnnxModelProbe.writeProbedOutputType(app, modelPath, output, inputTypes, expected);
+
+            // if 'vespa-analyze-onnx-model' was unavailable, specifically cache expected type
+            if (outputType.equals(TensorType.empty)) {
+                OnnxModelProbe.writeProbedOutputType(app, modelPath, output, inputTypes, expected);
+            } else {
+                assertEquals(outputType, expected);
+            }
 
             // Test loading from generated info
             storedAppDir.toFile().mkdirs();