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
62 changes: 55 additions & 7 deletions apps/sim/lib/execution/remote-sandbox/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,10 @@ import {
SANDBOX_OUTPUT_DIR_MAX_DEPTH,
SANDBOX_OUTPUT_DIR_SENTINEL,
} from '@/lib/execution/remote-sandbox/paths'
import {
createSandboxPhaseTimer,
type SandboxPhaseTimer,
} from '@/lib/execution/remote-sandbox/phase-timing'
import { resolvePiSandboxLifetimeMs } from '@/lib/execution/remote-sandbox/pi-lifetime'
import { resolveProvider } from '@/lib/execution/remote-sandbox/provider'
import {
Expand Down Expand Up @@ -116,13 +120,16 @@ async function leaseSandbox(
selected: ResolvedSandbox | null,
signal: AbortSignal,
meterUsage: boolean | undefined,
session: SandboxSessionRequest | undefined
session: SandboxSessionRequest | undefined,
phases: SandboxPhaseTimer
): Promise<SandboxLease> {
const provider = resolveProvider()
const sessionCapable = Boolean(session && !meterUsage && provider.findSessionSandbox)

if (session && sessionCapable) {
const endAccess = phases.start('access')
await prepareSandboxSessionAccess(session.key, signal)
endAccess()
return withSandboxSessionLock(session.key, signal, async (leaseSignal) => {
const { created, status } = await ensureSessionSandbox({
provider,
Expand All @@ -132,6 +139,7 @@ async function leaseSandbox(
session,
signal: leaseSignal,
bootstrapTimeoutMs: remainingSandboxBudgetMs(signal),
phases,
})
return {
created,
Expand Down Expand Up @@ -214,6 +222,16 @@ function raceSandboxAbort<T>(operation: Promise<T>, signal: AbortSignal): Promis
})
}

/** Attaches the phase breakdown to whichever result the execution returns, after its release. */
async function withPhaseTimings(
execute: (phases: SandboxPhaseTimer) => Promise<SandboxExecutionResult>
): Promise<SandboxExecutionResult> {
const phases = createSandboxPhaseTimer()
const result = await execute(phases)
result.phaseMs = phases.timings
return result
}

async function withSandboxExecutionBudget<T>(
timeoutMs: number,
parentSignal: AbortSignal | undefined,
Expand Down Expand Up @@ -864,22 +882,26 @@ async function acceptSessionOutputHistory(
async function executeInSandboxWithinBudget(
// The budget wrapper always injects the signal; the required-signal type states that
// invariant instead of a cast hiding it.
req: SandboxExecutionRequest & { signal: AbortSignal }
req: SandboxExecutionRequest & { signal: AbortSignal },
phases: SandboxPhaseTimer
): Promise<SandboxExecutionResult> {
const { code, language, signal } = req
const kind = req.sandboxKind ?? 'code'
throwIfAborted(signal)

// Resolved before the sandbox is created so a selection that cannot be honored
// fails without spending a provider create.
const endResolve = phases.start('resolve')
const selected = await resolveWorkspaceSandbox({
kind,
language,
workspaceId: req.workspaceId,
sandboxId: req.sandboxId,
})
endResolve()
throwIfAborted(signal)

const endLease = phases.start('lease')
const lease = await leaseSandbox(
kind,
{
Expand All @@ -890,8 +912,10 @@ async function executeInSandboxWithinBudget(
selected,
signal,
req.meterUsage,
req.session
req.session,
phases
)
endLease()
const created = lease.created
const sandbox = created.sandbox
const sandboxId = sandbox.sandboxId
Expand All @@ -906,6 +930,7 @@ async function executeInSandboxWithinBudget(
// the finally below. Dependencies land before the inputs so user code and its
// mounts always see a complete environment.
//
const endPrepare = phases.start('prepare')
if (lease.session && req.session)
await recordSessionFileInput(
req.session.key,
Expand All @@ -926,6 +951,7 @@ async function executeInSandboxWithinBudget(
signal,
Boolean(lease.session)
)
endPrepare()
const executionEnvironment = {
...selected?.envs,
...req.session?.envs,
Expand All @@ -944,13 +970,15 @@ async function executeInSandboxWithinBudget(

let execution: SandboxCodeResult
try {
const endRun = phases.start('run')
execution = await sandbox.runCode(code, {
timeoutMs: codeBudgetMs(req, signal),
javascriptPreload: buildJavaScriptRuntimeBindingsSource(req.runtimeBindings ?? []),
maxOutputBytes: MAX_SANDBOX_PROCESS_OUTPUT_BYTES,
signal,
...(hasExecutionEnvironment ? { envs: executionEnvironment } : {}),
})
endRun()
} catch (error) {
throwIfAborted(signal)
throw error
Expand Down Expand Up @@ -1009,11 +1037,13 @@ async function executeInSandboxWithinBudget(
...sessionField,
}
try {
const endCollect = phases.start('collect')
const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles(
sandbox,
req,
{ signal }
)
endCollect()
throwIfAborted(signal)
billableResult.exportedFileContent = exportedFileContent
billableResult.exportedFiles = exportedFiles
Expand Down Expand Up @@ -1044,6 +1074,7 @@ async function executeInSandboxWithinBudget(
if (cost && billableOutputError) {
attachTrustedSandboxOutputCost(billableOutputError, cost)
}
const endRelease = phases.start('release')
try {
await privateInputFiles?.cleanup()
await lease.release()
Expand All @@ -1052,6 +1083,7 @@ async function executeInSandboxWithinBudget(
providerId: created.providerId,
sandboxId,
})
endRelease()
}
}
}
Expand All @@ -1061,34 +1093,41 @@ export function executeInSandbox(req: SandboxExecutionRequest): Promise<SandboxE
req.timeoutMs,
req.signal,
Boolean(req.session && !req.meterUsage),
(signal) => executeInSandboxWithinBudget({ ...req, signal })
(signal) =>
withPhaseTimings((phases) => executeInSandboxWithinBudget({ ...req, signal }, phases))
)
}

async function executeShellInSandboxWithinBudget(
req: SandboxShellExecutionRequest & { signal: AbortSignal }
req: SandboxShellExecutionRequest & { signal: AbortSignal },
phases: SandboxPhaseTimer
): Promise<SandboxExecutionResult> {
const { code, envs, signal } = req
const kind = req.sandboxKind ?? 'shell'
throwIfAborted(signal)

// No language is passed: a shell execution runs commands rather than a language
// runtime, so whichever language the sandbox carries is the one it installs.
const endResolve = phases.start('resolve')
const selected = await resolveWorkspaceSandbox({
kind,
workspaceId: req.workspaceId,
sandboxId: req.sandboxId,
})
endResolve()
throwIfAborted(signal)

const endLease = phases.start('lease')
const lease = await leaseSandbox(
kind,
{ imageRef: selected?.imageRef, lifetimeMs: remainingSandboxBudgetMs(signal) },
selected,
signal,
req.meterUsage,
req.session
req.session,
phases
)
endLease()
const created = lease.created
const sandbox = created.sandbox
const sandboxId = sandbox.sandboxId
Expand All @@ -1102,6 +1141,7 @@ async function executeShellInSandboxWithinBudget(
// Inside the try so a failed install or mount still releases the sandbox via
// the finally below. The install shares the caller's budget rather than adding
// to it — see the note in `executeInSandbox`.
const endPrepare = phases.start('prepare')
if (lease.session && req.session)
await recordSessionFileInput(
req.session.key,
Expand All @@ -1123,9 +1163,11 @@ async function executeShellInSandboxWithinBudget(
signal,
Boolean(lease.session)
)
endPrepare()

let result: SandboxCommandResult
try {
const endRun = phases.start('run')
result = await sandbox.runCommand(code, {
envs: {
...selected?.envs,
Expand All @@ -1141,6 +1183,7 @@ async function executeShellInSandboxWithinBudget(
rootUser: !lease.session,
atMostOnce: true,
})
endRun()
} catch (error) {
throwIfAborted(signal)
throw error
Expand Down Expand Up @@ -1185,11 +1228,13 @@ async function executeShellInSandboxWithinBudget(
...sessionField,
}
try {
const endCollect = phases.start('collect')
const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles(
sandbox,
req,
{ signal }
)
endCollect()
throwIfAborted(signal)
billableResult.exportedFileContent = exportedFileContent
billableResult.exportedFiles = exportedFiles
Expand Down Expand Up @@ -1220,6 +1265,7 @@ async function executeShellInSandboxWithinBudget(
if (cost && billableOutputError) {
attachTrustedSandboxOutputCost(billableOutputError, cost)
}
const endRelease = phases.start('release')
try {
await privateInputFiles?.cleanup()
await lease.release()
Expand All @@ -1228,6 +1274,7 @@ async function executeShellInSandboxWithinBudget(
providerId: created.providerId,
sandboxId,
})
endRelease()
}
}
}
Expand All @@ -1239,7 +1286,8 @@ export function executeShellInSandbox(
req.timeoutMs,
req.signal,
Boolean(req.session && !req.meterUsage),
(signal) => executeShellInSandboxWithinBudget({ ...req, signal })
(signal) =>
withPhaseTimings((phases) => executeShellInSandboxWithinBudget({ ...req, signal }, phases))
)
}

Expand Down
39 changes: 39 additions & 0 deletions apps/sim/lib/execution/remote-sandbox/phase-timing.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
/**
* Phases of one sandbox execution, in the order they run. `lease` spans the whole acquisition
* (session lock included); for a session, `lookup` and `cli` are the parts of it spent finding or
* creating the machine and verifying its deployment tooling.
*/
type SandboxPhase =
| 'resolve'
| 'access'
| 'lease'
| 'lookup'
| 'cli'
| 'prepare'
| 'run'
| 'collect'
| 'release'

/** Wall-clock milliseconds per phase; a phase that did not finish is absent. */
export type SandboxPhaseTimings = Partial<Record<SandboxPhase, number>>

/** Accumulates per-phase wall time for one execution's log line. */
export interface SandboxPhaseTimer {
readonly timings: SandboxPhaseTimings
/** Starts timing a phase; the returned function records its duration. */
start(phase: SandboxPhase): () => void
}

/** Creates a timer whose callbacks record rounded milliseconds; unfinished phases stay absent. */
export function createSandboxPhaseTimer(): SandboxPhaseTimer {
Comment thread
waleedlatif1 marked this conversation as resolved.
Resolved
const timings: SandboxPhaseTimings = {}
return {
timings,
start(phase) {
const startedAt = performance.now()
return () => {
timings[phase] = Math.round(performance.now() - startedAt)
}
},
}
}
40 changes: 40 additions & 0 deletions apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -821,6 +821,46 @@ describe('session sandbox lease', () => {
expect(writes).toHaveBeenCalledTimes(firstWrites)
})

it('attributes a warm shell call’s wall time to the phase that spent it', async () => {
const { handle } = fakeSandbox('phase-timing')
mockFindSessionSandbox.mockResolvedValue(handle)
const request = {
code: 'true',
envs: {},
timeoutMs: 10_000,
session: { key: 'phase-timing', cli: CLI },
}
await executeShellInSandbox(request)

let nowMs = 0
const clock = vi.spyOn(performance, 'now').mockImplementation(() => nowMs)
const writes = vi.fn(handle.writeFile.bind(handle))
handle.writeFile = writes
const runCommand = handle.runCommand.bind(handle)
handle.runCommand = async (command, options) => {
nowMs += options.envs?.SIM_CLI_PATH ? 30 : 60
return runCommand(command, options)
Comment thread
waleedlatif1 marked this conversation as resolved.
Resolved
}
mockFindSessionSandbox.mockImplementation(async () => {
nowMs += 40
return handle
})

try {
const result = await executeShellInSandbox(request)

const phases = result.phaseMs ?? {}
expect(result.sandboxSession).toBe('reused')
expect(writes).not.toHaveBeenCalled()
expect(phases.lookup).toBe(40)
expect(phases.cli).toBe(30)
expect(phases.run).toBe(60)
expect(phases.lease).toBeGreaterThanOrEqual(70)
} finally {
clock.mockRestore()
}
})

it('pins shell CLI resolution ahead of caller environment paths', async () => {
const { handle, calls } = fakeSandbox('shell-cli-path')
mockFindSessionSandbox.mockResolvedValue(handle)
Expand Down
6 changes: 6 additions & 0 deletions apps/sim/lib/execution/remote-sandbox/session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { getErrorMessage } from '@sim/utils/errors'
import { createSandboxPricing } from '@/lib/billing/sandbox-pricing'
import { mapSettledWithConcurrency } from '@/lib/core/utils/concurrency'
import { type CreatedSandbox, createSelectedSandbox } from '@/lib/execution/remote-sandbox/create'
import type { SandboxPhaseTimer } from '@/lib/execution/remote-sandbox/phase-timing'
import { resolveProvider } from '@/lib/execution/remote-sandbox/provider'
import type { ResolvedSandbox } from '@/lib/execution/remote-sandbox/resolve'
import { ensureSessionCli } from '@/lib/execution/remote-sandbox/session-cli'
Expand Down Expand Up @@ -56,12 +57,14 @@ export async function ensureSessionSandbox(args: {
session: SandboxSessionRequest
signal: AbortSignal
bootstrapTimeoutMs: number
phases?: SandboxPhaseTimer
}): Promise<SessionAcquisition> {
const { provider, kind, options, selected, session, signal } = args
signal.throwIfAborted()
if (!provider.findSessionSandbox) throw new Error('This deployment has no persistent workbench')
const lifetimeMs = SESSION_SANDBOX_IDLE_MS + (options.lifetimeMs ?? 0)
const requestedAtMs = Date.now()
const endLookup = args.phases?.start('lookup')
const existing = await provider.findSessionSandbox(session.key, {
...(options.language ? { language: options.language } : {}),
lifetimeMs,
Expand All @@ -83,6 +86,7 @@ export async function ensureSessionSandbox(args: {
true,
provider
)
endLookup?.()
// Before the abort check: a stopped caller still leaves this machine for the next call to
// reuse, and a reused machine without history can never be certified.
if (!existing)
Expand All @@ -96,7 +100,9 @@ export async function ensureSessionSandbox(args: {
await created.sandbox.extendLifetime?.(lifetimeMs)
signal.throwIfAborted()
if (session.cli) {
const endCli = args.phases?.start('cli')
await ensureSessionCli(created.sandbox, session.cli, signal, args.bootstrapTimeoutMs)
endCli?.()
}
signal.throwIfAborted()
return { created, status: existing ? 'reused' : 'created' }
Expand Down
Loading
Loading