Question about jax.distributed error

[shyams2]

I've worked with jax.distributed before and I'm working with code that has worked in the past on other machines. I'm doing a few runs on a new cluster this week and notice these errors:

2022-07-29 12:12:03.885056: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:460] ptxas returned an error during compilation of ptx to sass: 'INTERNAL: ptxas exited with non-zero error code 65280, output: ptxas fatal   : Could not open output file '/var/tmp/pbs.313487/tmpxft_000021e4_0000000a'
'  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
Fatal Python error: Aborted

2022-07-29 12:13:50.890606: F external/org_tensorflow/tensorflow/compiler/xla/pjrt/distributed/client.h:71] Terminating process because the coordinator detected missing heartbeats. This most likely indicates that another task died; see the other task logs for more details. Status: ABORTED: Shutting down due to missed heartbeat from task 1
Fatal Python error: Aborted

Does anyone know what could be causing the error? Is there any potential way I can resolve it? Or is it just a permissions issue?

[hawkinsp]

The first error related to ptxas is the one you need to look at. The second error is just a consequence of the first one: since one of your processes died with the ptxas error, the others will eventually notice and shut themselves down also.

For the ptxas error, is the guess in the error message correct? Do you have enough space on that filesystem?

[shyams2]

@hawkinsp thanks a lot for the response! So this is a shared cluster - while I'm certain I haven't surpassed any personal quotas, I'm not sure about the /var/tmp/ directory. Maybe I don't even have sufficient permission to do so. Is there anyway to override this write location? Let me try reaching out to the system admin as well.

[hawkinsp]

You should be able to override the temporary directory with the TMPDIR environment variable.

(Code that decides on the directory: https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/core/platform/default/env.cc;drc=5f6e4b06871ca70f60bf9a7a84e68124f5af68df;l=261 )

[shyams2]

@hawkinsp Yes, that was it! It works now. Thank you very much! :)