feat(draft): add report= argument for uproot.dask; trigger report collection (take 2!)

pyproject.toml

@@ -51,7 +51,7 @@ requires-python = ">=3.8"
 [project.optional-dependencies]
 dev = [
  "boost_histogram>=0.13",
- "dask-awkward>=2023.10.0",
+ "dask-awkward>=2023.12.1",
  "dask[array]",
  "hist>=1.2",
  "pandas",

src/uproot/_dask.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import functools
 import math
+import socket
+import time
 from collections.abc import Callable, Iterable, Mapping
 
 try:
@@ -38,6 +41,7 @@ def dask(
  allow_missing=False,
  open_files=True,
  form_mapping=None,
+ allow_read_errors_with_report=False,
  **options,
 ):
  """
@@ -91,6 +95,10 @@ def dask(
  form_mapping (Callable[awkward.forms.Form] -> awkward.forms.Form | None): If not none
  and library="ak" then apply this remapping function to the awkward form of the input
  data. The form keys of the desired form should be available data in the input form.
+ allow_read_errors_with_report (bool): If True, catch OSError exceptions and return an
+ empty array for these nodes in the task graph. The return of this function then
+ becomes a two element tuple, where the first return is the dask-awkward collection
+ of interest and the second return is a report dask-awkward collection.
  options: See below.
 
  Returns dask equivalents of the backends supported by uproot. If ``library='np'``,
@@ -259,6 +267,7 @@ def dask(
  interp_options,
  form_mapping,
  steps_per_file,
+ allow_read_errors_with_report,
  )
  else:
  return _get_dak_array_delay_open(
@@ -274,6 +283,7 @@ def dask(
  interp_options,
  form_mapping,
  steps_per_file,
+ allow_read_errors_with_report,
  )
  else:
  raise NotImplementedError()
@@ -890,6 +900,15 @@ class UprootReadMixin:
  form_mapping_info: ImplementsFormMappingInfo
  common_keys: frozenset[str]
  interp_options: dict[str, Any]
+ allow_read_errors_with_report: bool
+
+ @property
+ def allowed_exceptions(self):
+ return (OSError,)
+
+ @property
+ def return_report(self) -> bool:
+ return self.allow_read_errors_with_report
 
  def read_tree(self, tree: HasBranches, start: int, stop: int) -> AwkArray:
  assert start <= stop
@@ -946,6 +965,15 @@ def mock(self) -> AwkArray:
  behavior=self.form_mapping_info.behavior,
  )
 
+ def mock_empty(self, backend) -> AwkArray:
+ awkward = uproot.extras.awkward()
+ return awkward.to_backend(
+ self.expected_form.length_zero_array(highlevel=False),
+ backend=backend,
+ highlevel=True,
+ behavior=self.form_mapping_info.behavior,
+ )
+
  def prepare_for_projection(self) -> tuple[AwkArray, TypeTracerReport, dict]:
  awkward = uproot.extras.awkward()
  dask_awkward = uproot.extras.dask_awkward()
@@ -1011,6 +1039,53 @@ def project_keys(self: T, keys: frozenset[str]) -> T:
  raise NotImplementedError
 
 
+def _report_failure(exception, call_time, *args, **kwargs):
+ awkward = uproot.extras.awkward()
+ return awkward.Array(
+ [
+ {
+ "call_time": call_time,
+ "duration": None,
+ "args": [repr(a) for a in args],
+ "kwargs": [[k, repr(v)] for k, v in kwargs.items()],
+ "exception": type(exception).__name__,
+ "message": str(exception),
+ "fqdn": socket.getfqdn(),
+ "hostname": socket.gethostname(),
+ }
+ ]
+ )
+
+
+def _report_success(duration, *args, **kwargs):
+ awkward = uproot.extras.awkward()
+ return awkward.Array(
+ [
+ {
+ "call_time": None,
+ "duration": duration,
+ "args": [repr(a) for a in args],
+ "kwargs": [[k, repr(v)] for k, v in kwargs.items()],
+ "exception": None,
+ "message": None,
+ "fqdn": None,
+ "hostname": None,
+ }
+ ]
+ )
+
+
+def with_duration(f):
+ @functools.wraps(f)
+ def wrapper(*args, **kwargs):
+ start = time.monotonic()
+ result = f(*args, **kwargs)
+ stop = time.monotonic()
+ return result, (stop - start)
+
+ return wrapper
+
+
 class _UprootRead(UprootReadMixin):
  def __init__(
  self,
@@ -1020,13 +1095,15 @@ def __init__(
  base_form: Form,
  expected_form: Form,
  form_mapping_info: ImplementsFormMappingInfo,
+ allow_read_errors_with_report: bool,
  ) -> None:
  self.ttrees = ttrees
  self.common_keys = frozenset(common_keys)
  self.interp_options = interp_options
  self.base_form = base_form
  self.expected_form = expected_form
  self.form_mapping_info = form_mapping_info
+ self.allow_read_errors_with_report = allow_read_errors_with_report
 
  def project_keys(self: T, keys: frozenset[str]) -> T:
  return _UprootRead(
@@ -1036,11 +1113,39 @@ def project_keys(self: T, keys: frozenset[str]) -> T:
  self.base_form,
  self.expected_form,
  self.form_mapping_info,
+ self.allow_read_errors_with_report,
  )
 
- def __call__(self, i_start_stop) -> AwkArray:
+ def __call__(self, i_start_stop):
  i, start, stop = i_start_stop
+ if self.return_report:
+ call_time = time.time_ns()
+ try:
+ result, duration = with_duration(self._call_impl)(i, start, stop)
+ return (
+ result,
+ _report_success(
+ duration,
+ self.ttrees[i],
+ start,
+ stop,
+ ),
+ )
+ except self.allowed_exceptions as err:
+ return (
+ self.mock_empty(backend="cpu"),
+ _report_failure(
+ err,
+ call_time,
+ self.ttrees[i],
+ start,
+ stop,
+ ),
+ )
 
+ return self._call_impl(i, start, stop)
+
+ def _call_impl(self, i, start, stop):
  return self.read_tree(self.ttrees[i], start, stop)
 
 
@@ -1055,6 +1160,7 @@ def __init__(
  base_form: Form,
  expected_form: Form,
  form_mapping_info: ImplementsFormMappingInfo,
+ allow_read_errors_with_report: bool,
  ) -> None:
  self.custom_classes = custom_classes
  self.allow_missing = allow_missing
@@ -1064,15 +1170,11 @@ def __init__(
  self.base_form = base_form
  self.expected_form = expected_form
  self.form_mapping_info = form_mapping_info
+ self.allow_read_errors_with_report = allow_read_errors_with_report
 
- def __call__(self, blockwise_args) -> AwkArray:
- (
- file_path,
- object_path,
- i_step_or_start,
- n_steps_or_stop,
- is_chunk,
- ) = blockwise_args
+ def _call_impl(
+ self, file_path, object_path, i_step_or_start, n_steps_or_stop, is_chunk
+ ):
  ttree = uproot._util.regularize_object_path(
  file_path,
  object_path,
@@ -1105,6 +1207,50 @@ def __call__(self, blockwise_args) -> AwkArray:
 
  return self.read_tree(ttree, start, stop)
 
+ def __call__(self, blockwise_args):
+ (
+ file_path,
+ object_path,
+ i_step_or_start,
+ n_steps_or_stop,
+ is_chunk,
+ ) = blockwise_args
+
+ if self.return_report:
+ call_time = time.time_ns()
+ try:
+ result, duration = with_duration(self._call_impl)(
+ file_path, object_path, i_step_or_start, n_steps_or_stop, is_chunk
+ )
+ return (
+ result,
+ _report_success(
+ duration,
+ file_path,
+ object_path,
+ i_step_or_start,
+ n_steps_or_stop,
+ is_chunk,
+ ),
+ )
+ except self.allowed_exceptions as err:
+ return (
+ self.mock_empty(backend="cpu"),
+ _report_failure(
+ err,
+ call_time,
+ file_path,
+ object_path,
+ i_step_or_start,
+ n_steps_or_stop,
+ is_chunk,
+ ),
+ )
+
+ return self._call_impl(
+ file_path, object_path, i_step_or_start, n_steps_or_stop, is_chunk
+ )
+
  def project_keys(self: T, keys: frozenset[str]) -> T:
  return _UprootOpenAndRead(
  self.custom_classes,
@@ -1115,6 +1261,7 @@ def project_keys(self: T, keys: frozenset[str]) -> T:
  self.base_form,
  self.expected_form,
  self.form_mapping_info,
+ self.allow_read_errors_with_report,
  )
 
 
@@ -1151,6 +1298,7 @@ def _get_dak_array(
  interp_options,
  form_mapping,
  steps_per_file,
+ allow_read_errors_with_report,
 ):
  dask_awkward = uproot.extras.dask_awkward()
  awkward = uproot.extras.awkward()
@@ -1306,15 +1454,18 @@ def real_filter_branch(branch):
  else:
  expected_form, form_mapping_info = form_mapping(base_form)
 
+ fn = _UprootRead(
+ ttrees,
+ common_keys,
+ interp_options,
+ base_form=base_form,
+ expected_form=expected_form,
+ form_mapping_info=form_mapping_info,
+ allow_read_errors_with_report=allow_read_errors_with_report,
+ )
+
  return dask_awkward.from_map(
- _UprootRead(
- ttrees,
- common_keys,
- interp_options,
- base_form=base_form,
- expected_form=expected_form,
- form_mapping_info=form_mapping_info,
- ),
+ fn,
  partition_args,
  divisions=tuple(divisions),
  label="from-uproot",
@@ -1334,6 +1485,7 @@ def _get_dak_array_delay_open(
  interp_options,
  form_mapping,
  steps_per_file,
+ allow_read_errors_with_report,
 ):
  dask_awkward = uproot.extras.dask_awkward()
  awkward = uproot.extras.awkward()
@@ -1396,17 +1548,20 @@ def _get_dak_array_delay_open(
  else:
  expected_form, form_mapping_info = form_mapping(base_form)
 
+ fn = _UprootOpenAndRead(
+ custom_classes,
+ allow_missing,
+ real_options,
+ common_keys,
+ interp_options,
+ base_form=base_form,
+ expected_form=expected_form,
+ form_mapping_info=form_mapping_info,
+ allow_read_errors_with_report=allow_read_errors_with_report,
+ )
+
  return dask_awkward.from_map(
- _UprootOpenAndRead(
- custom_classes,
- allow_missing,
- real_options,
- common_keys,
- interp_options,
- base_form=base_form,
- expected_form=expected_form,
- form_mapping_info=form_mapping_info,
- ),
+ fn,
  partition_args,
  divisions=None if divisions is None else tuple(divisions),
  label="from-uproot",

tests/test_1058_dask_awkward_report.py

@@ -0,0 +1,25 @@
+import pytest
+import skhep_testdata
+
+import uproot
+
+
+dask = pytest.importorskip("dask")
+dask_awkward = pytest.importorskip("dask_awkward")
+
+
+def test_with_report():
+ test_path1 = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+ test_path2 = skhep_testdata.data_path("uproot-Zmumu-uncompressed.root") + ":events"
+ test_path3 = "/some/file/that/doesnt/exist"
+ files = [test_path1, test_path2, test_path3]
+ collection, report = uproot.dask(
+ files,
+ library="ak",
+ open_files=False,
+ allow_read_errors_with_report=True,
+ )
+ _, creport = dask.compute(collection, report)
+ assert creport[0].exception is None # test_path1 is good
+ assert creport[1].exception is None # test_path2 is good
+ assert creport[2].exception == "FileNotFoundError" # test_path3 is a bad file