mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Use functools.cache on has_efa (#163439)
Cache the result of `has_efa` by `functools.cache`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/163439 Approved by: https://github.com/janeyx99
This commit is contained in:
committed by
PyTorch MergeBot
parent
e3b392bdfd
commit
d3a1345ed8
@ -1,6 +1,7 @@
|
||||
# mypy: ignore-errors
|
||||
|
||||
import faulthandler
|
||||
import functools
|
||||
import itertools
|
||||
import logging
|
||||
import multiprocessing
|
||||
@ -1140,30 +1141,24 @@ def run_subtests(
|
||||
c10d.barrier()
|
||||
|
||||
|
||||
# Cannot use functools.cache as it requires python 3.9
|
||||
EFA_PROBE_RESULT = None
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_efa() -> bool:
|
||||
"""
|
||||
If shell command `fi_info -p efa -t FI_EP_RDM` returns exit code 0 then we assume that the machine has
|
||||
Libfabric EFA interfaces and EFA software components installed,
|
||||
see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start.html.
|
||||
"""
|
||||
global EFA_PROBE_RESULT
|
||||
if EFA_PROBE_RESULT is not None:
|
||||
return EFA_PROBE_RESULT
|
||||
|
||||
try:
|
||||
EFA_PROBE_RESULT = (
|
||||
return (
|
||||
subprocess.run(
|
||||
["fi_info", "-p", "efa", "-t", "FI_EP_RDM"], check=False
|
||||
).returncode
|
||||
== 0
|
||||
)
|
||||
except FileNotFoundError:
|
||||
EFA_PROBE_RESULT = False
|
||||
return EFA_PROBE_RESULT
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def tp_transports():
|
||||
|
Reference in New Issue
Block a user