hudi-agent commented on code in PR #19677:
URL: https://github.com/apache/hudi/pull/19677#discussion_r3846020390
##########
hudi-utilities/src/main/java/org/apache/hudi/utilities/transform/embedding/EmbeddingTransformer.java:
##########
@@ -262,16 +277,60 @@ private void submitUpToWindow() {
}
}
- private ExecutorService executor() {
+ private synchronized ExecutorService executor() {
if (executor == null) {
executor = Executors.newFixedThreadPool(maxInflight,
new CustomizedThreadFactory("embedding-transformer", true));
+ // end() runs only once the input drains normally. A task killed by an
embeddings
+ // failure would otherwise strand maxInflight threads on an executor
JVM that Spark
+ // goes on reusing, so release them on task completion however the
task ends.
+ TaskContext taskContext = TaskContext.get();
+ if (taskContext != null) {
+ taskContext.addTaskCompletionListener(
+ (TaskCompletionListener) context -> shutdownExecutor());
+ }
}
return executor;
}
+ private synchronized void shutdownExecutor() {
+ if (executor != null) {
+ executor.shutdownNow();
+ executor = null;
+ }
+ }
+
private List<float[]> embed(List<String> texts) {
- return providerInstance().embed(texts);
+ Semaphore permits = permits(maxConcurrent);
+ try {
+ permits.acquire();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new HoodieException("Interrupted waiting to call the embeddings
API", e);
+ }
+ try {
+ return providerInstance().embed(texts);
+ } finally {
+ permits.release();
+ }
+ }
+
+ /**
+ * Caps embedding requests in flight across the whole JVM. Each partition
runs its own worker
+ * pool, so without a shared limit the load offered to one endpoint is
+ * (concurrent tasks x max.inflight.requests) and grows with the cluster:
an oversubscribed
+ * endpoint then returns timeouts or 429s, which retry, which adds more
load. Queueing here
+ * instead is free, since a waiting thread holds no connection and nothing
can time out
+ * while it waits.
+ *
+ * <p>Every task configured alike shares one semaphore, which is what
makes the cap hold
+ * across partitions. Executors are long lived and reused, so keying by
permit count keeps a
+ * second job with a different setting on its own budget rather than
silently inheriting the
+ * first job's.
+ */
+ @VisibleForTesting
Review Comment:
🤖 Since the semaphore is keyed only by permit count, two embedding
transforms in the same executor JVM that share a `maxConcurrent` value (e.g.
both on the default 4) but hit different endpoints would draw on one semaphore
— a combined cap of 4, not 4 per endpoint. Is that intended, or should the
provider/endpoint identity be part of the key so the doc's "requests to the
endpoint" framing holds?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag
quality.</i></sub>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
