Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
152 changes: 128 additions & 24 deletions packages/databricks-vscode/src/test/startCluster.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,12 @@ describe(__filename, function () {
): compute.ClusterDetails =>
({cluster_id: clusterId, state, ...extra}) as compute.ClusterDetails;

const placementFailure = () =>
details("TERMINATED", {
state_message:
"Unexpected failure during launch. databricks_error_message: Timeout while placing nodes.",
});

const whenGet = () =>
when(
mockedClient.request(
Expand Down Expand Up @@ -100,24 +106,6 @@ describe(__filename, function () {
verifyStarted(1);
});

it("fails fast when the cluster returns to a terminal state after start", async () => {
whenGet().thenResolve(
details("TERMINATED"),
details("TERMINATED", {state_message: "bad spark config"})
);
whenStart().thenResolve({});

const startPromise = startCluster(instance(mockedClient), clusterId);
const rejection = assert.rejects(
startPromise,
(e: Error) =>
e instanceof ClusterStartError &&
/bad spark config/.test(e.message)
);
await fakeTimer.runToLastAsync();
await rejection;
});

it("tolerates a concurrent start race on the shared cluster", async () => {
// Initial TERMINATED -> our start() races a sibling and throws -> the
// re-check finds it already coming up (PENDING) -> RUNNING.
Expand All @@ -138,9 +126,10 @@ describe(__filename, function () {
verifyStarted(1);
});

it("propagates a non-race start error when the cluster stays stopped", async () => {
it("fails fast when start() is rejected and the cluster stays stopped", async () => {
// start() fails and the re-check shows the cluster still stopped, so the
// original (actionable) error surfaces rather than being masked.
// original (actionable) error surfaces immediately rather than being
// masked or retried — retrying the same rejected call can't help.
whenGet().thenResolve(details("TERMINATED"), details("TERMINATED"));
whenStart().thenReject(new Error("permission denied"));

Expand All @@ -150,6 +139,8 @@ describe(__filename, function () {
);
await fakeTimer.runToLastAsync();
await rejection;

verifyStarted(1);
});

it("waits for a TERMINATING cluster to stop, then starts it", async () => {
Expand All @@ -167,10 +158,28 @@ describe(__filename, function () {
verifyStarted(1);
});

it("fails fast when the cluster is UNKNOWN after start", async () => {
it("re-starts after a transient post-start terminal failure, then reaches RUNNING", async () => {
whenGet().thenResolve(
details("TERMINATED"),
details("UNKNOWN", {state_message: "lost the cluster"})
placementFailure(),
details("PENDING"),
details("RUNNING")
);
whenStart().thenResolve({});

const startPromise = startCluster(instance(mockedClient), clusterId);
await fakeTimer.runAllAsync();
await startPromise;

// start() issued once per attempt: the failed one plus the recovery.
verifyStarted(2);
});

it("retries a persistent terminal failure up to the max, then surfaces its reason", async () => {
// Retry is not gated on the reason string, so even a deterministic
// failure is re-attempted (best-effort warm-up) before surfacing.
whenGet().thenResolve(
details("TERMINATED", {state_message: "bad spark config"})
);
whenStart().thenResolve({});

Expand All @@ -179,9 +188,104 @@ describe(__filename, function () {
startPromise,
(e: Error) =>
e instanceof ClusterStartError &&
/lost the cluster/.test(e.message)
/bad spark config/.test(e.message)
);
await fakeTimer.runToLastAsync();
await fakeTimer.runAllAsync();
await rejection;

verifyStarted(3);
});

it("stops re-starting once the deadline is exhausted by backoff", async () => {
// A short timeout: after the first failure the jittered backoff (>=10s)
// is capped to the remaining deadline and consumes it, so the loop must
// not issue a second start().
whenGet().thenResolve(placementFailure());
whenStart().thenResolve({});

const startPromise = startCluster(
instance(mockedClient),
clusterId,
new Time(5, TimeUnits.seconds)
);
const rejection = assert.rejects(
startPromise,
(e: Error) =>
e instanceof ClusterStartError &&
/Timeout while placing nodes/.test(e.message)
);
await fakeTimer.runAllAsync();
await rejection;

verifyStarted(1);
});

it("re-checks through a start race on a retry, then reaches RUNNING", async () => {
// A post-start terminal triggers a retry; the retry's start() races a
// sibling and throws, the re-check finds it coming up (PENDING), and the
// poll reaches RUNNING.
whenGet().thenResolve(
details("TERMINATED"),
placementFailure(),
details("PENDING"),
details("RUNNING")
);
whenStart()
.thenResolve({})
.thenReject(
new Error(
`Cluster ${clusterId} is in unexpected state Pending.`
)
);

const startPromise = startCluster(instance(mockedClient), clusterId);
await fakeTimer.runAllAsync();
await startPromise;

verifyStarted(2);
});

it("surfaces a stuck-PENDING poll timeout as-is, without retrying", async () => {
// A cluster that never leaves PENDING is out of scope: the poll times
// out (a non-ClusterStartError), which must propagate without a re-start.
whenGet().thenResolve(details("TERMINATED"), details("PENDING"));
whenStart().thenResolve({});

const startPromise = startCluster(
instance(mockedClient),
clusterId,
new Time(30, TimeUnits.seconds)
);
const rejection = assert.rejects(
startPromise,
(e: Error) => !(e instanceof ClusterStartError)
);
await fakeTimer.runAllAsync();
await rejection;

// One start(), then the poll timed out — no retry.
verifyStarted(1);
});

it("rejects without starting when the deadline is already spent on entry", async () => {
// Degenerate timeout: the loop's first guard trips before any start(),
// so the fallback error surfaces rather than throwing an undefined.
whenGet().thenResolve(details("TERMINATED"));

const startPromise = startCluster(
instance(mockedClient),
clusterId,
new Time(0, TimeUnits.milliseconds)
);
const rejection = assert.rejects(
startPromise,
(e: Error) =>
e instanceof ClusterStartError &&
/did not reach RUNNING/.test(e.message)
);
await fakeTimer.runAllAsync();
await rejection;

verifyStarted(0);
});
});
158 changes: 100 additions & 58 deletions packages/databricks-vscode/src/test/startCluster.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,29 @@ export class ClusterStartError extends Error {}
const DEFAULT_START_TIMEOUT = new Time(60, TimeUnits.minutes);
const POLL_INTERVAL = new Time(10, TimeUnits.seconds);

// Warm-up is best-effort: the cold-started shared cluster can fail placement
// transiently, so retry start() a bounded number of times before surfacing.
const MAX_START_ATTEMPTS = 3;
const START_RETRY_BACKOFF = new Time(20, TimeUnits.seconds);

function isTerminal(state?: compute.State): boolean {
return state === "TERMINATED" || state === "ERROR" || state === "UNKNOWN";
}

// Include both fields when present so a machine-readable code (e.g.
// CLOUD_PROVIDER_LAUNCH_FAILURE) isn't lost when state_message is generic or
// absent.
function terminationReason(cluster: compute.ClusterDetails): string {
const parts: string[] = [];
if (cluster.state_message) {
parts.push(cluster.state_message);
}
if (cluster.termination_reason) {
parts.push(JSON.stringify(cluster.termination_reason));
}
return parts.length > 0 ? parts.join(" ") : "unknown reason";
}

export async function startCluster(
client: ApiClient,
clusterId: string,
Expand All @@ -31,9 +54,9 @@ export async function startCluster(
}`
);

// One deadline across both the shutdown wait and the start poll, so a slow
// TERMINATING phase can't hand the poll a fresh full timeout and let the
// total exceed the caller's bound.
// One deadline across the shutdown wait, all start attempts, and their
// polls, so a slow phase can't hand a later step a fresh full timeout and
// let the total exceed the caller's bound.
const deadline = Date.now() + timeout.toMillSeconds().value;
const remaining = () =>
new Time(Math.max(0, deadline - Date.now()), TimeUnits.milliseconds);
Expand All @@ -59,66 +82,85 @@ export async function startCluster(
});
}

if (
cluster.state === "TERMINATED" ||
cluster.state === "ERROR" ||
cluster.state === "UNKNOWN"
) {
let lastError: unknown;
for (let attempt = 1; attempt <= MAX_START_ATTEMPTS; attempt++) {
// Backoff may have consumed the remaining time; don't launch past the
// deadline.
if (remaining().toMillSeconds().value <= 0) {
break;
}

if (isTerminal(cluster.state)) {
try {
await clusterApi.start({cluster_id: clusterId});
} catch (e) {
// Shared cluster: a sibling may have raced this start() into an
// error. If it's coming up we just raced (poll below); if still
// terminal the start was genuinely rejected, so surface it.
cluster = await clusterApi.get({cluster_id: clusterId});
log(cluster);
if (isTerminal(cluster.state)) {
throw e;
}
}
}

// On this cold-started shared cluster a post-start terminal state is
// usually a transient placement failure, so retry a bounded number of
// times rather than failing the whole shard.
try {
await clusterApi.start({cluster_id: clusterId});
await retry<void>({
timeout: remaining(),
retryPolicy: new retries.LinearRetryPolicy(POLL_INTERVAL),
fn: async () => {
cluster = await clusterApi.get({cluster_id: clusterId});
log(cluster);
switch (cluster.state) {
case "RUNNING":
return;
case "TERMINATED":
case "ERROR":
case "UNKNOWN":
throw new ClusterStartError(
`Cluster ${clusterId} failed to start (${
cluster.state
}): ${terminationReason(cluster)}`
);
default:
throw new retries.RetriableError();
}
},
});
return;
} catch (e) {
// The cluster is shared across ~40 shards, so a sibling may have
// already started it, racing this call into an error. Re-check: if
// it's now coming up we merely raced, so poll below; if it's still
// stopped the start genuinely failed (auth, permissions, bad
// request), so surface that actionable error rather than mask it.
cluster = await clusterApi.get({cluster_id: clusterId});
log(cluster);
if (
cluster.state === "TERMINATED" ||
cluster.state === "ERROR" ||
cluster.state === "UNKNOWN"
) {
// Only a post-start terminal (ClusterStartError) is retryable; a
// stuck-PENDING poll timeout and the like are not.
if (!(e instanceof ClusterStartError)) {
throw e;
}
lastError = e;
if (attempt >= MAX_START_ATTEMPTS) {
break;
}
// Jittered backoff so ~40 shards don't re-issue start() in lockstep
// (thundering herd). Capped by the remaining deadline.
const backoffMs = Math.min(
Math.round(
START_RETRY_BACKOFF.toMillSeconds().value *
(0.5 + Math.random())
),
remaining().toMillSeconds().value
);
if (backoffMs > 0) {
await new Promise((resolve) => setTimeout(resolve, backoffMs));
}
Comment on lines +147 to +156

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.

The cap here is doing two different jobs, and it's only right for one of them:

backoffMs = min(jitter, remaining)
              |        |
              |        +-- "don't sleep past the deadline"  <- wrong response
              +-- "space out ~40 shards"

When remaining < jitter there isn't room for another attempt, so the loop sleeps out the clock and then discovers that.

The timeline from your own stops re-starting once the deadline is exhausted by backoff test (timeout = 5s, jitter in [10s, 30s]):

t=0ms     get() -> TERMINATED
          attempt 1: remaining 5000 > 0  OK
          start()  ->  poll  ->  TERMINATED  ->  ClusterStartError
          attempt 1 < MAX, so back off
          backoffMs = min(  ~18000,  5000  )  =  5000
                            ^jitter  ^remaining  <- already knows a retry is impossible
t=0ms     sleep(5000) ---------------------+
                                           |  nothing happens
t=5000ms  <--------------------------------+
          attempt 2: remaining 0  ->  break
          throw lastError

The sleep is decided after the value that makes it pointless is already in hand.

The cap can go away entirely, because the guard subsumes it: if there's room to sleep, sleeping the full jitter is safe; if there isn't, the answer is break. That also drops the now-dead backoffMs > 0 check.

Suggested change
const backoffMs = Math.min(
Math.round(
START_RETRY_BACKOFF.toMillSeconds().value *
(0.5 + Math.random())
),
remaining().toMillSeconds().value
);
if (backoffMs > 0) {
await new Promise((resolve) => setTimeout(resolve, backoffMs));
}
const backoffMs = Math.round(
START_RETRY_BACKOFF.toMillSeconds().value *
(0.5 + Math.random())
);
// No room to back off and still launch, so stop rather than
// sleeping out the remaining deadline.
if (backoffMs >= remaining().toMillSeconds().value) {
break;
}
await new Promise((resolve) => setTimeout(resolve, backoffMs));

This is free with respect to the existing test, which asserts verifyStarted(1) and a /Timeout while placing nodes/ message. Both hold either way:

verifyStarted(1) error message wall clock
current (sleep 5s) pass pass 5000ms
with early break pass pass 0ms

So the test never pinned the sleep, only the clock did.

One knock-on worth noting: after this change the remaining() <= 0 guard at the top of the loop is no longer reachable via the backoff path. It stays useful only for a caller-supplied already-expired deadline, which your 0 ms test covers.

Worst case in production is one jitter window (~30s) per failing shard, so this is a nit rather than a blocker.

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.

Thanks Anton — agreed, this is a clean simplification (break instead of sleeping out the deadline, and the backoffMs > 0 check does become dead code). Since you flagged it non-blocking and it's test-neutral, I'm intentionally leaving it as-is here to avoid re-spinning the ~45-min integration run and churning the approval — noting it as a small follow-up. Appreciate the careful read on the shared-deadline invariant.

}
}

// Poll to RUNNING under one deadline. A terminal state here is a real launch
// failure (the start was already issued), so fail fast with the cloud-side
// reason instead of burning the whole timeout.
await retry<void>({
timeout: remaining(),
retryPolicy: new retries.LinearRetryPolicy(POLL_INTERVAL),
fn: async () => {
cluster = await clusterApi.get({cluster_id: clusterId});
log(cluster);
switch (cluster.state) {
case "RUNNING":
return;
case "TERMINATED":
case "ERROR":
case "UNKNOWN": {
// state_message is a string; termination_reason is an
// object — stringify only the latter so a plain message
// isn't wrapped in quotes.
const reason =
cluster.state_message ??
cluster.termination_reason ??
"unknown reason";
throw new ClusterStartError(
`Cluster ${clusterId} failed to start (${
cluster.state
}): ${
typeof reason === "string"
? reason
: JSON.stringify(reason)
}`
);
}
default:
throw new retries.RetriableError();
}
},
});
throw (
lastError ??
new ClusterStartError(
`Cluster ${clusterId} did not reach RUNNING within the timeout`
)
);
}
Loading