Commit 452feb9a authored by Seema Mirchandaney's avatar Seema Mirchandaney
Browse files

psana2 + legion

parent 4267bada
# This toml file contains settings for the Cori quicktart with PSANA and Legion
[debug]
test = "Quickstart settings for Cori with Psana and Legion"
verbose = true
verbosity = 0
use_callmonitor = false
[data]
in_dir = "${CFS}/m2859/data/3iyf/clean"
out_dir = "${SCRATCH}/spinifel_output"
name = "3iyf_sim_400k.h5"
[detector]
shape = [1, 128, 128]
[psana]
enable = true
ps_smd_n_events = 1000
ps_eb_nodes = 1
ps_srv_nodes = 0
ps_batch_size=100
ps_dir = "/global/cfs/cdirs/m2859/data/3iyf/xtc2"
exp = "xpptut15"
runnum = 1
ps_parallel = "legion"
[runtime]
N_images_per_rank = 1000
use_cuda = false
use_cufinufft = false
use_cupy = false
use_single_prec = false
chk_convergence = false
[gpu]
devices_per_node = 1
[algorithm]
N_generations = 3
nER = 50
nHIO = 25
beta = 0.3
cutoff = 0.05
N_phase_loops = 10
N_clipping = 0
N_binning = 0
N_orientations = 10000 # model_slices
N_batch_size = 100
Subproject commit e4006374e32bfd8b2c9644e756aae3265a564c3f
Subproject commit 6886cd006a7a4368baf36b3a71d860265596ba76
......@@ -3,7 +3,7 @@ import PyNVTX as nvtx
import os
import pygion
from pygion import acquire, attach_hdf5, execution_fence, task, Partition, Region, R, Tunable, WD
from pygion import acquire, attach_hdf5, execution_fence, task, Partition, Region, R, Tunable, WD, RO
from spinifel import settings, utils, contexts, checkpoint
from spinifel.prep import save_mrc
......@@ -15,43 +15,34 @@ from .orientation_matching import match
from . import mapper
from . import checkpoint
@task(replicable=True)
@nvtx.annotate("legion/main.py", is_prefix=True)
def main():
timer = utils.Timer()
total_procs = Tunable.select(Tunable.GLOBAL_PYS).get()
# Reading input images from hdf5
N_images_per_rank = settings.N_images_per_rank
batch_size = min(N_images_per_rank, 100)
max_events = min(settings.N_images_max, total_procs*N_images_per_rank)
writer_rank = 0 # pick writer rank as core 0
# Reading input images using psana2
ds = None
if settings.use_psana:
# For now, we use one smd chunk per node just to keep things simple.
# os.environ['PS_SMD_N_EVENTS'] = str(N_images_per_rank)
settings.ps_smd_n_events = N_images_per_rank
from psana import DataSource
logger.log("Using psana")
ds = DataSource(exp=settings.exp, run=settings.runnum,
dir=settings.data_dir, batch_size=batch_size,
max_events=max_events)
# Setup logger after knowing the writer rank
def load_psana():
logger = utils.Logger(True)
logger.log("In Legion main")
assert settings.use_psana == True
# Reading input images using psana2
# For now, we use one smd chunk per node just to keep things simple.
assert settings.ps_smd_n_events == settings.N_images_per_rank
from psana.psexp.tools import mode
from psana import DataSource
total_procs = Tunable.select(Tunable.GLOBAL_PYS).get()
batch_size = min(settings.N_images_per_rank, 100)
max_events = min(settings.N_images_max, total_procs*settings.N_images_per_rank)
logger.log(f'Using psana: exp={settings.ps_exp}, run={settings.ps_runnum}, dir={settings.ps_dir}, batch_size={batch_size}, max_events={max_events}, mode={mode}')
assert mode == 'legion'
ds = DataSource(exp=settings.ps_exp, run=settings.ps_runnum,
dir=settings.ps_dir, batch_size=batch_size,
max_events=max_events)
# Load unique set of intensity slices for python process
(pixel_position,
pixel_distance,
pixel_index,
pixel_index,
slices, slices_p) = get_data(ds)
logger.log(f"Loaded in {timer.lap():.2f}s.")
return pixel_position, pixel_distance, pixel_index, slices, slices_p
@task(privileges=[RO, RO, RO, RO])
def main_task(pixel_position, pixel_distance, pixel_index, slices, slices_p):
logger = utils.Logger(True)
timer = utils.Timer()
curr_gen = 0
if settings.load_gen > 0: # Load input from previous generation
curr_gen = settings.load_gen
......@@ -61,11 +52,8 @@ def main():
logger.log(f"AC recovered in {timer.lap():.2f}s.")
phased = phase(0, solved)
logger.log(f"Problem phased in {timer.lap():.2f}s.")
rho = np.fft.ifftshift(phased.rho_)
print('rho =', rho)
logger.log(f"Problem phased in {timer.lap():.2f}s.")
save_mrc(settings.out_dir / f"ac-0.mrc", phased.ac)
save_mrc(settings.out_dir / f"rho-0.mrc", rho)
......@@ -107,8 +95,6 @@ def main():
break;
rho = np.fft.ifftshift(phased.rho_)
print('rho =', rho)
save_mrc(settings.out_dir / f"ac-{generation}.mrc", phased.ac)
save_mrc(settings.out_dir / f"rho-{generation}.mrc", rho)
......@@ -116,5 +102,26 @@ def main():
logger.log(f"Results saved in {settings.out_dir}")
logger.log(f"Successfully completed in {timer.total():.2f}s.")
# read the data and run the main algorithm. This can be repeated
@nvtx.annotate("legion/main.py", is_prefix=True)
def main():
logger = utils.Logger(True)
logger.log("In Legion main")
ds = None
timer = utils.Timer()
# Reading input images using psana2
if settings.use_psana:
(pixel_position, pixel_distance, pixel_index, slices, slices_p) = load_psana()
# Reading input images from hdf5
else:
# Load unique set of intensity slices for python process
(pixel_position,
pixel_distance,
pixel_index,
slices, slices_p) = get_data(ds)
logger.log(f"Loaded in {timer.lap():.2f}s.")
main_task(pixel_position, pixel_distance, pixel_index, slices, slices_p)
......@@ -16,8 +16,7 @@ psana = None
if settings.use_psana:
import psana
from psana.psexp.legion_node import smd_chunks, smd_batches, batch_events
from psana.psexp.legion_node import smd_batches_without_transitions, smd_chunks_steps
@task(privileges=[WD])
@nvtx.annotate("legion/prep.py", is_prefix=True)
......@@ -25,34 +24,56 @@ def load_pixel_position(pixel_position):
prep.load_pixel_position_reciprocal(pixel_position.reciprocal)
@task(privileges=[WD])
@nvtx.annotate("legion/prep.py", is_prefix=True)
def load_pixel_position_psana(pixel_position,run):
pixel_position_val = 0
if run.expt == "xpptut15":
prep.load_pixel_position_reciprocal_psana(run,
pixel_position_val,
pixel_position.reciprocal)
# TODO - this needs to be setup for psana
elif run.expt == "amo06516":
pygion.fill(pixel_position, 'reciprocal', 0)
else:
assert False
@nvtx.annotate("legion/prep.py", is_prefix=True)
def get_pixel_position():
def get_pixel_position(run=None):
pixel_position_type = getattr(pygion, settings.pixel_position_type_str)
pixel_position = Region(settings.pixel_position_shape,
{'reciprocal': pixel_position_type})
load_pixel_position(pixel_position)
if run == None:
load_pixel_position(pixel_position)
else:
load_pixel_position_psana(pixel_position, run)
return pixel_position
@task(privileges=[WD])
@nvtx.annotate("legion/prep.py", is_prefix=True)
def load_pixel_index(pixel_index):
prep.load_pixel_index_map(pixel_index.map)
@task(privileges=[WD])
@nvtx.annotate("legion/prep.py", is_prefix=True)
def load_pixel_index_psana(pixel_index, run):
pixel_index.map[:] = np.moveaxis(run.beginruns[0].scan[0].raw.pixel_index_map, -1, 0)
@nvtx.annotate("legion/prep.py", is_prefix=True)
def get_pixel_index():
def get_pixel_index(run=None):
pixel_index_type = getattr(pygion, settings.pixel_index_type_str)
pixel_index = Region(settings.pixel_index_shape,
{'map': pixel_index_type})
load_pixel_index(pixel_index)
# psana
if run:
load_pixel_index_psana(pixel_index, run)
else:
load_pixel_index(pixel_index)
return pixel_index
# this is the equivalent of big data (the loop around batch_events)
@task(privileges=[WD])
@nvtx.annotate("legion/prep.py", is_prefix=True)
def load_slices_psana(slices, rank, N_images_per_rank, smd_chunk, run):
......@@ -70,6 +91,24 @@ def load_slices_psana(slices, rank, N_images_per_rank, smd_chunk, run):
print(f"{socket.gethostname()} loaded slices.", flush=True)
@task(privileges=[WD])
@nvtx.annotate("legion/prep.py", is_prefix=True)
def load_slices_psana2(slices, rank, N_images_per_rank, smd_chunk, run):
i = 0
if settings.verbosity > 0:
print(f'{socket.gethostname()} load_slices_psana2, rank = {rank}',flush=True)
det = run.Detector("amopnccd")
for smd_batch in smd_batches_without_transitions(smd_chunk, run):
for evt in batch_events(smd_batch,run):
raw = det.raw.calib(evt)
try:
slices.data[i] = raw
except IndexError:
raise RuntimeError(
f"Task received too many events.")
i += 1
if settings.verbosity > 0:
print(f"{socket.gethostname()} loaded slices.", flush=True)
@task(privileges=[WD])
@nvtx.annotate("legion/prep.py", is_prefix=True)
......@@ -82,8 +121,6 @@ def load_slices_hdf5(slices, rank, N_images_per_rank):
if settings.verbosity > 0:
print(f"{socket.gethostname()} loaded slices.", flush=True)
@nvtx.annotate("legion/prep.py", is_prefix=True)
def get_slices(ds):
N_images_per_rank = settings.N_images_per_rank
......@@ -91,21 +128,29 @@ def get_slices(ds):
sec_shape = settings.det_shape
slices, slices_p = lgutils.create_distributed_region(
N_images_per_rank, fields_dict, sec_shape)
pixel_position = None
pixel_index = None
if ds is not None:
n_nodes = Tunable.select(Tunable.NODE_COUNT).get()
chunk_i = 0
runs = list(ds.runs())
pygion.execution_fence(block=True)
for run in runs:
for smd_chunk in smd_chunks(run):
val = 0
for run in ds.runs():
# load pixel index map and pixel position reciprocal only once
if val == 0:
# pixel index map
pixel_index = get_pixel_index(run)
# TODO pixel_position.reciprocal for "am06516"
pixel_position = get_pixel_position(run)
val= val+1
for smd_chunk, step_data in smd_chunks_steps(run):
i = chunk_i % n_nodes
load_slices_psana(slices_p[i], i, N_images_per_rank, smd_chunk, run, point=i)
load_slices_psana2(slices_p[i], i, N_images_per_rank, bytearray(smd_chunk), run, point=i)
chunk_i += 1
pygion.execution_fence(block=True)
else:
for i, slices_subr in enumerate(slices_p):
load_slices_hdf5(slices_subr, i, N_images_per_rank, point=i)
return slices, slices_p
return slices, slices_p, pixel_position, pixel_index
@task(privileges=[WD])
......@@ -229,7 +274,6 @@ def show_image(pixel_index, images, image_index, name):
image.show_image(pixel_index.map, images.data[image_index], name)
@task(privileges=[RO, RO])
@nvtx.annotate("legion/prep.py", is_prefix=True)
def export_saxs(pixel_distance, mean_image, name):
......@@ -242,9 +286,16 @@ def export_saxs(pixel_distance, mean_image, name):
@nvtx.annotate("legion/prep.py", is_prefix=True)
def get_data(ds):
pixel_position = get_pixel_position()
pixel_index = get_pixel_index()
slices, slices_p = get_slices(ds)
print(f"{socket.gethostname()} loading slices.", flush=True)
# if psana load pixel_position/pixel_index using first 'run'
slices, slices_p, pixel_position, pixel_index = get_slices(ds)
# if not psana - load pixel_postion/pixel_index from hdf5 file
if ds == None:
pixel_position = get_pixel_position(ds)
pixel_index = get_pixel_index()
mean_image = compute_mean_image(slices, slices_p)
show_image(pixel_index, slices_p[0], 0, "image_0.png")
show_image(pixel_index, mean_image, ..., "mean_image.png")
......
......@@ -192,3 +192,12 @@ def save_mrc(savename, data, voxel_size=None):
mrc.voxel_size = voxel_size
mrc.close()
return
@nvtx.annotate("prep.py", is_prefix=True)
def load_pixel_position_reciprocal_psana(run, pixel_position, pixel_position_reciprocal):
if run.expt == "xpptut15":
pixel_position_reciprocal[:] = np.moveaxis(run.beginruns[0].scan[0].raw.pixel_position_reciprocal[:], -1, 0)
elif run.expt == "amo06516":
pixel_position = run.beginruns[0].scan[0].pixel_position
else:
assert False
......@@ -255,6 +255,21 @@ class SpinifelSettings(metaclass=Singleton):
str, "xpptut1",
"PSANA experiment name"
),
"_ps_batch_size": (
"psana", "ps_batch_size",
int, 100,
"PSANA batch size"
),
"_ps_dir": (
"psana", "ps_dir",
str, "",
"PSANA xtc2 directory"
),
"_ps_parallel": (
"psana", "ps_parallel",
str, "legion",
"Use legion or mpi mode for PSANA"
),
"_ps_runnum": (
"psana", "runnum",
int, 1,
......@@ -415,6 +430,7 @@ class SpinifelSettings(metaclass=Singleton):
"PS_SMD_N_EVENTS": ("_ps_smd_n_events", get_int),
"PS_EB_NODES": ("_ps_eb_nodes", get_int),
"PS_SRV_NODES": ("_ps_srv_nodes", get_int),
"PS_PARALLEL": ("_ps_parallel", get_str),
"USE_CALLMONITOR": ("_use_callmonitor", get_bool),
"CHK_CONVERGENCE": ("_chk_convergence", get_bool)
}
......@@ -660,6 +676,20 @@ class SpinifelSettings(metaclass=Singleton):
# update derived environment variable
environ["PS_SRV_NODES"] = str(val)
@property
def ps_parallel(self):
"""
ps parallel mode
"""
return self._ps_parallel
@ps_parallel.setter
def ps_parallel(self, val):
"""
update derived environment variable
"""
self._ps_parallel = val
environ["PS_PARALLEL"] = val
@property
def pixel_position_shape(self):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment