diff --git a/rsopt/configuration/jobs.py b/rsopt/configuration/jobs.py index 57a1870..ce22be8 100644 --- a/rsopt/configuration/jobs.py +++ b/rsopt/configuration/jobs.py @@ -25,12 +25,15 @@ def create_executor_arguments(setup: Setup, is_parallel: bool) -> dict: # Really creates Executor.submit() arguments if is_parallel: args = { - 'num_procs': setup.setup.get('cores', 1), + 'num_procs': None if setup.setup.get('gpu', False) else setup.setup.get('cores', 1), 'num_nodes': None, # No user interface right now 'procs_per_node': None, # No user interface right now 'machinefile': None, # Add in setup.machinefile if user wants to control 'app_args': setup.format_task_string(is_parallel), 'hyperthreads': False, # Add in setup.hyperthreads if this is needed + 'wait_on_start': True, + 'auto_assign_gpus': setup.setup.get('gpu', False), + 'match_procs_to_gpus': setup.setup.get('gpu', False), # 'app_name': None, # Handled at optimizer setup # 'stdout': None, # Handled at optimizer setup # 'stderr': None, # Handled at optimizer setup @@ -42,6 +45,9 @@ def create_executor_arguments(setup: Setup, is_parallel: bool) -> dict: args = { 'app_args': setup.format_task_string(is_parallel) } + # TODO: Should there really be this kind of hidden interface to override executor setup? + # for key, value in args.items(): + # args[key] = setup.setup.get(key, value) # Cannot be overridden args['calc_type'] = 'sim' diff --git a/rsopt/package_data/setup_schema.yml b/rsopt/package_data/setup_schema.yml index 27e2522..5053986 100644 --- a/rsopt/package_data/setup_schema.yml +++ b/rsopt/package_data/setup_schema.yml @@ -12,3 +12,4 @@ - ignored_files - shifter_image - code_arguments +- gpu diff --git a/rsopt/simulation.py b/rsopt/simulation.py index 8418a50..5a96f9a 100644 --- a/rsopt/simulation.py +++ b/rsopt/simulation.py @@ -154,6 +154,11 @@ def __call__(self, H, persis_info, sim_specs, libE_info): # NOTE: Right now f is not passed to the objective function. Would need to go inside J. Or pass J into # function job.execute(**kwargs) + if job.executor_args.get('auto_assign_gpus'): + from libensemble.tools.test_support import check_gpu_setting + if self.log.level == logging.DEBUG: + self.log.debug(check_gpu_setting(task, assert_setting=False, print_setting=True)) + if halt_job_sequence: break