Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Correctly set the exit code when an exception is raised in the atexit callback _prepare_to_exit. #6383

Merged
merged 11 commits into from
Feb 9, 2024
11 changes: 8 additions & 3 deletions torch_xla/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,9 +147,14 @@ def _setup_tpu_vm_library_path() -> bool:


def _prepare_to_exit():
_XLAC._prepare_to_exit()
if int(os.environ.get('PT_XLA_DEBUG', '0')):
_summarize_fn_tracker()
try:
_XLAC._prepare_to_exit()
if int(os.environ.get('PT_XLA_DEBUG', '0')):
_summarize_fn_tracker()
except Exception as e:
logging.error(
"Caught an exception when exiting the process. Exception: ", exc_info=e)
os._exit(1)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sorry, one more micro-nit: make this exit(1)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried and it didn't set the exit code correctly. Per https://docs.python.org/3/library/constants.html#exit, calling exit(1) raises a SystemExit exception. But due to https://bugs.python.org/issue27035, raising an exception won't set the exit code correctly.

Also, per https://docs.python.org/3/library/constants.html#exit, exit or quit "are useful for the interactive interpreter shell and should not be used in programs.".

I'll add a comment.



def _init_xla_lazy_backend():
Expand Down
Loading