nvidia-nvshmem-cu12

import nvshmem.core as nvshmem import os def main(): # Initialize NVSHMEM. This is a collective operation. # In a real scenario, this script would be launched with `mpiexec` or `nvshmrun`. nvshmem.init() # Query PE information my_pe = nvshmem.my_pe() n_pes = nvshmem.n_pes() print(f"Hello from PE {my_pe} of {n_pes}") # Perform some simple collective (e.g., a barrier) # This ensures all PEs reach this point before proceeding nvshmem.barrier_all() # Finalize NVSHMEM. This is also a collective operation. nvshmem.finalize() if __name__ == '__main__': # Note: This script needs to be run using an MPI launcher (e.g., mpiexec -n 2 python your_script.py) # or NVSHMEM's own launcher (nvshmrun). Running directly 'python your_script.py' # will result in an error or hang if NVSHMEM expects multiple processes. try: main() except Exception as e: # Catch potential errors if not launched collectively, for a more graceful exit print(f"Error: {e}") print("Please ensure the script is launched collectively, e.g., 'mpiexec -n 2 python quickstart.py'")