Skip to content

Commit d147cd0

Browse files
lorenzotomadalorenzotomada
committed
Update memory profiling
1 parent ae0b9a7 commit d147cd0

File tree

2 files changed

+156
-149
lines changed

2 files changed

+156
-149
lines changed

‎scripts/profiling_memory.py

Lines changed: 152 additions & 145 deletions
Original file line numberDiff line numberDiff line change
@@ -6,186 +6,193 @@
66
profile_scipy_eigvals,
77
poisson_2d_structure,
88
)
9+
from pyclassify.parallel_tridiag_eigen import parallel_tridiag_eigen
10+
911
import numpy as np
1012
import scipy
1113
import argparse
12-
from mpi4py import MPI
13-
import scipy
1414
import scipy.sparse as sp
1515
import psutil
1616
import gc
1717
import os
1818
import csv
1919
import sys
20-
21-
sys.path.append("scripts")
22-
# from mpi_running import compute_eigvals
20+
from mpi4py import MPI
2321

2422

25-
# Seed for reproducibility
2623
seed = 8422
2724
np.random.seed(seed)
2825

2926

30-
parser = argparse.ArgumentParser()
31-
parser.add_argument("--config", type=str, required=False, help="config file:")
32-
args = parser.parse_args()
33-
config_file = args.config if args.config else "experiments/config"
27+
comm = MPI.COMM_WORLD
28+
rank = comm.Get_rank()
29+
size = comm.Get_size()
30+
3431

35-
kwargs = read_config(config_file)
32+
if rank == 0:
33+
# parse arguments
34+
parser = argparse.ArgumentParser()
35+
parser.add_argument("--config", type=str, required=False, help="config file:")
36+
args = parser.parse_args()
37+
config_file = args.config if args.config else "experiments/config"
38+
kwargs = read_config(config_file)
39+
dim = kwargs["dim"]
40+
density = kwargs["density"]
41+
n_procs = kwargs["n_processes"]
42+
plot = kwargs["plot"]
43+
else:
44+
kwargs = None
45+
46+
# Broadcast config to all ranks
47+
kwargs = comm.bcast(kwargs, root=0)
3648
dim = kwargs["dim"]
3749
density = kwargs["density"]
3850
n_procs = kwargs["n_processes"]
3951
plot = kwargs["plot"]
4052

53+
# Now we build the matrix on rank 0
54+
# It is a scipy sparse matrix with the structure of a 2D Poisson problem matrix obtained using finite differences
55+
if rank == 0:
56+
A = poisson_2d_structure(dim)
57+
A_np = A.toarray()
58+
else:
59+
A_np = None
4160

42-
def parallel_eig(diag, off_diag, nprocs):
43-
print("Spawning a communicator")
44-
comm = MPI.COMM_SELF.Spawn(sys.executable, args=["scripts/run.py"], maxprocs=nprocs)
45-
46-
print("Sending data to children")
47-
comm.send(diag, dest=0, tag=11)
48-
comm.send(off_diag, dest=0, tag=12)
61+
A_np = comm.bcast(A_np, root=0)
4962

50-
print("Waiting for results...")
51-
sys.stdout.flush()
5263

53-
eigvals = comm.recv(source=0, tag=22)
54-
eigvecs = comm.recv(source=0, tag=23)
55-
delta_t = comm.recv(source=0, tag=24)
56-
total_mem_children = comm.recv(source=0, tag=25)
57-
comm.Disconnect()
64+
# On rank 0, we use the Lanczos method
65+
# We actually call it twice: the first time to ensure that the function is JIT-compiled by Numba, the second one for memory profiling
66+
if rank == 0:
67+
print("Precompiling Lanczos...")
68+
Q, diag, off_diag = Lanczos_PRO(A_np, np.ones_like(np.diag(A_np)) * 1.0)
69+
print("Done. Now reducing using Lanczos...")
70+
gc.collect()
71+
proc = psutil.Process()
72+
mem_before_lanczos = proc.memory_info().rss / 1024 / 1024 # MB
5873

59-
print("Data recieved!")
60-
return eigvals, eigvecs, delta_t, total_mem_children
74+
Q, diag, off_diag = Lanczos_PRO(A_np, np.ones_like(np.diag(A_np)) * 1.0)
6175

76+
gc.collect()
77+
mem_after_lanczos = proc.memory_info().rss / 1024 / 1024 # MB
78+
delta_mem_lanczos = mem_after_lanczos - mem_before_lanczos
79+
print("Done. Now computing eigenvalues...")
80+
else:
81+
diag = off_diag = None
6282

63-
def compute_eigvals(A, n_procs):
64-
print("Reducing using Lanczos")
65-
Q, diag, off_diag = Lanczos_PRO(A_np, np.ones_like(np.diag(A_np)) * 1.0)
83+
diag = comm.bcast(diag, root=0)
84+
off_diag = comm.bcast(off_diag, root=0)
6685

