Report where each notebook cell's time goes in Managed Spark

[ajma]

What

Cells that touch the cluster log what they cost. These are real lines from real IPython cells against a live session, not examples:

Cell took 56.08s, 56.08s in Managed Spark (session creation 56.08s)
Cell took 38.34s, 38.33s in Managed Spark across 2 round trips (fetch 34.00s, fetch 4.33s; transport 9.1 KiB down in 36.41s)
Cell took  3.98s,  3.97s in Managed Spark across 2 round trips (fetch 2.53s, fetch 1.44s; transport 4.0 KiB down in 3.97s)
Cell took  1.92s,  1.92s in Managed Spark across 1 round trip  (analyze 1.92s; transport 1.1 KiB down in 1.92s)
Cell took 21.04s, 20.11s in Managed Spark across 1 round trip  (fetch 20.11s; transport 62.1 MiB down in 20.88s, 3.0 MiB/s)

A cell doing pure client-side work logs nothing.

Each round trip is named by kind (execute, fetch, analyze) and timed individually — the three slowest are listed, with +N more beyond that. Note line two: a single collect() is two fetch round trips, which an aggregate count would have hidden.

Why

A throwaway spike measured where the time actually goes before any of this was designed:

	Measured
Cold session creation	56-65s — most of a first cell
A cell reading only .schema	1.9s, previously invisible — _analyze was untimed
Pulling 62 MiB back as pandas	20.1s of a 21.0s cell, at 3.0 MiB/s

Transport was ~95% of the only scenario that hurt, and the bridge hex-encodes (proxy.py), so 62 MiB of Arrow crossed the wire as 124 MiB of hex text. Throughput held at 2.5-3.0 MiB/s across four independent sessions. Without these counters a report like that just looks like "Spark is slow".

Jupyter, Colab and %%time already give you the first number in each line. None of them can give you the rest.

How

execution_timer.py — module-level state, no classes. A list of (kind, seconds) round trips plus three counters, all behind one lock, fed from three places and read by the IPython cell hooks.

• Dropped when no cell is in progress. Outside IPython nothing calls start_cell(), so the state stays bounded for the life of a batch job.
• Clipped to the cell window, so a generator built in an earlier cell and drained in this one can't report time from before the cell began.
• Clamped to the cell duration — both the Managed Spark total and the transport blocked time. Concurrent round trips and concurrent bridge threads each accumulate into one counter. Makes the reported MiB/s a lower bound, the safe direction for a diagnostic.
• register_cell_timing() is idempotent; it's called per session creation and IPython's events.register accepts duplicates silently.

session.py wraps _analyze, _execute and _execute_and_fetch_as_iterator, and feeds session creation. proxy.py counts payload bytes and blocked time in bridged_socket.recv/send.

Three bugs found by running it for real

The first cell reported nothing — the one paying the 56s. IPython fires pre_run_cell for cell 1 before anything has registered a handler, so no cell was marked started and record_session_creation hit the drop-when-no-cell rule. Every unit test passed while the feature was inert in its most important case. Fixed by registering at the top of getOrCreate() and having registration start a cell when none is in progress. Caveat: cell 1's cell_seconds is then measured from registration rather than true cell start, so it slightly under-reports its own wall time.

transport 0.0 MiB down in 43.24s, 0.0 MiB/s — sub-MiB control traffic printed as two zeros and read as broken. Now reports KiB and drops the meaningless rate, keeping the real signal that 43s went by blocked while moving almost nothing (a cold cluster).

0.00s in Dataproc Serverless across 0 operations — the original first-cell line was self-contradictory: it claimed zero time on a service where it had just spent a minute. Session creation now counts toward the headline, "operations" became "round trips", and the count is omitted rather than printed as zero.

The subtle bit

_execute_and_fetch_as_iterator is a generator function — calling it returns a generator without sending anything, so timing the call records ~0s for every df.collect(). The wrapper returns a wrapping generator, and because _timed_iterator is itself a generator, start isn't read until the first next(). There's a comment saying the generator-ness is load-bearing, and a regression test asserting zero round trips are recorded when the wrapper is called but not consumed.

time.monotonic() throughout, deviating from time.time() elsewhere in session.py — those are wall-clock deadlines, these are durations.

Decisions worth a second opinion

• Analyze is a kind of round trip, not a separate concept. record() takes a kind; there's no record_analyze(). A schema-only cell must count as a round trip, or it stays silent and the blind spot survives the fix.
• The logger.debug for session creation was promoted to logger.info. basicConfig(level=logging.INFO) was suppressing a 56-second event, and outside IPython the cell line never fires, so info is the only way a batch user learns of it.
• Nothing is emitted outside IPython. A "cell" doesn't exist in a script or batch job.
• Per-stage num_bytes_read was considered and dropped. The spike read 0.0 MiB on every stage — but every scenario used spark.range(), which reads no files, so that's untested rather than disproven.

Testing

pytest tests/unit → 227 passed, 13 subtests passed, 1 failed. The failure is pre-existing and environmental: test_create_session_without_application_default_credentials fails on any machine where ADC are configured. pyink --check clean.

Plus four live runs against real sessions, which is how all three bugs above were found.

[gemini-code-assist]

Code Review

This pull request introduces per-cell execution timing for Dataproc Serverless operations in IPython notebooks. It adds an ExecutionTimer class to track and merge round-trip intervals, registers IPython pre- and post-run cell hooks to log timing summaries, and wraps execution methods in ManagedSparkSession to record active durations. Comprehensive unit tests are also added. The feedback suggests improving _timed_iterator to measure only the active waiting time of individual next() calls on the underlying iterator, preventing client-side processing from inflating the reported Dataproc Serverless execution time.

[gemini-code-assist]

The current implementation of _timed_iterator measures the entire duration from the first next() call until the generator is fully consumed or closed. However, because this is a generator, any client-side processing (such as iterating over the rows, performing local computations, or plotting) that occurs between pulling chunks from the iterator will be included in the measured duration. This inflates the reported Dataproc Serverless time with client-side execution time, defeating the purpose of separating the two.

To accurately measure only the time spent waiting on the server, we should time each individual next() call on the underlying iterator and accumulate the active waiting time, then record a single interval representing the total active time.

Suggested change

	def _timed_iterator(iterator, timer):
	# This being a generator is load-bearing: `start` is not read until
	# the first next(), so the clock runs from when the caller begins
	# driving the RPC rather than from when the generator was built.
	# Hoisting it out of the generator body would time every fetch as 0s.
	start = time.monotonic()
	try:
	yield from iterator
	finally:
	timer.record(start, time.monotonic())
	def _timed_iterator(iterator, timer):
	# This being a generator is load-bearing: `start` is not read until
	# the first next(), so the clock runs from when the caller begins
	# driving the RPC rather than from when the generator was built.
	# Hoisting it out of the generator body would time every fetch as 0s.
	active_time = 0.0
	try:
	while True:
	start = time.monotonic()
	try:
	val = next(iterator)
	except StopIteration:
	break
	finally:
	active_time += time.monotonic() - start
	yield val
	finally:
	end = time.monotonic()
	timer.record(end - active_time, end)