Add runnable GPU main-path prototype
This commit is contained in:
@@ -9,6 +9,7 @@
|
||||
|
||||
|
||||
import AMSS_NCKU_Input as input_data
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
@@ -57,6 +58,48 @@ BUILD_JOBS = 64
|
||||
##################################################################
|
||||
|
||||
|
||||
##################################################################
|
||||
|
||||
def prepare_gpu_runtime_env():
|
||||
"""
|
||||
Create a user-private CUDA MPS environment for GPU runs.
|
||||
|
||||
On shared machines another user's daemon may already occupy the default
|
||||
/tmp/nvidia-mps pipe directory, which makes plain cudaSetDevice/cudaMalloc
|
||||
fail with cudaErrorMpsConnectionFailed. Binding AMSS-NCKU to a private
|
||||
pipe directory avoids cross-user interference.
|
||||
"""
|
||||
env = os.environ.copy()
|
||||
|
||||
pipe_dir = env.get("CUDA_MPS_PIPE_DIRECTORY", f"/tmp/amss-ncku-mps-{os.getuid()}")
|
||||
log_dir = env.get("CUDA_MPS_LOG_DIRECTORY", f"/tmp/amss-ncku-mps-log-{os.getuid()}")
|
||||
|
||||
os.makedirs(pipe_dir, exist_ok=True)
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
|
||||
env["CUDA_MPS_PIPE_DIRECTORY"] = pipe_dir
|
||||
env["CUDA_MPS_LOG_DIRECTORY"] = log_dir
|
||||
|
||||
control_socket = os.path.join(pipe_dir, "control")
|
||||
if not os.path.exists(control_socket):
|
||||
start = subprocess.run(
|
||||
["nvidia-cuda-mps-control", "-d"],
|
||||
env=env,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
if start.returncode != 0:
|
||||
print(f" Warning: failed to start private CUDA MPS daemon in {pipe_dir}")
|
||||
else:
|
||||
print(f" Using private CUDA MPS pipe directory: {pipe_dir}")
|
||||
else:
|
||||
print(f" Using existing private CUDA MPS pipe directory: {pipe_dir}")
|
||||
|
||||
return env
|
||||
|
||||
##################################################################
|
||||
|
||||
|
||||
|
||||
##################################################################
|
||||
|
||||
@@ -146,16 +189,29 @@ def run_ABE():
|
||||
|
||||
## Define the command to run; cast other values to strings as needed
|
||||
|
||||
run_env = None
|
||||
|
||||
if (input_data.GPU_Calculation == "no"):
|
||||
mpi_command = NUMACTL_CPU_BIND + " mpirun -np " + str(input_data.MPI_processes) + " ./ABE"
|
||||
#mpi_command = " mpirun -np " + str(input_data.MPI_processes) + " ./ABE"
|
||||
mpi_command_outfile = "ABE_out.log"
|
||||
elif (input_data.GPU_Calculation == "yes"):
|
||||
mpi_command = NUMACTL_CPU_BIND + " mpirun -np " + str(input_data.MPI_processes) + " ./ABEGPU"
|
||||
run_env = prepare_gpu_runtime_env()
|
||||
if int(input_data.MPI_processes) == 1:
|
||||
mpi_command = "./ABEGPU"
|
||||
else:
|
||||
mpi_command = NUMACTL_CPU_BIND + " mpirun -np " + str(input_data.MPI_processes) + " ./ABEGPU"
|
||||
mpi_command_outfile = "ABEGPU_out.log"
|
||||
|
||||
## Execute the MPI command and stream output
|
||||
mpi_process = subprocess.Popen(mpi_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
|
||||
mpi_process = subprocess.Popen(
|
||||
mpi_command,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
env=run_env,
|
||||
)
|
||||
|
||||
## Write ABE run output to file while printing to stdout
|
||||
with open(mpi_command_outfile, 'w') as file0:
|
||||
|
||||
Reference in New Issue
Block a user