67-
print("Done. Now computing eigenvalues.")
68-
eigvals, eigvecs, delta_t, total_mem_children = parallel_eig(
69-
diag, off_diag, n_procs
70-
)
86+
gc.collect()
87+
proc = psutil.Process()
88+
mem_before = proc.memory_info().rss / 1024 / 1024 # MB
7189

72-
print("Eigenvalues computed")
73-
return eigvals, eigvecs, delta_t, total_mem_children
90+
eigvals, eigvecs = parallel_tridiag_eigen(
91+
diag, off_diag, comm=comm, min_size=1, tol_factor=1e-10
92+
)
7493

94+
gc.collect()
95+
mem_after = proc.memory_info().rss / 1024 / 1024
96+
delta_mem = mem_after - mem_before
97+
98+
total_mem_children = comm.reduce(delta_mem, op=MPI.SUM, root=0)
99+
100+
if rank == 0:
101+
total_mem_all = delta_mem_lanczos
102+
print("Eigenvalues computed.")
103+
process = psutil.Process()
104+
105+
print(f"Total memory across all processes: {total_mem_all:.2f} MB")
106+
107+
mem_np = profile_numpy_eigvals(A_np)
108+
print(f"NumPy eig memory usage: {mem_np:.2f} MB")
109+
110+
mem_sp = profile_scipy_eigvals(A_np)
111+
print(f"SciPy eig memory usage: {mem_sp:.2f} MB")
112+
113+
os.makedirs("logs", exist_ok=True)
114+
log_file = "logs/memory_profile.csv"
115+
fieldnames = [
116+
"matrix_size",
117+
"n_processes",
118+
"mem_lanzos_mb",
119+
"mem_tridiag_mb",
120+
"mem_total_mb",
121+
"mem_numpy_mb",
122+
"mem_scipy_mb",
123+
]
124+
125+
write_header = not os.path.exists(log_file)
126+
with open(log_file, mode="a", newline="") as f:
127+
writer = csv.DictWriter(f, fieldnames=fieldnames)
128+
if write_header:
129+
writer.writeheader()
130+
writer.writerow(
131+
{
132+
"matrix_size": dim,
133+
"n_processes": size,
134+
"mem_lanzos_mb": round(delta_mem_lanczos, 2),
135+
"mem_tridiag_mb": round(total_mem_children, 2),
136+
"mem_total_mb": round(total_mem_all, 2),
137+
"mem_numpy_mb": round(mem_np, 2),
138+
"mem_scipy_mb": round(mem_sp, 2),
139+
}
140+
)
75141

76-
A = poisson_2d_structure(dim)
77-
A_np = A.toarray()
142+
if plot:
143+
import matplotlib.pyplot as plt
144+
import pandas as pd
78145

79-
Q, diag, off_diag = Lanczos_PRO(
80-
A_np, np.ones_like(np.diag(A_np)) * 1.0
81-
) # To compile using numba
146+
df = pd.read_csv("logs/memory_profile.csv")
147+
nproc_values = sorted(df["n_processes"].unique())
82148

83-
gc.collect()
84-
process = psutil.Process()
85-
mem_before = process.memory_info().rss / 1024 / 1024
149+
plt.figure(figsize=(10, 6))
86150

87-
eigvals, eigvecs, delta_t, total_mem_children = compute_eigvals(A_np, n_procs)
151+
numpy_avg = df.groupby("matrix_size")["mem_numpy_mb"].mean()
152+
plt.plot(
153+
numpy_avg.index,
154+
numpy_avg.values,
155+
color="green",
156+
marker="x",
157+
linestyle="--",
158+
label="NumPy",
159+
)
88160

