Provenance tracking for cuNumeric operators #596
Conversation
magnatelee
added
the
category:new-feature
cunumeric/coverage.py
Outdated
| @wraps(func) | ||
| def wrapper(*args: Any, **kwargs: Any) -> Any: | ||
| return func(*args, **kwargs) | ||
| if set_provenance := (runtime.legate_context.provenance is None): |
There was a problem hiding this comment.
This is to activate provenance tracking only for top-level invocations; some cuNumeric operators internally invoke other cuNumeric operators and because all of them share the same provenance, we need to set the provenance only once for the very first invocation made by the user program. (and extracting the stackframe isn't cheap, so we don't want to do it more than we need.)
There was a problem hiding this comment.
This seems fine but a context manager could be slightly nicer to automatically handle the call the reset_provenance when necessary. BTW is there any reason not to just call reset_provenance unconditionally? If legate_context.provenance was set, we need to reset it, but if legate_context.provenance was not set, then presumably resetting it is a no-op?
There was a problem hiding this comment.
This seems fine but a context manager could be slightly nicer to automatically handle the call the reset_provenance when necessary.
I guess I could do something like this:
func = runtime.legate_context.track_provenance(func, callback=find_last_user_frames)
and internally it calls the callback only when the provenance hasn't been set. Let me push changes for this.
BTW is there any reason not to just call reset_provenance unconditionally? If legate_context.provenance was set, we need to reset it, but if legate_context.provenance was not set, then presumably resetting it is a no-op?
No you don't always want to reset the provenance if you're already within a cunumeric API, otherwise, subsequent cuNumeric calls will lose the provenance. For example, in case you have a cuNumeric op A that internally calls cuNumeric ops B and C, unconditional resetting will reset the provenance as soon as it returns from B, which might not launch anything (e.g., if it was ndarray.__init__), thereby not tracking provenance for C.
There was a problem hiding this comment.
I meant something like this just to wrap up the code that is repeated here a few places (sketch, untested):
@contextmanager
def track_provenance():
# import runtime singleton here, probably, or pass it as an arg
if set_provenance := (runtime.legate_context.provenance is None):
runtime.legate_context.set_provenance(find_last_user_frames())
yield set_provenance
if set_provenance:
runtime.legate_context.reset_provenance()Then
with track_provenance():
return func()No you don't always want to reset the provenance if you're already within a cunumeric API, otherwise, subsequent cuNumeric calls will lose the provenance.
got it
There was a problem hiding this comment.
alright, I added a decorator for provenance tracking and it supports both top-level only tracking and nested tracking: 8448aec please also take a look at the companion PR: nv-legate/legate#370
* Use legate from Anaconda instead of artifacts * Call setup_conda_download * cleanup cupynumeric version check * Update legate_version * Reformat the meta * Set cpu_gpu_tag earlier * Define legate in run as well * Typo + Use uploaded legate version * Strip __cpu to _cpu * Cleanup-1 * Add a common script. Should be transferred to legate-gh-ci * Cleanup-2 * Check * Min. change and test * chk2 * chk3 * Enable all builds * Define and use legate label from versions.json * Placed label at the end * Syntax/Typo
2 participants
This PR activates provenance tracking for cuNumeric operators. This won't work until nv-legate/legate#370 gets merged.