From 62bf0837b503736a5c514584f56824778a18b784 Mon Sep 17 00:00:00 2001 From: Jan Kieseler Date: Fri, 13 Dec 2024 11:13:54 +0100 Subject: [PATCH] changed underlying geant encapsulation completely --- bind/G4Calo.py | 342 ++++++++++++++++++----------------------- bind/G4Calo_exec.py | 93 +++++++++++ bind/minicalo_tools.py | 177 +++++++++++++++++++++ docker/Dockerfile | 4 +- 4 files changed, 424 insertions(+), 192 deletions(-) create mode 100755 bind/G4Calo_exec.py create mode 100644 bind/minicalo_tools.py diff --git a/bind/G4Calo.py b/bind/G4Calo.py index fb99a49..f78d3d3 100644 --- a/bind/G4Calo.py +++ b/bind/G4Calo.py @@ -1,93 +1,186 @@ -import multiprocessing +import threading import warnings import contextlib import sys +from minicalo_tools import TempFileManager +from concurrent.futures import ThreadPoolExecutor + +''' +give G4Calo package an executable that takes input and output json files names + +This calls the actual Geant stuff. +The input are all the arguments and a geometry descriptor +The output is the data frame as well as the changed geometry descriptor + +The user only interacts with a wrapper that calls this executable through a system call +''' -def set_start_method(): - try: - multiprocessing.set_start_method('spawn', force=True) - except RuntimeError: - warnings.warn("Failed to set start method to 'spawn'. It may have already been set.") - -set_start_method() from minicalo import GeometryDescriptor -from minicalo import G4System as _G4System import plotly.graph_objects as go import pandas as pd import numpy as np import os -import glob import subprocess -import uproot -import time +import tempfile from IPython.display import Image, display -class __G4System(_G4System): - +def _run_mini_batch( + cw : GeometryDescriptor, + nEvents: int, + particleSpec , + minEnergy_GeV: float, + maxEnergy_GeV: float = -1.0, + seed : int = 0, + tmp_fname : str = "", + silent : bool = True, + return_geometry : bool = False): - def run_batch( - self, + to_process = {'geometry': cw, + 'nEvents': nEvents, + 'particleSpec': particleSpec, + 'minEnergy_GeV': minEnergy_GeV, + 'maxEnergy_GeV': maxEnergy_GeV, + 'temp_filename': tmp_fname, + 'seed': seed} + + # silent = False #DEBUG + + with TempFileManager() as tfm: + tfm.dump_input(to_process) + if silent: + subprocess.run(["python", "G4Calo_exec.py", tfm.input_file, tfm.output_file], + stdout=subprocess.DEVNULL)#, stderr=subprocess.DEVNULL) + else: + subprocess.run(["python", "G4Calo_exec.py", tfm.input_file, tfm.output_file]) + data = tfm.read_output() + if return_geometry: + return data['df'], data['geometry'] + else: + return data['df'] + +def _run_mini_batch_wrapper(all_args): + return _run_mini_batch(*all_args) + +def create_temp_root_fname(): + #check if we can put the temporary files in /dev/shm, otherwise use /tmp + if os.path.exists('/dev/shm'): + tmpdir = '/dev/shm' + else: + tmpdir = '/tmp' + #create a temp file name + tmpfilename = tempfile.NamedTemporaryFile(delete=False, dir=tmpdir, suffix=".root").name + #delete file as we only need the name + os.remove(tmpfilename) + return tmpfilename + +def run_batch( + gd : GeometryDescriptor, nEvents: int, particleSpec , minEnergy_GeV: float, maxEnergy_GeV: float = -1.0, filename: str = "", - ): + no_mp: bool = False, + manual_seed : int = -1): + ''' + splits the batch in jobs depending on how many cores are available and runs mini batches in parallel + this will create threads. + If no_mp is set to True, it will run the batch in a single thread + The seeding is not thread safe (unless provided by hand) + ''' + assert nEvents > 0 + assert minEnergy_GeV > 0 - if maxEnergy_GeV < 0: - maxEnergy_GeV = minEnergy_GeV + if not isinstance(particleSpec, str): + assert isinstance(particleSpec, list), "particleSpec must be a string or a list of strings" + for p in particleSpec: + assert isinstance(p, str), "particleSpec must be a string or a list of strings" - - #check if particleSpec is a list of strings or a single string and assert - if not isinstance(particleSpec, str): - assert isinstance(particleSpec, list), "particleSpec must be a string or a list of strings" - for p in particleSpec: - assert isinstance(p, str), "particleSpec must be a string or a list of strings" + nCores = os.cpu_count() + #make sure to adjust cores such that at least 80 events are run per core + nCores = min(nCores, nEvents // 80 + 1) + if no_mp: + nCores = 1 - if isinstance(particleSpec, str): - particleSpec = [particleSpec] + print(f"G4Calo: Running on {nCores} cores") + nEventsPerCore = nEvents // nCores + #print(f"Running {nEventsPerCore} events per core") + nEventsLastCore = nEvents - nEventsPerCore * (nCores - 1) + #print(f"Running {nEventsLastCore} events on last core") + + nevents = [nEventsPerCore if i < nCores - 1 else nEventsLastCore for i in range(nCores)] - save_file = len(filename) > 0 - # filename without file ending(!) - filename = "._" + str(time.perf_counter_ns()) + ".root" - _G4System.run_batch(self, nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename) - - # TO FIX: Geant4 adds "t" to the filename, circumvent this for one thread, but this is not a good solution - file = glob.glob(filename.replace(".root", "*.root")) - print(file) - if len(file)>1: - raise Exception("More than one file found! - Check mulithreading") + if manual_seed >= 0: + seed = manual_seed + else: + if os.path.exists(os.path.expanduser("~/.g4calo_seeds.txt")): + with open(os.path.expanduser("~/.g4calo_seeds.txt"), "r") as f: + seed = int(f.read()) else: - filename = file[0] + seed = 1 + with open(os.path.expanduser("~/.g4calo_seeds.txt"), "w") as f: + f.write(str(seed + nCores)) - # conversion from root to pandas dataframe - - # TO FIX: circumvent index out of bounds error - ttree = uproot.open(filename) - try: - df = ttree["Hits;1"].arrays(library="pd") - except: - df = index_out_of_bounds_workaround(ttree["Hits;1"]) + print(f"G4Calo: Batch seed: {seed}") - # save to pickle and delete root file - if save_file: - df.to_pickle(filename.replace(".root", ".pkl")) - os.system("rm "+filename) - - return df + + + #create a list of temporary file names, one for each core, no full path, just the names. + tmpfile = [f"{create_temp_root_fname()}_g4calo_{i}.root" for i in range(nCores)] + + args = [(gd, + nevents[i], + particleSpec, + minEnergy_GeV, + maxEnergy_GeV, + seed+i, + tmpfile[i], + True, + False) for i in range(nCores)] + + try: + if no_mp: + dfs = [] + for i in range(nCores): + dfs.append(_run_mini_batch(*args[i])) + else: + + with ThreadPoolExecutor(max_workers=nCores) as executor: + dfs = list(executor.map(_run_mini_batch_wrapper, args)) + + except Exception as e: + print(e) + raise e + finally: #make sure to clean up the temporary files in case of an exception + for f in tmpfile: + if os.path.exists(f): + os.remove(f) + + print('G4Calo: simulation finished, concatenating dataframes') + + alldf = pd.concat(dfs) + alldf.reset_index(drop=True, inplace=True) + + if len(filename): + alldf.to_pickle(filename) + + return alldf + +def _fill_event(gd : GeometryDescriptor, + particleSpec, + energy: float, + seed: int = -1): + + return _run_mini_batch(gd, 1, particleSpec, energy, return_geometry=True, seed=seed, tmp_fname=create_temp_root_fname())[1] - def displayEvent(self, logE = False, renderer=None): - raise NotImplementedError("This method is not implemented anymore. Use the function display_event instead.") - - - -# +# +## below only for visualization # some helpers # @@ -131,143 +224,12 @@ col_dict = { "G4_BRASS": 'slategrey', } -def index_out_of_bounds_workaround(tbranch): - # ugly workaround for index out of bounds error - dfs = [] - entries = tbranch['sensor_energy'].num_entries - - # check each entry - for entry in range(entries): - try: - df = tbranch.arrays(entry_start=entry,entry_stop=entry+1, library='pd') - dfs.append(df) - except: - continue - - df = pd.concat(dfs) - return df.reset_index(drop=True) - - - -def _run_mini_batch( - cw : GeometryDescriptor, - nEvents: int, - particleSpec , - minEnergy_GeV: float, - maxEnergy_GeV: float = -1.0, - batch_seed : int = 0, - counter : int = 0): - - - #this is now encapsuled - from G4Calo import __G4System - G4System = __G4System() - seed = int(batch_seed + counter) - - print(f"Running mini batch with seed {seed}") - G4System.init(cw, seed) #counter gives random seed offset - - df = G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV,"") - return df - -def _run_mini_batch_silent(*args, **kwargs): - #I tried - return _run_mini_batch(*args, **kwargs) - -_g4calo_threading_lock = None - -def init_g4calo_threading_lock(): - global _g4calo_threading_lock - _g4calo_threading_lock = multiprocessing.Lock() - - -def run_batch( - gd : GeometryDescriptor, - nEvents: int, - particleSpec , - minEnergy_GeV: float, - maxEnergy_GeV: float = -1.0, - filename: str = "", - no_mp: bool = False, - manual_seed : int = -1): - ''' - splits the batch in jobs depending on how many cores are available and runs mini batches in parallel - ''' - assert nEvents > 0 - assert minEnergy_GeV > 0 - - if not isinstance(particleSpec, str): - assert isinstance(particleSpec, list), "particleSpec must be a string or a list of strings" - for p in particleSpec: - assert isinstance(p, str), "particleSpec must be a string or a list of strings" - - nCores = multiprocessing.cpu_count() - #make sure to adjust cores such that at least 80 events are run per core - nCores = min(nCores, nEvents // 80 + 1) - - if no_mp: - nCores = 1 - - print(f"Running on {nCores} cores") - nEventsPerCore = nEvents // nCores - print(f"Running {nEventsPerCore} events per core") - nEventsLastCore = nEvents - nEventsPerCore * (nCores - 1) - print(f"Running {nEventsLastCore} events on last core") - - nevents = [nEventsPerCore if i < nCores - 1 else nEventsLastCore for i in range(nCores)] - - if manual_seed >= 0: - seed = manual_seed - else: - with _g4calo_threading_lock if _g4calo_threading_lock is not None else contextlib.nullcontext(): - # check if file exists, if so, read last seed. If not create it - if os.path.exists(os.path.expanduser("~/.g4calo_seeds.txt")): - with open(os.path.expanduser("~/.g4calo_seeds.txt"), "r") as f: - seed = int(f.read()) - else: - seed = 0 - with open(os.path.expanduser("~/.g4calo_seeds.txt"), "w") as f: - f.write(str(seed + nCores)) - - seed += 1 - print(f"Batch seed: {seed}") - - #if no_mp: - # dfs = [] - # for i in range(nCores): - # dfs.append(_run_mini_batch_silent(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, seed, i)) - #else: - with multiprocessing.Pool(nCores) as pool: - dfs = pool.starmap(_run_mini_batch_silent, [(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, seed, i) for i in range(nCores)]) - - print('simulation finished') - alldf = pd.concat(dfs) - alldf.reset_index(drop=True, inplace=True) - - if len(filename): - alldf.to_pickle(filename) - - return alldf - -def _fill_event(gd : GeometryDescriptor, - particleSpec, - energy: float, - seed: int = -1): - - from G4Calo import __G4System - G4System = __G4System() - G4System.init(gd, seed) - G4System.run_batch(1, particleSpec, energy, energy,"") - return gd - def display_event(gd : GeometryDescriptor, particleSpec, energy: float, logE = False, renderer=None, seed = -1, outfile : str = ""): - #run _fill_event in forked mode using 1-core multiprocessing to avoid G4 singletons to interfere - with multiprocessing.Pool(1) as pool: - gd = pool.apply(_fill_event, (gd, particleSpec, energy)) + gd = _fill_event(gd, particleSpec, energy) print('creating plot') #use gd to plot diff --git a/bind/G4Calo_exec.py b/bind/G4Calo_exec.py new file mode 100755 index 0000000..9c9b678 --- /dev/null +++ b/bind/G4Calo_exec.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +from minicalo import GeometryDescriptor +from minicalo import G4System +from minicalo_tools import TempFileManager +import uproot +import glob +import pandas as pd +import os + + +def index_out_of_bounds_workaround(tbranch): + # ugly workaround for index out of bounds error + dfs = [] + entries = tbranch['sensor_energy'].num_entries + + # check each entry + for entry in range(entries): + try: + df = tbranch.arrays(entry_start=entry,entry_stop=entry+1, library='pd') + dfs.append(df) + except: + continue + + df = pd.concat(dfs) + return df.reset_index(drop=True) + +def run_batch( + geometry: GeometryDescriptor, + nEvents: int, + particleSpec , + minEnergy_GeV: float, + maxEnergy_GeV: float, + temp_filename: str, + seed: int + ): + + if maxEnergy_GeV < 0: + maxEnergy_GeV = minEnergy_GeV + + #check if particleSpec is a list of strings or a single string and assert + if not isinstance(particleSpec, str): + assert isinstance(particleSpec, list), "particleSpec must be a string or a list of strings" + for p in particleSpec: + assert isinstance(p, str), "particleSpec must be a string or a list of strings" + if isinstance(particleSpec, str): + particleSpec = [particleSpec] + + # filename without file ending(!) + if not temp_filename.endswith(".root"): + filename = temp_filename+ ".root" + else: + filename = temp_filename + + _G4System = G4System() # create instance of G4System only here + _G4System.init(geometry, seed) + + try: + _G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename) + + # TO FIX: Geant4 adds "t" to the filename, circumvent this for one thread, but this is not a good solution + file = glob.glob(filename.replace(".root", "*.root")) + if len(file)>1: + raise Exception("More than one file found! - Check mulithreading") + else: + filename = file[0] + # conversion from root to pandas dataframe + + # TO FIX: circumvent index out of bounds error + ttree = uproot.open(filename) + try: + df = ttree["Hits;1"].arrays(library="pd") + except: + df = index_out_of_bounds_workaround(ttree["Hits;1"]) + # save to pickle and delete root file + except Exception as e: + raise e + finally: + os.remove(filename) + + return {'df':df, 'geometry':geometry} + +if __name__ == '__main__': + # do stuff + # get the first two arguments that correspond to input and output file + import sys + assert len(sys.argv) == 3, "Please provide input and output file" + + io = TempFileManager.slave_init(sys.argv[1], sys.argv[2]) + input_data = io.read_input() + output_data = run_batch(**input_data) + io.dump_output(output_data) diff --git a/bind/minicalo_tools.py b/bind/minicalo_tools.py new file mode 100644 index 0000000..c1524b5 --- /dev/null +++ b/bind/minicalo_tools.py @@ -0,0 +1,177 @@ +# this will be shared between isolated parts and exposed parts + +import os +import tempfile +import pickle + + +class TempFileManager: + def __init__(self, input_file=None, output_file=None, use_shm=True, file_permissions=0o644): + """ + Initialize the TempFileManager. + :param input_file: Path to the input file (optional). + :param output_file: Path to the output file (optional). + :param use_shm: Whether to attempt using /dev/shm for temporary files. + """ + self.input_file = input_file + self.output_file = output_file + self.use_shm = use_shm and os.path.exists('/dev/shm') + self.file_permissions = file_permissions + + @staticmethod + def slave_init(input_file, output_file): + """ + Initialize a TempFileManager instance for use with external filenames. + :param input_file: Path to the input file. + :param output_file: Path to the output file. + :return: A TempFileManager instance with the provided filenames. + """ + return TempFileManager(input_file=input_file, output_file=output_file) + + def __enter__(self): + """ + Create temporary files when entering the context, if not already provided. + """ + if not self.input_file: + base_dir = '/dev/shm' if self.use_shm else None + self.input_file = tempfile.NamedTemporaryFile(delete=False, dir=base_dir, suffix=".pkl").name + if not self.output_file: + base_dir = '/dev/shm' if self.use_shm else None + self.output_file = tempfile.NamedTemporaryFile(delete=False, dir=base_dir, suffix=".pkl").name + return self + + def dump_input(self, obj): + """ + Serialize and write the object to the input file. + :param obj: The object to serialize and write. + """ + with open(self.input_file, 'wb') as f: + pickle.dump(obj, f) + os.chmod(self.input_file, self.file_permissions) + + def read_input(self): + """ + Read and deserialize the input file. + :return: The deserialized object. + """ + with open(self.input_file, 'rb') as f: + return pickle.load(f) + + def dump_output(self, obj): + """ + Serialize and write the object to the output file. + :param obj: The object to serialize and write. + """ + with open(self.output_file, 'wb') as f: + pickle.dump(obj, f) + os.chmod(self.output_file, self.file_permissions) + + def read_output(self): + """ + Read and deserialize the output file. + :return: The deserialized object. + """ + with open(self.output_file, 'rb') as f: + return pickle.load(f) + + def __exit__(self, exc_type, exc_val, exc_tb): + """ + Cleanup temporary files when exiting the context. + """ + if os.path.exists(self.input_file): + os.remove(self.input_file) + if os.path.exists(self.output_file): + os.remove(self.output_file) + + + + +import unittest +import os +import json +import stat + + +class TestTempFileManager(unittest.TestCase): + + def test_file_creation_and_cleanup(self): + """Test that temporary files are created and cleaned up properly.""" + with TempFileManager() as tfm: + # Check if input and output files are created + self.assertTrue(os.path.exists(tfm.input_file)) + self.assertTrue(os.path.exists(tfm.output_file)) + input_path = tfm.input_file + output_path = tfm.output_file + + # After context exit, files should be deleted + self.assertFalse(os.path.exists(input_path)) + self.assertFalse(os.path.exists(output_path)) + + def test_dump_and_read(self): + """Test dumping an object to input and reading it back.""" + obj = {"key": "value", "number": 42} + with TempFileManager() as tfm: + # Dump object to input file + tfm.dump_input(obj) + + # Read it back to ensure correctness + with open(tfm.input_file, 'r') as f: + data = json.load(f) + self.assertEqual(data, obj) + + def test_slave_init(self): + """Test the slave_init static method for external file names.""" + input_file = "/tmp/test_input.json" + output_file = "/tmp/test_output.json" + + try: + # Create test input file + obj = {"key": "test"} + with open(input_file, 'w') as f: + json.dump(obj, f) + + # Initialize TempFileManager with slave_init + tfm = TempFileManager.slave_init(input_file, output_file) + + # Read the input file and verify content + data = tfm.read_input() + self.assertEqual(data, obj) + + # Write to the output file and verify + updated_obj = {"key": "updated"} + tfm.dump_output(updated_obj) + with open(output_file, 'r') as f: + output_data = json.load(f) + self.assertEqual(output_data, updated_obj) + finally: + # Cleanup test files + if os.path.exists(input_file): + os.unlink(input_file) + if os.path.exists(output_file): + os.unlink(output_file) + + def test_with_dev_shm(self): + """Test that /dev/shm is used if available.""" + if os.path.exists('/dev/shm'): + with TempFileManager() as tfm: + self.assertTrue(tfm.input_file.startswith('/dev/shm')) + self.assertTrue(tfm.output_file.startswith('/dev/shm')) + + def test_without_dev_shm(self): + """Test fallback to default location if /dev/shm is unavailable.""" + with TempFileManager(use_shm=False) as tfm: + self.assertFalse(tfm.input_file.startswith('/dev/shm')) + self.assertFalse(tfm.output_file.startswith('/dev/shm')) + + def test_permissions_after_dump(self): + """Test that file permissions are correctly set after dumping content.""" + with TempFileManager(file_permissions=0o600) as tfm: + # Dump content to the input file + tfm.dump_input({"key": "value"}) + + # Check that permissions are applied after dumping + permissions = stat.S_IMODE(os.stat(tfm.input_file).st_mode) + self.assertEqual(permissions, 0o600) + +#if __name__ == "__main__": +# unittest.main() diff --git a/docker/Dockerfile b/docker/Dockerfile index a9a8c6b..cb1ee7e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -95,9 +95,9 @@ ADD minicalosim /root/minicalosim RUN cd /root/minicalosim && git checkout $COMMIT && \ mkdir -p build && cd build && rm -rf * && cmake ../ && make -j4 &&\ - cp minicalo* ../bind/G4Calo.py /usr/local/lib/python3.8/dist-packages/ - + cp minicalo* ../bind/G4Calo.py ../bind/minicalo_tools.py /usr/local/lib/python3.8/dist-packages/ +RUN cp /root/minicalosim/bind/G4Calo_exec.py /usr/local/bin/ RUN cp /root/minicalosim/docker/start-notebook.py /usr/local/bin/ RUN cp /root/minicalosim/docker/start-notebook.sh /usr/local/bin/