Yicong-Huang commented on code in PR #59246:
URL: https://github.com/apache/spark/pull/59246#discussion_r4192113625
##########
python/pyspark/worker.py:
##########
@@ -128,21 +128,88 @@
)
-def report_metrics(
- outfile, boot, init, finish, execution_duration_ms, memory_bytes_spilled,
disk_bytes_spilled
-):
+class _WorkerTimer:
+ """Measure one named duration and add each interval to the collector's
task total.
+
+ The collector owns the duration dictionary. This scope keeps a reference
to it and a
+ start time for its current with block. Reuse the scope across sequential
batches.
+ """
+
+ def __init__(self, shared_duration_totals_ns: dict[str, int], metric_name:
str) -> None:
+ """Bind one metric name to the collector's shared duration totals."""
+ # Keep the collector's dictionary by reference so updates reach its
task totals.
+ self._shared_duration_totals_ns = shared_duration_totals_ns
+ self._metric_name = metric_name
+ self._block_start_ns: Optional[int] = None
+
+ def __enter__(self) -> None:
+ """Start the clock when Python enters the with block."""
+ # Re-entering this instance would overwrite an unfinished
measurement's start time.
+ if self._block_start_ns is not None:
+ raise RuntimeError("Worker timer is already running")
+ self._block_start_ns = time.perf_counter_ns()
+
+ def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+ """Accumulate elapsed time when Python leaves the block, including on
an exception."""
+ assert self._block_start_ns is not None
+ elapsed_ns = time.perf_counter_ns() - self._block_start_ns
+ self._block_start_ns = None
+ self._shared_duration_totals_ns[self._metric_name] += elapsed_ns
+ # Python supplies the exception arguments; returning None lets the
error propagate.
+
+
+class WorkerMetrics:
+ """Collect values and duration totals for one worker task.
+
+ set/add store values in their reporting units: counts, byte counts, or
epoch milliseconds.
+ measure creates timer scopes that accumulate nanoseconds in a separate
shared dictionary.
+ to_report combines both dictionaries, converting only durations to
milliseconds.
+ """
+
+ def __init__(self) -> None:
+ """Create fresh task state so a reused worker does not retain the
previous totals."""
+ # Counters, spill byte counts, and epoch timestamps need no conversion
at report time.
+ self._values_in_report_units: dict[str, int] = {}
+ # These totals belong to the collector; timer scopes share and update
them.
+ self._duration_totals_ns: dict[str, int] = {}
+
+ def set(self, name: str, value: int) -> None:
+ """Store a value in its reporting unit, replacing any earlier value.
+
+ Used for initial counter values, epoch timestamps, and final spill
byte counts.
+ Elapsed durations are accumulated through measure instead.
+ """
+ self._values_in_report_units[name] = value
+
+ def add(self, name: str, value: int = 1) -> None:
+ """Add to a counter in its reporting unit, starting from zero.
+
+ The default increment is one, as used for the number of timed batches.
+ """
+ self._values_in_report_units[name] =
self._values_in_report_units.get(name, 0) + value
+
+ def measure(self, name: str) -> _WorkerTimer:
+ """Create a named timing scope; entering it starts the clock."""
+ # Register zero so an empty supported task can still report this
timing.
+ self._duration_totals_ns.setdefault(name, 0)
+ return _WorkerTimer(self._duration_totals_ns, name)
+
+ def to_report(self) -> dict[str, int]:
Review Comment:
`to_dict` would be better? report is a noun and has a semantic. you can do
`report_metrics(metrics.to_dict())`
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
