Visitar URL original
fix: preserve attachment cancellation through storage errors · ranadeepsingh/SynapseML@eae6ccc · GitHub
Skip to content

Commit eae6ccc

Browse files
ranadeepsinghCopilot
ranadeepsingh
and
Copilot
committed
fix: preserve attachment cancellation through storage errors
Summary: - Check task/thread cancellation before and after attachment IO. - Recognize wrapped and suppressed interruption, including cyclic exception graphs, without treating storage timeouts as cancellation or replacing fatal errors. - Cover both legacy chat-completions and responses paths with local filesystem and HTTP fixtures; no live service calls. Prompting intent: Preserve reviewed follow-up fixes after the prerequisite PRs were merged, and complete current-source validation and review. Linked sources: Follow-up to microsoft#2807 and microsoft#2801. Rationale: Attachment storage can wrap or clear interruption signals, causing a canceled Spark task to continue or return an ordinary row error. This follow-up preserves cancellation through file reads and cleanup while keeping storage timeouts sanitized and row-local. Local validation: sbt "set cognitive / Test / fork := false" "cognitive/compile" "cognitive/Test/compile" "cognitive/testOnly com.microsoft.azure.synapse.ml.services.openai.OpenAIPromptMultimodalRequestSuite" "cognitive/scalastyle" "cognitive/Test/scalastyle" Passed: 69 tests, 0 skipped, compile/Test compile, both scalastyle gates. Publication: Explicitly authorized draft-first bootstrap; formal reviews and current-head CI remain pending. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5ba39157-c8dd-4adb-90e0-ff963a1e6733
1 parent 37948d7 commit eae6ccc

3 files changed

Lines changed: 389 additions & 94 deletions

File tree

‎cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIAttachmentUtils.scala‎

Lines changed: 65 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,40 +3,90 @@
33

44
package com.microsoft.azure.synapse.ml.services.openai
55

6-
import com.microsoft.azure.synapse.ml.io.binary.BinaryFileReader
6+
import org.apache.commons.io.IOUtils
7+
import org.apache.hadoop.conf.Configuration
78
import org.apache.hadoop.fs.{Path => HPath}
89
import org.apache.spark.{TaskContext, TaskKilledException}
910

10-
import java.io.ByteArrayInputStream
11+
import java.io.{ByteArrayInputStream, InterruptedIOException}
1112
import java.net.{URI, URLConnection}
1213
import java.nio.channels.ClosedByInterruptException
1314
import java.nio.charset.StandardCharsets
1415
import java.util.concurrent.CancellationException
15-
import java.util.{Base64, Locale}
16-
import scala.util.Try
16+
import java.util.{Base64, Collections, IdentityHashMap, Locale}
17+
import scala.util.{Try, Using}
1718
import scala.util.control.NonFatal
1819