89-
gc.collect()
90-
mem_after = process.memory_info().rss / 1024 / 1024
91-
delta_mem_parent = mem_after - mem_before
92-
93-
total_mem_all = delta_mem_parent + total_mem_children
94-
95-
print(f"Total memory across all processes: {total_mem_all:.2f} MB")
96-
97-
mem_np = profile_numpy_eigvals(A_np)
98-
print(f"NumPy eig memory usage: {mem_np:.2f} MB")
99-
100-
mem_sp = profile_scipy_eigvals(A_np)
101-
print(f"SciPy eig memory usage: {mem_sp:.2f} MB")
102-
103-
os.makedirs("logs", exist_ok=True)
104-
105-
log_file = "logs/memory_profile.csv"
106-
fieldnames = [
107-
"matrix_size",
108-
"n_processes",
109-
"mem_parent_mb",
110-
"mem_children_mb",
111-
"mem_total_mb",
112-
"mem_numpy_mb",
113-
"mem_scipy_mb",
114-
]
115-
116-
write_header = not os.path.exists(log_file)
117-
118-
with open(log_file, mode="a", newline="") as f:
119-
writer = csv.DictWriter(f, fieldnames=fieldnames)
120-
if write_header:
121-
writer.writeheader()
122-
writer.writerow(
123-
{
124-
"matrix_size": dim,
125-
"n_processes": n_procs,
126-
"mem_parent_mb": round(delta_mem_parent, 2),
127-
"mem_children_mb": round(total_mem_children, 2),
128-
"mem_total_mb": round(total_mem_all, 2),
129-
"mem_numpy_mb": round(mem_np, 2),
130-
"mem_scipy_mb": round(mem_sp, 2),
131-
}
132-
)
133-
134-
if plot:
135-
import matplotlib.pyplot as plt
136-
import pandas as pd
137-
import matplotlib.colors as mcolors
138-
import numpy as np
139-
140-
df = pd.read_csv("logs/memory_profile.csv")
141-
142-
nproc_values = sorted(df["n_processes"].unique())
143-
144-
plt.figure(figsize=(10, 6))
145-
146-
numpy_avg = df.groupby("matrix_size")["mem_numpy_mb"].mean()
147-
plt.plot(
148-
numpy_avg.index,
149-
numpy_avg.values,
150-
color="green",
151-
marker="x",
152-
linestyle="--",
153-
label="NumPy",
154-
)
155-
156-
scipy_avg = df.groupby("matrix_size")["mem_scipy_mb"].mean()
157-
plt.plot(
158-
scipy_avg.index,
159-
scipy_avg.values,
160-
color="red",
161-
marker="^",
162-
linestyle=":",
163-
label="SciPy",
164-
)
165-
166-
for nproc in nproc_values:
167-
subset = df[df["n_processes"] == nproc].sort_values("matrix_size")
168-
label = f"Divide and Conquer ({nproc} proc{'s' if nproc > 1 else ''})"
161+
scipy_avg = df.groupby("matrix_size")["mem_scipy_mb"].mean()
169162
plt.plot(
170-
subset["matrix_size"],
171-
subset["mem_total_mb"],
172-
marker="o",
173-
linestyle="-",
174-
label=label,
163+
scipy_avg.index,
164+
scipy_avg.values,
165+
color="red",
166+
marker="^",
167+
linestyle=":",
168+
label="SciPy",
169+
)
170+
171+
for nproc in nproc_values:
172+
subset = df[df["n_processes"] == nproc].sort_values("matrix_size")
173+
label = f"Divide et impera ({nproc} proc{'s' if nproc > 1 else ''})"
174+
plt.plot(
175+
subset["matrix_size"],
176+
subset["mem_total_mb"],
177+
marker="o",
178+
linestyle="-",
179+
label=label,
180+
)
181+
182+
plt.xlabel("Matrix size")
183+
plt.ylabel("Total memory (MB)")
184+
plt.xscale("log")
185+
plt.title("Memory usage vs. Matrix size")
186+
plt.grid(True)
187+
plt.tight_layout()
188+
189+
plt.legend(
190+
bbox_to_anchor=(1.05, 1),
191+
loc="upper left",
192+
borderaxespad=0.0,
193+
title="Method",
175194
)
195+
plt.subplots_adjust(right=0.75)
176196

177-
plt.xlabel("Matrix size")
178-
plt.ylabel("Total memory (MB)")
179-
plt.xscale("log")
180-
# plt.yscale("log")
181-
plt.title("Memory usage vs. Matrix size")
182-
plt.grid(True)
183-
plt.tight_layout()
184-
185-
plt.legend(
186-
bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0, title="Method"
187-
)
188-
plt.subplots_adjust(right=0.75)
189-
190-
plt.savefig("logs/mem_vs_size_all_methods.png", bbox_inches="tight")
191-
plt.show()
197+
plt.savefig("logs/mem_vs_size_all_methods.png", bbox_inches="tight")
198+
plt.show()

‎shell/submit.sh

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
#!/bin/bash
22

33
# Ranges over which we iterate
4-
n_processes=(1 2 4 8)
5-
matrix_sizes=(10 100 500 1000)
4+
n_processes=(1 2)
5+
matrix_sizes=(10 15 20)
66

77
last_dim="${matrix_sizes[-1]}"
88
last_nproc="${n_processes[-1]}"
@@ -31,11 +31,11 @@ for dim in "${matrix_sizes[@]}"; do
3131
sed -i "s/^plot: .*/plot: true/" $CONFIG_FILE
3232
fi
3333

34-
python scripts/profiling_memory.py
34+
mpirun -np ${n_p} python scripts/profiling_memory.py
3535
done
3636
done
3737

3838
# Restore the original config file
3939
mv ${CONFIG_FILE}.bak $CONFIG_FILE
4040

41-
echo "Experiment completed!"
41+
echo "Experiment completed!"

0 commit comments

Comments
 (0)