Skip to content

[Spark] Make cancel() cancel the Spark jobs and stop only a session the runner created #40103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tkaymak wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[Spark] Make cancel() cancel the Spark jobs and stop only a session the runner created #40103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 70 additions & 30 deletions ...apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,35 +25,49 @@
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;
import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
import org.apache.beam.runners.spark.structuredstreaming.translation.EvaluationContext;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.metrics.MetricResults;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
import org.apache.spark.SparkException;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
* Result of a pipeline submitted to the {@link SparkStructuredStreamingRunner}. The pipeline runs
* on a dedicated thread, {@link #cancel()} stops it and joins that thread.
*/
public class SparkStructuredStreamingPipelineResult implements PipelineResult {

private static final Logger LOG =
LoggerFactory.getLogger(SparkStructuredStreamingPipelineResult.class);

private final Future<?> pipelineExecution;
// Supplies the context of the translated pipeline, null until translation has completed.
private final Supplier<? extends @Nullable EvaluationContext> evaluationContext;
private final MetricsAccumulator metrics;
private final @Nullable Runnable onTerminalState;
private PipelineResult.State state;
private final AtomicBoolean cancelRequested;
private final Runnable cancelSparkJobs;
private volatile PipelineResult.State state;

SparkStructuredStreamingPipelineResult(
Future<?> pipelineExecution,
Supplier<? extends @Nullable EvaluationContext> evaluationContext,
MetricsAccumulator metrics,
final @Nullable Runnable onTerminalState) {
AtomicBoolean cancelRequested,
Runnable cancelSparkJobs) {
this.pipelineExecution = pipelineExecution;
this.evaluationContext = evaluationContext;
this.metrics = metrics;
this.onTerminalState = onTerminalState;
this.cancelRequested = cancelRequested;
this.cancelSparkJobs = cancelSparkJobs;
// pipelineExecution is expected to have started executing eagerly.
this.state = State.RUNNING;
}
Expand All @@ -77,13 +91,6 @@ private static RuntimeException unwrapCause(Throwable exception) {
: new Pipeline.PipelineExecutionException(firstNonNull(next, exception));
}

private State awaitTermination(Duration duration)
throws TimeoutException, ExecutionException, InterruptedException {
pipelineExecution.get(duration.getMillis(), TimeUnit.MILLISECONDS);
// Throws an exception if the job is not finished successfully in the given time.
return PipelineResult.State.DONE;
}

@Override
public PipelineResult.State getState() {
return state;
Expand All @@ -94,18 +101,33 @@ public PipelineResult.State waitUntilFinish() {
return waitUntilFinish(Duration.millis(Long.MAX_VALUE));
}

/**
* Waits up to {@code duration} for the execution thread. A pipeline that ends after {@link
* #cancel()} is CANCELLED, any other failure is rethrown and the pipeline is FAILED.
*/
@Override
public State waitUntilFinish(final Duration duration) {
try {
State finishState = awaitTermination(duration);
offerNewState(finishState);
pipelineExecution.get(duration.getMillis(), TimeUnit.MILLISECONDS);
state = cancelRequested.get() ? State.CANCELLED : State.DONE;
} catch (final TimeoutException e) {
// ignore.
} catch (final ExecutionException e) {
offerNewState(PipelineResult.State.FAILED);
if (cancelRequested.get()) {
LOG.info(
"Pipeline execution ended with an exception after cancel: {}",
String.valueOf(Throwables.getRootCause(e).getMessage()));
state = State.CANCELLED;
return state;
}
state = State.FAILED;
throw unwrapCause(firstNonNull(e.getCause(), e));
} catch (final InterruptedException e) {
Thread.currentThread().interrupt();
state = State.FAILED;
throw unwrapCause(e);
} catch (final Exception e) {
offerNewState(PipelineResult.State.FAILED);
state = State.FAILED;
throw unwrapCause(e);
}

Expand All @@ -117,26 +139,44 @@ public MetricResults metrics() {
return asAttemptedOnlyMetricResults(metrics.value());
}

/**
* Cancels the Spark jobs of the pipeline and blocks until the execution thread has ended. An
* execution that already ended keeps its state. An interrupted caller keeps the interrupt flag
* and gets the current state.
*/
@Override
public PipelineResult.State cancel() throws IOException {
public synchronized PipelineResult.State cancel() throws IOException {
if (state.isTerminal()) {
return state;
}
if (pipelineExecution.isDone()) {
try {
pipelineExecution.get();
state = cancelRequested.get() ? State.CANCELLED : State.DONE;
} catch (ExecutionException e) {
state = cancelRequested.get() ? State.CANCELLED : State.FAILED;
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
return state;
}
cancelRequested.set(true);
EvaluationContext ctx = evaluationContext.get();
if (ctx != null) {
ctx.stop();
}
pipelineExecution.cancel(true);
offerNewState(PipelineResult.State.CANCELLED);
return state;
}

private void offerNewState(State newState) {
State oldState = this.state;
this.state = newState;
if (!oldState.isTerminal() && newState.isTerminal() && onTerminalState != null) {
try {
onTerminalState.run();
} catch (Exception e) {
throw unwrapCause(e);
}
cancelSparkJobs.run();
try {
pipelineExecution.get();
} catch (ExecutionException e) {
LOG.info(
"Pipeline execution ended with an exception after cancel: {}",
String.valueOf(Throwables.getRootCause(e).getMessage()));
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
return state;
}
state = State.CANCELLED;
return state;
}
}
62 changes: 46 additions & 16 deletions ...ava/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,13 @@
*/
package org.apache.beam.runners.spark.structuredstreaming;

import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.metrics.MetricsPusher;
import org.apache.beam.runners.core.metrics.NoOpMetricsSink;
import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
Expand All @@ -42,6 +43,7 @@
import org.apache.beam.sdk.util.construction.SplittableParDo;
import org.apache.beam.sdk.util.construction.graph.ProjectionPushdownOptimizer;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.spark.SparkContext;
import org.apache.spark.SparkEnv$;
import org.apache.spark.metrics.MetricsSystem;
import org.apache.spark.sql.SparkSession;
Expand Down Expand Up @@ -145,27 +147,59 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) {

PipelineTranslator.detectStreamingMode(pipeline, options);

final SparkSession sparkSession = SparkSessionFactory.getOrCreateSession(options);
final MetricsAccumulator metrics = MetricsAccumulator.getInstance(sparkSession);
final boolean releaseSession = !options.getUseActiveSparkSession();
final SparkSession sparkSession = SparkSessionFactory.acquire(options);
final SparkContext sc = sparkSession.sparkContext();
final MetricsAccumulator metrics;
try {
metrics = MetricsAccumulator.getInstance(sparkSession);
} catch (RuntimeException e) {
if (releaseSession) {
SparkSessionFactory.release(sparkSession);
}
throw e;
}

// Set once the pipeline is translated, so the result can stop an ongoing (streaming)
// evaluation on cancel. Remains null until translation completes.
// Null until translation completes.
final AtomicReference<EvaluationContext> ctxRef = new AtomicReference<>();
final AtomicBoolean cancelRequested = new AtomicBoolean(false);

final String jobName = options.getJobName();
final String jobGroupId = "beam-" + jobName + "-" + UUID.randomUUID();
final Runnable cancelSparkJobs =
() -> {
try {
if (!sc.isStopped()) {
sc.cancelJobGroup(jobGroupId);
}
} catch (IllegalStateException e) {
// Context stopped concurrently.
}
};

final Future<?> submissionFuture =
runAsync(
() -> {
EvaluationContext ctx = translatePipeline(sparkSession, pipeline);
ctxRef.set(ctx);
ctx.evaluate();
try {
sc.setJobGroup(jobGroupId, "Beam " + jobName, true);
EvaluationContext ctx = translatePipeline(sparkSession, pipeline);
ctxRef.set(ctx);
if (!cancelRequested.get()) {
ctx.evaluate();
}
} finally {
if (!sc.isStopped()) {
sc.clearJobGroup();
}
if (releaseSession) {
SparkSessionFactory.release(sparkSession);
}
}
});

final SparkStructuredStreamingPipelineResult result =
new SparkStructuredStreamingPipelineResult(
submissionFuture,
ctxRef::get,
metrics,
sparkStopFn(sparkSession, options.getUseActiveSparkSession()));
submissionFuture, ctxRef::get, metrics, cancelRequested, cancelSparkJobs);

if (options.getEnableSparkMetricSinks()) {
registerMetricsSource(options.getAppName(), metrics);
Expand Down Expand Up @@ -228,8 +262,4 @@ private static Future<?> runAsync(Runnable task) {
execService.shutdown();
return future;
}

private static @Nullable Runnable sparkStopFn(SparkSession session, boolean isProvided) {
return !isProvided ? () -> session.stop() : null;
}
}
16 changes: 11 additions & 5 deletions ...java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ public interface NamedDataset<T> {

private final Collection<? extends NamedDataset<?>> leaves;
private final SparkSession session;
private volatile boolean stopped = false;

protected EvaluationContext(Collection<? extends NamedDataset<?>> leaves, SparkSession session) {
this.leaves = leaves;
Expand All @@ -63,9 +64,13 @@ protected Collection<? extends NamedDataset<?>> leaves() {
return leaves;
}

/** Trigger evaluation of all leaf datasets. */
/** Trigger evaluation of all leaf datasets. Returns early once {@link #stop()} was called. */
public void evaluate() {
for (NamedDataset<?> ds : leaves) {
if (stopped) {
LOG.info("Evaluation stopped, skipping remaining datasets");
return;
}
final Dataset<?> dataset = ds.dataset();
if (dataset == null) {
continue;
Expand Down Expand Up @@ -119,11 +124,12 @@ public static <T> void evaluate(String name, Dataset<T> ds) {
}

/**
* Stops any ongoing streaming execution triggered by this context.
*
* <p>This is a no-op for batch pipelines.
* Stops the evaluation after the current leaf dataset. Streaming contexts override this to stop
* their queries.
*/
public void stop() {}
public void stop() {
stopped = true;
}

public SparkSession getSparkSession() {
return session;
Expand Down
42 changes: 38 additions & 4 deletions ...va/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import javax.annotation.Nullable;
import org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
Expand Down Expand Up @@ -90,6 +91,7 @@
import org.apache.spark.sql.execution.datasources.v2.DataWritingSparkTaskResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Option;

public class SparkSessionFactory {

Expand All @@ -113,15 +115,47 @@ public class SparkSessionFactory {
"/com.esotericsoftware/kryo-shaded",
"/com/esotericsoftware/kryo-shaded");

// Users per session created here, guarded by the class lock.
private static final Map<SparkSession, Integer> OWNED_SESSIONS = new HashMap<>();

/**
* Gets active {@link SparkSession} or creates one using {@link
* SparkStructuredStreamingPipelineOptions}.
* Returns the {@link SparkSession} for a pipeline, pair with {@link #release}. Without {@code
* useActiveSparkSession} the session is created here unless a usable one exists, only sessions
* created here are stopped on release.
*/
public static SparkSession getOrCreateSession(SparkStructuredStreamingPipelineOptions options) {
public static synchronized SparkSession acquire(SparkStructuredStreamingPipelineOptions options) {
if (options.getUseActiveSparkSession()) {
return SparkSession.active();
}
return sessionBuilder(options.getSparkMaster(), options).getOrCreate();
// Spark 3 also returns stopped sessions.

@Abacn Abacn Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a few comments noting why we need to manage active sessions ourselves now.

@tkaymak tkaymak Sep 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will do: getOrCreate adopts an existing session, a pipeline must not stop one it did not create, and the next pipeline needs the previous one's session gone to get its own config.

boolean noUsableSession =
!isUsable(SparkSession.getActiveSession()) && !isUsable(SparkSession.getDefaultSession());
SparkSession session = sessionBuilder(options.getSparkMaster(), options).getOrCreate();
if (noUsableSession) {
OWNED_SESSIONS.put(session, 1);
} else {
OWNED_SESSIONS.computeIfPresent(session, (unused, count) -> count + 1);
}
return session;
}

/** Releases a session from {@link #acquire}, stops it when it was created here and unused. */
public static synchronized void release(SparkSession session) {
Integer count = OWNED_SESSIONS.get(session);
if (count == null) {
return;
}
if (count > 1) {
OWNED_SESSIONS.put(session, count - 1);
return;
}
OWNED_SESSIONS.remove(session);
LOG.info("Stopping SparkSession created by the runner");
session.stop();
}

private static boolean isUsable(Option<SparkSession> session) {
return session.isDefined() && !session.get().sparkContext().isStopped();
}

/** Creates Spark session builder with some optimizations for local mode, e.g. in tests. */
Expand Down
Loading
Loading