Add plot functionality for new future selector #820
Conversation
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
WalkthroughAdds FutureSelector utilities (split_future, get_item_from_future), exports them via executorlib. Updates interactive plotting to accept FutureSelector and relocates plotting imports. Adjusts tests to new import path and adds tests for the new FutureSelector behavior with executors and manual futures. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor U as User
participant E as Executor (SingleNode)
participant F as Future
participant S1 as FutureSelector[0]
participant S2 as FutureSelector[1]
participant S3 as FutureSelector[2]
U->>E: submit(fn)-> Future
E-->>U: F
U->>F: split_future(F, 3)
Note over S1,S3: Each selector proxies F and selects item by index
U-->>U: Await S1.result(), S2.result(), S3.result()
E-->>F: set_result((v0, v1, v2))
F-->>S1: result()-> v0
F-->>S2: result()-> v1
F-->>S3: result()-> v2
sequenceDiagram
autonumber
participant DP as dependency_plot
participant Arg as FutureSelector
participant UF as Underlying Future
DP->>DP: convert_arg(Arg)
DP->>Arg: read _future, _selector
DP-->>DP: serialize {future, selector, fn_name}
DP->>DP: add_element(start=UF, end=node, label="[selector]")
Note right of DP: Supports both int index and str key
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Pre-merge checks (2 passed, 1 warning)❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
Poem
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. ✨ Finishing touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #820 +/- ##
==========================================
+ Coverage 98.12% 98.17% +0.05%
==========================================
Files 33 34 +1
Lines 1650 1698 +48
==========================================
+ Hits 1619 1667 +48
Misses 31 31 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
for more information, see https://pre-commit.ci
# Conflicts: # executorlib/standalone/plot.py
for more information, see https://pre-commit.ci
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
executorlib/task_scheduler/interactive/dependency_plot.py (1)
52-61: Handle FutureSelector inside containers (lists/dicts) for correct plottingLists/dicts that contain
FutureSelectorelements are currently not recognized (onlyFutureis checked), so edges for nested selectors are missed.- elif isinstance(arg, list) and any(isinstance(a, Future) for a in arg): + elif isinstance(arg, list) and any(isinstance(a, (Future, FutureSelector)) for a in arg): @@ - for i, a in enumerate(arg): - if isinstance(a, Future): - add_element(arg=a, link_to=node_id, label="ind: " + str(i)) + for i, a in enumerate(arg): + if isinstance(a, (Future, FutureSelector)): + add_element(arg=a, link_to=node_id, label="ind: " + str(i)) @@ - elif isinstance(arg, dict) and any(isinstance(a, Future) for a in arg.values()): + elif isinstance(arg, dict) and any(isinstance(a, (Future, FutureSelector)) for a in arg.values()): @@ - for kt, vt in arg.items(): - if isinstance(vt, Future): - add_element(arg=vt, link_to=node_id, label="key: " + kt) + for kt, vt in arg.items(): + if isinstance(vt, (Future, FutureSelector)): + add_element(arg=vt, link_to=node_id, label="key: " + str(kt))Also applies to: 62-74
🧹 Nitpick comments (6)
executorlib/standalone/select.py (2)
23-25: Pass the FutureSelector itself to callbacksForwarding the underlying Future to
add_done_callbackcan surprise users (callback receives the parent Future instead of the selector). Wrap the callback so it getsself.- def add_done_callback(self, fn) -> None: - return self._future.add_done_callback(fn=fn) + def add_done_callback(self, fn) -> None: + def _wrap(_): + fn(self) + return self._future.add_done_callback(_wrap)
11-16: Clarify/codify cancellation semantics
cancel()on a selector cancels the underlying parent future (affecting all selectors). If this is intended, document it in the class docstring; otherwise, consider making selector-level cancellation a no-op.executorlib/task_scheduler/interactive/dependency_plot.py (1)
36-43: Disambiguate selector edge labelsFor dict selectors the edge label is just the key string (e.g., "b"), while list/tuple selectors show no index prefix. Use consistent prefixes to match existing dict/list labeling patterns.
- edge_lst.append( - { - "start": hash_id_dict[future_hash_inverse_dict[arg._future]], - "end": link_to, - "label": label + str(arg._selector), - } - ) + edge_lst.append( + { + "start": hash_id_dict[future_hash_inverse_dict[arg._future]], + "end": link_to, + "label": label + + ("key: " if isinstance(arg._selector, str) else "ind: ") + + str(arg._selector), + } + )tests/test_singlenodeexecutor_plot_dependency.py (1)
307-330: Nice executable examples for selectors in plotting modeCovers both tuple split and dict item selection with
plot_dependency_graph=True. Consider adding one case where a selector is nested inside a list/dict argument to exercise the container handling (after the suggested fix).tests/test_standalone_select.py (2)
16-18: Silence ARG001: unused parameterRename the unused argument to
_ito satisfy linters.-def function_with_exception(i): +def function_with_exception(_i): raise RuntimeError()
60-66: Optional: strengthen callback assertion (if callback wraps selector)If you adopt the
add_done_callbackwrapper suggestion, assert that the callback sees the selected value (not the full result).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
executorlib/__init__.py(2 hunks)executorlib/standalone/select.py(1 hunks)executorlib/task_scheduler/interactive/dependency.py(1 hunks)executorlib/task_scheduler/interactive/dependency_plot.py(3 hunks)tests/test_fluxjobexecutor_plot.py(1 hunks)tests/test_singlenodeexecutor_plot_dependency.py(3 hunks)tests/test_standalone_select.py(1 hunks)tests/test_testclusterexecutor.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
tests/test_fluxjobexecutor_plot.py (1)
executorlib/task_scheduler/interactive/dependency_plot.py (1)
generate_nodes_and_edges_for_plotting(10-91)
executorlib/task_scheduler/interactive/dependency_plot.py (1)
executorlib/standalone/select.py (1)
FutureSelector(5-43)
tests/test_testclusterexecutor.py (1)
executorlib/task_scheduler/interactive/dependency_plot.py (1)
generate_nodes_and_edges_for_plotting(10-91)
executorlib/task_scheduler/interactive/dependency.py (1)
executorlib/task_scheduler/base.py (1)
TaskSchedulerBase(17-198)
tests/test_standalone_select.py (3)
executorlib/executor/single.py (1)
SingleNodeExecutor(20-190)executorlib/standalone/select.py (13)
split_future(46-57)get_item_from_future(60-72)FutureSelector(5-43)result(26-31)done(20-21)add_done_callback(23-24)set_running_or_notify_cancel(36-37)running(17-18)set_result(39-40)cancel(11-12)cancelled(14-15)set_exception(42-43)exception(33-34)executorlib/standalone/serialize.py (1)
cloudpickle_register(9-28)
executorlib/__init__.py (1)
executorlib/standalone/select.py (2)
get_item_from_future(60-72)split_future(46-57)
tests/test_singlenodeexecutor_plot_dependency.py (3)
executorlib/standalone/select.py (2)
split_future(46-57)get_item_from_future(60-72)executorlib/task_scheduler/interactive/dependency_plot.py (1)
generate_nodes_and_edges_for_plotting(10-91)executorlib/task_scheduler/interactive/dependency.py (1)
submit(113-153)
🪛 Ruff (0.12.2)
tests/test_standalone_select.py
16-16: Unused function argument: i
(ARG001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: uml
- GitHub Check: notebooks_integration
🔇 Additional comments (5)
executorlib/task_scheduler/interactive/dependency.py (1)
14-19: LGTM: import relocationImporting plotting utilities from
interactive.dependency_plotaligns with the refactor; no behavior change here.executorlib/__init__.py (2)
28-28: Re-exporting selector utilities via top-level API is goodSurface area looks right and matches tests.
70-71: all updated appropriatelyExports include
get_item_from_futureandsplit_future.tests/test_fluxjobexecutor_plot.py (1)
6-6: Import path update looks correctMatches the new module location.
tests/test_testclusterexecutor.py (1)
7-7: Import path update looks correctConsistent with the plotting refactor.
Example 1:
Example 2:
Summary by CodeRabbit
New Features
Refactor