restage 0.2.4__py3-none-any.whl → 0.3.1__py3-none-any.whl

restage/cache.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from mccode_antlr.instr import Instr
 from .tables import InstrEntry, SimulationTableEntry, SimulationEntry
+from mccode_antlr.compiler.c import CBinaryTarget
 
 
 def setup_database(named: str):
@@ -37,40 +38,41 @@ def directory_under_module_data_path(sub: str, prefix=None, suffix=None, name=No
     return Path(mkdtemp(dir=under, prefix=prefix or '', suffix=suffix or ''))
 
 
-def _compile_instr(entry: InstrEntry, instr: Instr, config: dict | None = None, target=None, generator=None):
+def _compile_instr(entry: InstrEntry, instr: Instr, config: dict | None = None,
+                   mpi: bool = False, acc: bool = False,
+                   target=None, generator=None):
+    from tempfile import mkdtemp
     from mccode_antlr import __version__
     from mccode_antlr.compiler.c import compile_instrument, CBinaryTarget
     if config is None:
         config = dict(default_main=True, enable_trace=False, portable=False, include_runtime=True,
                       embed_instrument_file=False, verbose=False)
     if target is None:
-        target = CBinaryTarget(mpi=False, acc=False, count=1, nexus=False)
+        target = CBinaryTarget(mpi=mpi or False, acc=acc or False, count=1, nexus=False)
     if generator is None:
         from mccode_antlr.translators.target import MCSTAS_GENERATOR
         generator = MCSTAS_GENERATOR
 
-    output = module_data_path('bin').joinpath(entry.id)
-    if not output.exists():
-        output.mkdir(parents=True)
-
+    output = directory_under_module_data_path('bin')
     binary_path = compile_instrument(instr, target, output, generator=generator, config=config)
     entry.mccode_version = __version__
     entry.binary_path = str(binary_path)
     return entry
 
 
-def cache_instr(instr: Instr, mccode_version=None, binary_path=None, **kwargs) -> InstrEntry:
+def cache_instr(instr: Instr, mpi: bool = False, acc: bool = False, mccode_version=None, binary_path=None, **kwargs) -> InstrEntry:
     instr_contents = str(instr)
-    query = DATABASE.query_instr_file(search={'file_contents': instr_contents})  # returns a list[InstrTableEntry]
+    # the query returns a list[InstrTableEntry]
+    query = DATABASE.query_instr_file(search={'file_contents': instr_contents, 'mpi': mpi, 'acc': acc})
     if len(query) > 1:
         raise RuntimeError(f"Multiple entries for {instr_contents} in {DATABASE.instr_file_table}")
     elif len(query) == 1:
         return query[0]
 
-    instr_file_entry = InstrEntry(file_contents=instr_contents, binary_path=binary_path or '',
+    instr_file_entry = InstrEntry(file_contents=instr_contents, mpi=mpi, acc=acc, binary_path=binary_path or '',
                                   mccode_version=mccode_version or 'NONE')
     if binary_path is None:
-        instr_file_entry = _compile_instr(instr_file_entry, instr, **kwargs)
+        instr_file_entry = _compile_instr(instr_file_entry, instr, mpi=mpi, acc=acc, **kwargs)
 
     DATABASE.insert_instr_file(instr_file_entry)
     return instr_file_entry

restage/mcpl.py

@@ -4,11 +4,11 @@ from pathlib import Path
 def mcpl_real_filename(filename: Path) -> Path:
     """MCPL_output from McCode instruments has the bad habit of changing the output file name silently.
     Find the _real_ output file name by looking for the expected variants"""
-    if filename.exists():
+    if filename.exists() and filename.is_file():
         return filename
-    if filename.with_suffix('.mcpl').exists():
+    if filename.with_suffix('.mcpl').exists() and filename.with_suffix('.mcpl').is_file():
         return filename.with_suffix('.mcpl')
-    if filename.with_suffix('.mcpl.gz').exists():
+    if filename.with_suffix('.mcpl.gz').exists() and filename.with_suffix('.mcpl.gz').is_file():
         return filename.with_suffix('.mcpl.gz')
     raise FileNotFoundError(f'Could not find MCPL file {filename}')
 
@@ -37,7 +37,7 @@ def mcpl_particle_count(filename):
     return int(m.group(1))
 
 
-def mcpl_merge_files(files: list[Path], filename: str, keep_originals: bool = False):
+def mcpl_merge_files(files: list[Path], filepath: Path, keep_originals: bool = False):
     """Merge a list of MCPL files into a single file using mcpltool.
 
     :param files: The list of files to merge.
@@ -51,6 +51,14 @@ def mcpl_merge_files(files: list[Path], filename: str, keep_originals: bool = Fa
     """
     from subprocess import run
     real_filenames = [mcpl_real_filename(f) for f in files]
+    # if the real filenames have .mcpl or .mcpl.gz, the merged filename should too
+    ext = ''
+    if real_filenames[0].name.endswith('.mcpl.gz'):
+        ext = '.mcpl.gz'
+    elif real_filenames[0].name.endswith('.mcpl'):
+        ext = '.mcpl'
+    filename = filepath.with_suffix(ext).as_posix()
+
     command = ['mcpltool', '--merge', filename] + [str(f) for f in real_filenames]
     result = run(command)
     if result.returncode != 0:
@@ -58,3 +66,21 @@ def mcpl_merge_files(files: list[Path], filename: str, keep_originals: bool = Fa
     if not keep_originals:
         for file in real_filenames:
             file.unlink()
+
+
+def mcpl_rename_file(source: Path, dest: Path, strict: bool = False):
+    filepath = mcpl_real_filename(source)
+    filename = filepath.name  # this could be '{name}', '{name}.mcpl', or '{name}.mcpl.gz'
+    ext = ''
+    if filepath.name.endswith('.mcpl.gz'):
+        ext = '.mcpl.gz'
+    elif filepath.name.endswith('.mcpl'):
+        ext = '.mcpl'
+
+    if not dest.name.endswith(ext):
+        if strict:
+            raise RuntimeError(f"Destination {dest} does not have extension matching {source}")
+        dest = dest.with_suffix(ext)
+
+    filepath.rename(dest)
+    return dest

restage/splitrun.py

@@ -4,6 +4,7 @@ from typing import Union
 from pathlib import Path
 from .range import Singular, MRange
 from .tables import SimulationEntry, InstrEntry
+from mccode_antlr.compiler.c import CBinaryTarget
 
 
 def make_splitrun_parser():
@@ -30,6 +31,12 @@ def make_splitrun_parser():
        help='Maximum number of particles to simulate during first instrument simulations')
     aa('--dryrun', action='store_true', default=False,
        help='Do not run any simulations, just print the commands')
+    aa('--parallel', action='store_true', default=False,
+       help='Use MPI multi-process parallelism (primary instrument only at the moment)')
+    aa('--gpu', action='store_true', default=False,
+       help='Use GPU OpenACC parallelism (primary instrument only at the moment)')
+    aa('--process-count',  nargs=1, type=int, default=0,
+       help='MPI process count, 0 == System Default')
     # splitrun controlling parameters
     aa('--split-at', nargs=1, type=str, default=['mcpl_split'],
        help='Component at which to split -- DEFAULT: mcpl_split')
@@ -153,7 +160,10 @@ def splitrun_from_file(args, parameters, precision):
              format=args.format[0] if args.format is not None else None,
              minimum_particle_count=args.nmin[0] if args.nmin is not None else None,
              maximum_particle_count=args.nmax[0] if args.nmax is not None else None,
-             dry_run=args.dryrun
+             dry_run=args.dryrun,
+             parallel=args.parallel,
+             gpu=args.gpu,
+             process_count=args.process_count,
              )
 
 
@@ -161,6 +171,7 @@ def splitrun(instr, parameters, precision: dict[str, float], split_at=None, grid
              minimum_particle_count=None,
              maximum_particle_count=None,
              dry_run=False,
+             parallel=False, gpu=False, process_count=0,
              callback=None, callback_arguments: dict[str, str] | None = None,
              output_split_instrs=True,
              **runtime_arguments):
@@ -190,7 +201,7 @@ def splitrun(instr, parameters, precision: dict[str, float], split_at=None, grid
     pre_entry = splitrun_pre(pre, pre_parameters, grid, precision, **runtime_arguments,
                              minimum_particle_count=minimum_particle_count,
                              maximum_particle_count=maximum_particle_count,
-                             dry_run=dry_run)
+                             dry_run=dry_run, parallel=parallel, gpu=gpu, process_count=process_count)
 
     splitrun_combined(pre_entry, pre, post, pre_parameters, post_parameters, grid, precision,
                       dry_run=dry_run, callback=callback, callback_arguments=callback_arguments, **runtime_arguments)
@@ -198,16 +209,16 @@ def splitrun(instr, parameters, precision: dict[str, float], split_at=None, grid
 
 def splitrun_pre(instr, parameters, grid, precision: dict[str, float],
                  minimum_particle_count=None, maximum_particle_count=None, dry_run=False,
+                 parallel=False, gpu=False, process_count=0,
                  **runtime_arguments):
 
     from functools import partial
     from .cache import cache_instr
     from .energy import energy_to_chopper_translator
     from .range import parameters_to_scan
-
     # check if this instr is already represented in the module's cache database
     # if not, it is compiled and added to the cache with (hopefully sensible) defaults specified
-    entry = cache_instr(instr)
+    entry = cache_instr(instr, mpi=parallel, acc=gpu)
     # get the function with converts energy parameters to chopper parameters:
     translate = energy_to_chopper_translator(instr.name)
     # determine the scan in the user-defined parameters!
@@ -216,7 +227,7 @@ def splitrun_pre(instr, parameters, grid, precision: dict[str, float],
     sit_kw = {'seed': args.get('seed'), 'ncount': args.get('ncount'), 'gravitation': args.get('gravitation', False)}
 
     step = partial(_pre_step, instr, entry, names, precision, translate, sit_kw, minimum_particle_count,
-                   maximum_particle_count, dry_run)
+                   maximum_particle_count, dry_run, process_count)
 
     # this does not work due to the sqlite database being locked by the parallel processes
     # from joblib import Parallel, delayed
@@ -230,7 +241,7 @@ def splitrun_pre(instr, parameters, grid, precision: dict[str, float],
     return entry
 
 
-def _pre_step(instr, entry, names, precision, translate, kw, min_pc, max_pc, dry_run, values):
+def _pre_step(instr, entry, names, precision, translate, kw, min_pc, max_pc, dry_run, process_count, values):
     """The per-step function for the primary instrument simulation. Broken out for parallelization"""
     from .instr import collect_parameter_dict
     from .cache import cache_has_simulation, cache_simulation, cache_get_simulation
@@ -240,7 +251,8 @@ def _pre_step(instr, entry, names, precision, translate, kw, min_pc, max_pc, dry
         sim.output_path = do_primary_simulation(sim, entry, nv, kw,
                                                 minimum_particle_count=min_pc,
                                                 maximum_particle_count=max_pc,
-                                                dry_run=dry_run)
+                                                dry_run=dry_run,
+                                                process_count=process_count)
         cache_simulation(entry, sim)
     return cache_get_simulation(entry, sim)
 
@@ -255,7 +267,7 @@ def splitrun_combined(pre_entry, pre, post, pre_parameters, post_parameters, gri
     from .instr import collect_parameter_dict
     from .tables import best_simulation_entry_match
     from .emulate import mccode_sim_io, mccode_dat_io, mccode_dat_line
-    instr_entry = cache_instr(post)
+    instr_entry = cache_instr(post, mpi=False, acc=False)
     args = regular_mccode_runtime_dict(runtime_arguments)
     sit_kw = {'seed': args.get('seed'), 'ncount': args.get('ncount'), 'gravitation': args.get('gravitation', False)}
     # recombine the parameters to ensure the 'correct' scan is performed
@@ -341,7 +353,8 @@ def do_primary_simulation(sit: SimulationEntry,
                           args: dict,
                           minimum_particle_count: int | None = None,
                           maximum_particle_count: int | None = None,
-                          dry_run: bool = False
+                          dry_run: bool = False,
+                          process_count: int = 0,
                           ):
     from zenlog import log
     from pathlib import Path
@@ -349,10 +362,11 @@ def do_primary_simulation(sit: SimulationEntry,
     from mccode_antlr.compiler.c import run_compiled_instrument, CBinaryTarget
     from .cache import directory_under_module_data_path
     # create a directory for this simulation based on the uuid generated for the simulation entry
-    work_dir = directory_under_module_data_path('sim', prefix=f'p_{instr_file_entry.id}_')
+    work_dir = directory_under_module_data_path('sim', prefix=f'{Path(instr_file_entry.binary_path).parent.stem}_')
 
     binary_at = Path(instr_file_entry.binary_path)
-    target = CBinaryTarget(mpi=False, acc=False, count=1, nexus=False)
+    # process_count == 0 --> system MPI default process count (# physical cores, typically)
+    target = CBinaryTarget(mpi=instr_file_entry.mpi, acc=instr_file_entry.acc, count=process_count, nexus=False)
 
     # ensure the primary spectrometer uses our output directory
     args_dict = {k: v for k, v in args.items() if k != 'dir'}
@@ -368,6 +382,13 @@ def do_primary_simulation(sit: SimulationEntry,
         log.info('Expected mcpl_filename parameter in primary simulation, using default')
         mcpl_filename = f'{sit.id}.mcpl'
         sit.parameter_values['mcpl_filename'] = Value.str(mcpl_filename)
+
+    # strip the extension from the filename passed into the runner:
+    if mcpl_filename.endswith('.gz'):
+        mcpl_filename = mcpl_filename[:-3]
+    if mcpl_filename.endswith('.mcpl'):
+        mcpl_filename = mcpl_filename[:-5]
+
     mcpl_filepath = work_dir.joinpath(mcpl_filename)
     runner = partial(run_compiled_instrument, binary_at, target, capture=False, dry_run=dry_run)
     if dry_run or args.get('ncount') is None:
@@ -392,7 +413,7 @@ def repeat_simulation_until(count, runner, args: dict, parameters, work_dir: Pat
     from functools import partial
     from zenlog import log
     from .emulate import combine_mccode_dats_in_directories, combine_mccode_sims_in_directories
-    from .mcpl import mcpl_particle_count, mcpl_merge_files
+    from .mcpl import mcpl_particle_count, mcpl_merge_files, mcpl_rename_file
     goal, latest_result, one_trillion = count, -1, 1_000_000_000_000
     # avoid looping for too long by limiting the minimum number of particles to simulate
     minimum_particle_count = _clamp(1, one_trillion, minimum_particle_count or count)
@@ -420,21 +441,23 @@ def repeat_simulation_until(count, runner, args: dict, parameters, work_dir: Pat
             args['seed'] = random.randint(1, 2 ** 32 - 1)
 
         outputs.append(work_dir.joinpath(f'{len(files)}'))
-        files.append(work_dir.joinpath(f'part_{len(files)}.mcpl'))
+        files.append(work_dir.joinpath(f'part_{len(files)}'))  # appending the extension here breaks MCPL+MPI?
         args['dir'] = outputs[-1]
         # adjust our guess for how many particles to simulate : how many we need divided by the last transmission
         args['ncount'] = clamp(((goal - sum(counts)) * args['ncount']) // counts[-1] if len(counts) else goal)
-        runner(_args_pars_mcpl(args, parameters, files[-1]))
-        counts.append(mcpl_particle_count(files[-1]))
+        # recycle the intended-output mcpl filename to avoid breaking mcpl file-merging
+        runner(_args_pars_mcpl(args, parameters, mcpl_filepath))
+        counts.append(mcpl_particle_count(mcpl_filepath))
+        # rename the outputfile to this run's filename
+        files[-1] = mcpl_rename_file(mcpl_filepath, files[-1])
 
     # now we need to concatenate the mcpl files, and combine output (.dat and .sim) files
-    mcpl_merge_files(files, str(mcpl_filepath))
+    mcpl_merge_files(files, mcpl_filepath)
     combine_mccode_dats_in_directories(outputs, work_dir)
     combine_mccode_sims_in_directories(outputs, work_dir)
 
 
-def do_secondary_simulation(p_sit: SimulationEntry, entry: InstrEntry, pars: dict, args: dict,
-                            dry_run: bool = False):
+def do_secondary_simulation(p_sit: SimulationEntry, entry: InstrEntry, pars: dict, args: dict, dry_run: bool = False):
     from zenlog import log
     from pathlib import Path
     from shutil import copy
@@ -452,7 +475,7 @@ def do_secondary_simulation(p_sit: SimulationEntry, entry: InstrEntry, pars: dic
 
     mcpl_path = mcpl_real_filename(Path(p_sit.output_path).joinpath(mcpl_filename))
     executable = Path(entry.binary_path)
-    target = CBinaryTarget(mpi=False, acc=False, count=1, nexus=False)
+    target = CBinaryTarget(mpi=entry.mpi, acc=entry.acc, count=1 if not entry.mpi else 0, nexus=False)
     run_compiled_instrument(executable, target, _args_pars_mcpl(args, pars, mcpl_path), capture=False, dry_run=dry_run)
 
     if not dry_run:

restage/tables.py

@@ -316,6 +316,8 @@ class NexusStructureEntry:
 class InstrEntry:
     """A class to represent the instrument file and its compiled binary when stored as an entry in a table"""
     file_contents: str
+    mpi: bool
+    acc: bool
     binary_path: str
     mccode_version: str = field(default_factory=str)
     id: str = field(default_factory=uuid)
@@ -324,8 +326,8 @@ class InstrEntry:
 
     @classmethod
     def from_query_result(cls, values):
-        fid, file_contents, binary_path, mccode_version, creation, last_access = values
-        return cls(file_contents, binary_path, mccode_version, fid, creation, last_access)
+        fid, file_contents, mpi, acc, binary_path, mccode_version, creation, last_access = values
+        return cls(file_contents, mpi != 0, acc != 0, binary_path, mccode_version, fid, creation, last_access)
 
     def __post_init__(self):
         if len(self.mccode_version) == 0:
@@ -334,12 +336,14 @@ class InstrEntry:
 
     @staticmethod
     def columns():
-        return ['id', 'file_contents', 'binary_path', 'mccode_version', 'creation', 'last_access']
+        return ['id', 'file_contents', 'mpi', 'acc', 'binary_path', 'mccode_version', 'creation', 'last_access']
 
     def values(self):
         str_values = [f"'{x}'" for x in (self.id, self.file_contents, self.binary_path, self.mccode_version)]
+        int_values = [f'{x}' for x in (self.mpi, self.acc)]
         flt_values = [f'{self.creation}', f'{self.last_access}']
-        return str_values + flt_values
+        # matches id, file_contents, mpi, acc, binary_path, mccode_version, creation, last_access order
+        return str_values[:2] + int_values + str_values[2:] + flt_values
 
     @classmethod
     def create_sql_table(cls, table_name: str = 'instr_files'):

{restage-0.2.4.dist-info → restage-0.3.1.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: restage
-Version: 0.2.4
+Version: 0.3.1
 Author-email: Gregory Tucker <gregory.tucker@ess.eu>
 Classifier: License :: OSI Approved :: BSD License
 Description-Content-Type: text/markdown
-Requires-Dist: zenlog ==1.1
-Requires-Dist: platformdirs ==3.11
-Requires-Dist: psutil ==5.9.6
-Requires-Dist: mccode-antlr[hdf5] ==0.5.3
+Requires-Dist: zenlog >=1.1
+Requires-Dist: platformdirs >=3.11
+Requires-Dist: psutil >=5.9.6
+Requires-Dist: mccode-antlr[hdf5] ==0.7.0
 Requires-Dist: importlib-metadata ; python_version < "3.8"
 Provides-Extra: test
 Requires-Dist: pytest ; extra == 'test'

{restage-0.2.4.dist-info → restage-0.3.1.dist-info}/RECORD

@@ -1,19 +1,19 @@
 restage/__init__.py,sha256=v0dAUYVkvzjd3j6gjFdbunV_P8U9XxsGgLFGwbxBy6E,787
 restage/bifrost_choppers.py,sha256=aeAw4JgkGPSewU-mqGlGqo1AmN2dmf6FP9u4VDx-eLo,6338
-restage/cache.py,sha256=OAJCOzm6cgPm-SEoG53hzdS7rWUG-qebKZxw_n_xIfQ,8157
+restage/cache.py,sha256=Z60tn372Z_KlXC4e0UqMPPuM0NR8e6eD3aNbdR_q4rM,8374
 restage/cspec_choppers.py,sha256=ZWxyCcwYn4z9ZNqj_r6RC9ImbhVjYc1fmv-Ijm8A2Yk,206
 restage/database.py,sha256=pblHu8hCV5u3uyE8aUrnBSsfjDLYrxy9JRtnRuOvTXQ,9152
 restage/emulate.py,sha256=VrhfZJIbECdbDS-MHklqRuAIy9cRkjZkwPBTKQSQoe0,6164
 restage/energy.py,sha256=ltuMfByx3uwU8IBBriENbIH5EEfKdZ1-w84Urkcs0xM,2664
 restage/instr.py,sha256=G8iaXbgN0JKgv37S1oDuf0mEl2WP9VFuwAGeA0yw0R4,2412
-restage/mcpl.py,sha256=0F9A_edwSR9awQipicI3ipsMAJFk9zCspJhCOrmhkDg,2407
+restage/mcpl.py,sha256=BZYxBytughjc8slR6gUaBy3D7gzo7Yl3ACXrXhWgagI,3403
 restage/range.py,sha256=TjOf4DSKfgoAIcrWQvv6MrtksQpnGJHdsEjVI5K-UfI,8116
 restage/run.py,sha256=nk8d7cIyIqSt-5pyGm68Zak5H1a-fbo_z2_36eN-08E,1481
 restage/scan.py,sha256=Yx8OQSBG6I2_64sW0LIDb0glVKwWoxUQQznASXgDZFQ,1432
-restage/splitrun.py,sha256=QXmQOZ14EUdhBOzLwX1djTg9kPmxqi0bZJHQWsb3mhs,23339
-restage/tables.py,sha256=VMApcu4KtsBlUvAqCSKHDZ58FXhNPihr90OU12YAejw,15764
-restage-0.2.4.dist-info/METADATA,sha256=cv4eP6UBB5bMu0xFnwgKBTf-oUqWXLP9s7K304ETjNM,4773
-restage-0.2.4.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-restage-0.2.4.dist-info/entry_points.txt,sha256=dY3ZanXW6rBLd2vBo-WMcbgO7E10GVDE0t9GASuVjgY,90
-restage-0.2.4.dist-info/top_level.txt,sha256=iM_pb-taTZ0S2WMoDnt_qDMZoNMjmM19z3tTCuVm1IE,8
-restage-0.2.4.dist-info/RECORD,,
+restage/splitrun.py,sha256=PGWfiyVMDLpHkbc_L78sutBhunAN3e5bGijm_rWuDfk,24938
+restage/tables.py,sha256=FJ-8syvzgjdLvvJW9hVK4VV_rNiZhW34-4BMKQsGjtU,16034
+restage-0.3.1.dist-info/METADATA,sha256=pvRu9I6iVkvMP99RXZG7uXEzNdIFXdis3_TB8Epv29k,4773
+restage-0.3.1.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
+restage-0.3.1.dist-info/entry_points.txt,sha256=dY3ZanXW6rBLd2vBo-WMcbgO7E10GVDE0t9GASuVjgY,90
+restage-0.3.1.dist-info/top_level.txt,sha256=iM_pb-taTZ0S2WMoDnt_qDMZoNMjmM19z3tTCuVm1IE,8
+restage-0.3.1.dist-info/RECORD,,

{restage-0.2.4.dist-info → restage-0.3.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.42.0)
+Generator: setuptools (74.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any