From 69b45e8e6cc137228a557a0d191714a756d63fb9 Mon Sep 17 00:00:00 2001 From: fangyuedong Date: Tue, 11 Jun 2024 20:16:24 +0800 Subject: [PATCH] enable array job submission --- observation_sim/ObservationSim.py | 7 +++++-- observation_sim/_util.py | 4 +++- run_sim.py | 3 ++- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/observation_sim/ObservationSim.py b/observation_sim/ObservationSim.py index 1e82700..c230757 100755 --- a/observation_sim/ObservationSim.py +++ b/observation_sim/ObservationSim.py @@ -131,7 +131,7 @@ class Observation(object): chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024))) del chip.img - def runExposure_MPI_PointingList(self, pointing_list, chips=None): + def runExposure_MPI_PointingList(self, pointing_list, chips=None, task_id=None): comm = MPI.COMM_WORLD ind_thread = comm.Get_rank() num_thread = comm.Get_size() @@ -184,7 +184,10 @@ class Observation(object): for ichip in range(nchips_per_fp): i_process = process_counter + ichip - if i_process % num_thread != ind_thread: + if task_id is not None: + if i_process != task_id: + continue + elif i_process % num_thread != ind_thread: continue pid = os.getpid() diff --git a/observation_sim/_util.py b/observation_sim/_util.py index 47c91ce..6dcc91f 100755 --- a/observation_sim/_util.py +++ b/observation_sim/_util.py @@ -16,13 +16,15 @@ def parse_args(): parser.add_argument('--config_file', type=str, required=True, help='.yaml config file for simulation settings.') parser.add_argument('--catalog', type=str, - help='name of the catalog interface class to be loaded.') + help='Name of the catalog interface class to be loaded.') parser.add_argument('-c', '--config_dir', type=str, help='Directory that houses the .yaml config file.') parser.add_argument('-d', '--data_dir', type=str, help='Directory that houses the input data.') parser.add_argument('-w', '--work_dir', type=str, help='The path for output.') + parser.add_argument('--task_id', type=int, + help='The task id to run a particular chip, used when running with job array') return parser.parse_args() diff --git a/run_sim.py b/run_sim.py index f117985..594f040 100755 --- a/run_sim.py +++ b/run_sim.py @@ -103,7 +103,8 @@ def run_sim(): work_dir=config['work_dir'], data_dir=config['data_dir']) # Run simulation - obs.runExposure_MPI_PointingList(pointing_list=pointing_list) + obs.runExposure_MPI_PointingList( + pointing_list=pointing_list, task_id=args.task_id) if __name__ == '__main__': -- GitLab