INNER CODE UNIT · Python

get_cuda_version

awslabs/awsome-distributed-ai · validation/efa-versions.py:63

def get_cuda_version(container=[]):
    try:
        version = subprocess.check_output(container + ['nvcc', '--version'])
        version = version.decode('utf-8')
        version = re.search(r', V(\d+\.\d+\.\d+)', version).group(1)
        return version
    except Exception as e:
        print(f'Error: {e}')
        return None

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='This script checks the versions of NCCL, EFA, Libfabric and CUDA. ')
    parser.add_argument('-c', '--container-image', type=str, help='Container image to get versions from.')
    args = parser.parse_args()

    if args.container_image:
        table = PrettyTable(["Package", "Local", "Container"])

View source record →