fix: Retry OnnxHub call to improve test reliability (#1889)

microsoft · Mar 24, 2023 · 52919ce · 52919ce
1 parent 979c629
commit 52919ce
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/deep-learning/src/main/scala/com/microsoft/azure/synapse/ml/onnx/ONNXHub.scala b/deep-learning/src/main/scala/com/microsoft/azure/synapse/ml/onnx/ONNXHub.scala
@@ -139,10 +139,12 @@ class ONNXHub(val modelCacheDir: Path,
   }
 
   private def toStream(url: URL) = {
-    val urlCon = url.openConnection()
-    urlCon.setConnectTimeout(connectTimeout)
-    urlCon.setReadTimeout(readTimeout)
-    new BufferedInputStream(urlCon.getInputStream)
+    FaultToleranceUtils.retryWithTimeout(retryCount, Duration.apply(retryTimeoutInSeconds, "sec")) {
+      val urlCon = url.openConnection()
+      urlCon.setConnectTimeout(connectTimeout)
+      urlCon.setReadTimeout(readTimeout)
+      new BufferedInputStream(urlCon.getInputStream)
+    }
   }
 
   def listModels(repo: String = ONNXHub.DefaultRepo,