1920
private[openai] object OpenAIAttachmentUtils {
2021

2122
sealed class AttachmentReadException(fileName: String)
2223
extends IllegalArgumentException(s"Unable to read attachment '$fileName'")
2324

25+
/** Storage cancellation, not task interruption; legacy APIs keep this sanitized error row-local. */
2426
final class CancelledAttachmentReadException(fileName: String) extends AttachmentReadException(fileName)
2527

2628
private def interruptedRead(): Nothing = {
2729
Thread.currentThread().interrupt()
2830
throw new InterruptedException("Attachment read interrupted")
2931
}
3032

31-
private[openai] def readAttachment(fileName: String)(read: => Array[Byte]): Array[Byte] = try {
32-
read
33-
} catch {
34-
case _: InterruptedException | _: ClosedByInterruptException => interruptedRead()
35-
case NonFatal(_) if Thread.currentThread().isInterrupted => interruptedRead()
36-
case NonFatal(_) if Option(TaskContext.get()).exists(_.isInterrupted()) =>
37-
throw new TaskKilledException("Attachment read canceled")
38-
case _: CancellationException => throw new CancelledAttachmentReadException(fileName)
39-
case NonFatal(_) => throw new AttachmentReadException(fileName)
33+
private def taskCancelled: Boolean = Option(TaskContext.get()).exists(_.isInterrupted())
34+
35+
private def taskKilled(): Nothing = throw new TaskKilledException("Attachment read canceled")
36+
37+
private def checkCancellation(): Unit = {
38+
if (Thread.currentThread().isInterrupted) interruptedRead()
39+
if (taskCancelled) taskKilled()
40+
}
41+
42+
private def isTimeout(error: Throwable): Boolean = error match {
43+
case io: InterruptedIOException =>
44+
// Storage SDKs relocate HTTP timeout classes, including their subclasses.
45+
@scala.annotation.tailrec
46+
def timeoutClass(clazz: Class[_]): Boolean =
47+
clazz != null && (clazz.getName.endsWith("TimeoutException") || timeoutClass(clazz.getSuperclass))
48+
timeoutClass(io.getClass)
49+
case _ => false
50+
}
51+
52+
private def isInterruption(error: Throwable): Boolean = error match {
53+
case _: InterruptedException | _: ClosedByInterruptException => true
54+
case io: InterruptedIOException => !isTimeout(io)
55+
case _ => false
56+
}
57+
58+
private def isInterruptedRead(error: Throwable): Boolean = {
59+
if (!NonFatal(error)) isInterruption(error)
60+
else {
61+
// Storage wrappers and resource cleanup can hide interruption beneath an ordinary read failure.
62+
val visited = Collections.newSetFromMap(new IdentityHashMap[Throwable, java.lang.Boolean]())
63+
@scala.annotation.tailrec
64+
def inspect(pending: List[Throwable]): Boolean = pending match {
65+
case Nil => false
66+
case current :: remaining =>
67+
if (!visited.add(current)) inspect(remaining)
68+
else if (isInterruption(current)) true
69+
else {
70+
// A timeout owns its cause classification; cleanup can still add a separate interruption.
71+
val causes = if (isTimeout(current)) Nil else Option(current.getCause).toList
72+
inspect(causes ::: current.getSuppressed.toList ::: remaining)
73+
}
74+
}
75+
inspect(List(error))
76+
}
77+
}
78+
79+
private[openai] def readAttachment(fileName: String)(read: => Array[Byte]): Array[Byte] = {
80+
checkCancellation()
81+
val content = try read catch {
82+
case error if isInterruptedRead(error) => interruptedRead()
83+
case NonFatal(_) if Thread.currentThread().isInterrupted => interruptedRead()
84+
case NonFatal(_) if taskCancelled => taskKilled()
85+
case _: CancellationException => throw new CancelledAttachmentReadException(fileName)
86+
case NonFatal(_) => throw new AttachmentReadException(fileName)
87+
}
88+
checkCancellation()
89+
content
4090
}
4191

4292
private case class DataUrlAttachment(
@@ -249,7 +299,8 @@ private[openai] object OpenAIAttachmentUtils {
249299
} else {
250300
val fileName = attachmentFilename(filePathStr)
251301
val fileBytes = readAttachment(fileName) {
252-
BinaryFileReader.readSingleFileBytes(new HPath(filePathStr))
302+
val path = new HPath(filePathStr)
303+
Using.resource(path.getFileSystem(new Configuration()).open(path))(IOUtils.toByteArray)
253304
}
254305
(fileName, fileBytes, None)
255306
}

‎cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/AttachmentReadFailureFileSystem.scala‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
// Copyright (C) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License. See LICENSE in project root for information.
3+
4+
package com.microsoft.azure.synapse.ml.services.openai
5+
6+
import org.apache.hadoop.fs.{FSDataInputStream, FSInputStream, Path, RawLocalFileSystem}
7+
import org.apache.http.conn.ConnectTimeoutException
8+
import org.apache.spark.TaskContext
9+
10+
import java.io.{ByteArrayInputStream, InterruptedIOException, IOException}
11+
import java.net.{SocketTimeoutException, URI}
12+
import java.nio.channels.ClosedByInterruptException
13+
import java.util.concurrent.{CancellationException, ConcurrentHashMap, ConcurrentLinkedQueue, CountDownLatch, TimeUnit}
14+
import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
15+
16+
private[openai] object AttachmentReadFailures {
17+
val Location: URI = URI.create("attachmentfailure:///")
18+
19+
class Gate(val content: Option[Array[Byte]] = None) {
20+
val entered = new CountDownLatch(1)
21+
val release = new CountDownLatch(1)
22+
val completed = new CountDownLatch(1)
23+
val context = new AtomicReference[TaskContext]()
24+
val taskFailure = new AtomicReference[Throwable]()
25+
val completionInterrupted = new AtomicBoolean()
26+
}
27+
28+
val Active = new ConcurrentHashMap[String, Gate]()
29+
30+
val TimeoutNames = Seq("socket-timeout.png", "connect-timeout.png",
31+
"close-socket-timeout.png", "close-connect-timeout.png", "wrapped-socket-timeout.png",
32+
"wrapped-connect-timeout.png", "relocated-connect-timeout.png", "relocated-pool-timeout.png",
33+
"nested-connect-timeout.png", "nested-interrupted-timeout.png", "httpclient5-timeout.png")
34+
35+
private object RelocatedHttpClient {
36+
class ConnectTimeoutException extends InterruptedIOException("Synthetic private connect timeout")
37+
class ConnectionPoolTimeoutException extends ConnectTimeoutException
38+
}
39+
40+
private def timeoutCausedBy(cause: Throwable): SocketTimeoutException = {
41+
val error = new SocketTimeoutException("Synthetic private storage timeout")
42+
error.initCause(new IOException("Synthetic private storage wrapper", cause))
43+
error
44+
}
45+
46+
def cyclicExceptionGraph(includeCauses: Boolean): (IOException, IOException) = {
47+
val root = new IOException("Synthetic private path")
48+
val last = (1 to 2048).foldLeft(root) { (previous, index) =>
49+
val next = new IOException("Synthetic private cleanup failure")
50+
if (includeCauses && index % 2 == 0) previous.initCause(next) else previous.addSuppressed(next)
51+
next
52+
}
53+
last.addSuppressed(root)
54+
root -> last
55+
}
56+
57+
private def throwWithClearedInterrupt(error: IOException): Nothing = {
58+
Thread.currentThread().interrupt()
59+
Thread.interrupted()
60+
throw error
61+
}
62+
63+
private def cleanupFailure(id: String): Option[IOException] = id match {
64+
case "close-interrupted-io.png" => Some(new InterruptedIOException("Synthetic private close failure"))
65+
case "close-socket-timeout.png" => Some(new SocketTimeoutException("Synthetic private close failure"))
66+
case "close-connect-timeout.png" => Some(new ConnectTimeoutException("Synthetic private close failure"))
67+
case _ => None
68+
}
69+
70+
private def storageTimeout(id: String): Option[IOException] = id match {
71+
case "relocated-connect-timeout.png" => Some(new RelocatedHttpClient.ConnectTimeoutException)
72+
case "relocated-pool-timeout.png" => Some(new RelocatedHttpClient.ConnectionPoolTimeoutException {})
73+
case "nested-connect-timeout.png" =>
74+
Some(timeoutCausedBy(new RelocatedHttpClient.ConnectTimeoutException))
75+
case "nested-interrupted-timeout.png" =>
76+
Some(timeoutCausedBy(new InterruptedIOException("Synthetic private timeout cause")))
77+
case "httpclient5-timeout.png" =>
78+
Some(new org.apache.hc.client5.http.ConnectTimeoutException("Synthetic private connect timeout"))
79+
case "timeout-close-interrupted.png" =>
80+
val error = timeoutCausedBy(new RelocatedHttpClient.ConnectTimeoutException)
81+
error.addSuppressed(new InterruptedIOException("Synthetic private cleanup interruption"))
82+
Some(error)
83+
case _ => None
84+
}
85+
86+
private def wrappedFailure(id: String): Option[IOException] = {
87+
val cause = id match {
88+
case "wrapped-interrupted.png" => Some(new InterruptedException("Synthetic private path"))
89+
case "wrapped-interrupted-io.png" => Some(new InterruptedIOException("Synthetic private path"))
90+
case "wrapped-closed-interrupt.png" => Some(new ClosedByInterruptException())
91+
case "wrapped-socket-timeout.png" => Some(new SocketTimeoutException("Synthetic private path"))
92+
case "wrapped-connect-timeout.png" => Some(new ConnectTimeoutException("Synthetic private path"))
93+
case _ => None
94+
}
95+
cause.orElse(storageTimeout(id)).map(new IOException("Synthetic private wrapped failure", _))
96+
}
97+
98+
def open(id: String): FSDataInputStream = {
99+
val active = Option(Active.get(id))
100+
active.foreach { gate =>
101+
val context = TaskContext.get()
102+
gate.context.set(context)
103+
context.addTaskFailureListener((_: TaskContext, error: Throwable) => gate.taskFailure.set(error))
104+
context.addTaskCompletionListener[Unit] { _ =>
105+
gate.completionInterrupted.set(Thread.currentThread().isInterrupted)
106+
gate.completed.countDown()
107+
}
108+
gate.entered.countDown()
109+
require(gate.release.await(30, TimeUnit.SECONDS), "Attachment cancellation test gate timed out")
110+
}
111+
wrappedFailure(id).foreach(throwWithClearedInterrupt)
112+
if (id == "interrupted.png") throw new ClosedByInterruptException()
113+
if (id == "plain-interrupted.png") throw new InterruptedException("Synthetic private path")
114+
if (id == "interrupted-io.png") {
115+
throwWithClearedInterrupt(new InterruptedIOException("Synthetic private path"))
116+
}
117+
if (id == "socket-timeout.png") throw new SocketTimeoutException("Synthetic private path")
118+
if (id == "connect-timeout.png") throw new ConnectTimeoutException("Synthetic private path")
119+
if (id == "flagged.png") Thread.currentThread().interrupt()
120+
if (id == "cancelled.png") throw new CancellationException("Synthetic private path")
121+
val cleanupError = cleanupFailure(id)
122+
val streamContent = active.flatMap(_.content).orElse(cleanupError.map(_ => Array.emptyByteArray))
123+
if (id == "read-close.png" || streamContent.isDefined) {
124+
val content = streamContent.getOrElse(Array.emptyByteArray)
125+
val bytes = new ByteArrayInputStream(content)
126+
new FSDataInputStream(new FSInputStream {
127+
override def getPos: Long = content.length - bytes.available()
128+
override def seek(position: Long): Unit = throw new UnsupportedOperationException("Test stream is not seekable")
129+
override def seekToNewSource(position: Long): Boolean = false
130+
override def read(): Int = {
131+
if (id == "read-close.png") throw new InterruptedException("Synthetic private path")
132+
if (cleanupError.isDefined) throw new IOException("Synthetic private read failure")
133+
bytes.read()
134+
}
135+
override def close(): Unit = {
136+
bytes.close()
137+
if (id == "read-close.png") throw new IOException("Synthetic private close failure")
138+
cleanupError.foreach(throwWithClearedInterrupt)
139+
}
140+
})
141+
} else {
142+
throw new IOException("Synthetic attachment read failure")
143+
}
144+
}
145+
}
146+
147+
class AttachmentReadFailureFileSystem extends RawLocalFileSystem {
148+
val opened = new ConcurrentLinkedQueue[String]()
149+
150+
override def getUri: URI = AttachmentReadFailures.Location
151+
152+
override def open(path: Path, bufferSize: Int): FSDataInputStream = {
153+
opened.add(path.getName)
154+
AttachmentReadFailures.open(path.getName)
155+
}
156+
}

0 commit comments

Comments
 (0)