Call the given method within the invocation context corresponding to `module` and passing it `args` and `kwargs`. Args: module: The module whose context the mnethod should be called in. *args: Positional arguments to `method_fn`. method_fn: The method to call.
(
module: "Module", *args, method_fn: Callable, method_name: str, **kwargs
)
| 649 | "Each thread should call install_context_stack() to install its own stack." |
| 650 | ) |
| 651 | current = current_context() |
| 652 | assert current is not None |
| 653 | context = current.add_child(name, **kwargs) |
| 654 | with set_current_context(context) as c: |
| 655 | yield c |
| 656 | |
| 657 | |
| 658 | @traceback_util.wrap |
| 659 | @no_stack_summary |
| 660 | def _call_method_in_context( |
| 661 | module: "Module", *args, method_fn: Callable, method_name: str, **kwargs |
| 662 | ): |
| 663 | """Call the given method within the invocation context corresponding to `module` and passing |
| 664 | it `args` and `kwargs`. |
| 665 | |
| 666 | Args: |
| 667 | module: The module whose context the mnethod should be called in. |
| 668 | *args: Positional arguments to `method_fn`. |
| 669 | method_fn: The method to call. |
| 670 | method_name: The name of the method to call. |
| 671 | **kwargs: Keyword arguments to `method_fn`. |
| 672 | |
| 673 | Returns: |
| 674 | The output of `method_fn(*args, **kwawrgs)` when called from within the invocation context |
| 675 | of `module`. |
| 676 | """ |
| 677 | if len(args) > 1: |
| 678 | logging.log_first_n( |
| 679 | logging.WARNING, |
| 680 | "Multiple positional arguments for %s.%s. Consider using keyword arguments instead.", |
| 681 | 3, |
| 682 | type(module), |
| 683 | method_name, |
| 684 | ) |
| 685 | |
| 686 | # Use ExitStack since we need to repeatedly enter a context in a loop. |
| 687 | # This cannot be done with a parenthesized context manager since, confusingly, |
| 688 | # even though you can do something like `with (mgr1, mgr2)`, |
| 689 | # it is not allowed to do `z = (mgr1, mgr2)` and then `with z`. |
| 690 | # We prefer the ExitStack() approach over recursion since it does not add unnecessary frames to |
| 691 | # the stack, which can make it harder to use a debugger with AXLearn. |
| 692 | with contextlib.ExitStack() as stack: |
| 693 | context = current_context() |
| 694 | if context is not None: |
| 695 | try: |
| 696 | # Enter context for descendant module if not already in it. |
| 697 | reversed_path_to_descendant = list( |
| 698 | reversed(context.module.path_to_descendant_module(module)) |
| 699 | ) |
| 700 | while reversed_path_to_descendant: |
| 701 | stack.enter_context(child_context(reversed_path_to_descendant.pop())) |
| 702 | except InvalidDescendantError as e: |
| 703 | # If an ancestor shared this module, use the shared module context since the module |
| 704 | # may not be a descendant of the current module. |
| 705 | try: |
| 706 | shared_module = context.module.get_shared_module(module) |
| 707 | stack.enter_context(child_context(**shared_module._asdict())) |
| 708 | except InvalidDescendantError: |
nothing calls this directly
no test coverage detected